{"version":1,"extractor":"496c9d83e6eb3668","files":{"scripts/build_data.py":{"size":2901,"mtimeMs":1789275830304.2024,"hash":"d0873fc2ab964d945c3ca9108ef62c4c71542f63bfa04ff6dfa512f34c148c97","nodes":[{"id":"scripts/build_data.py","name":"build_data.py","kind":"file","path":"scripts/build_data.py","span":"L1-L72","signature":null,"exported":true,"origin":"ast","body_hash":"d0873fc2ab964d945c3ca9108ef62c4c71542f63bfa04ff6dfa512f34c148c97","chars":2901,"summary_state":"pending","summary":null,"crux":null,"body_text":"\"\"\"Build the paper's NQ + HotpotQA training data from Search-R1-compatible JSONL. Expected JSONL fields: ``question`` and ``golden_answers``. The input directory must contain ``nq/{train,test}.jsonl`` and ``hotpotqa/{train,dev}.jsonl``. \"\"\" import argparse from pathlib import Path import datasets as hf_datasets SPLITS = { \"nq\": (\"nq\", \"train.jsonl\", \"test.jsonl\"), \"hotpotqa\": (\"hotpotqa\", \"train.jsonl\", \"dev.jsonl\"), } if __name__ == \"__main__\": main()"},{"id":"scripts/build_data.py#make_prompt","name":"make_prompt","kind":"function","path":"scripts/build_data.py","span":"L18-L29","signature":"def make_prompt(question: str) -> str","exported":true,"origin":"ast","body_hash":"7f871974966d4eb477938a9641af876b27bce5279aaba506b4ba520375d1aba0","body_text":"def make_prompt(question: str) -> str: return ( \"Answer the given question. You can search for external knowledge using \" \"<search> query </search>. Search results will be returned between \" \"<information> and </information>. You can search as many times as you want. \" \"After receiving search results, first output your current confidence in \" \"<confidence> X </confidence> (integer 1-10), then either search again or \" \"give your answer. When giving your final answer, write \" \"<answer> your answer </answer> followed by <final-confidence> X \" \"</final-confidence> (integer 1-10). \" f\"Question: {question}\" )","summary_state":"pending","summary":null,"crux":null},{"id":"scripts/build_data.py#load_and_convert","name":"load_and_convert","kind":"function","path":"scripts/build_data.py","span":"L32-L47","signature":"def load_and_convert(source: Path, dataset_name: str, split: str)","exported":true,"origin":"ast","body_hash":"22135b108020b0ecd42a33759f1eca74f7870e8bbf22ee87c71c155a1d7f0240","body_text":"def load_and_convert(source: Path, dataset_name: str, split: str): dataset = hf_datasets.load_dataset(\"json\", data_files=str(source), split=\"train\") def convert(example, index): question = example[\"question\"].strip() if not question.endswith(\"?\"): question += \"?\" return { \"data_source\": dataset_name, \"prompt\": [{\"role\": \"user\", \"content\": make_prompt(question)}], \"ability\": \"fact-reasoning\", \"reward_model\": {\"style\": \"rule\", \"ground_truth\": {\"target\": example[\"golden_answers\"]}}, \"extra_info\": {\"split\": split, \"index\": index, \"dataset\": dataset_name}, } return dataset.map(convert, with_indices=True, remove_columns=dataset.column_names)","summary_state":"pending","summary":null,"crux":null},{"id":"scripts/build_data.py#load_and_convert.convert","name":"convert","kind":"function","path":"scripts/build_data.py","span":"L35-L45","signature":"def convert(example, index)","exported":true,"origin":"ast","body_hash":"62c6a7e8a0dd606b144f6d591104784f4b0618dec64a07ef04a3249f5705db2b","body_text":"def convert(example, index): question = example[\"question\"].strip() if not question.endswith(\"?\"): question += \"?\" return { \"data_source\": dataset_name, \"prompt\": [{\"role\": \"user\", \"content\": make_prompt(question)}], \"ability\": \"fact-reasoning\", \"reward_model\": {\"style\": \"rule\", \"ground_truth\": {\"target\": example[\"golden_answers\"]}}, \"extra_info\": {\"split\": split, \"index\": index, \"dataset\": dataset_name}, }","summary_state":"pending","summary":null,"crux":null},{"id":"scripts/build_data.py#main","name":"main","kind":"function","path":"scripts/build_data.py","span":"L50-L67","signature":"def main() -> None","exported":true,"origin":"ast","body_hash":"9416f796320ba5e26c26820381b4f5b2a102e2dae355fb87429387573fd5f585","body_text":"def main() -> None: parser = argparse.ArgumentParser() parser.add_argument(\"--raw-root\", type=Path, required=True, help=\"Directory containing NQ and HotpotQA JSONL files.\") parser.add_argument(\"--output-dir\", type=Path, default=Path(\"data/nq_hotpotqa\")) args = parser.parse_args() converted = {\"train\": [], \"test\": []} for dataset_name, (subdir, train_name, test_name) in SPLITS.items(): for split, filename in ((\"train\", train_name), (\"test\", test_name)): source = args.raw_root / subdir / filename if not source.is_file(): raise FileNotFoundError(f\"Missing input file: {source}\") converted[split].append(load_and_convert(source, dataset_name, split)) args.output_dir.mkdir(parents=True, exist_ok=True) for split, parts in converted.items(): combined = hf_datasets.concatenate_datasets(parts) combined.to_parquet(str(args.output_dir / f\"{split}.parquet\"))","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"scripts/build_data.py","relation":"imports","specifier":"argparse","file":"scripts/build_data.py"},{"source":"scripts/build_data.py","relation":"imports","specifier":"pathlib","file":"scripts/build_data.py"},{"source":"scripts/build_data.py","relation":"contains","targetId":"scripts/build_data.py#make_prompt","file":"scripts/build_data.py"},{"source":"scripts/build_data.py","relation":"contains","targetId":"scripts/build_data.py#load_and_convert","file":"scripts/build_data.py"},{"source":"scripts/build_data.py#load_and_convert","relation":"calls","name":"load_dataset","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#load_and_convert","relation":"calls","name":"str","viaMember":false,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#load_and_convert","relation":"contains","targetId":"scripts/build_data.py#load_and_convert.convert","file":"scripts/build_data.py"},{"source":"scripts/build_data.py#load_and_convert.convert","relation":"calls","name":"strip","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#load_and_convert.convert","relation":"calls","name":"endswith","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#load_and_convert.convert","relation":"calls","name":"make_prompt","viaMember":false,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#load_and_convert","relation":"calls","name":"map","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py","relation":"contains","targetId":"scripts/build_data.py#main","file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"Path","viaMember":false,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"parse_args","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"items","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"is_file","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"FileNotFoundError","viaMember":false,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"append","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"load_and_convert","viaMember":false,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"mkdir","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"items","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"concatenate_datasets","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"to_parquet","viaMember":true,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py#main","relation":"calls","name":"str","viaMember":false,"file":"scripts/build_data.py"},{"source":"scripts/build_data.py","relation":"calls","name":"main","viaMember":false,"file":"scripts/build_data.py"}]},"scripts/calc_last_step_conf.py":{"size":6960,"mtimeMs":1789275830304.2024,"hash":"bda4281f0b370aff1639de51c865002f930271ee0ecd2d5c68195303e4075305","nodes":[{"id":"scripts/calc_last_step_conf.py","name":"calc_last_step_conf.py","kind":"file","path":"scripts/calc_last_step_conf.py","span":"L1-L181","signature":null,"exported":true,"origin":"ast","body_hash":"bda4281f0b370aff1639de51c865002f930271ee0ecd2d5c68195303e4075305","chars":6426,"summary_state":"pending","summary":null,"crux":null,"body_text":"\"\"\" 计算 final confidence 和 last-step confidence 的 ECE / AUROC / Brier Score。 过滤条件： - final confidence 有效：confidence >= 1（排除 confidence=0，即模型未输出 <final-confidence> 标签的样本） - last-step confidence 有效：per_step_conf_values 中 answer 步骤有非 None 值 - 只保留两者同时有效的样本（交集） 格式正确（has_valid_format=1）的子集单独统计。 \"\"\" import argparse import ast import sys from pathlib import Path import numpy as np import pandas as pd sys.path.insert(0, str(Path(__file__).resolve().parent)) from compute_metrics_ext import compute_ece, compute_auroc, compute_brier_score, print_bin_detail if __name__ == \"__main__\": main()"},{"id":"scripts/calc_last_step_conf.py#get_last_step_conf","name":"get_last_step_conf","kind":"function","path":"scripts/calc_last_step_conf.py","span":"L23-L46","signature":"def get_last_step_conf(row)","exported":true,"origin":"ast","body_hash":"aa34e3da9e7f2192972f1da1dbd037334c0b7fe4151bcb1ba8626dd9753f1c46","body_text":"def get_last_step_conf(row): \"\"\" 从 per_step_conf_values 和 per_step_action_types 中提取 answer 步骤的置信度。 优先取最后一个 action_type=='answer' 且 conf 非 None 的步骤； 若无，取最后一个非 None 的置信度。 \"\"\" try: conf_values = ast.literal_eval(str(row[\"per_step_conf_values\"])) action_types = ast.literal_eval(str(row[\"per_step_action_types\"])) except Exception: return None if not conf_values or not action_types: return None for i in range(len(action_types) - 1, -1, -1): if i < len(conf_values) and action_types[i] == \"answer\" and conf_values[i] is not None: return conf_values[i] for i in range(len(conf_values) - 1, -1, -1): if conf_values[i] is not None: return conf_values[i] return None","summary_state":"pending","summary":null,"crux":null},{"id":"scripts/calc_last_step_conf.py#print_metrics","name":"print_metrics","kind":"function","path":"scripts/calc_last_step_conf.py","span":"L49-L80","signature":"def print_metrics(label, confs, is_correct, n_bins, detail=False)","exported":true,"origin":"ast","body_hash":"dfb09f8b386827e14def92aa2e87419540020e0abcab6bfae68252cb4ce0ff0f","body_text":"def print_metrics(label, confs, is_correct, n_bins, detail=False): \"\"\"打印一组样本的校准指标。\"\"\" n = len(confs) if n == 0: print(f\"[{label}] 无有效样本，跳过。\") return overall_acc = is_correct.mean() avg_conf = confs.mean() ece = compute_ece(confs, is_correct, n_bins=n_bins) auroc = compute_auroc(confs, is_correct) brier = compute_brier_score(confs, is_correct) print(f\"\\n{'='*55}\") print(f\" {label} (n={n})\") print(f\"{'='*55}\") print(f\" 准确率: {overall_acc:.4f}\") print(f\" 平均置信度(归一化): {avg_conf:.4f}\") print(f\" 置信度范围: [{confs.min():.2f}, {confs.max():.2f}]\") print(f\" ECE (↓ 20桶): {ece:.4f}\") print(f\" AUROC (↑): {auroc:.4f}\") print(f\" Brier Score (↓): {brier:.4f}\") correct_confs = confs[is_correct == 1] incorrect_confs = confs[is_correct == 0] if len(correct_confs) > 0: print(f\" 答对 {len(correct_confs)} 样本，平均置信度: {correct_confs.mean():.4f}\") if len(incorrect_confs) > 0: print(f\" 答错 {len(incorrect_confs)} 样本，平均置信度: {incorrect_confs.mean():.4f}\") if detail: print_bin_detail(confs, is_correct, n_bins=n_bins)","summary_state":"pending","summary":null,"crux":null},{"id":"scripts/calc_last_step_conf.py#main","name":"main","kind":"function","path":"scripts/calc_last_step_conf.py","span":"L83-L176","signature":"def main()","exported":true,"origin":"ast","body_hash":"287c5c42f9869eb8e07a3628c8fc2d38c007559bf342680017b9754a175623d0","body_text":"def main(): parser = argparse.ArgumentParser() parser.add_argument(\"--csv\", type=str, required=True) parser.add_argument(\"--n_bins\", type=int, default=20) parser.add_argument(\"--detail\", action=\"store_true\") args = parser.parse_args() print(f\"\\n读取文件：{args.csv}\") df = pd.read_csv(args.csv) n_total = len(df) print(f\"总样本数: {n_total}\") print(f\"has_valid_format 分布: {dict(df['has_valid_format'].value_counts().sort_index())}\") print(f\"confidence 分布: {dict(df['confidence'].value_counts().sort_index())}\") # ------------------------------------------------------------------ # # 1. 提取 last-step confidence # ------------------------------------------------------------------ # df[\"last_step_conf\"] = df.apply(get_last_step_conf, axis=1) df[\"last_step_conf\"] = pd.to_numeric(df[\"last_step_conf\"], errors=\"coerce\") # ------------------------------------------------------------------ # # 2. 过滤：同时有效的样本（final conf >= 1 AND last_step_conf 非空） # ------------------------------------------------------------------ # mask_final = df[\"confidence\"] >= 1 # 排除 confidence=0（未输出标签） mask_last = df[\"last_step_conf\"].notna() # last-step 有值 mask_both = mask_final & mask_last df_both = df[mask_both].copy() n_both = len(df_both) print(f\"\\n同时有 final conf(>=1) 且有 last-step conf 的样本: {n_both} / {n_total} \" f\"({n_both/n_total*100:.1f}%)\") print(f\" 其中 has_valid_format=1: {df_both['has_valid_format'].sum()} / {n_both}\") # ------------------------------------------------------------------ # # 3. 归一化 /10 # ------------------------------------------------------------------ # df_both[\"final_conf_norm\"] = df_both[\"confidence\"].astype(float) / 10.0 df_both[\"last_conf_norm\"] = df_both[\"last_step_conf\"].astype(float) / 10.0 # ------------------------------------------------------------------ # # 4. 分两组：全部 & 格式正确 # ------------------------------------------------------------------ # groups = { \"【全部有效样本】\": df_both, \"【格式正确(has_valid_format=1)】\": df_both[df_both[\"has_valid_format\"] == 1], } print(\"\\n\\n\" + \"#\"*60) print(\"# Final Confidence (answer步骤的 <final-confidence> 标签)\") print(\"#\"*60) for label, sub in groups.items(): print_metrics( label, sub[\"final_conf_norm\"].values, sub[\"is_correct\"].values.astype(float), args.n_bins, args.detail, ) print(\"\\n\\n\" + \"#\"*60) print(\"# Last-Step Confidence (per_step_conf_values 中 answer 步骤的 <confidence> 标签)\") print(\"#\"*60) for label, sub in groups.items(): print_metrics( label, sub[\"last_conf_norm\"].values, sub[\"is_correct\"].values.astype(float), args.n_bins, args.detail, ) # ------------------------------------------------------------------ # # 5. 汇总对比表 # ------------------------------------------------------------------ # print(\"\\n\\n\" + \"=\"*70) print(\"Summary comparison\") print(\"=\"*70) header = f\"{'子集':<28} {'类型':<12} {'ECE':>7} {'AUROC':>7} {'Brier':>7} {'n':>6}\" print(header) print(\"-\"*70) for label, sub in groups.items(): n = len(sub) if n == 0: continue for conf_type, col in [(\"Final Conf\", \"final_conf_norm\"), (\"Last-Step\", \"last_conf_norm\")]: c = sub[col].values y = sub[\"is_correct\"].values.astype(float) ece = compute_ece(c, y, n_bins=args.n_bins) auroc = compute_auroc(c, y) brier = compute_brier_score(c, y) short_label = label.replace(\"【\", \"\").replace(\"】\", \"\") print(f\"{short_label:<28} {conf_type:<12} {ece:>7.4f} {auroc:>7.4f} {brier:>7.4f} {n:>6}\") print()","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"scripts/calc_last_step_conf.py","relation":"imports","specifier":"argparse","file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"imports","specifier":"ast","file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"imports","specifier":"sys","file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"imports","specifier":"pathlib","file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"calls","name":"insert","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"calls","name":"str","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"calls","name":"resolve","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"calls","name":"Path","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"imports","specifier":"compute_metrics_ext","file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"contains","targetId":"scripts/calc_last_step_conf.py#get_last_step_conf","file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#get_last_step_conf","relation":"calls","name":"literal_eval","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#get_last_step_conf","relation":"calls","name":"str","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#get_last_step_conf","relation":"calls","name":"literal_eval","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#get_last_step_conf","relation":"calls","name":"str","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#get_last_step_conf","relation":"calls","name":"range","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#get_last_step_conf","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#get_last_step_conf","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#get_last_step_conf","relation":"calls","name":"range","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#get_last_step_conf","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"contains","targetId":"scripts/calc_last_step_conf.py#print_metrics","file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"mean","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"mean","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"compute_ece","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"compute_auroc","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"compute_brier_score","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"min","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"max","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"mean","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"mean","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#print_metrics","relation":"calls","name":"print_bin_detail","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"contains","targetId":"scripts/calc_last_step_conf.py#main","file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"parse_args","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"read_csv","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"dict","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"sort_index","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"value_counts","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"dict","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"sort_index","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"value_counts","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"apply","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"to_numeric","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"notna","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"copy","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"sum","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"astype","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"astype","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"items","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print_metrics","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"astype","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"items","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print_metrics","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"astype","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"items","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"len","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"astype","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"compute_ece","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"compute_auroc","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"compute_brier_score","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"replace","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"replace","viaMember":true,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/calc_last_step_conf.py"},{"source":"scripts/calc_last_step_conf.py","relation":"calls","name":"main","viaMember":false,"file":"scripts/calc_last_step_conf.py"}]},"scripts/compute_metrics_ext.py":{"size":8550,"mtimeMs":1789275830304.2024,"hash":"edb48fbecb6ec160c2d4fe7f7ea62b9b33d128d7ed3ff67a2f62203f22257df7","nodes":[{"id":"scripts/compute_metrics_ext.py","name":"compute_metrics_ext.py","kind":"file","path":"scripts/compute_metrics_ext.py","span":"L1-L220","signature":null,"exported":true,"origin":"ast","body_hash":"edb48fbecb6ec160c2d4fe7f7ea62b9b33d128d7ed3ff67a2f62203f22257df7","chars":7933,"summary_state":"pending","summary":null,"crux":null,"body_text":"\"\"\" compute_metrics_ext.py 从 evaluate_batch_conf_ext.sh 生成的评估产物目录中，计算： 1. bamboogle：读取 4 次 rep1~rep4 的 CSV，合并后整体计算指标（等效取均值） 2. simpleqa ：读取单次 CSV，直接计算指标 每个数据集输出与 train.md 表格一致的指标： acc | final-conf ECE↓ | final-conf AUROC↑ | final-conf BS↓ | last-step-conf ECE↓ | last-step-conf AUROC↑ | last-step-conf BS↓ 用法： # 计算单个模型产物目录 python scripts/compute_metrics_ext.py --eval_dir evaluation/<model_name>_ext_<timestamp> # 批量计算所有 *_ext_* 目录 python scripts/compute_metrics_ext.py --eval_root evaluation --all \"\"\" import argparse import ast import os import sys import glob import numpy as np import pandas as pd sys.path.insert(0, os.path.join(os.path.dirname(__file__), '../statistics')) from confidence_metrics import compute_ece, compute_auroc, compute_brier_score # ============================================================ # last-step confidence 提取（同 calc_last_step_conf.py） # ============================================================ # ============================================================ # 从 sample_statistics 目录中找 CSV # ============================================================ # ============================================================ # 计算一个 DataFrame 的全套指标 # ============================================================ # ============================================================ # 处理单个 eval_dir（一个模型的扩展评估产物目录） # ============================================================ # ============================================================ # 主程序 # ============================================================ if __name__ == \"__main__\": main()"},{"id":"scripts/compute_metrics_ext.py#get_last_step_conf","name":"get_last_step_conf","kind":"function","path":"scripts/compute_metrics_ext.py","span":"L35-L49","signature":"def get_last_step_conf(row)","exported":true,"origin":"ast","body_hash":"99bb16e6a45f5babdc87f8283fe66fed38ce3450d2b7c049d4df270d39228340","body_text":"def get_last_step_conf(row): try: conf_values = ast.literal_eval(str(row[\"per_step_conf_values\"])) action_types = ast.literal_eval(str(row[\"per_step_action_types\"])) except Exception: return None if not conf_values or not action_types: return None for i in range(len(action_types) - 1, -1, -1): if i < len(conf_values) and action_types[i] == \"answer\" and conf_values[i] is not None: return conf_values[i] for i in range(len(conf_values) - 1, -1, -1): if conf_values[i] is not None: return conf_values[i] return None","summary_state":"pending","summary":null,"crux":null},{"id":"scripts/compute_metrics_ext.py#find_csv","name":"find_csv","kind":"function","path":"scripts/compute_metrics_ext.py","span":"L55-L65","signature":"def find_csv(stats_dir: str) -> str | None","exported":true,"origin":"ast","body_hash":"378f3c1b0834e5f3817cfdff370bc47079b09cf99169fbcd86a8e3ce1bbff43f","body_text":"def find_csv(stats_dir: str) -> str | None: \"\"\"在 stats_dir/sample_statistics/ 下找到第一个 *_val.csv 文件。\"\"\" pattern = os.path.join(stats_dir, \"sample_statistics\", \"*_val.csv\") files = glob.glob(pattern) if not files: # 有时直接在 stats_dir 下 pattern2 = os.path.join(stats_dir, \"sample_statistics\", \"*.csv\") files = glob.glob(pattern2) if files: return sorted(files)[0] return None","summary_state":"pending","summary":null,"crux":null},{"id":"scripts/compute_metrics_ext.py#compute_metrics_df","name":"compute_metrics_df","kind":"function","path":"scripts/compute_metrics_ext.py","span":"L71-L115","signature":"def compute_metrics_df(df: pd.DataFrame, n_bins: int = 20) -> dict","exported":true,"origin":"ast","body_hash":"3a83e84eb458894ccfb5d418e44c2e98f589381eb87afb3387c1ce077d0a7a59","body_text":"def compute_metrics_df(df: pd.DataFrame, n_bins: int = 20) -> dict: \"\"\" 返回字典： acc, n, fc_ece, fc_auroc, fc_bs, # final-conf lsc_ece, lsc_auroc, lsc_bs # last-step-conf \"\"\" # last-step conf df = df.copy() df[\"last_step_conf\"] = df.apply(get_last_step_conf, axis=1) df[\"last_step_conf\"] = pd.to_numeric(df[\"last_step_conf\"], errors=\"coerce\") # 过滤：final conf valid (confidence_valid==1 且 confidence>=1) mask_fc = (df[\"confidence_valid\"] == 1) & (pd.to_numeric(df[\"confidence\"], errors=\"coerce\") >= 1) mask_lsc = df[\"last_step_conf\"].notna() mask_both = mask_fc & mask_lsc df_v = df[mask_both].copy() df_v[\"confidence\"] = pd.to_numeric(df_v[\"confidence\"], errors=\"coerce\") df_v[\"fc_norm\"] = df_v[\"confidence\"].astype(float) / 10.0 df_v[\"lsc_norm\"] = df_v[\"last_step_conf\"].astype(float) / 10.0 n = len(df_v) if n == 0: nan = float(\"nan\") return dict(acc=nan, n=0, fc_ece=nan, fc_auroc=nan, fc_bs=nan, lsc_ece=nan, lsc_auroc=nan, lsc_bs=nan) y = df_v[\"is_correct\"].values.astype(float) fc = df_v[\"fc_norm\"].values lc = df_v[\"lsc_norm\"].values acc = df[\"is_correct\"].values.astype(float).mean() # 整体 acc（不过滤） return dict( acc = float(acc), n = n, fc_ece = compute_ece(fc, y, n_bins), fc_auroc= compute_auroc(fc, y), fc_bs = compute_brier_score(fc, y), lsc_ece = compute_ece(lc, y, n_bins), lsc_auroc= compute_auroc(lc, y), lsc_bs = compute_brier_score(lc, y), )","summary_state":"pending","summary":null,"crux":null},{"id":"scripts/compute_metrics_ext.py#process_eval_dir","name":"process_eval_dir","kind":"function","path":"scripts/compute_metrics_ext.py","span":"L121-L178","signature":"def process_eval_dir(eval_dir: str, n_bins: int = 20)","exported":true,"origin":"ast","body_hash":"f6e5c54f72db369b14cb00d9f2b49261cb4b9016b514a89b86fb8e84ee4c5059","body_text":"def process_eval_dir(eval_dir: str, n_bins: int = 20): model_name = os.path.basename(eval_dir) print(f\"\\n{'='*70}\") print(f\"模型目录: {model_name}\") print(f\"{'='*70}\") results = {} # ---- bamboogle: 收集所有 rep 的 CSV ---- bamboogle_csvs = [] for rep in range(1, 10): rep_dir = os.path.join(eval_dir, \"bamboogle\", f\"rep{rep}\") if not os.path.isdir(rep_dir): break csv_path = find_csv(rep_dir) if csv_path: bamboogle_csvs.append(csv_path) else: print(f\" ⚠ bamboogle rep{rep}: 未找到 CSV（{rep_dir}）\") if bamboogle_csvs: n_reps = len(bamboogle_csvs) print(f\"\\nbamboogle: 找到 {n_reps} 次重复\") dfs = [pd.read_csv(f) for f in bamboogle_csvs] df_all = pd.concat(dfs, ignore_index=True) m = compute_metrics_df(df_all, n_bins) results[\"bamboogle\"] = m print(f\" 合并样本数 (valid conf): {m['n']} (每次 {len(dfs[0])} 条 × {n_reps} 次)\") _print_row(\"bamboogle\", m) else: print(\" ⚠ bamboogle: 无 CSV，跳过\") # ---- simpleqa: 单次 ---- simpleqa_dir = os.path.join(eval_dir, \"simpleqa\") if os.path.isdir(simpleqa_dir): csv_path = find_csv(simpleqa_dir) if csv_path: print(f\"\\nsimpleqa: {csv_path}\") df = pd.read_csv(csv_path) m = compute_metrics_df(df, n_bins) results[\"simpleqa\"] = m print(f\" 总样本数: {len(df)} 有效 conf: {m['n']}\") _print_row(\"simpleqa\", m) else: print(f\" ⚠ simpleqa: 未找到 CSV（{simpleqa_dir}）\") else: print(f\" ⚠ simpleqa 目录不存在: {simpleqa_dir}\") # ---- Markdown 格式输出 ---- print(f\"\\n--- Markdown 表格行（可直接贴入 train.md）---\") print(f\"| | acc | final-conf ECE↓ | final-conf AUROC↑ | final-conf BS↓ | last-step-conf ECE↓ | last-step-conf AUROC↑ | last-step-conf BS↓ |\") print(f\"|---|---|---|---|---|---|---|---|\") for ds, m in results.items(): n_label = f\"{m['n']}\" if m['n'] else \"?\" label = f\"Bamboogle ({len(bamboogle_csvs)}次均)\" if ds == \"bamboogle\" else \"SimpleQA\" print(f\"| {label} | {m['acc']:.3f} | {m['fc_ece']:.4f} | {m['fc_auroc']:.4f} | {m['fc_bs']:.4f} | {m['lsc_ece']:.4f} | {m['lsc_auroc']:.4f} | {m['lsc_bs']:.4f} |\") return results","summary_state":"pending","summary":null,"crux":null},{"id":"scripts/compute_metrics_ext.py#_print_row","name":"_print_row","kind":"function","path":"scripts/compute_metrics_ext.py","span":"L181-L184","signature":"def _print_row(ds, m)","exported":false,"origin":"ast","body_hash":"8cc7892e35a379ac8663b203abbbe83286af8625e5717f729ade2dbda6740cd0","body_text":"def _print_row(ds, m): print(f\" acc={m['acc']:.3f} \" f\"fc_ECE={m['fc_ece']:.4f} fc_AUROC={m['fc_auroc']:.4f} fc_BS={m['fc_bs']:.4f} \" f\"lsc_ECE={m['lsc_ece']:.4f} lsc_AUROC={m['lsc_auroc']:.4f} lsc_BS={m['lsc_bs']:.4f}\")","summary_state":"pending","summary":null,"crux":null},{"id":"scripts/compute_metrics_ext.py#main","name":"main","kind":"function","path":"scripts/compute_metrics_ext.py","span":"L190-L215","signature":"def main()","exported":true,"origin":"ast","body_hash":"0288a0e2b8e2fa34cb7b3e78387efc3eb44bd760b55f46eeb074185f9fd83032","body_text":"def main(): parser = argparse.ArgumentParser(description=\"计算 bamboogle/simpleqa 扩展评估指标\") group = parser.add_mutually_exclusive_group(required=True) group.add_argument(\"--eval_dir\", type=str, help=\"单个模型的 evaluation 产物目录（如 evaluation/way1-stake_ext_20260515_120000）\") group.add_argument(\"--eval_root\", type=str, help=\"evaluation 根目录，配合 --all 批量处理所有 *_ext_* 子目录\") parser.add_argument(\"--all\", action=\"store_true\", help=\"批量处理 eval_root 下所有 *_ext_* 目录\") parser.add_argument(\"--n_bins\", type=int, default=20) args = parser.parse_args() if args.eval_dir: process_eval_dir(args.eval_dir, args.n_bins) else: if not args.all: print(\"请指定 --all 以批量处理，或直接用 --eval_dir 指定单目录\") sys.exit(1) pattern = os.path.join(args.eval_root, \"*_ext_*\") dirs = sorted(glob.glob(pattern)) if not dirs: print(f\"未找到 *_ext_* 目录 in {args.eval_root}\") sys.exit(1) for d in dirs: if os.path.isdir(d): process_eval_dir(d, args.n_bins)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"scripts/compute_metrics_ext.py","relation":"imports","specifier":"argparse","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"imports","specifier":"ast","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"imports","specifier":"os","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"imports","specifier":"sys","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"imports","specifier":"glob","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"calls","name":"insert","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"calls","name":"join","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"calls","name":"dirname","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"imports","specifier":"confidence_metrics","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"contains","targetId":"scripts/compute_metrics_ext.py#get_last_step_conf","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#get_last_step_conf","relation":"calls","name":"literal_eval","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#get_last_step_conf","relation":"calls","name":"str","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#get_last_step_conf","relation":"calls","name":"literal_eval","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#get_last_step_conf","relation":"calls","name":"str","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#get_last_step_conf","relation":"calls","name":"range","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#get_last_step_conf","relation":"calls","name":"len","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#get_last_step_conf","relation":"calls","name":"len","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#get_last_step_conf","relation":"calls","name":"range","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#get_last_step_conf","relation":"calls","name":"len","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"contains","targetId":"scripts/compute_metrics_ext.py#find_csv","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#find_csv","relation":"calls","name":"join","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#find_csv","relation":"calls","name":"glob","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#find_csv","relation":"calls","name":"join","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#find_csv","relation":"calls","name":"glob","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#find_csv","relation":"calls","name":"sorted","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"contains","targetId":"scripts/compute_metrics_ext.py#compute_metrics_df","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"copy","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"apply","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"to_numeric","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"to_numeric","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"notna","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"copy","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"to_numeric","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"astype","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"astype","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"len","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"float","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"dict","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"astype","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"mean","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"astype","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"dict","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"float","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"compute_ece","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"compute_auroc","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"compute_brier_score","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"compute_ece","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"compute_auroc","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#compute_metrics_df","relation":"calls","name":"compute_brier_score","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"contains","targetId":"scripts/compute_metrics_ext.py#process_eval_dir","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"basename","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"range","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"join","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"isdir","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"find_csv","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"append","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"len","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"read_csv","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"concat","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"compute_metrics_df","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"len","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"_print_row","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"join","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"isdir","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"find_csv","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"read_csv","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"compute_metrics_df","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"len","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"_print_row","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"items","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"len","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#process_eval_dir","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"contains","targetId":"scripts/compute_metrics_ext.py#_print_row","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#_print_row","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"contains","targetId":"scripts/compute_metrics_ext.py#main","file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"add_mutually_exclusive_group","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"parse_args","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"process_eval_dir","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"exit","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"join","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"sorted","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"glob","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"print","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"exit","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"isdir","viaMember":true,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py#main","relation":"calls","name":"process_eval_dir","viaMember":false,"file":"scripts/compute_metrics_ext.py"},{"source":"scripts/compute_metrics_ext.py","relation":"calls","name":"main","viaMember":false,"file":"scripts/compute_metrics_ext.py"}]},"search_r1/__init__.py":{"size":0,"mtimeMs":1789275830304.2024,"hash":"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855","nodes":[{"id":"search_r1/__init__.py","name":"__init__.py","kind":"file","path":"search_r1/__init__.py","span":"L1-L1","signature":null,"exported":true,"origin":"ast","body_hash":"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855","chars":0,"summary_state":"pending","summary":null,"crux":null,"body_text":""}],"rawEdges":[]},"search_r1/llm_agent/__init__.py":{"size":0,"mtimeMs":1789275830304.2024,"hash":"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855","nodes":[{"id":"search_r1/llm_agent/__init__.py","name":"__init__.py","kind":"file","path":"search_r1/llm_agent/__init__.py","span":"L1-L1","signature":null,"exported":true,"origin":"ast","body_hash":"e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855","chars":0,"summary_state":"pending","summary":null,"crux":null,"body_text":""}],"rawEdges":[]},"search_r1/llm_agent/generation.py":{"size":32961,"mtimeMs":1789275830305.2024,"hash":"8b26cc6df6c2d0ffa17d3757916aa5392d881b14962567362f93a1599166447a","nodes":[{"id":"search_r1/llm_agent/generation.py","name":"generation.py","kind":"file","path":"search_r1/llm_agent/generation.py","span":"L1-L692","signature":null,"exported":true,"origin":"ast","body_hash":"8b26cc6df6c2d0ffa17d3757916aa5392d881b14962567362f93a1599166447a","chars":32948,"summary_state":"pending","summary":null,"crux":null,"body_text":"import torch import re import numpy as np from collections import defaultdict import os from typing import List, Dict, Any, Tuple from dataclasses import dataclass from .tensor_helper import TensorHelper, TensorConfig from verl import DataProto from verl.utils.tracking import Tracking import shutil import requests @dataclass"},{"id":"search_r1/llm_agent/generation.py#GenerationConfig","name":"GenerationConfig","kind":"class","path":"search_r1/llm_agent/generation.py","span":"L15-L24","signature":"class GenerationConfig","exported":true,"origin":"ast","body_hash":"ed9bbddf644ef1654d25979be5cb6162da26580c83fdefa3ee4ed92fae774e08","body_text":"class GenerationConfig: max_turns: int max_start_length: int max_prompt_length: int max_response_length: int max_obs_length: int num_gpus: int no_think_rl: bool=False search_url: str = None topk: int = 3","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager","name":"LLMGenerationManager","kind":"class","path":"search_r1/llm_agent/generation.py","span":"L26-L692","signature":"class LLMGenerationManager","exported":true,"origin":"ast","body_hash":"feac04afdab1326d60033994c91f0ff10e52475cf7b179b76a74daeeae8ff467","body_text":"class LLMGenerationManager: def __init__( self, tokenizer, actor_rollout_wg, config: GenerationConfig, is_validation: bool = False, ): self.tokenizer = tokenizer self.actor_rollout_wg = actor_rollout_wg self.config = config self.is_validation = is_validation self.tensor_fn = TensorHelper(TensorConfig( pad_token_id=tokenizer.pad_token_id, max_prompt_length=config.max_prompt_length, max_obs_length=config.max_obs_length, max_start_length=config.max_start_length )) def _batch_tokenize(self, responses: List[str]) -> torch.Tensor: \"\"\"Tokenize a batch of responses.\"\"\" return self.tokenizer( responses, add_special_tokens=False, return_tensors='pt', padding=\"longest\" )['input_ids'] def _postprocess_responses(self, responses: torch.Tensor, is_final: bool = False) -> Tuple[torch.Tensor, List[str]]: \"\"\"Process responses to stop at the appropriate tag. Normal turns (is_final=False): 1. </search> found → truncate at </search> (per-step <confidence> is naturally included before) 2. </final-confidence> found → truncate at </final-confidence> 3. Neither → no truncation (invalid action, error prompt injected next turn) Final turn (is_final=True): 1. </final-confidence> found → truncate at </final-confidence> 2. Not found → no truncation (invalid action) Note: </search> is completely ignored in the final turn. \"\"\" responses_str = self.tokenizer.batch_decode( responses, skip_special_tokens=True ) new_responses_str = [] for resp in responses_str: if not is_final: # Normal turn: prefer search, then final-confidence if '</search>' in resp: new_responses_str.append(resp.split('</search>')[0] + '</search>') elif '</final-confidence>' in resp: new_responses_str.append(resp.split('</final-confidence>')[0] + '</final-confidence>') else: new_responses_str.append(resp) else: # Final turn: only truncate at final-confidence, ignore search if '</final-confidence>' in resp: new_responses_str.append(resp.split('</final-confidence>')[0] + '</final-confidence>') else: new_responses_str.append(resp) if self.config.no_think_rl: raise ValueError('stop') responses = self._batch_tokenize(new_responses_str) return responses, new_responses_str def _process_next_obs(self, next_obs: List[str]) -> torch.Tensor: \"\"\"Process next observations from environment.\"\"\" next_obs_ids = self.tokenizer( next_obs, padding='longest', return_tensors='pt', add_special_tokens=False, )['input_ids'] if next_obs_ids.shape[1] > self.config.max_obs_length: print(f\"[WARNING] OBSERVATION TOO LONG, CONSIDER CHANGING YOUR CONFIG, {next_obs_ids.shape[1]} & {self.config.max_obs_length}\") next_obs_ids = next_obs_ids[:, :self.config.max_obs_length] return next_obs_ids def _update_rolling_state(self, rollings: DataProto, cur_responses: torch.Tensor, next_obs_ids: torch.Tensor) -> Dict: \"\"\"Update rolling state with new responses and observations.\"\"\" new_input_ids = self.tensor_fn.concatenate_with_padding([ rollings.batch['input_ids'], cur_responses, next_obs_ids ]) new_attention_mask = self.tensor_fn.create_attention_mask(new_input_ids) new_position_ids = self.tensor_fn.create_position_ids(new_attention_mask) effective_len = new_attention_mask.sum(dim=1).max() max_len = min(self.config.max_prompt_length, effective_len) new_rollings = DataProto.from_dict({ 'input_ids': new_input_ids[:, -max_len:], 'position_ids': new_position_ids[:, -max_len:], 'attention_mask': new_attention_mask[:, -max_len:] }) new_rollings.meta_info.update(rollings.meta_info) return new_rollings def _info_masked_concatenate_with_padding(self, prompt: torch.Tensor, prompt_with_mask: torch.Tensor, response: torch.Tensor, info: torch.Tensor = None, pad_to_left: bool = True ) -> torch.Tensor: \"\"\"Concatenate tensors and handle padding. Additionally, create a mask (info_mask) to cover the information block if it exists.\"\"\" pad_id = self.tokenizer.pad_token_id tensors = [prompt, response] tensors_with_mask = [prompt_with_mask, response] if info is not None: tensors.append(info) info_mask = torch.full(info.size(), pad_id, dtype=info.dtype, device=info.device) tensors_with_mask.append(info_mask) concatenated = torch.cat(tensors, dim=1) concatenated_with_info = torch.cat(tensors_with_mask, dim=1) mask = concatenated != pad_id if pad_to_left else concatenated == pad_id sorted_indices = mask.to(torch.int64).argsort(dim=1, stable=True) padded_tensor = concatenated.gather(1, sorted_indices) padded_tensor_with_info = concatenated_with_info.gather(1, sorted_indices) return padded_tensor, padded_tensor_with_info def _update_right_side(self, right_side: Dict, cur_responses: torch.Tensor, next_obs_ids: torch.Tensor = None) -> Dict: \"\"\"Update right side state.\"\"\" if next_obs_ids != None: responses, responses_with_info_mask = self._info_masked_concatenate_with_padding( right_side['responses'], right_side['responses_with_info_mask'], cur_responses, next_obs_ids, pad_to_left=False ) else: responses, responses_with_info_mask = self._info_masked_concatenate_with_padding( right_side['responses'], right_side['responses_with_info_mask'], cur_r","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager.__init__","name":"__init__","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L27-L44","signature":"def __init__( self, tokenizer, actor_rollout_wg, config: GenerationConfig, is_validation: bool = False, )","exported":false,"origin":"ast","body_hash":"1a9c949b6bacb1f9d52eedf9a13d0fdbe4cf527cc543c65914d679e6eff2f3f8","body_text":"def __init__( self, tokenizer, actor_rollout_wg, config: GenerationConfig, is_validation: bool = False, ): self.tokenizer = tokenizer self.actor_rollout_wg = actor_rollout_wg self.config = config self.is_validation = is_validation self.tensor_fn = TensorHelper(TensorConfig( pad_token_id=tokenizer.pad_token_id, max_prompt_length=config.max_prompt_length, max_obs_length=config.max_obs_length, max_start_length=config.max_start_length ))","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager._batch_tokenize","name":"_batch_tokenize","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L46-L53","signature":"def _batch_tokenize(self, responses: List[str]) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"c95fca88e11c937ce3356b7065c1945d6a00e52e487761801514f64544a8d5a9","body_text":"def _batch_tokenize(self, responses: List[str]) -> torch.Tensor: \"\"\"Tokenize a batch of responses.\"\"\" return self.tokenizer( responses, add_special_tokens=False, return_tensors='pt', padding=\"longest\" )['input_ids']","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","name":"_postprocess_responses","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L55-L94","signature":"def _postprocess_responses(self, responses: torch.Tensor, is_final: bool = False) -> Tuple[torch.Tensor, List[str]]","exported":false,"origin":"ast","body_hash":"85b7d5b3af878196cda4ed7fae75690dd3d5f1bd386e68ce65b8279cedd7a6d8","body_text":"def _postprocess_responses(self, responses: torch.Tensor, is_final: bool = False) -> Tuple[torch.Tensor, List[str]]: \"\"\"Process responses to stop at the appropriate tag. Normal turns (is_final=False): 1. </search> found → truncate at </search> (per-step <confidence> is naturally included before) 2. </final-confidence> found → truncate at </final-confidence> 3. Neither → no truncation (invalid action, error prompt injected next turn) Final turn (is_final=True): 1. </final-confidence> found → truncate at </final-confidence> 2. Not found → no truncation (invalid action) Note: </search> is completely ignored in the final turn. \"\"\" responses_str = self.tokenizer.batch_decode( responses, skip_special_tokens=True ) new_responses_str = [] for resp in responses_str: if not is_final: # Normal turn: prefer search, then final-confidence if '</search>' in resp: new_responses_str.append(resp.split('</search>')[0] + '</search>') elif '</final-confidence>' in resp: new_responses_str.append(resp.split('</final-confidence>')[0] + '</final-confidence>') else: new_responses_str.append(resp) else: # Final turn: only truncate at final-confidence, ignore search if '</final-confidence>' in resp: new_responses_str.append(resp.split('</final-confidence>')[0] + '</final-confidence>') else: new_responses_str.append(resp) if self.config.no_think_rl: raise ValueError('stop') responses = self._batch_tokenize(new_responses_str) return responses, new_responses_str","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager._process_next_obs","name":"_process_next_obs","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L96-L109","signature":"def _process_next_obs(self, next_obs: List[str]) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"9aaf6b736e2221a980b03e51389257b8d54e23c639b7506249ea4472be2b689d","body_text":"def _process_next_obs(self, next_obs: List[str]) -> torch.Tensor: \"\"\"Process next observations from environment.\"\"\" next_obs_ids = self.tokenizer( next_obs, padding='longest', return_tensors='pt', add_special_tokens=False, )['input_ids'] if next_obs_ids.shape[1] > self.config.max_obs_length: print(f\"[WARNING] OBSERVATION TOO LONG, CONSIDER CHANGING YOUR CONFIG, {next_obs_ids.shape[1]} & {self.config.max_obs_length}\") next_obs_ids = next_obs_ids[:, :self.config.max_obs_length] return next_obs_ids","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_rolling_state","name":"_update_rolling_state","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L111-L133","signature":"def _update_rolling_state(self, rollings: DataProto, cur_responses: torch.Tensor, next_obs_ids: torch.Tensor) -> Dict","exported":false,"origin":"ast","body_hash":"d2b8a454470d80d9d13564d66543ec967580e5e97f23b320200150b426073766","body_text":"def _update_rolling_state(self, rollings: DataProto, cur_responses: torch.Tensor, next_obs_ids: torch.Tensor) -> Dict: \"\"\"Update rolling state with new responses and observations.\"\"\" new_input_ids = self.tensor_fn.concatenate_with_padding([ rollings.batch['input_ids'], cur_responses, next_obs_ids ]) new_attention_mask = self.tensor_fn.create_attention_mask(new_input_ids) new_position_ids = self.tensor_fn.create_position_ids(new_attention_mask) effective_len = new_attention_mask.sum(dim=1).max() max_len = min(self.config.max_prompt_length, effective_len) new_rollings = DataProto.from_dict({ 'input_ids': new_input_ids[:, -max_len:], 'position_ids': new_position_ids[:, -max_len:], 'attention_mask': new_attention_mask[:, -max_len:] }) new_rollings.meta_info.update(rollings.meta_info) return new_rollings","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","name":"_info_masked_concatenate_with_padding","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L135-L158","signature":"def _info_masked_concatenate_with_padding(self, prompt: torch.Tensor, prompt_with_mask: torch.Tensor, response: torch.Tensor, info: torch.Tensor = None, pad_to_left: bool = True ) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"76aaf87efda6142dec0d8522a5ce505e3fb7b52630717481d76b473404b27bd9","body_text":"def _info_masked_concatenate_with_padding(self, prompt: torch.Tensor, prompt_with_mask: torch.Tensor, response: torch.Tensor, info: torch.Tensor = None, pad_to_left: bool = True ) -> torch.Tensor: \"\"\"Concatenate tensors and handle padding. Additionally, create a mask (info_mask) to cover the information block if it exists.\"\"\" pad_id = self.tokenizer.pad_token_id tensors = [prompt, response] tensors_with_mask = [prompt_with_mask, response] if info is not None: tensors.append(info) info_mask = torch.full(info.size(), pad_id, dtype=info.dtype, device=info.device) tensors_with_mask.append(info_mask) concatenated = torch.cat(tensors, dim=1) concatenated_with_info = torch.cat(tensors_with_mask, dim=1) mask = concatenated != pad_id if pad_to_left else concatenated == pad_id sorted_indices = mask.to(torch.int64).argsort(dim=1, stable=True) padded_tensor = concatenated.gather(1, sorted_indices) padded_tensor_with_info = concatenated_with_info.gather(1, sorted_indices) return padded_tensor, padded_tensor_with_info","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_right_side","name":"_update_right_side","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L160-L182","signature":"def _update_right_side(self, right_side: Dict, cur_responses: torch.Tensor, next_obs_ids: torch.Tensor = None) -> Dict","exported":false,"origin":"ast","body_hash":"332430f9f020e34af67f9778c967004ce56dffe8d5e022eb2e370a81cd38dc4e","body_text":"def _update_right_side(self, right_side: Dict, cur_responses: torch.Tensor, next_obs_ids: torch.Tensor = None) -> Dict: \"\"\"Update right side state.\"\"\" if next_obs_ids != None: responses, responses_with_info_mask = self._info_masked_concatenate_with_padding( right_side['responses'], right_side['responses_with_info_mask'], cur_responses, next_obs_ids, pad_to_left=False ) else: responses, responses_with_info_mask = self._info_masked_concatenate_with_padding( right_side['responses'], right_side['responses_with_info_mask'], cur_responses, pad_to_left=False ) effective_len = self.tensor_fn.create_attention_mask(responses).sum(dim=1).max() max_len = min(self.config.max_prompt_length, effective_len) return {'responses': responses[:, :max_len], 'responses_with_info_mask': responses_with_info_mask[:, :max_len]}","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","name":"_generate_with_gpu_padding","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L184-L218","signature":"def _generate_with_gpu_padding(self, active_batch: DataProto) -> DataProto","exported":false,"origin":"ast","body_hash":"f525246a4b63d532823d8fc40e1023226a2be05a8caf4472b60b0d4797c2e9e0","body_text":"def _generate_with_gpu_padding(self, active_batch: DataProto) -> DataProto: \"\"\"Wrapper for generation that handles multi-GPU padding requirements.\"\"\" num_gpus = self.config.num_gpus if num_gpus <= 1: return self.actor_rollout_wg.generate_sequences(active_batch) batch_size = active_batch.batch['input_ids'].shape[0] remainder = batch_size % num_gpus if remainder == 0: return self.actor_rollout_wg.generate_sequences(active_batch) padding_size = num_gpus - remainder padded_batch = {} for k, v in active_batch.batch.items(): pad_sequence = v[0:1].repeat(padding_size, *[1] * (len(v.shape) - 1)) padded_batch[k] = torch.cat([v, pad_sequence], dim=0) padded_active_batch = DataProto.from_dict(padded_batch) padded_output = self.actor_rollout_wg.generate_sequences(padded_active_batch) trimmed_batch = {k: v[:-padding_size] for k, v in padded_output.batch.items()} if hasattr(padded_output, 'meta_info') and padded_output.meta_info: trimmed_meta = {} for k, v in padded_output.meta_info.items(): if isinstance(v, torch.Tensor): trimmed_meta[k] = v[:-padding_size] else: trimmed_meta[k] = v padded_output.meta_info = trimmed_meta padded_output.batch = trimmed_batch return padded_output","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","name":"postprocess_predictions","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L220-L293","signature":"def postprocess_predictions(self, predictions: List[Any], is_final: bool = False) -> Tuple[List[str], List[str], List[Any], List[Any]]","exported":true,"origin":"ast","body_hash":"9763761058f219b486421607fba6c15481915eac76903e40f4a30af4e15bdb65","body_text":"def postprocess_predictions(self, predictions: List[Any], is_final: bool = False) -> Tuple[List[str], List[str], List[Any], List[Any]]: \"\"\"Process predictions from llm into actions, contents, final_confidences and step_confidences. Returns: actions: list of 'search' / 'answer' / None contents: list of extracted content strings final_confidences: list of <final-confidence> values (float or None); only set when action='answer' step_confidences: list of per-step <confidence> values (float or None); only valid when the tag appears BEFORE the action tag \"\"\" actions = [] contents = [] final_confidences = [] step_confidences = [] for prediction in predictions: if isinstance(prediction, str): # ------ 1. Extract action and content ------ answer_match = re.search(r'<answer>(.*?)</answer>', prediction, re.DOTALL) if answer_match: content = answer_match.group(1).strip() action = 'answer' else: search_match = re.search(r'<search>(.*?)</search>', prediction, re.DOTALL) if search_match: content = search_match.group(1).strip() action = 'search' else: content = '' action = None # ------ 2. Extract final-confidence (only when action='answer') ------ final_conf = None if action == 'answer': fc_match = re.search(r'<final-confidence>(.*?)</final-confidence>', prediction, re.DOTALL) if fc_match: try: val = float(fc_match.group(1).strip()) if 0 <= val <= 10: final_conf = val except ValueError: pass # ------ 3. Extract per-step confidence (position must be BEFORE the action tag) ------ step_conf = None conf_match = re.search(r'<confidence>(.*?)</confidence>', prediction, re.DOTALL) if conf_match: # Determine the start position of the action tag action_tag_start = len(prediction) # default: no action tag, no position limit if action == 'search': sm = re.search(r'<search>', prediction) if sm: action_tag_start = sm.start() elif action == 'answer': am = re.search(r'<answer>', prediction) if am: action_tag_start = am.start() # Only accept if confidence tag ends before the action tag starts if conf_match.end() <= action_tag_start: try: val = float(conf_match.group(1).strip()) if 0 <= val <= 10: step_conf = val except ValueError: pass else: raise ValueError(f\"Invalid prediction type: {type(prediction)}\") actions.append(action) contents.append(content) final_confidences.append(final_conf) step_confidences.append(step_conf) return actions, contents, final_confidences, step_confidences","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","name":"execute_predictions","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L295-L339","signature":"def execute_predictions(self, predictions: List[str], pad_token: str, active_mask=None, do_search=True, is_final=False) -> Tuple[List[str], List[int], List[int], List[int], List[Any], List[Any]]","exported":true,"origin":"ast","body_hash":"5af4ab4eb6e492d7b0399ecb6da7ef7f3246492b706d6365c761ff6c278957cb","body_text":"def execute_predictions(self, predictions: List[str], pad_token: str, active_mask=None, do_search=True, is_final=False) -> Tuple[List[str], List[int], List[int], List[int], List[Any], List[Any]]: \"\"\"Execute predictions and return observations, dones, validity, search flags, final_confidences, step_confidences.\"\"\" cur_actions, contents, final_confidences, step_confidences = self.postprocess_predictions(predictions, is_final=is_final) next_obs, dones, valid_action, is_search = [], [], [], [] search_queries = [content for action, content in zip(cur_actions, contents) if action == 'search'] if do_search: search_results = self.batch_search(search_queries) assert len(search_results) == sum([1 for action in cur_actions if action == 'search']) else: search_results = [''] * sum([1 for action in cur_actions if action == 'search']) for i, (action, active) in enumerate(zip(cur_actions, active_mask)): if not active: next_obs.append('') dones.append(1) valid_action.append(0) is_search.append(0) else: if action == 'answer': next_obs.append('') dones.append(1) valid_action.append(1) is_search.append(0) elif action == 'search': next_obs.append(f'\\n\\n<information>{search_results.pop(0).strip()}</information>\\n\\n') dones.append(0) valid_action.append(1) is_search.append(1) else: next_obs.append( '\\nMy previous action is invalid. ' 'If I want to search, I should output <confidence>X</confidence> first (if I have received search results), ' 'then <search> query </search>. ' 'If I want to give the final answer, I should output <confidence>X</confidence> first (if I have received search results), ' 'then <answer> my answer </answer>, then <final-confidence>X</final-confidence>. ' 'Let me try again.\\n' ) dones.append(0) valid_action.append(0) is_search.append(0) assert len(search_results) == 0 return next_obs, dones, valid_action, is_search, final_confidences, step_confidences","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager.batch_search","name":"batch_search","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L341-L344","signature":"def batch_search(self, queries: List[str] = None) -> str","exported":true,"origin":"ast","body_hash":"8b70444002dbfe9bcea334133812a84e87aca99d27bd1f0b581e261e91b0a560","body_text":"def batch_search(self, queries: List[str] = None) -> str: \"\"\"Batchified search for queries.\"\"\" results = self._batch_search(queries)['result'] return [self._passages2string(result) for result in results]","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager._batch_search","name":"_batch_search","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L346-L352","signature":"def _batch_search(self, queries)","exported":false,"origin":"ast","body_hash":"6b8bb803bb9a81dbeb906bd62656b104804226b4e61a1c0eb53f77cfb6bf44c2","body_text":"def _batch_search(self, queries): payload = { \"queries\": queries, \"topk\": self.config.topk, \"return_scores\": True } return requests.post(self.config.search_url, json=payload).json()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager._passages2string","name":"_passages2string","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L354-L361","signature":"def _passages2string(self, retrieval_result)","exported":false,"origin":"ast","body_hash":"108ef56e7beaef356fc6f070cb24c9064e5922a758d88cdc718d61d8c15a5b10","body_text":"def _passages2string(self, retrieval_result): format_reference = '' for idx, doc_item in enumerate(retrieval_result): content = doc_item['document']['contents'] title = content.split(\"\\n\")[0] text = \"\\n\".join(content.split(\"\\n\")[1:]) format_reference += f\"Doc {idx+1}(Title: {title}) {text}\\n\" return format_reference","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","name":"run_llm_loop","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L363-L662","signature":"def run_llm_loop(self, gen_batch, initial_input_ids: torch.Tensor) -> Tuple[Dict, Dict]","exported":true,"origin":"ast","body_hash":"341cd9b45a7eedd9363a5af80b8525c13b73777c442d5e271246ea79468204db","body_text":"def run_llm_loop(self, gen_batch, initial_input_ids: torch.Tensor) -> Tuple[Dict, Dict]: \"\"\"Run main LLM generation loop with statistics and confidence collection.\"\"\" batch_size = gen_batch.batch['input_ids'].shape[0] # ========== Initialise statistics and confidence accumulators ========== generation_statistics = { 'llm_output_tokens': [0] * batch_size, 'information_tokens': [0] * batch_size, 'total_output_tokens': [0] * batch_size, 'reasoning_tokens': [0] * batch_size, 'search_count': [0] * batch_size, 'turn_count': [0] * batch_size, } # Final confidence (scalar per sample, written to batch['confidences']) batch_confidences_tensor = torch.full( (batch_size,), fill_value=-1, dtype=torch.float ) # Per-step confidence: list of dicts per sample per_step_confidences = [[] for _ in range(batch_size)] # Whether each sample has received at least one <information> block has_retrieval = [False] * batch_size # ========== End of initialisation ========== original_left_side = {'input_ids': initial_input_ids[:, -self.config.max_start_length:]} original_right_side = {'responses': initial_input_ids[:, []], 'responses_with_info_mask': initial_input_ids[:, []]} active_mask = torch.ones(batch_size, dtype=torch.bool) turns_stats = torch.ones(batch_size, dtype=torch.int) valid_action_stats = torch.zeros(batch_size, dtype=torch.int) valid_search_stats = torch.zeros(batch_size, dtype=torch.int) active_num_list = [active_mask.sum().item()] rollings = gen_batch # ==================== Main generation loop ==================== for step in range(self.config.max_turns): if not active_mask.sum(): break rollings.batch = self.tensor_fn.cut_to_effective_len( rollings.batch, keys=['input_ids', 'attention_mask', 'position_ids'] ) rollings_active = DataProto.from_dict({ k: v[active_mask] for k, v in rollings.batch.items() }) gen_output = self._generate_with_gpu_padding(rollings_active) meta_info = gen_output.meta_info responses_ids, responses_str = self._postprocess_responses(gen_output.batch['responses'], is_final=False) responses_ids, responses_str = self.tensor_fn._example_level_pad(responses_ids, responses_str, active_mask) # Token statistics for this turn current_response_tokens = (responses_ids != self.tokenizer.pad_token_id).sum(dim=1).tolist() for idx in range(batch_size): if active_mask[idx]: generation_statistics['llm_output_tokens'][idx] += current_response_tokens[idx] generation_statistics['turn_count'][idx] += 1 # Execute predictions (now returns 6 values) next_obs, dones, valid_action, is_search, final_confs, step_confs = self.execute_predictions( responses_str, self.tokenizer.pad_token, active_mask, is_final=False ) # Collect final confidence into tensor for idx in range(batch_size): if active_mask[idx] and final_confs[idx] is not None: batch_confidences_tensor[idx] = final_confs[idx] # Collect per-step confidence record for idx in range(batch_size): if active_mask[idx]: if is_search[idx]: action_type = 'search' elif dones[idx] and valid_action[idx]: action_type = 'answer' else: action_type = 'invalid' per_step_confidences[idx].append({ 'turn': step + 1, # 1-indexed 'confidence': step_confs[idx], # None if missing or malformed 'had_information': has_retrieval[idx], 'action_type': action_type, }) # Update has_retrieval AFTER recording this step # (search result arrives in the NEXT turn as <information>) for idx in range(batch_size): if active_mask[idx] and is_search[idx]: has_retrieval[idx] = True # Search count statistics for idx in range(batch_size): if is_search[idx]: generation_statistics['search_count'][idx] += 1 curr_active_mask = torch.tensor([not done for done in dones], dtype=torch.bool) active_mask = active_mask * curr_active_mask active_num_list.append(active_mask.sum().item()) turns_stats[curr_active_mask] += 1 valid_action_stats += torch.tensor(valid_action, dtype=torch.int) valid_search_stats += torch.tensor(is_search, dtype=torch.int) next_obs_ids = self._process_next_obs(next_obs) # Information token statistics info_token_counts = (next_obs_ids != self.tokenizer.pad_token_id).sum(dim=1).tolist() for idx in range(batch_size): generation_statistics['information_tokens'][idx] += info_token_counts[idx] # Update rolling state and right side rollings = self._update_rolling_state(rollings, responses_ids, next_obs_ids) original_right_side = self._update_right_side(original_right_side, responses_ids, next_obs_ids) # Ensure meta_info fields always exist (even if the loop runs zero times) meta_info['turns_stats'] = turns_stats.tolist() meta_info['active_mask'] = active_mask.tolist() meta_info['valid_action_stats'] = valid_action_stats.tolist() meta_info['valid_search_stats'] = valid_search_stats.tolist() # ==================== Final rollout (with FINAL stage prompt) ==================== if active_mask.sum(): # ------ Inject FINAL stage prompt (rolling_state only, NOT right_side) ------ final_stage_prompts = [] for idx in range(batch_size): if active_mask[idx]: if has_retrieval[idx]: prompt = ( '\\nThis is your final turn. ","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"},{"id":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","name":"_compose_final_output","kind":"method","path":"search_r1/llm_agent/generation.py","span":"L664-L692","signature":"def _compose_final_output(self, left_side: Dict, right_side: Dict, meta_info: Dict) -> Tuple[Dict, Dict]","exported":false,"origin":"ast","body_hash":"6c7839b96b5e2e4cfcce94386ff8c73d231ceb9f1a714ef94c1aa7ac9a6ed6d2","body_text":"def _compose_final_output(self, left_side: Dict, right_side: Dict, meta_info: Dict) -> Tuple[Dict, Dict]: \"\"\"Compose final generation output.\"\"\" final_output = right_side.copy() final_output['prompts'] = left_side['input_ids'] final_output['input_ids'] = torch.cat([ left_side['input_ids'], right_side['responses'] ], dim=1) final_output['attention_mask'] = torch.cat([ self.tensor_fn.create_attention_mask(left_side['input_ids']), self.tensor_fn.create_attention_mask(final_output['responses']) ], dim=1) final_output['info_mask'] = torch.cat([ self.tensor_fn.create_attention_mask(left_side['input_ids']), self.tensor_fn.create_attention_mask(final_output['responses_with_info_mask']) ], dim=1) final_output['position_ids'] = self.tensor_fn.create_position_ids( final_output['attention_mask'] ) final_output = DataProto.from_dict(final_output) final_output.meta_info.update(meta_info) return final_output","summary_state":"pending","summary":null,"crux":null,"owner":"LLMGenerationManager"}],"rawEdges":[{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":"torch","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":"re","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":"collections","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":"os","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":"typing","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":"dataclasses","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":".tensor_helper","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":"verl","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":"verl.utils.tracking","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":"shutil","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"imports","specifier":"requests","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"contains","targetId":"search_r1/llm_agent/generation.py#GenerationConfig","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager.__init__","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.__init__","relation":"calls","name":"TensorHelper","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.__init__","relation":"calls","name":"TensorConfig","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager._batch_tokenize","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._batch_tokenize","relation":"calls","name":"tokenizer","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"batch_decode","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"split","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"split","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"split","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"ValueError","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._postprocess_responses","relation":"calls","name":"_batch_tokenize","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager._process_next_obs","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._process_next_obs","relation":"calls","name":"tokenizer","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._process_next_obs","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_rolling_state","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_rolling_state","relation":"calls","name":"concatenate_with_padding","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_rolling_state","relation":"calls","name":"create_attention_mask","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_rolling_state","relation":"calls","name":"create_position_ids","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_rolling_state","relation":"calls","name":"max","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_rolling_state","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_rolling_state","relation":"calls","name":"min","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_rolling_state","relation":"calls","name":"from_dict","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_rolling_state","relation":"calls","name":"update","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","relation":"calls","name":"full","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","relation":"calls","name":"size","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","relation":"calls","name":"cat","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","relation":"calls","name":"cat","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","relation":"calls","name":"argsort","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","relation":"calls","name":"to","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","relation":"calls","name":"gather","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._info_masked_concatenate_with_padding","relation":"calls","name":"gather","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_right_side","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_right_side","relation":"calls","name":"_info_masked_concatenate_with_padding","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_right_side","relation":"calls","name":"_info_masked_concatenate_with_padding","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_right_side","relation":"calls","name":"max","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_right_side","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_right_side","relation":"calls","name":"create_attention_mask","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._update_right_side","relation":"calls","name":"min","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"generate_sequences","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"generate_sequences","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"items","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"repeat","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"len","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"cat","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"from_dict","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"generate_sequences","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"items","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"hasattr","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"items","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._generate_with_gpu_padding","relation":"calls","name":"isinstance","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"isinstance","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"search","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"strip","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"group","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"search","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"strip","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"group","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"search","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"float","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"strip","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"group","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"search","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"len","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"search","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"start","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"search","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"start","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"end","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"float","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"strip","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"group","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"ValueError","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"type","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.postprocess_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"postprocess_predictions","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"zip","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"batch_search","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"len","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"sum","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"sum","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"enumerate","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"zip","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"strip","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"pop","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.execute_predictions","relation":"calls","name":"len","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager.batch_search","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.batch_search","relation":"calls","name":"_batch_search","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.batch_search","relation":"calls","name":"_passages2string","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager._batch_search","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._batch_search","relation":"calls","name":"json","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._batch_search","relation":"calls","name":"post","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager._passages2string","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._passages2string","relation":"calls","name":"enumerate","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._passages2string","relation":"calls","name":"split","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._passages2string","relation":"calls","name":"join","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._passages2string","relation":"calls","name":"split","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"full","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"ones","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"ones","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"zeros","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"zeros","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"item","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"cut_to_effective_len","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"from_dict","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"items","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_generate_with_gpu_padding","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_postprocess_responses","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_example_level_pad","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"execute_predictions","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tensor","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"item","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tensor","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tensor","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_process_next_obs","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_update_rolling_state","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_update_right_side","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_process_next_obs","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"zeros","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_update_rolling_state","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"cut_to_effective_len","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"from_dict","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"items","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_generate_with_gpu_padding","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"update","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_postprocess_responses","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_example_level_pad","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"execute_predictions","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tensor","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"append","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"item","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tensor","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tensor","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_update_right_side","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"full","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"range","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"reversed","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"get","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"get","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"len","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"mean","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"min","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"max","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"len","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"len","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"sum","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"mean","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"len","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"mean","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"min","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"max","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"mean","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"mean","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"mean","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"mean","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"mean","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"print","viaMember":false,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager.run_llm_loop","relation":"calls","name":"_compose_final_output","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"LLMGenerationManager"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager","relation":"contains","targetId":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"copy","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"Dict"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"cat","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"cat","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"create_attention_mask","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"create_attention_mask","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"cat","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"create_attention_mask","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"create_attention_mask","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"create_position_ids","viaMember":true,"file":"search_r1/llm_agent/generation.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"from_dict","viaMember":true,"file":"search_r1/llm_agent/generation.py"},{"source":"search_r1/llm_agent/generation.py#LLMGenerationManager._compose_final_output","relation":"calls","name":"update","viaMember":true,"file":"search_r1/llm_agent/generation.py"}]},"search_r1/llm_agent/tensor_helper.py":{"size":3148,"mtimeMs":1789275830305.2024,"hash":"b000f3b61c5a2be13dfcf64e1ce9091ce89dda99ed3b331fb24ab1f43b4dce69","nodes":[{"id":"search_r1/llm_agent/tensor_helper.py","name":"tensor_helper.py","kind":"file","path":"search_r1/llm_agent/tensor_helper.py","span":"L1-L75","signature":null,"exported":true,"origin":"ast","body_hash":"b000f3b61c5a2be13dfcf64e1ce9091ce89dda99ed3b331fb24ab1f43b4dce69","chars":3148,"summary_state":"pending","summary":null,"crux":null,"body_text":"import torch from typing import Dict, Tuple, List from dataclasses import dataclass @dataclass"},{"id":"search_r1/llm_agent/tensor_helper.py#TensorConfig","name":"TensorConfig","kind":"class","path":"search_r1/llm_agent/tensor_helper.py","span":"L6-L10","signature":"class TensorConfig","exported":true,"origin":"ast","body_hash":"9b95415b18b58a663b50726bc8c20672bda2417dca98d179e149eb2e43291131","body_text":"class TensorConfig: pad_token_id: int max_prompt_length: int max_obs_length: int max_start_length: int","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/llm_agent/tensor_helper.py#TensorHelper","name":"TensorHelper","kind":"class","path":"search_r1/llm_agent/tensor_helper.py","span":"L12-L75","signature":"class TensorHelper","exported":true,"origin":"ast","body_hash":"325674ed1de50bab13f7c8b0afd5d8610e98ea54dd8e3abbad9bde1555bf1c25","body_text":"class TensorHelper: def __init__(self, config: TensorConfig): self.config = config def cut_to_effective_len(self, tensor_dict: Dict[str, torch.Tensor], keys: List[str], cut_left: bool = True) -> Dict[str, torch.Tensor]: \"\"\"Cut tensors to their effective length based on attention mask.\"\"\" effective_len = tensor_dict['attention_mask'].sum(dim=1).max() result = tensor_dict.copy() for key in keys: if cut_left: result[key] = tensor_dict[key][:, -effective_len:] else: result[key] = tensor_dict[key][:, :effective_len] return result def convert_pad_structure(self, tensor: torch.Tensor, pad_to_left: bool = True) -> Tuple[torch.Tensor, torch.Tensor]: \"\"\"Convert padding structure and return sorted tensor with indices.\"\"\" mask = tensor != self.config.pad_token_id if pad_to_left else tensor == self.config.pad_token_id sorted_indices = mask.to(torch.int64).argsort(dim=1, stable=True) return tensor.gather(1, sorted_indices), sorted_indices def create_attention_mask(self, input_ids: torch.Tensor) -> torch.Tensor: \"\"\"Create attention mask from input ids.\"\"\" return torch.where(input_ids != self.config.pad_token_id, 1, 0) def create_position_ids(self, attention_mask: torch.Tensor) -> torch.Tensor: \"\"\"Create position ids from attention mask.\"\"\" return (torch.cumsum(attention_mask, dim=1) - 1) * attention_mask def concatenate_with_padding(self, tensors: List[torch.Tensor], pad_to_left: bool = True) -> torch.Tensor: \"\"\"Concatenate tensors and handle padding.\"\"\" concatenated = torch.cat(tensors, dim=1) padded_tensor, _ = self.convert_pad_structure(concatenated, pad_to_left) return padded_tensor def _example_level_pad(self, responses: torch.Tensor, responses_str: List[str], active_mask: torch.Tensor) -> Tuple[torch.Tensor, List[str]]: \"\"\" Pad responses for non-active examples with pad tokens. \"\"\" assert active_mask.sum() == responses.shape[0] # Create masked responses tensor batch_size = active_mask.shape[0] seq_len = responses.shape[1] padded_responses = torch.full( (batch_size, seq_len), self.config.pad_token_id, dtype=responses.dtype, device=responses.device ) padded_responses[active_mask] = responses # Create masked response strings padded_responses_str = [\"\"] * batch_size s = 0 for i, is_active in enumerate(active_mask): if is_active: padded_responses_str[i] = responses_str[s] s += 1 return padded_responses, padded_responses_str","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/llm_agent/tensor_helper.py#TensorHelper.__init__","name":"__init__","kind":"method","path":"search_r1/llm_agent/tensor_helper.py","span":"L13-L14","signature":"def __init__(self, config: TensorConfig)","exported":false,"origin":"ast","body_hash":"9db37495552a6b6d162672bfcd40830dcbc96bcf9910f8da8b8180ba069b6bb1","body_text":"def __init__(self, config: TensorConfig): self.config = config","summary_state":"pending","summary":null,"crux":null,"owner":"TensorHelper"},{"id":"search_r1/llm_agent/tensor_helper.py#TensorHelper.cut_to_effective_len","name":"cut_to_effective_len","kind":"method","path":"search_r1/llm_agent/tensor_helper.py","span":"L16-L27","signature":"def cut_to_effective_len(self, tensor_dict: Dict[str, torch.Tensor], keys: List[str], cut_left: bool = True) -> Dict[str, torch.Tensor]","exported":true,"origin":"ast","body_hash":"4707a563a97c83087c8bafd83ecfb7014afdba12aabf09af282f40811013b04d","body_text":"def cut_to_effective_len(self, tensor_dict: Dict[str, torch.Tensor], keys: List[str], cut_left: bool = True) -> Dict[str, torch.Tensor]: \"\"\"Cut tensors to their effective length based on attention mask.\"\"\" effective_len = tensor_dict['attention_mask'].sum(dim=1).max() result = tensor_dict.copy() for key in keys: if cut_left: result[key] = tensor_dict[key][:, -effective_len:] else: result[key] = tensor_dict[key][:, :effective_len] return result","summary_state":"pending","summary":null,"crux":null,"owner":"TensorHelper"},{"id":"search_r1/llm_agent/tensor_helper.py#TensorHelper.convert_pad_structure","name":"convert_pad_structure","kind":"method","path":"search_r1/llm_agent/tensor_helper.py","span":"L29-L33","signature":"def convert_pad_structure(self, tensor: torch.Tensor, pad_to_left: bool = True) -> Tuple[torch.Tensor, torch.Tensor]","exported":true,"origin":"ast","body_hash":"a604ff450f720b5f1a8c143ceedd44723325ba4fa6250ad2ae4403bf79282157","body_text":"def convert_pad_structure(self, tensor: torch.Tensor, pad_to_left: bool = True) -> Tuple[torch.Tensor, torch.Tensor]: \"\"\"Convert padding structure and return sorted tensor with indices.\"\"\" mask = tensor != self.config.pad_token_id if pad_to_left else tensor == self.config.pad_token_id sorted_indices = mask.to(torch.int64).argsort(dim=1, stable=True) return tensor.gather(1, sorted_indices), sorted_indices","summary_state":"pending","summary":null,"crux":null,"owner":"TensorHelper"},{"id":"search_r1/llm_agent/tensor_helper.py#TensorHelper.create_attention_mask","name":"create_attention_mask","kind":"method","path":"search_r1/llm_agent/tensor_helper.py","span":"L35-L37","signature":"def create_attention_mask(self, input_ids: torch.Tensor) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"b3d8ef40037b2e737f2e74def851d345b4225ad3e69e76f4499a856ecce40083","body_text":"def create_attention_mask(self, input_ids: torch.Tensor) -> torch.Tensor: \"\"\"Create attention mask from input ids.\"\"\" return torch.where(input_ids != self.config.pad_token_id, 1, 0)","summary_state":"pending","summary":null,"crux":null,"owner":"TensorHelper"},{"id":"search_r1/llm_agent/tensor_helper.py#TensorHelper.create_position_ids","name":"create_position_ids","kind":"method","path":"search_r1/llm_agent/tensor_helper.py","span":"L39-L41","signature":"def create_position_ids(self, attention_mask: torch.Tensor) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"83996e2b2d50fe42b07b5f97562054ac05ab1b172a61e56ec4c552e79553f348","body_text":"def create_position_ids(self, attention_mask: torch.Tensor) -> torch.Tensor: \"\"\"Create position ids from attention mask.\"\"\" return (torch.cumsum(attention_mask, dim=1) - 1) * attention_mask","summary_state":"pending","summary":null,"crux":null,"owner":"TensorHelper"},{"id":"search_r1/llm_agent/tensor_helper.py#TensorHelper.concatenate_with_padding","name":"concatenate_with_padding","kind":"method","path":"search_r1/llm_agent/tensor_helper.py","span":"L43-L48","signature":"def concatenate_with_padding(self, tensors: List[torch.Tensor], pad_to_left: bool = True) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"149ccfc9f3861eb09e2360cfa106d79b961b318b775bd2713363a1cad150dfd9","body_text":"def concatenate_with_padding(self, tensors: List[torch.Tensor], pad_to_left: bool = True) -> torch.Tensor: \"\"\"Concatenate tensors and handle padding.\"\"\" concatenated = torch.cat(tensors, dim=1) padded_tensor, _ = self.convert_pad_structure(concatenated, pad_to_left) return padded_tensor","summary_state":"pending","summary":null,"crux":null,"owner":"TensorHelper"},{"id":"search_r1/llm_agent/tensor_helper.py#TensorHelper._example_level_pad","name":"_example_level_pad","kind":"method","path":"search_r1/llm_agent/tensor_helper.py","span":"L50-L75","signature":"def _example_level_pad(self, responses: torch.Tensor, responses_str: List[str], active_mask: torch.Tensor) -> Tuple[torch.Tensor, List[str]]","exported":false,"origin":"ast","body_hash":"c8c58743f044fd65474170c0d9053e841a449dc5dcf45e43c260da1ac882a4e6","body_text":"def _example_level_pad(self, responses: torch.Tensor, responses_str: List[str], active_mask: torch.Tensor) -> Tuple[torch.Tensor, List[str]]: \"\"\" Pad responses for non-active examples with pad tokens. \"\"\" assert active_mask.sum() == responses.shape[0] # Create masked responses tensor batch_size = active_mask.shape[0] seq_len = responses.shape[1] padded_responses = torch.full( (batch_size, seq_len), self.config.pad_token_id, dtype=responses.dtype, device=responses.device ) padded_responses[active_mask] = responses # Create masked response strings padded_responses_str = [\"\"] * batch_size s = 0 for i, is_active in enumerate(active_mask): if is_active: padded_responses_str[i] = responses_str[s] s += 1 return padded_responses, padded_responses_str","summary_state":"pending","summary":null,"crux":null,"owner":"TensorHelper"}],"rawEdges":[{"source":"search_r1/llm_agent/tensor_helper.py","relation":"imports","specifier":"torch","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py","relation":"imports","specifier":"typing","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py","relation":"imports","specifier":"dataclasses","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py","relation":"contains","targetId":"search_r1/llm_agent/tensor_helper.py#TensorConfig","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py","relation":"contains","targetId":"search_r1/llm_agent/tensor_helper.py#TensorHelper","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper","relation":"contains","targetId":"search_r1/llm_agent/tensor_helper.py#TensorHelper.__init__","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper","relation":"contains","targetId":"search_r1/llm_agent/tensor_helper.py#TensorHelper.cut_to_effective_len","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper.cut_to_effective_len","relation":"calls","name":"max","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper.cut_to_effective_len","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper.cut_to_effective_len","relation":"calls","name":"copy","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper","relation":"contains","targetId":"search_r1/llm_agent/tensor_helper.py#TensorHelper.convert_pad_structure","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper.convert_pad_structure","relation":"calls","name":"argsort","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper.convert_pad_structure","relation":"calls","name":"to","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper.convert_pad_structure","relation":"calls","name":"gather","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper","relation":"contains","targetId":"search_r1/llm_agent/tensor_helper.py#TensorHelper.create_attention_mask","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper.create_attention_mask","relation":"calls","name":"where","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper","relation":"contains","targetId":"search_r1/llm_agent/tensor_helper.py#TensorHelper.create_position_ids","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper.create_position_ids","relation":"calls","name":"cumsum","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper","relation":"contains","targetId":"search_r1/llm_agent/tensor_helper.py#TensorHelper.concatenate_with_padding","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper.concatenate_with_padding","relation":"calls","name":"cat","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper.concatenate_with_padding","relation":"calls","name":"convert_pad_structure","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py","recvType":"TensorHelper"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper","relation":"contains","targetId":"search_r1/llm_agent/tensor_helper.py#TensorHelper._example_level_pad","file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper._example_level_pad","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper._example_level_pad","relation":"calls","name":"full","viaMember":true,"file":"search_r1/llm_agent/tensor_helper.py"},{"source":"search_r1/llm_agent/tensor_helper.py#TensorHelper._example_level_pad","relation":"calls","name":"enumerate","viaMember":false,"file":"search_r1/llm_agent/tensor_helper.py"}]},"search_r1/search/index_builder.py":{"size":12495,"mtimeMs":1789275830306.2024,"hash":"320d7076a3b9cf05feeddc464616d5f0b37e3d5ed70f5dafb034bdf14743359c","nodes":[{"id":"search_r1/search/index_builder.py","name":"index_builder.py","kind":"file","path":"search_r1/search/index_builder.py","span":"L1-L349","signature":null,"exported":true,"origin":"ast","body_hash":"320d7076a3b9cf05feeddc464616d5f0b37e3d5ed70f5dafb034bdf14743359c","chars":12495,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import faiss import json import warnings import numpy as np from typing import cast, List, Dict import shutil import subprocess import argparse import torch from tqdm import tqdm # from LongRAG.retriever.utils import load_model, load_corpus, pooling import datasets from transformers import AutoTokenizer, AutoModel, AutoConfig MODEL2POOLING = { \"e5\": \"mean\", \"bge\": \"cls\", \"contriever\": \"mean\", 'jina': 'mean' } if __name__ == \"__main__\": main()"},{"id":"search_r1/search/index_builder.py#load_model","name":"load_model","kind":"function","path":"search_r1/search/index_builder.py","span":"L17-L29","signature":"def load_model( model_path: str, use_fp16: bool = False )","exported":true,"origin":"ast","body_hash":"2d4e1f08f900b69638da6c1525eb3006b46ed0822e5e3f12fb86d8faa60ab34b","body_text":"def load_model( model_path: str, use_fp16: bool = False ): model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) model = AutoModel.from_pretrained(model_path, trust_remote_code=True) model.eval() model.cuda() if use_fp16: model = model.half() tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, trust_remote_code=True) return model, tokenizer","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/index_builder.py#pooling","name":"pooling","kind":"function","path":"search_r1/search/index_builder.py","span":"L32-L46","signature":"def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = \"mean\" )","exported":true,"origin":"ast","body_hash":"6e356a5fbac440ef6226d2119161f4cf7a669ceb3505baddff6331a4a50030ec","body_text":"def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = \"mean\" ): if pooling_method == \"mean\": last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0) return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None] elif pooling_method == \"cls\": return last_hidden_state[:, 0] elif pooling_method == \"pooler\": return pooler_output else: raise NotImplementedError(\"Pooling method not implemented!\")","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/index_builder.py#load_corpus","name":"load_corpus","kind":"function","path":"search_r1/search/index_builder.py","span":"L49-L55","signature":"def load_corpus(corpus_path: str)","exported":true,"origin":"ast","body_hash":"63642a3531f6506e9fae80592a2170b344a1b6179a3d27a3002081529a9e3d23","body_text":"def load_corpus(corpus_path: str): corpus = datasets.load_dataset( 'json', data_files=corpus_path, split=\"train\", num_proc=4) return corpus","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/index_builder.py#Index_Builder","name":"Index_Builder","kind":"class","path":"search_r1/search/index_builder.py","span":"L58-L285","signature":"class Index_Builder","exported":true,"origin":"ast","body_hash":"725abc2e713a922925432014acf86aecb268c47847921f93d8840dbda661630a","body_text":"class Index_Builder: r\"\"\"A tool class used to build an index used in retrieval. \"\"\" def __init__( self, retrieval_method, model_path, corpus_path, save_dir, max_length, batch_size, use_fp16, pooling_method, faiss_type=None, embedding_path=None, save_embedding=False, faiss_gpu=False ): self.retrieval_method = retrieval_method.lower() self.model_path = model_path self.corpus_path = corpus_path self.save_dir = save_dir self.max_length = max_length self.batch_size = batch_size self.use_fp16 = use_fp16 self.pooling_method = pooling_method self.faiss_type = faiss_type if faiss_type is not None else 'Flat' self.embedding_path = embedding_path self.save_embedding = save_embedding self.faiss_gpu = faiss_gpu self.gpu_num = torch.cuda.device_count() # prepare save dir print(self.save_dir) if not os.path.exists(self.save_dir): os.makedirs(self.save_dir) else: if not self._check_dir(self.save_dir): warnings.warn(\"Some files already exists in save dir and may be overwritten.\", UserWarning) self.index_save_path = os.path.join(self.save_dir, f\"{self.retrieval_method}_{self.faiss_type}.index\") self.embedding_save_path = os.path.join(self.save_dir, f\"emb_{self.retrieval_method}.memmap\") self.corpus = load_corpus(self.corpus_path) print(\"Finish loading...\") @staticmethod def _check_dir(dir_path): r\"\"\"Check if the dir path exists and if there is content. \"\"\" if os.path.isdir(dir_path): if len(os.listdir(dir_path)) > 0: return False else: os.makedirs(dir_path, exist_ok=True) return True def build_index(self): r\"\"\"Constructing different indexes based on selective retrieval method. \"\"\" if self.retrieval_method == \"bm25\": self.build_bm25_index() else: self.build_dense_index() def build_bm25_index(self): \"\"\"Building BM25 index based on Pyserini library. Reference: https://github.com/castorini/pyserini/blob/master/docs/usage-index.md#building-a-bm25-index-direct-java-implementation \"\"\" # to use pyserini pipeline, we first need to place jsonl file in the folder self.save_dir = os.path.join(self.save_dir, \"bm25\") os.makedirs(self.save_dir, exist_ok=True) temp_dir = self.save_dir + \"/temp\" temp_file_path = temp_dir + \"/temp.jsonl\" os.makedirs(temp_dir) # if self.have_contents: # shutil.copyfile(self.corpus_path, temp_file_path) # else: # with open(temp_file_path, \"w\") as f: # for item in self.corpus: # f.write(json.dumps(item) + \"\\n\") shutil.copyfile(self.corpus_path, temp_file_path) print(\"Start building bm25 index...\") pyserini_args = [\"--collection\", \"JsonCollection\", \"--input\", temp_dir, \"--index\", self.save_dir, \"--generator\", \"DefaultLuceneDocumentGenerator\", \"--threads\", \"1\"] subprocess.run([\"python\", \"-m\", \"pyserini.index.lucene\"] + pyserini_args) shutil.rmtree(temp_dir) print(\"Finish!\") def _load_embedding(self, embedding_path, corpus_size, hidden_size): all_embeddings = np.memmap( embedding_path, mode=\"r\", dtype=np.float32 ).reshape(corpus_size, hidden_size) return all_embeddings def _save_embedding(self, all_embeddings): memmap = np.memmap( self.embedding_save_path, shape=all_embeddings.shape, mode=\"w+\", dtype=all_embeddings.dtype ) length = all_embeddings.shape[0] # add in batch save_batch_size = 10000 if length > save_batch_size: for i in tqdm(range(0, length, save_batch_size), leave=False, desc=\"Saving Embeddings\"): j = min(i + save_batch_size, length) memmap[i: j] = all_embeddings[i: j] else: memmap[:] = all_embeddings def encode_all(self): if self.gpu_num > 1: print(\"Use multi gpu!\") self.encoder = torch.nn.DataParallel(self.encoder) self.batch_size = self.batch_size * self.gpu_num all_embeddings = [] for start_idx in tqdm(range(0, len(self.corpus), self.batch_size), desc='Inference Embeddings:'): batch_data_title = self.corpus[start_idx:start_idx+self.batch_size]['title'] batch_data_text = self.corpus[start_idx:start_idx+self.batch_size]['text'] batch_data = ['\"' + title + '\"\\n' + text for title, text in zip(batch_data_title, batch_data_text)] if self.retrieval_method == \"e5\": batch_data = [f\"passage: {doc}\" for doc in batch_data] inputs = self.tokenizer( batch_data, padding=True, truncation=True, return_tensors='pt', max_length=self.max_length, ).to('cuda') inputs = {k: v.cuda() for k, v in inputs.items()} #TODO: support encoder-only T5 model if \"T5\" in type(self.encoder).__name__: # T5-based retrieval model decoder_input_ids = torch.zeros( (inputs['input_ids'].shape[0], 1), dtype=torch.long ).to(inputs['input_ids'].device) output = self.encoder( **inputs, decoder_input_ids=decoder_input_ids, return_dict=True ) embeddings = output.last_hidden_state[:, 0, :] else: output = self.encoder(**inputs, return_dict=True) embeddings = pooling(output.pooler_output, output.last_hidden_state, inputs['attention_mask'], self.pooling_method) if \"dpr\" not in self.retrieval_method: embeddings = torch.nn.functional.normalize(embeddings, dim=-1) embeddings = cast(torch.Tensor, embeddings) embeddings = embeddings.detach().cpu().numpy() all_embeddings.append(embeddings) all_embeddings = np.concatenate(all_embeddings, axis=0) all_embeddings","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/index_builder.py#Index_Builder.__init__","name":"__init__","kind":"method","path":"search_r1/search/index_builder.py","span":"L62-L106","signature":"def __init__( self, retrieval_method, model_path, corpus_path, save_dir, max_length, batch_size, use_fp16, pooling_method, faiss_type=None, embedding_path=None, save_embedding=False, faiss_gpu=False )","exported":false,"origin":"ast","body_hash":"24bd9895a6b17014d34182c8189174889d6c68ec55cded5a407f9e8c6afbf4f7","body_text":"def __init__( self, retrieval_method, model_path, corpus_path, save_dir, max_length, batch_size, use_fp16, pooling_method, faiss_type=None, embedding_path=None, save_embedding=False, faiss_gpu=False ): self.retrieval_method = retrieval_method.lower() self.model_path = model_path self.corpus_path = corpus_path self.save_dir = save_dir self.max_length = max_length self.batch_size = batch_size self.use_fp16 = use_fp16 self.pooling_method = pooling_method self.faiss_type = faiss_type if faiss_type is not None else 'Flat' self.embedding_path = embedding_path self.save_embedding = save_embedding self.faiss_gpu = faiss_gpu self.gpu_num = torch.cuda.device_count() # prepare save dir print(self.save_dir) if not os.path.exists(self.save_dir): os.makedirs(self.save_dir) else: if not self._check_dir(self.save_dir): warnings.warn(\"Some files already exists in save dir and may be overwritten.\", UserWarning) self.index_save_path = os.path.join(self.save_dir, f\"{self.retrieval_method}_{self.faiss_type}.index\") self.embedding_save_path = os.path.join(self.save_dir, f\"emb_{self.retrieval_method}.memmap\") self.corpus = load_corpus(self.corpus_path) print(\"Finish loading...\")","summary_state":"pending","summary":null,"crux":null,"owner":"Index_Builder"},{"id":"search_r1/search/index_builder.py#Index_Builder._check_dir","name":"_check_dir","kind":"method","path":"search_r1/search/index_builder.py","span":"L108-L118","signature":"def _check_dir(dir_path)","exported":false,"origin":"ast","body_hash":"a5ad6d853672e30fb2c257dcf191693abc7f1180f4dc684b7ef544cf9cd0612d","body_text":"def _check_dir(dir_path): r\"\"\"Check if the dir path exists and if there is content. \"\"\" if os.path.isdir(dir_path): if len(os.listdir(dir_path)) > 0: return False else: os.makedirs(dir_path, exist_ok=True) return True","summary_state":"pending","summary":null,"crux":null,"owner":"Index_Builder"},{"id":"search_r1/search/index_builder.py#Index_Builder.build_index","name":"build_index","kind":"method","path":"search_r1/search/index_builder.py","span":"L120-L127","signature":"def build_index(self)","exported":true,"origin":"ast","body_hash":"071d7c3ac8faf081890da7b74ce59e7b4b58300c235ba66a00c38cdd179b5725","body_text":"def build_index(self): r\"\"\"Constructing different indexes based on selective retrieval method. \"\"\" if self.retrieval_method == \"bm25\": self.build_bm25_index() else: self.build_dense_index()","summary_state":"pending","summary":null,"crux":null,"owner":"Index_Builder"},{"id":"search_r1/search/index_builder.py#Index_Builder.build_bm25_index","name":"build_bm25_index","kind":"method","path":"search_r1/search/index_builder.py","span":"L129-L161","signature":"def build_bm25_index(self)","exported":true,"origin":"ast","body_hash":"9a5cc51a0c64fd8e0cc3b12cd68460d1084ca1709ac6d10457fcafacd526c824","body_text":"def build_bm25_index(self): \"\"\"Building BM25 index based on Pyserini library. Reference: https://github.com/castorini/pyserini/blob/master/docs/usage-index.md#building-a-bm25-index-direct-java-implementation \"\"\" # to use pyserini pipeline, we first need to place jsonl file in the folder self.save_dir = os.path.join(self.save_dir, \"bm25\") os.makedirs(self.save_dir, exist_ok=True) temp_dir = self.save_dir + \"/temp\" temp_file_path = temp_dir + \"/temp.jsonl\" os.makedirs(temp_dir) # if self.have_contents: # shutil.copyfile(self.corpus_path, temp_file_path) # else: # with open(temp_file_path, \"w\") as f: # for item in self.corpus: # f.write(json.dumps(item) + \"\\n\") shutil.copyfile(self.corpus_path, temp_file_path) print(\"Start building bm25 index...\") pyserini_args = [\"--collection\", \"JsonCollection\", \"--input\", temp_dir, \"--index\", self.save_dir, \"--generator\", \"DefaultLuceneDocumentGenerator\", \"--threads\", \"1\"] subprocess.run([\"python\", \"-m\", \"pyserini.index.lucene\"] + pyserini_args) shutil.rmtree(temp_dir) print(\"Finish!\")","summary_state":"pending","summary":null,"crux":null,"owner":"Index_Builder"},{"id":"search_r1/search/index_builder.py#Index_Builder._load_embedding","name":"_load_embedding","kind":"method","path":"search_r1/search/index_builder.py","span":"L163-L169","signature":"def _load_embedding(self, embedding_path, corpus_size, hidden_size)","exported":false,"origin":"ast","body_hash":"bc3be3444bdac9c1a4eb9cb4c5ba41d398d887d80f03c6ff2c2429237d9fa215","body_text":"def _load_embedding(self, embedding_path, corpus_size, hidden_size): all_embeddings = np.memmap( embedding_path, mode=\"r\", dtype=np.float32 ).reshape(corpus_size, hidden_size) return all_embeddings","summary_state":"pending","summary":null,"crux":null,"owner":"Index_Builder"},{"id":"search_r1/search/index_builder.py#Index_Builder._save_embedding","name":"_save_embedding","kind":"method","path":"search_r1/search/index_builder.py","span":"L171-L186","signature":"def _save_embedding(self, all_embeddings)","exported":false,"origin":"ast","body_hash":"e4aeb62957e99dc7205e0de5e791434605b0d3101637e8606befc4ee1fc6057f","body_text":"def _save_embedding(self, all_embeddings): memmap = np.memmap( self.embedding_save_path, shape=all_embeddings.shape, mode=\"w+\", dtype=all_embeddings.dtype ) length = all_embeddings.shape[0] # add in batch save_batch_size = 10000 if length > save_batch_size: for i in tqdm(range(0, length, save_batch_size), leave=False, desc=\"Saving Embeddings\"): j = min(i + save_batch_size, length) memmap[i: j] = all_embeddings[i: j] else: memmap[:] = all_embeddings","summary_state":"pending","summary":null,"crux":null,"owner":"Index_Builder"},{"id":"search_r1/search/index_builder.py#Index_Builder.encode_all","name":"encode_all","kind":"method","path":"search_r1/search/index_builder.py","span":"L188-L242","signature":"def encode_all(self)","exported":true,"origin":"ast","body_hash":"66b326b53707ef442cf20ed0dd2db96ad739819e6897d848aeac86d566e39ad1","body_text":"def encode_all(self): if self.gpu_num > 1: print(\"Use multi gpu!\") self.encoder = torch.nn.DataParallel(self.encoder) self.batch_size = self.batch_size * self.gpu_num all_embeddings = [] for start_idx in tqdm(range(0, len(self.corpus), self.batch_size), desc='Inference Embeddings:'): batch_data_title = self.corpus[start_idx:start_idx+self.batch_size]['title'] batch_data_text = self.corpus[start_idx:start_idx+self.batch_size]['text'] batch_data = ['\"' + title + '\"\\n' + text for title, text in zip(batch_data_title, batch_data_text)] if self.retrieval_method == \"e5\": batch_data = [f\"passage: {doc}\" for doc in batch_data] inputs = self.tokenizer( batch_data, padding=True, truncation=True, return_tensors='pt', max_length=self.max_length, ).to('cuda') inputs = {k: v.cuda() for k, v in inputs.items()} #TODO: support encoder-only T5 model if \"T5\" in type(self.encoder).__name__: # T5-based retrieval model decoder_input_ids = torch.zeros( (inputs['input_ids'].shape[0], 1), dtype=torch.long ).to(inputs['input_ids'].device) output = self.encoder( **inputs, decoder_input_ids=decoder_input_ids, return_dict=True ) embeddings = output.last_hidden_state[:, 0, :] else: output = self.encoder(**inputs, return_dict=True) embeddings = pooling(output.pooler_output, output.last_hidden_state, inputs['attention_mask'], self.pooling_method) if \"dpr\" not in self.retrieval_method: embeddings = torch.nn.functional.normalize(embeddings, dim=-1) embeddings = cast(torch.Tensor, embeddings) embeddings = embeddings.detach().cpu().numpy() all_embeddings.append(embeddings) all_embeddings = np.concatenate(all_embeddings, axis=0) all_embeddings = all_embeddings.astype(np.float32) return all_embeddings","summary_state":"pending","summary":null,"crux":null,"owner":"Index_Builder"},{"id":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","name":"build_dense_index","kind":"method","path":"search_r1/search/index_builder.py","span":"L245-L285","signature":"def build_dense_index(self)","exported":true,"origin":"ast","body_hash":"96cbca4b67e6832e65ca3d82bc06022d21d67494951939e20d12f6da4997d9bc","body_text":"def build_dense_index(self): \"\"\"Obtain the representation of documents based on the embedding model(BERT-based) and construct a faiss index. \"\"\" if os.path.exists(self.index_save_path): print(\"The index file already exists and will be overwritten.\") self.encoder, self.tokenizer = load_model(model_path = self.model_path, use_fp16 = self.use_fp16) if self.embedding_path is not None: hidden_size = self.encoder.config.hidden_size corpus_size = len(self.corpus) all_embeddings = self._load_embedding(self.embedding_path, corpus_size, hidden_size) else: all_embeddings = self.encode_all() if self.save_embedding: self._save_embedding(all_embeddings) del self.corpus # build index print(\"Creating index\") dim = all_embeddings.shape[-1] faiss_index = faiss.index_factory(dim, self.faiss_type, faiss.METRIC_INNER_PRODUCT) if self.faiss_gpu: co = faiss.GpuMultipleClonerOptions() co.useFloat16 = True co.shard = True faiss_index = faiss.index_cpu_to_all_gpus(faiss_index, co) if not faiss_index.is_trained: faiss_index.train(all_embeddings) faiss_index.add(all_embeddings) faiss_index = faiss.index_gpu_to_cpu(faiss_index) else: if not faiss_index.is_trained: faiss_index.train(all_embeddings) faiss_index.add(all_embeddings) faiss.write_index(faiss_index, self.index_save_path) print(\"Finish!\")","summary_state":"pending","summary":null,"crux":null,"owner":"Index_Builder"},{"id":"search_r1/search/index_builder.py#main","name":"main","kind":"function","path":"search_r1/search/index_builder.py","span":"L296-L344","signature":"def main()","exported":true,"origin":"ast","body_hash":"a1bed8de3a8e93a31af6a93435d9b615b3f47ff0a4d8a199443f772b34e4b687","body_text":"def main(): parser = argparse.ArgumentParser(description = \"Creating index.\") # Basic parameters parser.add_argument('--retrieval_method', type=str) parser.add_argument('--model_path', type=str, default=None) parser.add_argument('--corpus_path', type=str) parser.add_argument('--save_dir', default= 'indexes/',type=str) # Parameters for building dense index parser.add_argument('--max_length', type=int, default=180) parser.add_argument('--batch_size', type=int, default=512) parser.add_argument('--use_fp16', default=False, action='store_true') parser.add_argument('--pooling_method', type=str, default=None) parser.add_argument('--faiss_type',default=None,type=str) parser.add_argument('--embedding_path', default=None, type=str) parser.add_argument('--save_embedding', action='store_true', default=False) parser.add_argument('--faiss_gpu', default=False, action='store_true') args = parser.parse_args() if args.pooling_method is None: pooling_method = 'mean' for k,v in MODEL2POOLING.items(): if k in args.retrieval_method.lower(): pooling_method = v break else: if args.pooling_method not in ['mean','cls','pooler']: raise NotImplementedError else: pooling_method = args.pooling_method index_builder = Index_Builder( retrieval_method = args.retrieval_method, model_path = args.model_path, corpus_path = args.corpus_path, save_dir = args.save_dir, max_length = args.max_length, batch_size = args.batch_size, use_fp16 = args.use_fp16, pooling_method = pooling_method, faiss_type = args.faiss_type, embedding_path = args.embedding_path, save_embedding = args.save_embedding, faiss_gpu = args.faiss_gpu ) index_builder.build_index()","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"os","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"faiss","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"json","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"warnings","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"typing","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"shutil","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"subprocess","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"argparse","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"torch","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"tqdm","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"datasets","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"imports","specifier":"transformers","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"contains","targetId":"search_r1/search/index_builder.py#load_model","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#load_model","relation":"calls","name":"eval","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#load_model","relation":"calls","name":"cuda","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#load_model","relation":"calls","name":"half","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"contains","targetId":"search_r1/search/index_builder.py#pooling","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#pooling","relation":"calls","name":"masked_fill","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#pooling","relation":"calls","name":"bool","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#pooling","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#pooling","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#pooling","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"contains","targetId":"search_r1/search/index_builder.py#load_corpus","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#load_corpus","relation":"calls","name":"load_dataset","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"contains","targetId":"search_r1/search/index_builder.py#Index_Builder","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder","relation":"contains","targetId":"search_r1/search/index_builder.py#Index_Builder.__init__","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"lower","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"device_count","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"exists","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"makedirs","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"_check_dir","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"Index_Builder"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"warn","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"join","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"join","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"load_corpus","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.__init__","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder","relation":"contains","targetId":"search_r1/search/index_builder.py#Index_Builder._check_dir","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder._check_dir","relation":"calls","name":"isdir","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder._check_dir","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder._check_dir","relation":"calls","name":"listdir","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder._check_dir","relation":"calls","name":"makedirs","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder","relation":"contains","targetId":"search_r1/search/index_builder.py#Index_Builder.build_index","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_index","relation":"calls","name":"build_bm25_index","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"Index_Builder"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_index","relation":"calls","name":"build_dense_index","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"Index_Builder"},{"source":"search_r1/search/index_builder.py#Index_Builder","relation":"contains","targetId":"search_r1/search/index_builder.py#Index_Builder.build_bm25_index","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_bm25_index","relation":"calls","name":"join","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_bm25_index","relation":"calls","name":"makedirs","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_bm25_index","relation":"calls","name":"makedirs","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_bm25_index","relation":"calls","name":"copyfile","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_bm25_index","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_bm25_index","relation":"calls","name":"run","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_bm25_index","relation":"calls","name":"rmtree","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_bm25_index","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder","relation":"contains","targetId":"search_r1/search/index_builder.py#Index_Builder._load_embedding","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder._load_embedding","relation":"calls","name":"reshape","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder._load_embedding","relation":"calls","name":"memmap","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder","relation":"contains","targetId":"search_r1/search/index_builder.py#Index_Builder._save_embedding","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder._save_embedding","relation":"calls","name":"memmap","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder._save_embedding","relation":"calls","name":"tqdm","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder._save_embedding","relation":"calls","name":"range","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder._save_embedding","relation":"calls","name":"min","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder","relation":"contains","targetId":"search_r1/search/index_builder.py#Index_Builder.encode_all","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"DataParallel","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"tqdm","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"range","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"zip","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"to","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"tokenizer","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"Index_Builder"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"cuda","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"items","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"type","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"to","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"zeros","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"encoder","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"Index_Builder"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"encoder","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"Index_Builder"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"pooling","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"normalize","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"cast","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"numpy","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"cpu","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"detach","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"cast"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"concatenate","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.encode_all","relation":"calls","name":"astype","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder","relation":"calls","name":"no_grad","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder","relation":"contains","targetId":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"exists","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"load_model","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"_load_embedding","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"Index_Builder"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"encode_all","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"Index_Builder"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"_save_embedding","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"Index_Builder"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"index_factory","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"GpuMultipleClonerOptions","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"index_cpu_to_all_gpus","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"train","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"add","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"index_gpu_to_cpu","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"train","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"add","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"write_index","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#Index_Builder.build_dense_index","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py","relation":"contains","targetId":"search_r1/search/index_builder.py#main","file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"parse_args","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"items","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"lower","viaMember":true,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"Index_Builder","viaMember":false,"file":"search_r1/search/index_builder.py"},{"source":"search_r1/search/index_builder.py#main","relation":"calls","name":"build_index","viaMember":true,"file":"search_r1/search/index_builder.py","recvType":"Index_Builder"},{"source":"search_r1/search/index_builder.py","relation":"calls","name":"main","viaMember":false,"file":"search_r1/search/index_builder.py"}]},"search_r1/search/retrieval.py":{"size":12700,"mtimeMs":1789275830306.2024,"hash":"bab4cff67fb67a034dad561c5fef6b4017f0f3d612a9fcbd65f630e5ec4efe62","nodes":[{"id":"search_r1/search/retrieval.py","name":"retrieval.py","kind":"file","path":"search_r1/search/retrieval.py","span":"L1-L369","signature":null,"exported":true,"origin":"ast","body_hash":"bab4cff67fb67a034dad561c5fef6b4017f0f3d612a9fcbd65f630e5ec4efe62","chars":12700,"summary_state":"pending","summary":null,"crux":null,"body_text":"import json import os import warnings from typing import List, Dict import functools from tqdm import tqdm from multiprocessing import Pool import faiss import torch import numpy as np from transformers import AutoConfig, AutoTokenizer, AutoModel import argparse import datasets if __name__ == '__main__': parser = argparse.ArgumentParser(description = \"Retrieval\") # Basic parameters parser.add_argument('--retrieval_method', type=str) parser.add_argument('--retrieval_topk', type=int, default=10) parser.add_argument('--index_path', type=str, default=None) parser.add_argument('--corpus_path', type=str) parser.add_argument('--dataset_path', default=None, type=str) parser.add_argument('--faiss_gpu', default=True, type=bool) parser.add_argument('--data_split', default=\"train\", type=str) parser.add_argument('--retrieval_model_path', type=str, default=None) parser.add_argument('--retrieval_pooling_method', default='mean', type=str) parser.add_argument('--retrieval_query_max_length', default=256, type=str) parser.add_argument('--retrieval_use_fp16', action='store_true', default=False) parser.add_argument('--retrieval_batch_size', default=512, type=int) args = parser.parse_args() args.index_path = os.path.join(args.index_path, f'{args.retrieval_method}_Flat.index') if args.retrieval_method != 'bm25' else os.path.join(args.index_path, 'bm25') # load dataset all_split = get_dataset(args) input_query = [sample['question'] for sample in all_split[:512]] # initialize the retriever and conduct retrieval retriever = get_retriever(args) print('Start Retrieving ...') results, scores = retriever.batch_search(input_query, return_score=True) # from IPython import embed # embed()"},{"id":"search_r1/search/retrieval.py#load_corpus","name":"load_corpus","kind":"function","path":"search_r1/search/retrieval.py","span":"L16-L22","signature":"def load_corpus(corpus_path: str)","exported":true,"origin":"ast","body_hash":"63642a3531f6506e9fae80592a2170b344a1b6179a3d27a3002081529a9e3d23","body_text":"def load_corpus(corpus_path: str): corpus = datasets.load_dataset( 'json', data_files=corpus_path, split=\"train\", num_proc=4) return corpus","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#read_jsonl","name":"read_jsonl","kind":"function","path":"search_r1/search/retrieval.py","span":"L25-L32","signature":"def read_jsonl(file_path)","exported":true,"origin":"ast","body_hash":"4c47880b2c2ac231addf857ba43f05bdfc1be3f19de4f48067418ca65a407f1d","body_text":"def read_jsonl(file_path): data = [] with open(file_path, \"r\") as f: readin = f.readlines() for line in readin: data.append(json.loads(line)) return data","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#load_docs","name":"load_docs","kind":"function","path":"search_r1/search/retrieval.py","span":"L35-L38","signature":"def load_docs(corpus, doc_idxs)","exported":true,"origin":"ast","body_hash":"6391624ce3919bef69c4bd7c04e58574246a106871a10aa208bcde98337a2922","body_text":"def load_docs(corpus, doc_idxs): results = [corpus[int(idx)] for idx in doc_idxs] return results","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#load_model","name":"load_model","kind":"function","path":"search_r1/search/retrieval.py","span":"L41-L53","signature":"def load_model( model_path: str, use_fp16: bool = False )","exported":true,"origin":"ast","body_hash":"2d4e1f08f900b69638da6c1525eb3006b46ed0822e5e3f12fb86d8faa60ab34b","body_text":"def load_model( model_path: str, use_fp16: bool = False ): model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) model = AutoModel.from_pretrained(model_path, trust_remote_code=True) model.eval() model.cuda() if use_fp16: model = model.half() tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, trust_remote_code=True) return model, tokenizer","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#pooling","name":"pooling","kind":"function","path":"search_r1/search/retrieval.py","span":"L56-L70","signature":"def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = \"mean\" )","exported":true,"origin":"ast","body_hash":"6e356a5fbac440ef6226d2119161f4cf7a669ceb3505baddff6331a4a50030ec","body_text":"def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = \"mean\" ): if pooling_method == \"mean\": last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0) return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None] elif pooling_method == \"cls\": return last_hidden_state[:, 0] elif pooling_method == \"pooler\": return pooler_output else: raise NotImplementedError(\"Pooling method not implemented!\")","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#Encoder","name":"Encoder","kind":"class","path":"search_r1/search/retrieval.py","span":"L73-L129","signature":"class Encoder","exported":true,"origin":"ast","body_hash":"a513ae91bd7e48309b1f39780b3beac631d16168574d922913ff8662dd199b11","body_text":"class Encoder: def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16): self.model_name = model_name self.model_path = model_path self.pooling_method = pooling_method self.max_length = max_length self.use_fp16 = use_fp16 self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16) @torch.no_grad() def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders if isinstance(query_list, str): query_list = [query_list] if \"e5\" in self.model_name.lower(): if is_query: query_list = [f\"query: {query}\" for query in query_list] else: query_list = [f\"passage: {query}\" for query in query_list] if \"bge\" in self.model_name.lower(): if is_query: query_list = [f\"Represent this sentence for searching relevant passages: {query}\" for query in query_list] inputs = self.tokenizer(query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors=\"pt\" ) inputs = {k: v.cuda() for k, v in inputs.items()} if \"T5\" in type(self.model).__name__: # T5-based retrieval model decoder_input_ids = torch.zeros( (inputs['input_ids'].shape[0], 1), dtype=torch.long ).to(inputs['input_ids'].device) output = self.model( **inputs, decoder_input_ids=decoder_input_ids, return_dict=True ) query_emb = output.last_hidden_state[:, 0, :] else: output = self.model(**inputs, return_dict=True) query_emb = pooling(output.pooler_output, output.last_hidden_state, inputs['attention_mask'], self.pooling_method) if \"dpr\" not in self.model_name.lower(): query_emb = torch.nn.functional.normalize(query_emb, dim=-1) query_emb = query_emb.detach().cpu().numpy() query_emb = query_emb.astype(np.float32, order=\"C\") return query_emb","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#Encoder.__init__","name":"__init__","kind":"method","path":"search_r1/search/retrieval.py","span":"L74-L82","signature":"def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16)","exported":false,"origin":"ast","body_hash":"91ad5eec7d331441c6f1b920061f211d1e755e85e3509cf6d2b1b47a3215c1dd","body_text":"def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16): self.model_name = model_name self.model_path = model_path self.pooling_method = pooling_method self.max_length = max_length self.use_fp16 = use_fp16 self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16)","summary_state":"pending","summary":null,"crux":null,"owner":"Encoder"},{"id":"search_r1/search/retrieval.py#Encoder.encode","name":"encode","kind":"method","path":"search_r1/search/retrieval.py","span":"L85-L129","signature":"def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders","exported":true,"origin":"ast","body_hash":"6a6511e1300b1aa62cedb981fcdf443894d0ee025705cf692573931f36fd8dc0","body_text":"def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders if isinstance(query_list, str): query_list = [query_list] if \"e5\" in self.model_name.lower(): if is_query: query_list = [f\"query: {query}\" for query in query_list] else: query_list = [f\"passage: {query}\" for query in query_list] if \"bge\" in self.model_name.lower(): if is_query: query_list = [f\"Represent this sentence for searching relevant passages: {query}\" for query in query_list] inputs = self.tokenizer(query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors=\"pt\" ) inputs = {k: v.cuda() for k, v in inputs.items()} if \"T5\" in type(self.model).__name__: # T5-based retrieval model decoder_input_ids = torch.zeros( (inputs['input_ids'].shape[0], 1), dtype=torch.long ).to(inputs['input_ids'].device) output = self.model( **inputs, decoder_input_ids=decoder_input_ids, return_dict=True ) query_emb = output.last_hidden_state[:, 0, :] else: output = self.model(**inputs, return_dict=True) query_emb = pooling(output.pooler_output, output.last_hidden_state, inputs['attention_mask'], self.pooling_method) if \"dpr\" not in self.model_name.lower(): query_emb = torch.nn.functional.normalize(query_emb, dim=-1) query_emb = query_emb.detach().cpu().numpy() query_emb = query_emb.astype(np.float32, order=\"C\") return query_emb","summary_state":"pending","summary":null,"crux":null,"owner":"Encoder"},{"id":"search_r1/search/retrieval.py#BaseRetriever","name":"BaseRetriever","kind":"class","path":"search_r1/search/retrieval.py","span":"L132-L162","signature":"class BaseRetriever","exported":true,"origin":"ast","body_hash":"1ce69eef8b2362b2435db716c275651dd2e03bf182aa743b91f8230edd7fabef","body_text":"class BaseRetriever: \"\"\"Base object for all retrievers.\"\"\" def __init__(self, config): self.config = config self.retrieval_method = config.retrieval_method self.topk = config.retrieval_topk self.index_path = config.index_path self.corpus_path = config.corpus_path # self.cache_save_path = os.path.join(config.save_dir, 'retrieval_cache.json') def _search(self, query: str, num: int, return_score:bool) -> List[Dict[str, str]]: r\"\"\"Retrieve topk relevant documents in corpus. Return: list: contains information related to the document, including: contents: used for building index title: (if provided) text: (if provided) \"\"\" pass def _batch_search(self, query_list, num, return_score): pass def search(self, *args, **kwargs): return self._search(*args, **kwargs) def batch_search(self, *args, **kwargs): return self._batch_search(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#BaseRetriever.__init__","name":"__init__","kind":"method","path":"search_r1/search/retrieval.py","span":"L135-L143","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"1592cd9ea0215fbf4afe28ec72ebc8027028107edbd1f2bfb86b7b06a1204680","body_text":"def __init__(self, config): self.config = config self.retrieval_method = config.retrieval_method self.topk = config.retrieval_topk self.index_path = config.index_path self.corpus_path = config.corpus_path # self.cache_save_path = os.path.join(config.save_dir, 'retrieval_cache.json')","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"search_r1/search/retrieval.py#BaseRetriever._search","name":"_search","kind":"method","path":"search_r1/search/retrieval.py","span":"L145-L153","signature":"def _search(self, query: str, num: int, return_score:bool) -> List[Dict[str, str]]","exported":false,"origin":"ast","body_hash":"a1863af52b58a6982768a8c05eb0ceae5946ccc8ab516251836104adf4c6d0a3","body_text":"def _search(self, query: str, num: int, return_score:bool) -> List[Dict[str, str]]: r\"\"\"Retrieve topk relevant documents in corpus. Return: list: contains information related to the document, including: contents: used for building index title: (if provided) text: (if provided) \"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"search_r1/search/retrieval.py#BaseRetriever._batch_search","name":"_batch_search","kind":"method","path":"search_r1/search/retrieval.py","span":"L155-L156","signature":"def _batch_search(self, query_list, num, return_score)","exported":false,"origin":"ast","body_hash":"3a4873d6a4beb109c53552d834cfc803b0c161cbc908217f0aea5c67d4667fb1","body_text":"def _batch_search(self, query_list, num, return_score): pass","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"search_r1/search/retrieval.py#BaseRetriever.search","name":"search","kind":"method","path":"search_r1/search/retrieval.py","span":"L158-L159","signature":"def search(self, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"3d2bd614b882a4bd4196108a68168e90abd4d275d61e4c7e493e3800af3888a1","body_text":"def search(self, *args, **kwargs): return self._search(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"search_r1/search/retrieval.py#BaseRetriever.batch_search","name":"batch_search","kind":"method","path":"search_r1/search/retrieval.py","span":"L161-L162","signature":"def batch_search(self, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"d37a6c77b70ed1c7532c9a1a534aa2bb6e6e942eeacc42bc02faa2e163da8629","body_text":"def batch_search(self, *args, **kwargs): return self._batch_search(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"search_r1/search/retrieval.py#BM25Retriever","name":"BM25Retriever","kind":"class","path":"search_r1/search/retrieval.py","span":"L165-L224","signature":"class BM25Retriever(BaseRetriever)","exported":true,"origin":"ast","body_hash":"18bc06875d301abb97a7e3e5cc3ffaace07af3b8d85938f73ab0b7322bdb9f35","body_text":"class BM25Retriever(BaseRetriever): r\"\"\"BM25 retriever based on pre-built pyserini index.\"\"\" def __init__(self, config): super().__init__(config) from pyserini.search.lucene import LuceneSearcher self.searcher = LuceneSearcher(self.index_path) self.contain_doc = self._check_contain_doc() if not self.contain_doc: self.corpus = load_corpus(self.corpus_path) self.max_process_num = 8 def _check_contain_doc(self): r\"\"\"Check if the index contains document content \"\"\" return self.searcher.doc(0).raw() is not None def _search(self, query: str, num: int = None, return_score = False) -> List[Dict[str, str]]: if num is None: num = self.topk hits = self.searcher.search(query, num) if len(hits) < 1: if return_score: return [],[] else: return [] scores = [hit.score for hit in hits] if len(hits) < num: warnings.warn('Not enough documents retrieved!') else: hits = hits[:num] if self.contain_doc: all_contents = [json.loads(self.searcher.doc(hit.docid).raw())['contents'] for hit in hits] results = [{'title': content.split(\"\\n\")[0].strip(\"\\\"\"), 'text': \"\\n\".join(content.split(\"\\n\")[1:]), 'contents': content} for content in all_contents] else: results = load_docs(self.corpus, [hit.docid for hit in hits]) if return_score: return results, scores else: return results def _batch_search(self, query_list, num: int = None, return_score = False): # TODO: modify batch method results = [] scores = [] for query in query_list: item_result, item_score = self._search(query, num,True) results.append(item_result) scores.append(item_score) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#BM25Retriever.__init__","name":"__init__","kind":"method","path":"search_r1/search/retrieval.py","span":"L168-L175","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"b9c8fcd86e95a5d5fb5a36d4aaddf2661d2b39e1467fc208c07f078d58d1626b","body_text":"def __init__(self, config): super().__init__(config) from pyserini.search.lucene import LuceneSearcher self.searcher = LuceneSearcher(self.index_path) self.contain_doc = self._check_contain_doc() if not self.contain_doc: self.corpus = load_corpus(self.corpus_path) self.max_process_num = 8","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"search_r1/search/retrieval.py#BM25Retriever._check_contain_doc","name":"_check_contain_doc","kind":"method","path":"search_r1/search/retrieval.py","span":"L177-L180","signature":"def _check_contain_doc(self)","exported":false,"origin":"ast","body_hash":"f7a574ad883e23e81357c53f3a2947ddd8491adbc804954c0d23985aa507d3c4","body_text":"def _check_contain_doc(self): r\"\"\"Check if the index contains document content \"\"\" return self.searcher.doc(0).raw() is not None","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"search_r1/search/retrieval.py#BM25Retriever._search","name":"_search","kind":"method","path":"search_r1/search/retrieval.py","span":"L182-L210","signature":"def _search(self, query: str, num: int = None, return_score = False) -> List[Dict[str, str]]","exported":false,"origin":"ast","body_hash":"a07f442a9561ebd63dc4c1e3f79003636e2f38b206f554402e2ee3148fd9a28f","body_text":"def _search(self, query: str, num: int = None, return_score = False) -> List[Dict[str, str]]: if num is None: num = self.topk hits = self.searcher.search(query, num) if len(hits) < 1: if return_score: return [],[] else: return [] scores = [hit.score for hit in hits] if len(hits) < num: warnings.warn('Not enough documents retrieved!') else: hits = hits[:num] if self.contain_doc: all_contents = [json.loads(self.searcher.doc(hit.docid).raw())['contents'] for hit in hits] results = [{'title': content.split(\"\\n\")[0].strip(\"\\\"\"), 'text': \"\\n\".join(content.split(\"\\n\")[1:]), 'contents': content} for content in all_contents] else: results = load_docs(self.corpus, [hit.docid for hit in hits]) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"search_r1/search/retrieval.py#BM25Retriever._batch_search","name":"_batch_search","kind":"method","path":"search_r1/search/retrieval.py","span":"L212-L224","signature":"def _batch_search(self, query_list, num: int = None, return_score = False): # TODO: modify batch method","exported":false,"origin":"ast","body_hash":"b569a216808f8f2a47a87bac80d0ffb727e6e0c4e1cf265dcd4d35ed86f12b0a","body_text":"def _batch_search(self, query_list, num: int = None, return_score = False): # TODO: modify batch method results = [] scores = [] for query in query_list: item_result, item_score = self._search(query, num,True) results.append(item_result) scores.append(item_score) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"search_r1/search/retrieval.py#get_available_gpu_memory","name":"get_available_gpu_memory","kind":"function","path":"search_r1/search/retrieval.py","span":"L226-L233","signature":"def get_available_gpu_memory()","exported":true,"origin":"ast","body_hash":"05f2123a0d851b4102dc85306cc3b77385e3dcef549527e9d5557e96833eac8e","body_text":"def get_available_gpu_memory(): memory_info = [] for i in range(torch.cuda.device_count()): total_memory = torch.cuda.get_device_properties(i).total_memory allocated_memory = torch.cuda.memory_allocated(i) free_memory = total_memory - allocated_memory memory_info.append((i, free_memory / 1e9)) # Convert to GB return memory_info","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#DenseRetriever","name":"DenseRetriever","kind":"class","path":"search_r1/search/retrieval.py","span":"L236-L309","signature":"class DenseRetriever(BaseRetriever)","exported":true,"origin":"ast","body_hash":"80b057df367d81e26c6dba9eaa0fdd8d2e21d862b349a01b786202ad097b2a9b","body_text":"class DenseRetriever(BaseRetriever): r\"\"\"Dense retriever based on pre-built faiss index.\"\"\" def __init__(self, config: dict): super().__init__(config) self.index = faiss.read_index(self.index_path) if config.faiss_gpu: co = faiss.GpuMultipleClonerOptions() co.useFloat16 = True co.shard = True self.index = faiss.index_cpu_to_all_gpus(self.index, co=co) # self.index = faiss.index_cpu_to_all_gpus(self.index) self.corpus = load_corpus(self.corpus_path) self.encoder = Encoder( model_name = self.retrieval_method, model_path = config.retrieval_model_path, pooling_method = config.retrieval_pooling_method, max_length = config.retrieval_query_max_length, use_fp16 = config.retrieval_use_fp16 ) self.topk = config.retrieval_topk self.batch_size = self.config.retrieval_batch_size def _search(self, query: str, num: int = None, return_score = False): if num is None: num = self.topk query_emb = self.encoder.encode(query) scores, idxs = self.index.search(query_emb, k=num) idxs = idxs[0] scores = scores[0] results = load_docs(self.corpus, idxs) if return_score: return results, scores else: return results def _batch_search(self, query_list: List[str], num: int = None, return_score = False): if isinstance(query_list, str): query_list = [query_list] if num is None: num = self.topk batch_size = self.batch_size results = [] scores = [] for start_idx in tqdm(range(0, len(query_list), batch_size), desc='Retrieval process: '): query_batch = query_list[start_idx:start_idx + batch_size] # from time import time # a = time() batch_emb = self.encoder.encode(query_batch) # b = time() # print(f'################### encode time {b-a} #####################') batch_scores, batch_idxs = self.index.search(batch_emb, k=num) batch_scores = batch_scores.tolist() batch_idxs = batch_idxs.tolist() # print(f'################### search time {time()-b} #####################') # exit() flat_idxs = sum(batch_idxs, []) batch_results = load_docs(self.corpus, flat_idxs) batch_results = [batch_results[i*num : (i+1)*num] for i in range(len(batch_idxs))] scores.extend(batch_scores) results.extend(batch_results) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#DenseRetriever.__init__","name":"__init__","kind":"method","path":"search_r1/search/retrieval.py","span":"L239-L258","signature":"def __init__(self, config: dict)","exported":false,"origin":"ast","body_hash":"5cd69be146497d06ea80a44cf5eb03a9d37bcdc7eb4e1cf51cc83da3c421294b","body_text":"def __init__(self, config: dict): super().__init__(config) self.index = faiss.read_index(self.index_path) if config.faiss_gpu: co = faiss.GpuMultipleClonerOptions() co.useFloat16 = True co.shard = True self.index = faiss.index_cpu_to_all_gpus(self.index, co=co) # self.index = faiss.index_cpu_to_all_gpus(self.index) self.corpus = load_corpus(self.corpus_path) self.encoder = Encoder( model_name = self.retrieval_method, model_path = config.retrieval_model_path, pooling_method = config.retrieval_pooling_method, max_length = config.retrieval_query_max_length, use_fp16 = config.retrieval_use_fp16 ) self.topk = config.retrieval_topk self.batch_size = self.config.retrieval_batch_size","summary_state":"pending","summary":null,"crux":null,"owner":"DenseRetriever"},{"id":"search_r1/search/retrieval.py#DenseRetriever._search","name":"_search","kind":"method","path":"search_r1/search/retrieval.py","span":"L260-L272","signature":"def _search(self, query: str, num: int = None, return_score = False)","exported":false,"origin":"ast","body_hash":"fcfb5c0e8906a79e9da094fbea0615a2453c04c999b61a71f5d388d653eeb25f","body_text":"def _search(self, query: str, num: int = None, return_score = False): if num is None: num = self.topk query_emb = self.encoder.encode(query) scores, idxs = self.index.search(query_emb, k=num) idxs = idxs[0] scores = scores[0] results = load_docs(self.corpus, idxs) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"DenseRetriever"},{"id":"search_r1/search/retrieval.py#DenseRetriever._batch_search","name":"_batch_search","kind":"method","path":"search_r1/search/retrieval.py","span":"L274-L309","signature":"def _batch_search(self, query_list: List[str], num: int = None, return_score = False)","exported":false,"origin":"ast","body_hash":"ebf27a073c2ea006af0b6ce3bf9c3440f698d20995f31395b6d128966a00b635","body_text":"def _batch_search(self, query_list: List[str], num: int = None, return_score = False): if isinstance(query_list, str): query_list = [query_list] if num is None: num = self.topk batch_size = self.batch_size results = [] scores = [] for start_idx in tqdm(range(0, len(query_list), batch_size), desc='Retrieval process: '): query_batch = query_list[start_idx:start_idx + batch_size] # from time import time # a = time() batch_emb = self.encoder.encode(query_batch) # b = time() # print(f'################### encode time {b-a} #####################') batch_scores, batch_idxs = self.index.search(batch_emb, k=num) batch_scores = batch_scores.tolist() batch_idxs = batch_idxs.tolist() # print(f'################### search time {time()-b} #####################') # exit() flat_idxs = sum(batch_idxs, []) batch_results = load_docs(self.corpus, flat_idxs) batch_results = [batch_results[i*num : (i+1)*num] for i in range(len(batch_idxs))] scores.extend(batch_scores) results.extend(batch_results) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"DenseRetriever"},{"id":"search_r1/search/retrieval.py#get_retriever","name":"get_retriever","kind":"function","path":"search_r1/search/retrieval.py","span":"L311-L323","signature":"def get_retriever(config)","exported":true,"origin":"ast","body_hash":"95eccf1a0555d1c2cb970bddb44c981ee738b92f26a34fb6aa7a7efc653ea830","body_text":"def get_retriever(config): r\"\"\"Automatically select retriever class based on config's retrieval method Args: config (dict): configuration with 'retrieval_method' key Returns: Retriever: retriever instance \"\"\" if config.retrieval_method == \"bm25\": return BM25Retriever(config) else: return DenseRetriever(config)","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval.py#get_dataset","name":"get_dataset","kind":"function","path":"search_r1/search/retrieval.py","span":"L326-L330","signature":"def get_dataset(config)","exported":true,"origin":"ast","body_hash":"cb479cb13378fd0dfa376417ff3fd23c2f1f972d04a6f5e46780bc6a40bbf20d","body_text":"def get_dataset(config): \"\"\"Load dataset from config.\"\"\" split_path = os.path.join(config.dataset_path, f'{config.data_split}.jsonl') return read_jsonl(split_path)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"json","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"os","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"warnings","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"typing","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"functools","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"tqdm","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"multiprocessing","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"faiss","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"torch","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"transformers","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"argparse","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"datasets","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#load_corpus","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#load_corpus","relation":"calls","name":"load_dataset","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#read_jsonl","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#read_jsonl","relation":"calls","name":"open","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#read_jsonl","relation":"calls","name":"readlines","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#read_jsonl","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#read_jsonl","relation":"calls","name":"loads","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#load_docs","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#load_docs","relation":"calls","name":"int","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#load_model","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#load_model","relation":"calls","name":"eval","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#load_model","relation":"calls","name":"cuda","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#load_model","relation":"calls","name":"half","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#pooling","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#pooling","relation":"calls","name":"masked_fill","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#pooling","relation":"calls","name":"bool","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#pooling","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#pooling","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#pooling","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#Encoder","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder","relation":"contains","targetId":"search_r1/search/retrieval.py#Encoder.__init__","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.__init__","relation":"calls","name":"load_model","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder","relation":"calls","name":"no_grad","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder","relation":"contains","targetId":"search_r1/search/retrieval.py#Encoder.encode","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"isinstance","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"lower","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"lower","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"tokenizer","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"Encoder"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"cuda","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"items","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"type","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"to","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"zeros","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"model","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"Encoder"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"model","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"Encoder"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"pooling","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"lower","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"normalize","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"numpy","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"cpu","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"detach","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"pooling"},{"source":"search_r1/search/retrieval.py#Encoder.encode","relation":"calls","name":"astype","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"pooling"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#BaseRetriever","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BaseRetriever","relation":"contains","targetId":"search_r1/search/retrieval.py#BaseRetriever.__init__","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BaseRetriever","relation":"contains","targetId":"search_r1/search/retrieval.py#BaseRetriever._search","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BaseRetriever","relation":"contains","targetId":"search_r1/search/retrieval.py#BaseRetriever._batch_search","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BaseRetriever","relation":"contains","targetId":"search_r1/search/retrieval.py#BaseRetriever.search","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BaseRetriever.search","relation":"calls","name":"_search","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"BaseRetriever"},{"source":"search_r1/search/retrieval.py#BaseRetriever","relation":"contains","targetId":"search_r1/search/retrieval.py#BaseRetriever.batch_search","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BaseRetriever.batch_search","relation":"calls","name":"_batch_search","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"BaseRetriever"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#BM25Retriever","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever","relation":"extends","name":"BaseRetriever","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever","relation":"contains","targetId":"search_r1/search/retrieval.py#BM25Retriever.__init__","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever.__init__","relation":"calls","name":"super","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"imports","specifier":"pyserini.search.lucene","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever.__init__","relation":"calls","name":"LuceneSearcher","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever.__init__","relation":"calls","name":"_check_contain_doc","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"BM25Retriever"},{"source":"search_r1/search/retrieval.py#BM25Retriever.__init__","relation":"calls","name":"load_corpus","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever","relation":"contains","targetId":"search_r1/search/retrieval.py#BM25Retriever._check_contain_doc","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._check_contain_doc","relation":"calls","name":"raw","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._check_contain_doc","relation":"calls","name":"doc","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"LuceneSearcher"},{"source":"search_r1/search/retrieval.py#BM25Retriever","relation":"contains","targetId":"search_r1/search/retrieval.py#BM25Retriever._search","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"search","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"LuceneSearcher"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"warn","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"loads","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"raw","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"doc","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"LuceneSearcher"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"strip","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"split","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"join","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"split","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._search","relation":"calls","name":"load_docs","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever","relation":"contains","targetId":"search_r1/search/retrieval.py#BM25Retriever._batch_search","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._batch_search","relation":"calls","name":"_search","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"BM25Retriever"},{"source":"search_r1/search/retrieval.py#BM25Retriever._batch_search","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#BM25Retriever._batch_search","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#get_available_gpu_memory","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#get_available_gpu_memory","relation":"calls","name":"range","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#get_available_gpu_memory","relation":"calls","name":"device_count","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#get_available_gpu_memory","relation":"calls","name":"get_device_properties","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#get_available_gpu_memory","relation":"calls","name":"memory_allocated","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#get_available_gpu_memory","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#DenseRetriever","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever","relation":"extends","name":"BaseRetriever","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever","relation":"contains","targetId":"search_r1/search/retrieval.py#DenseRetriever.__init__","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever.__init__","relation":"calls","name":"super","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever.__init__","relation":"calls","name":"read_index","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever.__init__","relation":"calls","name":"GpuMultipleClonerOptions","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever.__init__","relation":"calls","name":"index_cpu_to_all_gpus","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever.__init__","relation":"calls","name":"load_corpus","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever.__init__","relation":"calls","name":"Encoder","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever","relation":"contains","targetId":"search_r1/search/retrieval.py#DenseRetriever._search","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._search","relation":"calls","name":"encode","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"Encoder"},{"source":"search_r1/search/retrieval.py#DenseRetriever._search","relation":"calls","name":"search","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._search","relation":"calls","name":"load_docs","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever","relation":"contains","targetId":"search_r1/search/retrieval.py#DenseRetriever._batch_search","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"isinstance","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"tqdm","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"range","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"encode","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"Encoder"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"search","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"sum","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"load_docs","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"range","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"extend","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#DenseRetriever._batch_search","relation":"calls","name":"extend","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#get_retriever","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#get_retriever","relation":"calls","name":"BM25Retriever","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#get_retriever","relation":"calls","name":"DenseRetriever","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"contains","targetId":"search_r1/search/retrieval.py#get_dataset","file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#get_dataset","relation":"calls","name":"join","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py#get_dataset","relation":"calls","name":"read_jsonl","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"parse_args","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"join","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"join","viaMember":true,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"get_dataset","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"get_retriever","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/retrieval.py"},{"source":"search_r1/search/retrieval.py","relation":"calls","name":"batch_search","viaMember":true,"file":"search_r1/search/retrieval.py","recvType":"get_retriever"}]},"search_r1/search/retrieval_request.py":{"size":500,"mtimeMs":1789275830306.2024,"hash":"469671b8625574a7fc75011c0fb566004a12c3208dcea4b2aba010fe975f6c6b","nodes":[{"id":"search_r1/search/retrieval_request.py","name":"retrieval_request.py","kind":"file","path":"search_r1/search/retrieval_request.py","span":"L1-L24","signature":null,"exported":true,"origin":"ast","body_hash":"469671b8625574a7fc75011c0fb566004a12c3208dcea4b2aba010fe975f6c6b","chars":500,"summary_state":"pending","summary":null,"crux":null,"body_text":"import requests # URL for your local FastAPI server url = \"http://127.0.0.1:8000/retrieve\" # Example payload payload = { \"queries\": [\"What is the capital of France?\", \"Explain neural networks.\"] * 200, \"topk\": 5, \"return_scores\": True } # Send POST request response = requests.post(url, json=payload) # Raise an exception if the request failed response.raise_for_status() # Get the JSON response retrieved_data = response.json() print(\"Response from server:\") print(retrieved_data)"}],"rawEdges":[{"source":"search_r1/search/retrieval_request.py","relation":"imports","specifier":"requests","file":"search_r1/search/retrieval_request.py"},{"source":"search_r1/search/retrieval_request.py","relation":"calls","name":"post","viaMember":true,"file":"search_r1/search/retrieval_request.py"},{"source":"search_r1/search/retrieval_request.py","relation":"calls","name":"raise_for_status","viaMember":true,"file":"search_r1/search/retrieval_request.py"},{"source":"search_r1/search/retrieval_request.py","relation":"calls","name":"json","viaMember":true,"file":"search_r1/search/retrieval_request.py"},{"source":"search_r1/search/retrieval_request.py","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/retrieval_request.py"},{"source":"search_r1/search/retrieval_request.py","relation":"calls","name":"print","viaMember":false,"file":"search_r1/search/retrieval_request.py"}]},"search_r1/search/retrieval_server.py":{"size":13482,"mtimeMs":1789275830306.2024,"hash":"a420181f59974cd463d864d6fadd217b85dcc2130c595c14d02b833c2584a3b6","nodes":[{"id":"search_r1/search/retrieval_server.py","name":"retrieval_server.py","kind":"file","path":"search_r1/search/retrieval_server.py","span":"L1-L383","signature":null,"exported":true,"origin":"ast","body_hash":"a420181f59974cd463d864d6fadd217b85dcc2130c595c14d02b833c2584a3b6","chars":13482,"summary_state":"pending","summary":null,"crux":null,"body_text":"import json import os import warnings from typing import List, Dict, Optional import argparse import faiss import torch import numpy as np from transformers import AutoConfig, AutoTokenizer, AutoModel from tqdm import tqdm import datasets import uvicorn from fastapi import FastAPI from pydantic import BaseModel parser = argparse.ArgumentParser(description=\"Launch the local faiss retriever.\") parser.add_argument(\"--index_path\", type=str, default=\"/home/peterjin/mnt/index/wiki-18/e5_Flat.index\", help=\"Corpus indexing file.\") parser.add_argument(\"--corpus_path\", type=str, default=\"/home/peterjin/mnt/data/retrieval-corpus/wiki-18.jsonl\", help=\"Local corpus file.\") parser.add_argument(\"--topk\", type=int, default=3, help=\"Number of retrieved passages for one query.\") parser.add_argument(\"--retriever_model\", type=str, default=\"intfloat/e5-base-v2\", help=\"Name of the retriever model.\") args = parser.parse_args() ##################################### # FastAPI server below ##################################### app = FastAPI() # 1) Build a config (could also parse from arguments). # In real usage, you'd parse your CLI arguments or environment variables. config = Config( retrieval_method = \"e5\", # or \"dense\" index_path=args.index_path, corpus_path=args.corpus_path, retrieval_topk=args.topk, faiss_gpu=True, retrieval_model_path=args.retriever_model, retrieval_pooling_method=\"mean\", retrieval_query_max_length=256, retrieval_use_fp16=True, retrieval_batch_size=512, ) # 2) Instantiate a global retriever so it is loaded once and reused. retriever = get_retriever(config) @app.post(\"/retrieve\") if __name__ == \"__main__\": # 3) Launch the server. By default, it listens on http://127.0.0.1:8000 uvicorn.run(app, host=\"0.0.0.0\", port=8000)"},{"id":"search_r1/search/retrieval_server.py#load_corpus","name":"load_corpus","kind":"function","path":"search_r1/search/retrieval_server.py","span":"L27-L34","signature":"def load_corpus(corpus_path: str)","exported":true,"origin":"ast","body_hash":"d9d3814371c2370fe29e0f7a30ae71eb08b596382c0f571a875ffb5fac036db8","body_text":"def load_corpus(corpus_path: str): corpus = datasets.load_dataset( 'json', data_files=corpus_path, split=\"train\", num_proc=4 ) return corpus","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#read_jsonl","name":"read_jsonl","kind":"function","path":"search_r1/search/retrieval_server.py","span":"L36-L41","signature":"def read_jsonl(file_path)","exported":true,"origin":"ast","body_hash":"acd404dc59d31053de24b5e6190c4a528f1eff1019a2a98ceb6d1779a380df66","body_text":"def read_jsonl(file_path): data = [] with open(file_path, \"r\") as f: for line in f: data.append(json.loads(line)) return data","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#load_docs","name":"load_docs","kind":"function","path":"search_r1/search/retrieval_server.py","span":"L43-L45","signature":"def load_docs(corpus, doc_idxs)","exported":true,"origin":"ast","body_hash":"81714fbc2d808de489a2631e4ffecc03e0004457c3c8fe0405a90988575e6231","body_text":"def load_docs(corpus, doc_idxs): results = [corpus[int(idx)] for idx in doc_idxs] return results","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#load_model","name":"load_model","kind":"function","path":"search_r1/search/retrieval_server.py","span":"L47-L55","signature":"def load_model(model_path: str, use_fp16: bool = False)","exported":true,"origin":"ast","body_hash":"824e2cdac2884a5194858ad8eebcb75a234255bd31f9b172f1cd92f66d43a5bd","body_text":"def load_model(model_path: str, use_fp16: bool = False): model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) model = AutoModel.from_pretrained(model_path, trust_remote_code=True) model.eval() model.cuda() if use_fp16: model = model.half() tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, trust_remote_code=True) return model, tokenizer","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#pooling","name":"pooling","kind":"function","path":"search_r1/search/retrieval_server.py","span":"L57-L71","signature":"def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = \"mean\" )","exported":true,"origin":"ast","body_hash":"6cca3097df29910b6a09b2ec4a1f24f674be73af9408895c52f6955fe7640d34","body_text":"def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = \"mean\" ): if pooling_method == \"mean\": last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0) return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None] elif pooling_method == \"cls\": return last_hidden_state[:, 0] elif pooling_method == \"pooler\": return pooler_output else: raise NotImplementedError(\"Pooling method not implemented!\")","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#Encoder","name":"Encoder","kind":"class","path":"search_r1/search/retrieval_server.py","span":"L73-L128","signature":"class Encoder","exported":true,"origin":"ast","body_hash":"1af11518030a2b67a59abd136cad4f0e77f365155de3c8f7be0e95812bd12e6c","body_text":"class Encoder: def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16): self.model_name = model_name self.model_path = model_path self.pooling_method = pooling_method self.max_length = max_length self.use_fp16 = use_fp16 self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16) self.model.eval() @torch.no_grad() def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders if isinstance(query_list, str): query_list = [query_list] if \"e5\" in self.model_name.lower(): if is_query: query_list = [f\"query: {query}\" for query in query_list] else: query_list = [f\"passage: {query}\" for query in query_list] if \"bge\" in self.model_name.lower(): if is_query: query_list = [f\"Represent this sentence for searching relevant passages: {query}\" for query in query_list] inputs = self.tokenizer(query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors=\"pt\" ) inputs = {k: v.cuda() for k, v in inputs.items()} if \"T5\" in type(self.model).__name__: # T5-based retrieval model decoder_input_ids = torch.zeros( (inputs['input_ids'].shape[0], 1), dtype=torch.long ).to(inputs['input_ids'].device) output = self.model( **inputs, decoder_input_ids=decoder_input_ids, return_dict=True ) query_emb = output.last_hidden_state[:, 0, :] else: output = self.model(**inputs, return_dict=True) query_emb = pooling(output.pooler_output, output.last_hidden_state, inputs['attention_mask'], self.pooling_method) if \"dpr\" not in self.model_name.lower(): query_emb = torch.nn.functional.normalize(query_emb, dim=-1) query_emb = query_emb.detach().cpu().numpy() query_emb = query_emb.astype(np.float32, order=\"C\") return query_emb","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#Encoder.__init__","name":"__init__","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L74-L82","signature":"def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16)","exported":false,"origin":"ast","body_hash":"2721883cf802b21052d15bfc422a4b89a0c191775cf0146e9e310beb69d67c3f","body_text":"def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16): self.model_name = model_name self.model_path = model_path self.pooling_method = pooling_method self.max_length = max_length self.use_fp16 = use_fp16 self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16) self.model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"Encoder"},{"id":"search_r1/search/retrieval_server.py#Encoder.encode","name":"encode","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L85-L128","signature":"def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders","exported":true,"origin":"ast","body_hash":"bdbdf0cb31dc98da40dc24ab4aafc75bec72106856942f4e7ba1b836fe3e5970","body_text":"def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders if isinstance(query_list, str): query_list = [query_list] if \"e5\" in self.model_name.lower(): if is_query: query_list = [f\"query: {query}\" for query in query_list] else: query_list = [f\"passage: {query}\" for query in query_list] if \"bge\" in self.model_name.lower(): if is_query: query_list = [f\"Represent this sentence for searching relevant passages: {query}\" for query in query_list] inputs = self.tokenizer(query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors=\"pt\" ) inputs = {k: v.cuda() for k, v in inputs.items()} if \"T5\" in type(self.model).__name__: # T5-based retrieval model decoder_input_ids = torch.zeros( (inputs['input_ids'].shape[0], 1), dtype=torch.long ).to(inputs['input_ids'].device) output = self.model( **inputs, decoder_input_ids=decoder_input_ids, return_dict=True ) query_emb = output.last_hidden_state[:, 0, :] else: output = self.model(**inputs, return_dict=True) query_emb = pooling(output.pooler_output, output.last_hidden_state, inputs['attention_mask'], self.pooling_method) if \"dpr\" not in self.model_name.lower(): query_emb = torch.nn.functional.normalize(query_emb, dim=-1) query_emb = query_emb.detach().cpu().numpy() query_emb = query_emb.astype(np.float32, order=\"C\") return query_emb","summary_state":"pending","summary":null,"crux":null,"owner":"Encoder"},{"id":"search_r1/search/retrieval_server.py#BaseRetriever","name":"BaseRetriever","kind":"class","path":"search_r1/search/retrieval_server.py","span":"L130-L149","signature":"class BaseRetriever","exported":true,"origin":"ast","body_hash":"2641a450bc83bd67357bb34f3ec02200e64dc2d5d1131d2d888438c0a6aa3f1d","body_text":"class BaseRetriever: def __init__(self, config): self.config = config self.retrieval_method = config.retrieval_method self.topk = config.retrieval_topk self.index_path = config.index_path self.corpus_path = config.corpus_path def _search(self, query: str, num: int, return_score: bool): raise NotImplementedError def _batch_search(self, query_list: List[str], num: int, return_score: bool): raise NotImplementedError def search(self, query: str, num: int = None, return_score: bool = False): return self._search(query, num, return_score) def batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): return self._batch_search(query_list, num, return_score)","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#BaseRetriever.__init__","name":"__init__","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L131-L137","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"0c431b7e5ce8a5ebda1442f79c5f24ee117298b8fc3f094542e74aa50395a7ca","body_text":"def __init__(self, config): self.config = config self.retrieval_method = config.retrieval_method self.topk = config.retrieval_topk self.index_path = config.index_path self.corpus_path = config.corpus_path","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"search_r1/search/retrieval_server.py#BaseRetriever._search","name":"_search","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L139-L140","signature":"def _search(self, query: str, num: int, return_score: bool)","exported":false,"origin":"ast","body_hash":"888c954f200f65e3aa486d1b7c2ab1333228397c10d5f5434fcaf026975ff5a3","body_text":"def _search(self, query: str, num: int, return_score: bool): raise NotImplementedError","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"search_r1/search/retrieval_server.py#BaseRetriever._batch_search","name":"_batch_search","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L142-L143","signature":"def _batch_search(self, query_list: List[str], num: int, return_score: bool)","exported":false,"origin":"ast","body_hash":"ee271dba633f341c0ecfd54febfbcd940453507aba1e572bbbfccf13ec11ce98","body_text":"def _batch_search(self, query_list: List[str], num: int, return_score: bool): raise NotImplementedError","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"search_r1/search/retrieval_server.py#BaseRetriever.search","name":"search","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L145-L146","signature":"def search(self, query: str, num: int = None, return_score: bool = False)","exported":true,"origin":"ast","body_hash":"f1686fdcde6b0a4a38918ae7792e7e057f384f5cefe3343baa1e2c6010b848f1","body_text":"def search(self, query: str, num: int = None, return_score: bool = False): return self._search(query, num, return_score)","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"search_r1/search/retrieval_server.py#BaseRetriever.batch_search","name":"batch_search","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L148-L149","signature":"def batch_search(self, query_list: List[str], num: int = None, return_score: bool = False)","exported":true,"origin":"ast","body_hash":"2e0ee2d7db8e8a7035ac5c089b034944ab6ffa092e7fa4d858c331ef8657bd04","body_text":"def batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): return self._batch_search(query_list, num, return_score)","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"search_r1/search/retrieval_server.py#BM25Retriever","name":"BM25Retriever","kind":"class","path":"search_r1/search/retrieval_server.py","span":"L151-L210","signature":"class BM25Retriever(BaseRetriever)","exported":true,"origin":"ast","body_hash":"4cc3b88f5678ce5794c817675c79f2fe58238e9a53465daf66571238a71819b6","body_text":"class BM25Retriever(BaseRetriever): def __init__(self, config): super().__init__(config) from pyserini.search.lucene import LuceneSearcher self.searcher = LuceneSearcher(self.index_path) self.contain_doc = self._check_contain_doc() if not self.contain_doc: self.corpus = load_corpus(self.corpus_path) self.max_process_num = 8 def _check_contain_doc(self): return self.searcher.doc(0).raw() is not None def _search(self, query: str, num: int = None, return_score: bool = False): if num is None: num = self.topk hits = self.searcher.search(query, num) if len(hits) < 1: if return_score: return [], [] else: return [] scores = [hit.score for hit in hits] if len(hits) < num: warnings.warn('Not enough documents retrieved!') else: hits = hits[:num] if self.contain_doc: all_contents = [ json.loads(self.searcher.doc(hit.docid).raw())['contents'] for hit in hits ] results = [ { 'title': content.split(\"\\n\")[0].strip(\"\\\"\"), 'text': \"\\n\".join(content.split(\"\\n\")[1:]), 'contents': content } for content in all_contents ] else: results = load_docs(self.corpus, [hit.docid for hit in hits]) if return_score: return results, scores else: return results def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): results = [] scores = [] for query in query_list: item_result, item_score = self._search(query, num, True) results.append(item_result) scores.append(item_score) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#BM25Retriever.__init__","name":"__init__","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L152-L159","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"b9c8fcd86e95a5d5fb5a36d4aaddf2661d2b39e1467fc208c07f078d58d1626b","body_text":"def __init__(self, config): super().__init__(config) from pyserini.search.lucene import LuceneSearcher self.searcher = LuceneSearcher(self.index_path) self.contain_doc = self._check_contain_doc() if not self.contain_doc: self.corpus = load_corpus(self.corpus_path) self.max_process_num = 8","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"search_r1/search/retrieval_server.py#BM25Retriever._check_contain_doc","name":"_check_contain_doc","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L161-L162","signature":"def _check_contain_doc(self)","exported":false,"origin":"ast","body_hash":"178b6eb860dc4f097984975673ec39f6db7807bc1c2e5deb037a1bf8624ab269","body_text":"def _check_contain_doc(self): return self.searcher.doc(0).raw() is not None","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"search_r1/search/retrieval_server.py#BM25Retriever._search","name":"_search","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L164-L198","signature":"def _search(self, query: str, num: int = None, return_score: bool = False)","exported":false,"origin":"ast","body_hash":"09b72557989c3f31687a416a1f18b7e970b44ea2f2185393bf40074ce242efdd","body_text":"def _search(self, query: str, num: int = None, return_score: bool = False): if num is None: num = self.topk hits = self.searcher.search(query, num) if len(hits) < 1: if return_score: return [], [] else: return [] scores = [hit.score for hit in hits] if len(hits) < num: warnings.warn('Not enough documents retrieved!') else: hits = hits[:num] if self.contain_doc: all_contents = [ json.loads(self.searcher.doc(hit.docid).raw())['contents'] for hit in hits ] results = [ { 'title': content.split(\"\\n\")[0].strip(\"\\\"\"), 'text': \"\\n\".join(content.split(\"\\n\")[1:]), 'contents': content } for content in all_contents ] else: results = load_docs(self.corpus, [hit.docid for hit in hits]) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"search_r1/search/retrieval_server.py#BM25Retriever._batch_search","name":"_batch_search","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L200-L210","signature":"def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False)","exported":false,"origin":"ast","body_hash":"c5de1b7501526e6fc77e0d484660d65e564e6c7876f1bbbe41e57b26eea8e66c","body_text":"def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): results = [] scores = [] for query in query_list: item_result, item_score = self._search(query, num, True) results.append(item_result) scores.append(item_score) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"search_r1/search/retrieval_server.py#DenseRetriever","name":"DenseRetriever","kind":"class","path":"search_r1/search/retrieval_server.py","span":"L212-L272","signature":"class DenseRetriever(BaseRetriever)","exported":true,"origin":"ast","body_hash":"aea7c46390c33b2d7b937621b31e11e92bdd97c4fe146143661dd4ed5c1279ab","body_text":"class DenseRetriever(BaseRetriever): def __init__(self, config): super().__init__(config) self.index = faiss.read_index(self.index_path) if config.faiss_gpu: co = faiss.GpuMultipleClonerOptions() co.useFloat16 = True co.shard = True self.index = faiss.index_cpu_to_all_gpus(self.index, co=co) self.corpus = load_corpus(self.corpus_path) self.encoder = Encoder( model_name = self.retrieval_method, model_path = config.retrieval_model_path, pooling_method = config.retrieval_pooling_method, max_length = config.retrieval_query_max_length, use_fp16 = config.retrieval_use_fp16 ) self.topk = config.retrieval_topk self.batch_size = config.retrieval_batch_size def _search(self, query: str, num: int = None, return_score: bool = False): if num is None: num = self.topk query_emb = self.encoder.encode(query) scores, idxs = self.index.search(query_emb, k=num) idxs = idxs[0] scores = scores[0] results = load_docs(self.corpus, idxs) if return_score: return results, scores.tolist() else: return results def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): if isinstance(query_list, str): query_list = [query_list] if num is None: num = self.topk results = [] scores = [] for start_idx in tqdm(range(0, len(query_list), self.batch_size), desc='Retrieval process: '): query_batch = query_list[start_idx:start_idx + self.batch_size] batch_emb = self.encoder.encode(query_batch) batch_scores, batch_idxs = self.index.search(batch_emb, k=num) batch_scores = batch_scores.tolist() batch_idxs = batch_idxs.tolist() # load_docs is not vectorized, but is a python list approach flat_idxs = sum(batch_idxs, []) batch_results = load_docs(self.corpus, flat_idxs) # chunk them back batch_results = [batch_results[i*num : (i+1)*num] for i in range(len(batch_idxs))] results.extend(batch_results) scores.extend(batch_scores) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#DenseRetriever.__init__","name":"__init__","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L213-L231","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"0d45ca695efc264e1f301b07f3d486a4768aa1396287f6cb83fcc7c925872142","body_text":"def __init__(self, config): super().__init__(config) self.index = faiss.read_index(self.index_path) if config.faiss_gpu: co = faiss.GpuMultipleClonerOptions() co.useFloat16 = True co.shard = True self.index = faiss.index_cpu_to_all_gpus(self.index, co=co) self.corpus = load_corpus(self.corpus_path) self.encoder = Encoder( model_name = self.retrieval_method, model_path = config.retrieval_model_path, pooling_method = config.retrieval_pooling_method, max_length = config.retrieval_query_max_length, use_fp16 = config.retrieval_use_fp16 ) self.topk = config.retrieval_topk self.batch_size = config.retrieval_batch_size","summary_state":"pending","summary":null,"crux":null,"owner":"DenseRetriever"},{"id":"search_r1/search/retrieval_server.py#DenseRetriever._search","name":"_search","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L233-L244","signature":"def _search(self, query: str, num: int = None, return_score: bool = False)","exported":false,"origin":"ast","body_hash":"6d0530829950c5765c63940f4de06bed3c2d3f47a7c0390022a29d35fb5b5fb4","body_text":"def _search(self, query: str, num: int = None, return_score: bool = False): if num is None: num = self.topk query_emb = self.encoder.encode(query) scores, idxs = self.index.search(query_emb, k=num) idxs = idxs[0] scores = scores[0] results = load_docs(self.corpus, idxs) if return_score: return results, scores.tolist() else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"DenseRetriever"},{"id":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","name":"_batch_search","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L246-L272","signature":"def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False)","exported":false,"origin":"ast","body_hash":"59a42e96e32f723d344f468090e47e43cdebfda11937097a76be3d3d1f082f12","body_text":"def _batch_search(self, query_list: List[str], num: int = None, return_score: bool = False): if isinstance(query_list, str): query_list = [query_list] if num is None: num = self.topk results = [] scores = [] for start_idx in tqdm(range(0, len(query_list), self.batch_size), desc='Retrieval process: '): query_batch = query_list[start_idx:start_idx + self.batch_size] batch_emb = self.encoder.encode(query_batch) batch_scores, batch_idxs = self.index.search(batch_emb, k=num) batch_scores = batch_scores.tolist() batch_idxs = batch_idxs.tolist() # load_docs is not vectorized, but is a python list approach flat_idxs = sum(batch_idxs, []) batch_results = load_docs(self.corpus, flat_idxs) # chunk them back batch_results = [batch_results[i*num : (i+1)*num] for i in range(len(batch_idxs))] results.extend(batch_results) scores.extend(batch_scores) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"DenseRetriever"},{"id":"search_r1/search/retrieval_server.py#get_retriever","name":"get_retriever","kind":"function","path":"search_r1/search/retrieval_server.py","span":"L274-L278","signature":"def get_retriever(config)","exported":true,"origin":"ast","body_hash":"9cb8b9010e6be3870a7606c4331ad08e38c2912ba26d967ceeea651ab195ea15","body_text":"def get_retriever(config): if config.retrieval_method == \"bm25\": return BM25Retriever(config) else: return DenseRetriever(config)","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#Config","name":"Config","kind":"class","path":"search_r1/search/retrieval_server.py","span":"L285-L316","signature":"class Config","exported":true,"origin":"ast","body_hash":"da7389089df79a5565aa5da08f2aca681165e890786dc5ec037d4621f54ac40c","body_text":"class Config: \"\"\" Minimal config class (simulating your argparse) Replace this with your real arguments or load them dynamically. \"\"\" def __init__( self, retrieval_method: str = \"bm25\", retrieval_topk: int = 10, index_path: str = \"./index/bm25\", corpus_path: str = \"./data/corpus.jsonl\", dataset_path: str = \"./data\", data_split: str = \"train\", faiss_gpu: bool = True, retrieval_model_path: str = \"./model\", retrieval_pooling_method: str = \"mean\", retrieval_query_max_length: int = 256, retrieval_use_fp16: bool = False, retrieval_batch_size: int = 128 ): self.retrieval_method = retrieval_method self.retrieval_topk = retrieval_topk self.index_path = index_path self.corpus_path = corpus_path self.dataset_path = dataset_path self.data_split = data_split self.faiss_gpu = faiss_gpu self.retrieval_model_path = retrieval_model_path self.retrieval_pooling_method = retrieval_pooling_method self.retrieval_query_max_length = retrieval_query_max_length self.retrieval_use_fp16 = retrieval_use_fp16 self.retrieval_batch_size = retrieval_batch_size","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#Config.__init__","name":"__init__","kind":"method","path":"search_r1/search/retrieval_server.py","span":"L290-L316","signature":"def __init__( self, retrieval_method: str = \"bm25\", retrieval_topk: int = 10, index_path: str = \"./index/bm25\", corpus_path: str = \"./data/corpus.jsonl\", dataset_path: str = \"./data\", data_split: str = \"train\", faiss_gpu: bool = True, retrieval_model_path: str = \"./model\", retrieval_pooling_method: str = \"mean\", retrieval_query_max_length: int = 256, retrieval_use_fp16: bool = False, retrieval_batch_size: int = 128 )","exported":false,"origin":"ast","body_hash":"dfb28fe95ac552021e70f0af878ead067d16e36a584cd1c3ef27e68de36477a8","body_text":"def __init__( self, retrieval_method: str = \"bm25\", retrieval_topk: int = 10, index_path: str = \"./index/bm25\", corpus_path: str = \"./data/corpus.jsonl\", dataset_path: str = \"./data\", data_split: str = \"train\", faiss_gpu: bool = True, retrieval_model_path: str = \"./model\", retrieval_pooling_method: str = \"mean\", retrieval_query_max_length: int = 256, retrieval_use_fp16: bool = False, retrieval_batch_size: int = 128 ): self.retrieval_method = retrieval_method self.retrieval_topk = retrieval_topk self.index_path = index_path self.corpus_path = corpus_path self.dataset_path = dataset_path self.data_split = data_split self.faiss_gpu = faiss_gpu self.retrieval_model_path = retrieval_model_path self.retrieval_pooling_method = retrieval_pooling_method self.retrieval_query_max_length = retrieval_query_max_length self.retrieval_use_fp16 = retrieval_use_fp16 self.retrieval_batch_size = retrieval_batch_size","summary_state":"pending","summary":null,"crux":null,"owner":"Config"},{"id":"search_r1/search/retrieval_server.py#QueryRequest","name":"QueryRequest","kind":"class","path":"search_r1/search/retrieval_server.py","span":"L319-L322","signature":"class QueryRequest(BaseModel)","exported":true,"origin":"ast","body_hash":"eb790423bbdb7e3a41e936eea701b96ce08c9938de69aa6bbee1083839a6d9eb","body_text":"class QueryRequest(BaseModel): queries: List[str] topk: Optional[int] = None return_scores: bool = False","summary_state":"pending","summary":null,"crux":null},{"id":"search_r1/search/retrieval_server.py#retrieve_endpoint","name":"retrieve_endpoint","kind":"function","path":"search_r1/search/retrieval_server.py","span":"L346-L377","signature":"def retrieve_endpoint(request: QueryRequest)","exported":true,"origin":"ast","body_hash":"12cccb4e450a190a80545e5c9698326b5df2dccceb057d3352ef4405eb9853b5","body_text":"def retrieve_endpoint(request: QueryRequest): \"\"\" Endpoint that accepts queries and performs retrieval. Input format: { \"queries\": [\"What is Python?\", \"Tell me about neural networks.\"], \"topk\": 3, \"return_scores\": true } \"\"\" if not request.topk: request.topk = config.retrieval_topk # fallback to default # Perform batch retrieval results, scores = retriever.batch_search( query_list=request.queries, num=request.topk, return_score=request.return_scores ) # Format response resp = [] for i, single_result in enumerate(results): if request.return_scores: # If scores are returned, combine them with results combined = [] for doc, score in zip(single_result, scores[i]): combined.append({\"document\": doc, \"score\": score}) resp.append(combined) else: resp.append(single_result) return {\"result\": resp}","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"json","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"os","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"warnings","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"typing","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"argparse","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"faiss","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"torch","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"transformers","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"tqdm","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"datasets","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"uvicorn","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"fastapi","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"pydantic","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"ArgumentParser","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"add_argument","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"parse_args","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#load_corpus","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#load_corpus","relation":"calls","name":"load_dataset","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#read_jsonl","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#read_jsonl","relation":"calls","name":"open","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#read_jsonl","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#read_jsonl","relation":"calls","name":"loads","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#load_docs","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#load_docs","relation":"calls","name":"int","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#load_model","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#load_model","relation":"calls","name":"eval","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#load_model","relation":"calls","name":"cuda","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#load_model","relation":"calls","name":"half","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#pooling","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#pooling","relation":"calls","name":"masked_fill","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#pooling","relation":"calls","name":"bool","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#pooling","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#pooling","relation":"calls","name":"sum","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#pooling","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#Encoder","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder","relation":"contains","targetId":"search_r1/search/retrieval_server.py#Encoder.__init__","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.__init__","relation":"calls","name":"load_model","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.__init__","relation":"calls","name":"eval","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder","relation":"calls","name":"no_grad","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder","relation":"contains","targetId":"search_r1/search/retrieval_server.py#Encoder.encode","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"isinstance","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"lower","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"lower","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"tokenizer","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"Encoder"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"cuda","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"items","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"type","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"to","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"zeros","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"model","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"Encoder"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"model","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"Encoder"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"pooling","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"lower","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"normalize","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"numpy","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"cpu","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"detach","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"pooling"},{"source":"search_r1/search/retrieval_server.py#Encoder.encode","relation":"calls","name":"astype","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"pooling"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BaseRetriever","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BaseRetriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BaseRetriever.__init__","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BaseRetriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BaseRetriever._search","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BaseRetriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BaseRetriever._batch_search","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BaseRetriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BaseRetriever.search","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BaseRetriever.search","relation":"calls","name":"_search","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"BaseRetriever"},{"source":"search_r1/search/retrieval_server.py#BaseRetriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BaseRetriever.batch_search","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BaseRetriever.batch_search","relation":"calls","name":"_batch_search","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"BaseRetriever"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BM25Retriever","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever","relation":"extends","name":"BaseRetriever","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BM25Retriever.__init__","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever.__init__","relation":"calls","name":"super","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"imports","specifier":"pyserini.search.lucene","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever.__init__","relation":"calls","name":"LuceneSearcher","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever.__init__","relation":"calls","name":"_check_contain_doc","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"BM25Retriever"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever.__init__","relation":"calls","name":"load_corpus","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BM25Retriever._check_contain_doc","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._check_contain_doc","relation":"calls","name":"raw","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._check_contain_doc","relation":"calls","name":"doc","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"LuceneSearcher"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BM25Retriever._search","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"search","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"LuceneSearcher"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"warn","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"loads","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"raw","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"doc","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"LuceneSearcher"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"strip","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"split","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"join","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"split","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._search","relation":"calls","name":"load_docs","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#BM25Retriever._batch_search","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._batch_search","relation":"calls","name":"_search","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"BM25Retriever"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._batch_search","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#BM25Retriever._batch_search","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#DenseRetriever","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever","relation":"extends","name":"BaseRetriever","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#DenseRetriever.__init__","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever.__init__","relation":"calls","name":"super","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever.__init__","relation":"calls","name":"read_index","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever.__init__","relation":"calls","name":"GpuMultipleClonerOptions","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever.__init__","relation":"calls","name":"index_cpu_to_all_gpus","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever.__init__","relation":"calls","name":"load_corpus","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever.__init__","relation":"calls","name":"Encoder","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#DenseRetriever._search","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._search","relation":"calls","name":"encode","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"Encoder"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._search","relation":"calls","name":"search","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._search","relation":"calls","name":"load_docs","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._search","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever","relation":"contains","targetId":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"isinstance","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"tqdm","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"range","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"encode","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"Encoder"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"search","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"tolist","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"sum","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"load_docs","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"range","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"len","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"extend","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#DenseRetriever._batch_search","relation":"calls","name":"extend","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#get_retriever","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#get_retriever","relation":"calls","name":"BM25Retriever","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#get_retriever","relation":"calls","name":"DenseRetriever","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#Config","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#Config","relation":"contains","targetId":"search_r1/search/retrieval_server.py#Config.__init__","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#QueryRequest","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#QueryRequest","relation":"extends","name":"BaseModel","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"FastAPI","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"Config","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"get_retriever","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"post","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"FastAPI"},{"source":"search_r1/search/retrieval_server.py","relation":"contains","targetId":"search_r1/search/retrieval_server.py#retrieve_endpoint","file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#retrieve_endpoint","relation":"calls","name":"batch_search","viaMember":true,"file":"search_r1/search/retrieval_server.py","recvType":"get_retriever"},{"source":"search_r1/search/retrieval_server.py#retrieve_endpoint","relation":"calls","name":"enumerate","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#retrieve_endpoint","relation":"calls","name":"zip","viaMember":false,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#retrieve_endpoint","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#retrieve_endpoint","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py#retrieve_endpoint","relation":"calls","name":"append","viaMember":true,"file":"search_r1/search/retrieval_server.py"},{"source":"search_r1/search/retrieval_server.py","relation":"calls","name":"run","viaMember":true,"file":"search_r1/search/retrieval_server.py"}]},"setup.py":{"size":1917,"mtimeMs":1789275830306.2024,"hash":"1f45bacb406604058229b111b07a226d8b1c0299132dfbafc2b28df011c8100e","nodes":[{"id":"setup.py","name":"setup.py","kind":"file","path":"setup.py","span":"L1-L54","signature":null,"exported":true,"origin":"ast","body_hash":"1f45bacb406604058229b111b07a226d8b1c0299132dfbafc2b28df011c8100e","chars":1917,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 from setuptools import setup, find_packages import os 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() with open('requirements.txt') as f: required = f.read().splitlines() install_requires = [item.strip() for item in required if item.strip()[0] != '#'] extras_require = { 'test': ['pytest', 'yapf'] } from pathlib import Path 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/volcengine/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'],}, include_package_data=True, long_description=long_description, long_description_content_type='text/markdown' )"}],"rawEdges":[{"source":"setup.py","relation":"imports","specifier":"setuptools","file":"setup.py"},{"source":"setup.py","relation":"imports","specifier":"os","file":"setup.py"},{"source":"setup.py","relation":"calls","name":"dirname","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"join","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"abspath","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"open","viaMember":false,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"join","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"strip","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"read","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"open","viaMember":false,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"splitlines","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"read","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"strip","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"strip","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"imports","specifier":"pathlib","file":"setup.py"},{"source":"setup.py","relation":"calls","name":"Path","viaMember":false,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"read_text","viaMember":true,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"setup","viaMember":false,"file":"setup.py"},{"source":"setup.py","relation":"calls","name":"find_packages","viaMember":false,"file":"setup.py"}]},"verl/__init__.py":{"size":927,"mtimeMs":1789275830307.2024,"hash":"41d6e41c33606205fb21af76c9ea180ade26413345a073b07c24848cb06b41f6","nodes":[{"id":"verl/__init__.py","name":"__init__.py","kind":"file","path":"verl/__init__.py","span":"L1-L28","signature":null,"exported":true,"origin":"ast","body_hash":"41d6e41c33606205fb21af76c9ea180ade26413345a073b07c24848cb06b41f6","chars":927,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 version_folder = os.path.dirname(os.path.join(os.path.abspath(__file__))) with open(os.path.join(version_folder, 'version/version')) as f: __version__ = f.read().strip() from .protocol import DataProto from .utils.logging_utils import set_basic_config import logging set_basic_config(level=logging.WARNING)"}],"rawEdges":[{"source":"verl/__init__.py","relation":"imports","specifier":"os","file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"calls","name":"dirname","viaMember":true,"file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"calls","name":"join","viaMember":true,"file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"calls","name":"abspath","viaMember":true,"file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"calls","name":"open","viaMember":false,"file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"calls","name":"join","viaMember":true,"file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"calls","name":"strip","viaMember":true,"file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"calls","name":"read","viaMember":true,"file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"imports","specifier":".protocol","file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"imports","specifier":".utils.logging_utils","file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"imports","specifier":"logging","file":"verl/__init__.py"},{"source":"verl/__init__.py","relation":"calls","name":"set_basic_config","viaMember":false,"file":"verl/__init__.py"}]},"verl/models/__init__.py":{"size":600,"mtimeMs":1789275830307.2024,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/models/__init__.py","name":"__init__.py","kind":"file","path":"verl/models/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/models/llama/__init__.py":{"size":600,"mtimeMs":1789275830307.2024,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/models/llama/__init__.py","name":"__init__.py","kind":"file","path":"verl/models/llama/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/models/llama/megatron/__init__.py":{"size":944,"mtimeMs":1789275830307.2024,"hash":"bacfa32a69412dc2a3d55134a1a5a287288d2179c76700b72eed4bc42ed94d53","nodes":[{"id":"verl/models/llama/megatron/__init__.py","name":"__init__.py","kind":"file","path":"verl/models/llama/megatron/__init__.py","span":"L1-L25","signature":null,"exported":true,"origin":"ast","body_hash":"bacfa32a69412dc2a3d55134a1a5a287288d2179c76700b72eed4bc42ed94d53","chars":944,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 .modeling_llama_megatron import ( # original model with megatron ParallelLlamaModel, ParallelLlamaForCausalLM, # rmpad with megatron ParallelLlamaForCausalLMRmPad, ParallelLlamaForValueRmPad, # rmpad with megatron and pipeline parallelism ParallelLlamaForCausalLMRmPadPP, ParallelLlamaForValueRmPadPP)"}],"rawEdges":[{"source":"verl/models/llama/megatron/__init__.py","relation":"imports","specifier":".modeling_llama_megatron","file":"verl/models/llama/megatron/__init__.py"}]},"verl/models/llama/megatron/checkpoint_utils/__init__.py":{"size":600,"mtimeMs":1789275830307.2024,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/models/llama/megatron/checkpoint_utils/__init__.py","name":"__init__.py","kind":"file","path":"verl/models/llama/megatron/checkpoint_utils/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/models/llama/megatron/checkpoint_utils/llama_loader.py":{"size":19665,"mtimeMs":1789275830307.2024,"hash":"8132d66d0756561bdf7bcc53a6b77d85bc67a38385f4eff8233377d5ad123a29","nodes":[{"id":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","name":"llama_loader.py","kind":"file","path":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","span":"L1-L447","signature":null,"exported":true,"origin":"ast","body_hash":"8132d66d0756561bdf7bcc53a6b77d85bc67a38385f4eff8233377d5ad123a29","chars":19665,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 import time from typing import Dict, Any, Callable, Optional import torch.distributed as dist"},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#_megatron_calc_layer_map","name":"_megatron_calc_layer_map","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","span":"L21-L48","signature":"def _megatron_calc_layer_map(config)","exported":false,"origin":"ast","body_hash":"c6c7e2d0f3c35feca085f495ed37ddbe27e0bef9d0608f4f544cd748c86a0d5d","body_text":"def _megatron_calc_layer_map(config): \"\"\"Calculate the mapping of global layer_idx to local layer_idx Returns: layer_map (Dict: int -> tuple(int, int, int)): mapping from the global layer index to a tuple of (pp_rank, virtual_pp_rank, layer_idx inside model) \"\"\" import megatron from megatron.core import mpu pp_size = mpu.get_pipeline_model_parallel_world_size() virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1 layer_map = dict() num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers for pp_rank_idx in range(pp_size): for virtual_pp_rank_idx in range(virtual_pp_size): layer_offset = (virtual_pp_rank_idx * (config.num_hidden_layers // virtual_pp_size) + pp_rank_idx * num_layers_per_model) for layer_idx in range(num_layers_per_model): layer_map[layer_offset + layer_idx] = ( pp_rank_idx, virtual_pp_rank_idx, layer_idx, ) return layer_map","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","name":"load_state_dict_to_megatron_llama","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","span":"L51-L446","signature":"def load_state_dict_to_megatron_llama(state_dict, wrapped_models, config, params_dtype, is_value_model=False)","exported":true,"origin":"ast","body_hash":"6b71c33de24eae027f69c9d90e14200f15bdb8ae1fee67951f83cd8b53992d19","body_text":"def load_state_dict_to_megatron_llama(state_dict, wrapped_models, config, params_dtype, is_value_model=False): \"\"\"Load merged state_dict to sharded Megatron module in training. \"\"\" import megatron from megatron.core import mpu from megatron.utils import print_rank_0, unwrap_model from megatron.core.transformer.module import Float16Module from megatron.core import DistributedDataParallel as LocalDDP from torch.nn.parallel import DistributedDataParallel as torchDDP start_time = time.time() def _get_gpt_model(model): return model def broadcast_params(module): for param in module.parameters(): torch.distributed.broadcast(param.data, src=mpu.get_data_parallel_src_rank(), group=mpu.get_data_parallel_group()) dp_rank = mpu.get_data_parallel_rank() pp_rank = mpu.get_pipeline_model_parallel_rank() pp_size = mpu.get_pipeline_model_parallel_world_size() virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1 mp_group = mpu.get_model_parallel_group() if torch.distributed.get_rank() == 0: assert mp_group.rank() == 0, f\"mp_rank:[{mp_group.rank}] != 0 on rank #0\" assert pp_rank == 0, f\"pp_rank:[{pp_rank}] != 0 on rank #0\" assert dp_rank == 0, f\"dp_rank:[{dp_rank}] != 0 on rank #0\" if not isinstance(wrapped_models, (list, tuple)): wrapped_models = list(wrapped_models) assert len(wrapped_models) == virtual_pp_size num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers models = [None] * len(wrapped_models) for i, wrapped_model in enumerate(wrapped_models): models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module)) gpt_model_module = _get_gpt_model(models[i]) assert len(gpt_model_module.model.layers) == num_layers_per_model def _broadcast_tensor(tensor, name) -> torch.Tensor: \"\"\"broadcast tensor from rank0 across mp_group\"\"\" nonlocal state_dict nonlocal mp_group if torch.distributed.get_rank() == 0: if name in state_dict: weight = state_dict[name] tensor_shape = weight.shape else: tensor_shape = None else: weight = None tensor_shape = None obj_list = [tensor_shape] dist.broadcast_object_list(obj_list, src=0, group=mp_group) tensor_shape = obj_list[0] if tensor_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tensor:[{name}] not in state_dict, skip load\") return if tensor is None: tensor = torch.empty( tensor_shape, dtype=params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) if torch.distributed.get_rank() == 0: tensor.data.copy_(weight) dist.broadcast(tensor, src=0, group=mp_group) def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor: \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() if torch.distributed.get_rank() == 0: if name in state_dict: full_weight = state_dict[name] if mutate_func is not None: full_weight = mutate_func(full_weight) tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim) chunk_shape = tensor_chunk[0].shape else: chunk_shape = None else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=0, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\") return if tensor is None: sync_tensor = torch.empty( chunk_shape, dtype=params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) else: assert (tensor.shape == chunk_shape ), f\"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}\" sync_tensor = torch.empty_like(tensor, device=torch.cuda.current_device(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: sync_tensor.data.copy_(tensor_chunk[i]) dist.broadcast(sync_tensor, src=0, group=mp_group) if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor) def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor: \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() if torch.distributed.get_rank() == 0: if name in state_dict: full_weight = state_dict[name] if mutate_func is not None: full_weight = mutate_func(full_weight) tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim) chunk_shape = tensor_chunk[0].shape else: chunk_shape = None else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=0, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\") return if tensor is None: sync_tensor = torch.empty( chunk_shape, dtype=","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._get_gpt_model","name":"_get_gpt_model","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","span":"L63-L64","signature":"def _get_gpt_model(model)","exported":false,"origin":"ast","body_hash":"c616fc9942bbf7a93cfd936b39e34055ee56d16434699c954dafff40af8155a0","body_text":"def _get_gpt_model(model): return model","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama.broadcast_params","name":"broadcast_params","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","span":"L66-L70","signature":"def broadcast_params(module)","exported":true,"origin":"ast","body_hash":"370662409e80a26a99a0658c3fba78fc5d602f866447b521de8781ff420e78dc","body_text":"def broadcast_params(module): for param in module.parameters(): torch.distributed.broadcast(param.data, src=mpu.get_data_parallel_src_rank(), group=mpu.get_data_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tensor","name":"_broadcast_tensor","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","span":"L97-L129","signature":"def _broadcast_tensor(tensor, name) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"64cf0bb00b9f907a7700f02eca115e6f09d21efa74522b91c206ace7d82975d2","body_text":"def _broadcast_tensor(tensor, name) -> torch.Tensor: \"\"\"broadcast tensor from rank0 across mp_group\"\"\" nonlocal state_dict nonlocal mp_group if torch.distributed.get_rank() == 0: if name in state_dict: weight = state_dict[name] tensor_shape = weight.shape else: tensor_shape = None else: weight = None tensor_shape = None obj_list = [tensor_shape] dist.broadcast_object_list(obj_list, src=0, group=mp_group) tensor_shape = obj_list[0] if tensor_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tensor:[{name}] not in state_dict, skip load\") return if tensor is None: tensor = torch.empty( tensor_shape, dtype=params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) if torch.distributed.get_rank() == 0: tensor.data.copy_(weight) dist.broadcast(tensor, src=0, group=mp_group)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","name":"_broadcast_tp_shard_tensor_vocab","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","span":"L131-L176","signature":"def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"59255bd48899004f650fa659a4d1363693799b3aa03d7c72097021507c3d293e","body_text":"def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor: \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() if torch.distributed.get_rank() == 0: if name in state_dict: full_weight = state_dict[name] if mutate_func is not None: full_weight = mutate_func(full_weight) tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim) chunk_shape = tensor_chunk[0].shape else: chunk_shape = None else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=0, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\") return if tensor is None: sync_tensor = torch.empty( chunk_shape, dtype=params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) else: assert (tensor.shape == chunk_shape ), f\"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}\" sync_tensor = torch.empty_like(tensor, device=torch.cuda.current_device(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: sync_tensor.data.copy_(tensor_chunk[i]) dist.broadcast(sync_tensor, src=0, group=mp_group) if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","name":"_broadcast_tp_shard_tensor","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","span":"L178-L222","signature":"def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"1e5c7126ab35f65111c322dad4d56bc64d026670c63a975331eb490640d04ecf","body_text":"def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor: \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() if torch.distributed.get_rank() == 0: if name in state_dict: full_weight = state_dict[name] if mutate_func is not None: full_weight = mutate_func(full_weight) tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim) chunk_shape = tensor_chunk[0].shape else: chunk_shape = None else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=0, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\") return if tensor is None: sync_tensor = torch.empty( chunk_shape, dtype=params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) else: assert (tensor.shape == chunk_shape ), f\"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}\" sync_tensor = torch.empty_like(tensor, device=torch.cuda.current_device(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: sync_tensor.data.copy_(tensor_chunk[i]) dist.broadcast(sync_tensor, src=0, group=mp_group) if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","name":"_broadcast_tp_shard_tensor_gate_up","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","span":"L224-L276","signature":"def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"3ea59bfc74724e6c5b5c521e314e13e168045e949b9b2c86609180b55376c72c","body_text":"def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor: \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() if torch.distributed.get_rank() == 0: gate_weight = state_dict[gate_name] up_weight = state_dict[up_name] new_gate_up_weight = torch.empty(config.intermediate_size * 2, config.hidden_size, dtype=params_dtype, device=torch.cuda.current_device()) for i in range(tp_size): intermediate_size_tp = config.intermediate_size // tp_size gate_weight_tp = gate_weight[i * intermediate_size_tp:(i + 1) * intermediate_size_tp] up_weight_tp = up_weight[i * intermediate_size_tp:(i + 1) * intermediate_size_tp] new_gate_up_weight[intermediate_size_tp * 2 * i:intermediate_size_tp * 2 * (i + 1)].copy_( torch.cat([gate_weight_tp, up_weight_tp], dim=0)) tensor_chunk = torch.chunk(new_gate_up_weight, tp_size, dim=0) chunk_shape = tensor_chunk[0].shape else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=0, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tp_shard tensor:[{gate_name, up_name}] not in state_dict, skip loading\") return if tensor is None: sync_tensor = torch.empty( chunk_shape, dtype=params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) else: assert ( tensor.shape == chunk_shape ), f\"rank #{torch.distributed.get_rank() == 0:} tensor {gate_name, up_name} shape {tensor.shape} != {chunk_shape}\" sync_tensor = torch.empty_like(tensor, device=torch.cuda.current_device(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: sync_tensor.data.copy_(tensor_chunk[i]) dist.broadcast(sync_tensor, src=0, group=mp_group) if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","name":"_broadcast_tp_shard_tensor_qkv","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","span":"L278-L355","signature":"def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"e85764a9f0b621be556ca2f965cbd16259fe679a347f9c54a92203dd5c37e84c","body_text":"def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor: \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() if torch.distributed.get_rank() == 0: assert (q_name in state_dict and k_name in state_dict and v_name in state_dict) full_weight_q = state_dict[q_name] full_weight_k = state_dict[k_name] full_weight_v = state_dict[v_name] hidden_size_per_head = config.hidden_size // config.num_attention_heads if config.num_key_value_heads >= tp_size: q_size_tp = config.hidden_size // tp_size kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size total_size = q_size_tp + 2 * kv_size_tp new_weight_qkv = torch.empty(total_size * tp_size, config.hidden_size, dtype=params_dtype, device=torch.cuda.current_device()) for i in range(tp_size): q_part = full_weight_q[i * q_size_tp:(i + 1) * q_size_tp] k_part = full_weight_k[i * kv_size_tp:(i + 1) * kv_size_tp] v_part = full_weight_v[i * kv_size_tp:(i + 1) * kv_size_tp] new_weight_qkv[i * total_size:(i + 1) * total_size].copy_(torch.cat([q_part, k_part, v_part], dim=0)) else: q_size_tp = config.hidden_size // tp_size kv_size_tp = hidden_size_per_head total_size = q_size_tp + 2 * kv_size_tp new_weight_qkv = torch.empty(total_size * tp_size, config.hidden_size, dtype=params_dtype, device=torch.cuda.current_device()) for i in range(tp_size): q_part = full_weight_q[i * q_size_tp:(i + 1) * q_size_tp] start_idx = i * config.num_key_value_heads // tp_size * hidden_size_per_head end_idx = (i * config.num_key_value_heads // tp_size + 1) * hidden_size_per_head k_part = full_weight_k[start_idx:end_idx] v_part = full_weight_v[start_idx:end_idx] new_weight_qkv[i * total_size:(i + 1) * total_size].copy_(torch.cat([q_part, k_part, v_part], dim=0)) tensor_chunk = torch.chunk(new_weight_qkv, tp_size, dim=0) chunk_shape = tensor_chunk[0].shape else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=0, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\") return if tensor is None: sync_tensor = torch.empty( chunk_shape, dtype=params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) else: assert (tensor.shape == chunk_shape ), f\"rank #{torch.distributed.get_rank()} tensor {q_name} shape {tensor.shape} != {chunk_shape}\" sync_tensor = torch.empty_like(tensor, device=torch.cuda.current_device(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: sync_tensor.data.copy_(tensor_chunk[i]) dist.broadcast(sync_tensor, src=0, group=mp_group) if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"time","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"typing","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#_megatron_calc_layer_map","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"megatron","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#_megatron_calc_layer_map","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#_megatron_calc_layer_map","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#_megatron_calc_layer_map","relation":"calls","name":"dict","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#_megatron_calc_layer_map","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#_megatron_calc_layer_map","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#_megatron_calc_layer_map","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"megatron","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"megatron.utils","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"megatron.core.transformer.module","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py","relation":"imports","specifier":"torch.nn.parallel","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"time","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._get_gpt_model","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama.broadcast_params","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama.broadcast_params","relation":"calls","name":"parameters","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama.broadcast_params","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama.broadcast_params","relation":"calls","name":"get_data_parallel_src_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama.broadcast_params","relation":"calls","name":"get_data_parallel_group","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"get_data_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"get_model_parallel_group","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"list","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"len","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"len","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"unwrap_model","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_get_gpt_model","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"len","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tensor","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tensor","relation":"calls","name":"broadcast_object_list","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tensor","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tensor","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tensor","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tensor","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tensor","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"mutate_func","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"chunk","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"broadcast_object_list","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_vocab","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"mutate_func","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"chunk","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"broadcast_object_list","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"chunk","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"broadcast_object_list","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"chunk","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"broadcast_object_list","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"copy_","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_get_gpt_model","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tp_shard_tensor_vocab","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_megatron_calc_layer_map","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_get_gpt_model","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tp_shard_tensor_qkv","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tp_shard_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tp_shard_tensor_gate_up","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tp_shard_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_get_gpt_model","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"getattr","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"_broadcast_tp_shard_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"barrier","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"broadcast_params","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py#load_state_dict_to_megatron_llama","relation":"calls","name":"time","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_loader.py"}]},"verl/models/llama/megatron/checkpoint_utils/llama_saver.py":{"size":18189,"mtimeMs":1789275830308.2024,"hash":"97d6ea7470f04372dac5e6e50b8f4febe1e6db3759706c7da8a2c3593e9017f6","nodes":[{"id":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","name":"llama_saver.py","kind":"file","path":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","span":"L1-L450","signature":null,"exported":true,"origin":"ast","body_hash":"97d6ea7470f04372dac5e6e50b8f4febe1e6db3759706c7da8a2c3593e9017f6","chars":18189,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 megatron from megatron.core import mpu from megatron.utils import print_rank_0, unwrap_model from megatron.model import Float16Module from megatron.model import DistributedDataParallel as LocalDDP from torch.nn.parallel import DistributedDataParallel as torchDDP import torch import time from typing import Optional import torch.distributed as dist from megatron import get_args"},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_global_rank","name":"_megatron_calc_global_rank","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","span":"L28-L42","signature":"def _megatron_calc_global_rank(tp_rank: int = 0, dp_rank: int = 0, pp_rank: int = 0)","exported":false,"origin":"ast","body_hash":"9b0c446f50064da1d9b1bf32de75cc134e3a4180e858fc9c7846c128883b3bdb","body_text":"def _megatron_calc_global_rank(tp_rank: int = 0, dp_rank: int = 0, pp_rank: int = 0): \"\"\"given TP,DP,PP rank to get the global rank.\"\"\" args = get_args() tp_size = mpu.get_tensor_model_parallel_world_size() dp_size = mpu.get_data_parallel_world_size() pp_size = mpu.get_pipeline_model_parallel_world_size() assert (tp_size * dp_size * pp_size == torch.distributed.get_world_size() ), f\"{tp_size} x {dp_size} x {pp_size} != {torch.distributed.get_world_size()}\" if args.switch_dp_and_pp_grouping: # TP-PP-DP grouping return (dp_rank * pp_size + pp_rank) * tp_size + tp_rank else: # TP-DP-PP grouping return (pp_rank * dp_size + dp_rank) * tp_size + tp_rank","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_layer_map","name":"_megatron_calc_layer_map","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","span":"L45-L73","signature":"def _megatron_calc_layer_map(config)","exported":false,"origin":"ast","body_hash":"dc7be979a73e481ce00f8b5bb77fd84d10449de0869f2f3dd7ab95a110fa9b2d","body_text":"def _megatron_calc_layer_map(config): \"\"\"Calculate the mapping of global layer_idx to local layer_idx Returns: layer_map (Dict: int -> tuple(int, int, int)): mapping from the global layer index to a tuple of (pp_rank, virtual_pp_rank, layer_idx inside model) \"\"\" import megatron from megatron.core import mpu pp_size = mpu.get_pipeline_model_parallel_world_size() virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1 args = megatron.get_args() layer_map = dict() num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers for pp_rank_idx in range(pp_size): for virtual_pp_rank_idx in range(virtual_pp_size): layer_offset = (virtual_pp_rank_idx * (config.num_hidden_layers // virtual_pp_size) + pp_rank_idx * num_layers_per_model) for layer_idx in range(num_layers_per_model): layer_map[layer_offset + layer_idx] = ( pp_rank_idx, virtual_pp_rank_idx, layer_idx, ) return layer_map","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","name":"merge_megatron_ckpt_llama","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","span":"L76-L449","signature":"def merge_megatron_ckpt_llama(wrapped_models, config, is_value_model=False, dtype='bf16')","exported":true,"origin":"ast","body_hash":"fe9498e38cda1e5a06f103cf0ca7833a84734e76d66a64317cbfb97e1e07c545","body_text":"def merge_megatron_ckpt_llama(wrapped_models, config, is_value_model=False, dtype='bf16'): \"\"\"Merge sharded parameters of a Megatron module into a merged checkpoint. Args: wrapped_modelss (list of megatron.model.DistributedDataParallel): The local DDP wrapped megatron modules. dtype (str or None): The data type of state_dict. if None, the data type of the original parameters is used. gpt_model_key: key to access model Returns: state_dict (dict): The merged state_dict in rank 0, and an empty dictionary in other ranks. \"\"\" start_time = time.time() args = megatron.get_args() def _get_gpt_model(model): return model dp_rank = mpu.get_data_parallel_rank() pp_size = mpu.get_pipeline_model_parallel_world_size() pp_rank = mpu.get_pipeline_model_parallel_rank() virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1 mp_group = mpu.get_model_parallel_group() if dist.get_rank() == 0: assert mp_group.rank() == 0, f\"mp_rank:[{mp_group.rank}] != 0 on rank #0\" assert pp_rank == 0, f\"pp_rank:[{pp_rank}] != 0 on rank #0\" assert dp_rank == 0, f\"dp_rank:[{dp_rank}] != 0 on rank #0\" if not isinstance(wrapped_models, (list, tuple)): wrapped_models = list(wrapped_models) assert len(wrapped_models) == virtual_pp_size num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers models = [None] * len(wrapped_models) for i, wrapped_model in enumerate(wrapped_models): models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module)) assert len(models[i].model.layers ) == num_layers_per_model, 'len model layers {} not equal to num_layers_per_model {}'.format( len(models[i].model.layers), num_layers_per_model) state_dict = dict() def _get_cpu_tensor(tensor: torch.Tensor): if tensor is None: return None if tensor.device == torch.device(\"cpu\"): return tensor.detach().clone() return tensor.detach().cpu() def _broadcast_tensor(tensor, name, src_pp_rank) -> torch.Tensor: \"\"\"broadcast tensor across mp_group\"\"\" nonlocal state_dict nonlocal mp_group src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) if torch.distributed.get_rank() == src_rank: if tensor is None: weight = None tensor_shape = None else: weight = tensor tensor_shape = weight.shape else: weight = None tensor_shape = None obj_list = [tensor_shape] dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group) tensor_shape = obj_list[0] if tensor_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tensor:[{name}] not exist, skip collect\") return if weight is None: weight = torch.empty( tensor_shape, dtype=args.params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) dist.broadcast(weight, src=src_rank, group=mp_group) if torch.distributed.get_rank() == 0: state_dict[name] = _get_cpu_tensor(weight) def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_func=None) -> torch.Tensor: \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) if torch.distributed.get_rank() == src_rank: chunk_shape = tensor.shape else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tp_shard tensor:[{name}] not exist, skip collecting\") return buffer_tensor = torch.empty( chunk_shape, dtype=args.params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) chunk_tensors = [None] * tp_size for i in range(tp_size): cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank) sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group) if torch.distributed.get_rank() == 0: chunk_tensors[i] = _get_cpu_tensor(sync_tensor) if torch.distributed.get_rank() == 0: full_tensor = torch.concat(chunk_tensors, dim=concat_dim) if mutate_func is not None: full_tensor = mutate_func(full_tensor) state_dict[name] = full_tensor def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) -> torch.Tensor: \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) if torch.distributed.get_rank() == src_rank: chunk_shape = tensor.shape else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._get_gpt_model","name":"_get_gpt_model","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","span":"L93-L94","signature":"def _get_gpt_model(model)","exported":false,"origin":"ast","body_hash":"c616fc9942bbf7a93cfd936b39e34055ee56d16434699c954dafff40af8155a0","body_text":"def _get_gpt_model(model): return model","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._get_cpu_tensor","name":"_get_cpu_tensor","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","span":"L124-L129","signature":"def _get_cpu_tensor(tensor: torch.Tensor)","exported":false,"origin":"ast","body_hash":"634438126ca52c64e5a6a3a05857c5225dd50a8a2ff864f8ff29af2740ef8836","body_text":"def _get_cpu_tensor(tensor: torch.Tensor): if tensor is None: return None if tensor.device == torch.device(\"cpu\"): return tensor.detach().clone() return tensor.detach().cpu()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","name":"_broadcast_tensor","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","span":"L131-L168","signature":"def _broadcast_tensor(tensor, name, src_pp_rank) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"f751f760ddbbb612007997767e0139d9182d9cf378e93bc462a911675818b339","body_text":"def _broadcast_tensor(tensor, name, src_pp_rank) -> torch.Tensor: \"\"\"broadcast tensor across mp_group\"\"\" nonlocal state_dict nonlocal mp_group src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) if torch.distributed.get_rank() == src_rank: if tensor is None: weight = None tensor_shape = None else: weight = tensor tensor_shape = weight.shape else: weight = None tensor_shape = None obj_list = [tensor_shape] dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group) tensor_shape = obj_list[0] if tensor_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tensor:[{name}] not exist, skip collect\") return if weight is None: weight = torch.empty( tensor_shape, dtype=args.params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) dist.broadcast(weight, src=src_rank, group=mp_group) if torch.distributed.get_rank() == 0: state_dict[name] = _get_cpu_tensor(weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","name":"_broadcast_tp_shard_tensor","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","span":"L170-L212","signature":"def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_func=None) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"99b7728f93aa86843e2769e9a7fa2039c043e054aa5cee1fb7445cf9ae2dd7db","body_text":"def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_func=None) -> torch.Tensor: \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) if torch.distributed.get_rank() == src_rank: chunk_shape = tensor.shape else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tp_shard tensor:[{name}] not exist, skip collecting\") return buffer_tensor = torch.empty( chunk_shape, dtype=args.params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) chunk_tensors = [None] * tp_size for i in range(tp_size): cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank) sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group) if torch.distributed.get_rank() == 0: chunk_tensors[i] = _get_cpu_tensor(sync_tensor) if torch.distributed.get_rank() == 0: full_tensor = torch.concat(chunk_tensors, dim=concat_dim) if mutate_func is not None: full_tensor = mutate_func(full_tensor) state_dict[name] = full_tensor","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","name":"_broadcast_tp_shard_tensor_gate_up","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","span":"L214-L265","signature":"def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) -> torch.Tensor","exported":false,"origin":"ast","body_hash":"9d071bedd414ad0e06c61ff3afed5382e0336d9afb46ee39e3141eba42db58de","body_text":"def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) -> torch.Tensor: \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) if torch.distributed.get_rank() == src_rank: chunk_shape = tensor.shape else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tp_shard tensor:[{gate_name, up_name}] not exist, skip collecting\") return buffer_tensor = torch.empty( chunk_shape, dtype=args.params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) chunk_tensors = [None] * tp_size for i in range(tp_size): cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank) sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group) if torch.distributed.get_rank() == 0: chunk_tensors[i] = _get_cpu_tensor(sync_tensor) if torch.distributed.get_rank() == 0: full_tensor = torch.concat(chunk_tensors, dim=0) intermediate_size_tp = config.intermediate_size // tp_size gate_weight_list = [] up_weight_list = [] for i in range(tp_size): gate_up_weight_tp = full_tensor[intermediate_size_tp * 2 * i:intermediate_size_tp * 2 * (i + 1)] gate_weight_tp = gate_up_weight_tp[:intermediate_size_tp] up_weight_tp = gate_up_weight_tp[intermediate_size_tp:] gate_weight_list.append(gate_weight_tp) up_weight_list.append(up_weight_tp) state_dict[gate_name] = torch.cat(gate_weight_list, dim=0) state_dict[up_name] = torch.cat(up_weight_list, dim=0)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","name":"_broadcast_tp_shard_tensor_qkv","kind":"function","path":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","span":"L267-L340","signature":"def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank)","exported":false,"origin":"ast","body_hash":"f2dbbc6cf55676942e81e9043e085873fed5977d432a3915507dcebcdc56dc01","body_text":"def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): \"\"\"broadcast tensor in tp shards across mp_group\"\"\" nonlocal state_dict nonlocal mp_group tp_rank = mpu.get_tensor_model_parallel_rank() tp_size = mpu.get_tensor_model_parallel_world_size() src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) if torch.distributed.get_rank() == src_rank: chunk_shape = tensor.shape else: chunk_shape = None obj_list = [chunk_shape] dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group) chunk_shape = obj_list[0] if chunk_shape is None: # all or none ranks in the mp_group should reach here print_rank_0(f\"tp_shard tensor:[{q_name}] not exist, skip collecting\") return buffer_tensor = torch.empty( chunk_shape, dtype=args.params_dtype, device=torch.cuda.current_device(), requires_grad=False, ) chunk_tensors = [None] * tp_size for i in range(tp_size): cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank) sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group) if torch.distributed.get_rank() == 0: chunk_tensors[i] = _get_cpu_tensor(sync_tensor) if torch.distributed.get_rank() == 0: full_tensor = torch.concat(chunk_tensors, dim=0) q_weight_list = [] k_weight_list = [] v_weight_list = [] hidden_size_per_head = config.hidden_size // config.num_attention_heads if config.num_key_value_heads >= tp_size: q_size_tp = config.hidden_size // tp_size kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size total_size = q_size_tp + 2 * kv_size_tp for i in range(tp_size): qkv_part = full_tensor[i * total_size:(i + 1) * total_size] q_part = qkv_part[:q_size_tp] k_part = qkv_part[q_size_tp:q_size_tp + kv_size_tp] v_part = qkv_part[q_size_tp + kv_size_tp:total_size] q_weight_list.append(q_part) k_weight_list.append(k_part) v_weight_list.append(v_part) else: q_size_tp = config.hidden_size // tp_size kv_size_tp = hidden_size_per_head total_size = q_size_tp + 2 * kv_size_tp for i in range(tp_size): qkv_part = full_tensor[i * total_size:(i + 1) * total_size] q_part = qkv_part[:q_size_tp] k_part = qkv_part[q_size_tp:q_size_tp + kv_size_tp] v_part = qkv_part[q_size_tp + kv_size_tp:total_size] q_weight_list.append(q_part) if i * config.num_key_value_heads % tp_size == 0: k_weight_list.append(k_part) v_weight_list.append(v_part) state_dict[q_name] = torch.cat(q_weight_list, dim=0) state_dict[k_name] = torch.cat(k_weight_list, dim=0) state_dict[v_name] = torch.cat(v_weight_list, dim=0)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"megatron","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"megatron.utils","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"megatron.model","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"megatron.model","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"torch.nn.parallel","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"time","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"typing","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"megatron","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_global_rank","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_global_rank","relation":"calls","name":"get_args","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_global_rank","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_global_rank","relation":"calls","name":"get_data_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_global_rank","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_global_rank","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_global_rank","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_layer_map","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"megatron","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_layer_map","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_layer_map","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_layer_map","relation":"calls","name":"get_args","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_layer_map","relation":"calls","name":"dict","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_layer_map","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_layer_map","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#_megatron_calc_layer_map","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"time","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"get_args","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._get_gpt_model","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"get_data_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"get_model_parallel_group","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"list","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"len","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"len","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"unwrap_model","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"len","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"format","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"len","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"dict","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._get_cpu_tensor","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._get_cpu_tensor","relation":"calls","name":"device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._get_cpu_tensor","relation":"calls","name":"clone","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._get_cpu_tensor","relation":"calls","name":"detach","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._get_cpu_tensor","relation":"calls","name":"cpu","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._get_cpu_tensor","relation":"calls","name":"detach","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","relation":"calls","name":"_megatron_calc_global_rank","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","relation":"calls","name":"broadcast_object_list","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tensor","relation":"calls","name":"_get_cpu_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"_megatron_calc_global_rank","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"broadcast_object_list","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"_megatron_calc_global_rank","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"_get_cpu_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"concat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor","relation":"calls","name":"mutate_func","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"_megatron_calc_global_rank","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"broadcast_object_list","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"_megatron_calc_global_rank","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"_get_cpu_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"concat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"append","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"append","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_gate_up","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"contains","targetId":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"_megatron_calc_global_rank","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"broadcast_object_list","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"empty","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"current_device","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"_megatron_calc_global_rank","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"_get_cpu_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"concat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"append","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"append","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"append","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"append","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"append","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"append","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama._broadcast_tp_shard_tensor_qkv","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_get_gpt_model","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_broadcast_tp_shard_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_megatron_calc_layer_map","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_get_gpt_model","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_broadcast_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_broadcast_tp_shard_tensor_qkv","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_broadcast_tp_shard_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_broadcast_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_broadcast_tp_shard_tensor_gate_up","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_broadcast_tp_shard_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_get_gpt_model","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_broadcast_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"getattr","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_broadcast_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"getattr","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"_broadcast_tp_shard_tensor","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"getattr","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"barrier","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"print","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"exit","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"items","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py","recvType":"dict"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"print_rank_0","viaMember":false,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"},{"source":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py#merge_megatron_ckpt_llama","relation":"calls","name":"time","viaMember":true,"file":"verl/models/llama/megatron/checkpoint_utils/llama_saver.py"}]},"verl/models/llama/megatron/layers/__init__.py":{"size":838,"mtimeMs":1789275830308.2024,"hash":"164f9b46a38b7c500c99f5a1f11a6b6d0a3160bda61296f8062d20c2720d5743","nodes":[{"id":"verl/models/llama/megatron/layers/__init__.py","name":"__init__.py","kind":"file","path":"verl/models/llama/megatron/layers/__init__.py","span":"L1-L19","signature":null,"exported":true,"origin":"ast","body_hash":"164f9b46a38b7c500c99f5a1f11a6b6d0a3160bda61296f8062d20c2720d5743","chars":838,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 .parallel_attention import ParallelLlamaAttention from .parallel_decoder import ParallelLlamaDecoderLayer, ParallelLlamaDecoderLayerRmPad from .parallel_mlp import ParallelLlamaMLP from .parallel_rmsnorm import ParallelLlamaRMSNorm"}],"rawEdges":[{"source":"verl/models/llama/megatron/layers/__init__.py","relation":"imports","specifier":".parallel_attention","file":"verl/models/llama/megatron/layers/__init__.py"},{"source":"verl/models/llama/megatron/layers/__init__.py","relation":"imports","specifier":".parallel_decoder","file":"verl/models/llama/megatron/layers/__init__.py"},{"source":"verl/models/llama/megatron/layers/__init__.py","relation":"imports","specifier":".parallel_mlp","file":"verl/models/llama/megatron/layers/__init__.py"},{"source":"verl/models/llama/megatron/layers/__init__.py","relation":"imports","specifier":".parallel_rmsnorm","file":"verl/models/llama/megatron/layers/__init__.py"}]},"verl/models/llama/megatron/layers/parallel_attention.py":{"size":20129,"mtimeMs":1789275830308.2024,"hash":"223115766f99e5c92baa9a2d6f2d0451c2f472ddc3e05da176399dbd438daa5e","nodes":[{"id":"verl/models/llama/megatron/layers/parallel_attention.py","name":"parallel_attention.py","kind":"file","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L1-L419","signature":null,"exported":true,"origin":"ast","body_hash":"223115766f99e5c92baa9a2d6f2d0451c2f472ddc3e05da176399dbd438daa5e","chars":20129,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX # and OPT implementations in this library. It has been modified from its # original forms to accommodate minor architectural differences compared # to GPT-NeoX and OPT used by the Meta AI team that trained the model. # # 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 math from typing import Optional, Tuple import torch from megatron.core import parallel_state as mpu from megatron.core import tensor_parallel from megatron.core import ModelParallelConfig from torch import nn from transformers import LlamaConfig from verl.models.llama.megatron.layers.parallel_linear import QKVParallelLinear from verl.utils.megatron import tensor_parallel as tp_utils \"\"\" Remove padding Attention - Using Flash-attn 2 - Compatible with sequence parallel \"\"\" from transformers.utils import is_flash_attn_2_available import torch.nn.functional as F from einops import rearrange if is_flash_attn_2_available(): from flash_attn import flash_attn_varlen_func from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa from flash_attn.layers.rotary import apply_rotary_emb # use flash-attn rotary embeddings with rmpad # cos/sin shoudl be: (seq_length, rotary_dim / 2)"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding","name":"LlamaRotaryEmbedding","kind":"class","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L35-L69","signature":"class LlamaRotaryEmbedding(nn.Module)","exported":true,"origin":"ast","body_hash":"f2845998eacac15267e954eb2c3c90b3c220c3bf35fdbf6e87aad53211b02dfb","body_text":"class LlamaRotaryEmbedding(nn.Module): def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (self.base**(torch.arange(0, self.dim, 2).float().to(device) / self.dim)) self.register_buffer(\"inv_freq\", inv_freq, persistent=False) # Build here to make `torch.jit.trace` work. self._set_cos_sin_cache(seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False) self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False) def forward(self, x, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size] if seq_len > self.max_seq_len_cached: self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) return ( self.cos_cached[:seq_len].to(dtype=x.dtype), self.sin_cached[:seq_len].to(dtype=x.dtype), )","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L37-L49","signature":"def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None)","exported":false,"origin":"ast","body_hash":"ce5a30b7a72ae26b73b6ddc92332d2fa42c49960c4d0a467c6750451fbf39020","body_text":"def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): super().__init__() self.dim = dim self.max_position_embeddings = max_position_embeddings self.base = base inv_freq = 1.0 / (self.base**(torch.arange(0, self.dim, 2).float().to(device) / self.dim)) self.register_buffer(\"inv_freq\", inv_freq, persistent=False) # Build here to make `torch.jit.trace` work. self._set_cos_sin_cache(seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype())","summary_state":"pending","summary":null,"crux":null,"owner":"LlamaRotaryEmbedding"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","name":"_set_cos_sin_cache","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L51-L59","signature":"def _set_cos_sin_cache(self, seq_len, device, dtype)","exported":false,"origin":"ast","body_hash":"a25fc5c56e8d3d18792e9162b55dc7e1f2f81e1114d49c9bb81da9a301e97b49","body_text":"def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False) self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)","summary_state":"pending","summary":null,"crux":null,"owner":"LlamaRotaryEmbedding"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L61-L69","signature":"def forward(self, x, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size]","exported":true,"origin":"ast","body_hash":"6df6ac35c10318f3f601a9fbd9724d72866380292e021604aaf0249e03d2b921","body_text":"def forward(self, x, seq_len=None): # x: [bs, num_attention_heads, seq_len, head_size] if seq_len > self.max_seq_len_cached: self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype) return ( self.cos_cached[:seq_len].to(dtype=x.dtype), self.sin_cached[:seq_len].to(dtype=x.dtype), )","summary_state":"pending","summary":null,"crux":null,"owner":"LlamaRotaryEmbedding"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding","name":"LlamaLinearScalingRotaryEmbedding","kind":"class","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L72-L88","signature":"class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding)","exported":true,"origin":"ast","body_hash":"8e5b43c49306cf379c31fabec41431de0b56421827e65e12f0ad0f6ae56b0a0e","body_text":"class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding): \"\"\"LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev\"\"\" def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): self.scaling_factor = scaling_factor super().__init__(dim, max_position_embeddings, base, device) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) t = t / self.scaling_factor freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False) self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L75-L77","signature":"def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0)","exported":false,"origin":"ast","body_hash":"d0313c98509e4d31d4a11532dd1321bd3a85291f672a65b70cda3341ee793e3b","body_text":"def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): self.scaling_factor = scaling_factor super().__init__(dim, max_position_embeddings, base, device)","summary_state":"pending","summary":null,"crux":null,"owner":"LlamaLinearScalingRotaryEmbedding"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","name":"_set_cos_sin_cache","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L79-L88","signature":"def _set_cos_sin_cache(self, seq_len, device, dtype)","exported":false,"origin":"ast","body_hash":"c5f8c29872a488225a6a17de383de4c60a7f66dbac9be5882d5cdbd23115fb59","body_text":"def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) t = t / self.scaling_factor freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False) self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)","summary_state":"pending","summary":null,"crux":null,"owner":"LlamaLinearScalingRotaryEmbedding"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding","name":"LlamaDynamicNTKScalingRotaryEmbedding","kind":"class","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L91-L113","signature":"class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding)","exported":true,"origin":"ast","body_hash":"eb738a4933524ec4b044a1f19783b4b776467e064d7bfe346123515f8173b276","body_text":"class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding): \"\"\"LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla\"\"\" def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): self.scaling_factor = scaling_factor super().__init__(dim, max_position_embeddings, base, device) def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len if seq_len > self.max_position_embeddings: base = self.base * ((self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1))**(self.dim / (self.dim - 2)) inv_freq = 1.0 / (base**(torch.arange(0, self.dim, 2).float().to(device) / self.dim)) self.register_buffer(\"inv_freq\", inv_freq, persistent=False) t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False) self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L94-L96","signature":"def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0)","exported":false,"origin":"ast","body_hash":"d0313c98509e4d31d4a11532dd1321bd3a85291f672a65b70cda3341ee793e3b","body_text":"def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0): self.scaling_factor = scaling_factor super().__init__(dim, max_position_embeddings, base, device)","summary_state":"pending","summary":null,"crux":null,"owner":"LlamaDynamicNTKScalingRotaryEmbedding"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","name":"_set_cos_sin_cache","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L98-L113","signature":"def _set_cos_sin_cache(self, seq_len, device, dtype)","exported":false,"origin":"ast","body_hash":"f941565f73f5eaaaddd0fb79b9e7f797ea51aecca7b244684f921d31c6208ff7","body_text":"def _set_cos_sin_cache(self, seq_len, device, dtype): self.max_seq_len_cached = seq_len if seq_len > self.max_position_embeddings: base = self.base * ((self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1))**(self.dim / (self.dim - 2)) inv_freq = 1.0 / (base**(torch.arange(0, self.dim, 2).float().to(device) / self.dim)) self.register_buffer(\"inv_freq\", inv_freq, persistent=False) t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype) freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq) # Different from paper, but it uses a different permutation in order to obtain the same calculation emb = torch.cat((freqs, freqs), dim=-1) self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False) self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)","summary_state":"pending","summary":null,"crux":null,"owner":"LlamaDynamicNTKScalingRotaryEmbedding"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#rotate_half","name":"rotate_half","kind":"function","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L116-L120","signature":"def rotate_half(x)","exported":true,"origin":"ast","body_hash":"ca1458f15a1ae0848a6ac905e0d26b8057b076202f76165f28c1b40e4c509b2f","body_text":"def rotate_half(x): \"\"\"Rotates half the hidden dims of the input.\"\"\" x1 = x[..., :x.shape[-1] // 2] x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb","name":"apply_rotary_pos_emb","kind":"function","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L123-L128","signature":"def apply_rotary_pos_emb(q, k, cos, sin, position_ids)","exported":true,"origin":"ast","body_hash":"433506a5e611643a00df552be6d569541eae9b46897327c856b7a8522e395e2b","body_text":"def apply_rotary_pos_emb(q, k, cos, sin, position_ids): cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim] sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim] q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) return q_embed, k_embed","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#repeat_kv","name":"repeat_kv","kind":"function","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L131-L140","signature":"def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"d7cb84ea104d6a1d59eef5013d518dd66a36183d083536ef5f46aa65a7143f49","body_text":"def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: \"\"\" This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) \"\"\" batch, num_key_value_heads, slen, head_dim = hidden_states.shape if n_rep == 1: return hidden_states hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim) return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention","name":"ParallelLlamaAttention","kind":"class","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L143-L280","signature":"class ParallelLlamaAttention(nn.Module)","exported":true,"origin":"ast","body_hash":"4348cbf9d8fe41615455fb270610a95cd6f658871b2afac5fa0a4e81374c4d31","body_text":"class ParallelLlamaAttention(nn.Module): \"\"\"Multi-headed attention from 'Attention Is All You Need' paper\"\"\" def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.config = config self.megatron_config = megatron_config self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.hidden_size // self.num_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta # assign values after tp tp_size = mpu.get_tensor_model_parallel_world_size() assert self.num_heads % tp_size == 0, f'num_head must be divisible by tp_size. Got num_head={self.num_heads}, tp_size={tp_size}' assert self.num_key_value_heads % tp_size == 0, \\ f'num_key_value_heads must be divisible by tp_size. Got num_key_value_heads={self.num_key_value_heads}, tp_size={tp_size}' self.num_heads_per_tp = self.num_heads // tp_size self.num_key_value_heads_per_tp = self.num_key_value_heads // tp_size self.hidden_size_per_tp = self.hidden_size // tp_size if (self.head_dim * self.num_heads) != self.hidden_size: raise ValueError(f\"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}\" f\" and `num_heads`: {self.num_heads}).\") column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() row_kwargs = tp_utils.get_default_kwargs_for_row_parallel_linear() if megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' assert row_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, megatron_config) tp_utils.update_kwargs_with_config(row_kwargs, megatron_config) # [self.q_size, self.k_size, self.v_size] self.qkv_proj = QKVParallelLinear(input_size=self.hidden_size, num_heads=self.num_heads, num_key_value_heads=self.num_key_value_heads, head_dim=self.head_dim, bias=config.attention_bias, gather_output=False, skip_bias_add=False, **column_kwargs) self.q_size = self.num_heads_per_tp * self.head_dim self.k_size = self.num_key_value_heads_per_tp * self.head_dim self.v_size = self.num_key_value_heads_per_tp * self.head_dim self.o_proj = tensor_parallel.RowParallelLinear(input_size=self.num_heads * self.head_dim, output_size=self.hidden_size, bias=config.attention_bias, input_is_parallel=True, skip_bias_add=False, **row_kwargs) self._init_rope() def _init_rope(self): if self.config.rope_scaling is None: self.rotary_emb = LlamaRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, ) else: scaling_type = self.config.rope_scaling[\"type\"] scaling_factor = self.config.rope_scaling[\"factor\"] if scaling_type == \"linear\": self.rotary_emb = LlamaLinearScalingRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor, base=self.rope_theta, ) elif scaling_type == \"dynamic\": self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor, base=self.rope_theta, ) else: raise ValueError(f\"Unknown RoPE scaling type {scaling_type}\") def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous() def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() qkv = self.qkv_proj(hidden_states)[0] query_states, key_states, value_states = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) query_states = query_states.view(bsz, q_len, self.num_heads_per_tp, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads_per_tp, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads_per_tp, self.head_dim).transpose(1, 2) kv_seq_len = key_states.shape[-2] cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) if attn_weights.size() != (bsz, self.num_heads_per_tp, q_len, kv_seq_len): raise ValueError( f\"Attention weights should be of size {(bsz, self.num_heads_per_tp, q_len, kv_seq_len)}, but is\" f\" {attn_weights.size()}\") if attention_mask is not None: if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): raise ValueError( f\"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mas","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L146-L202","signature":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig)","exported":false,"origin":"ast","body_hash":"28215bfd00e5c29e91f28f8136f7b46f2eb67df11950f9fc204560a49e98214d","body_text":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.config = config self.megatron_config = megatron_config self.hidden_size = config.hidden_size self.num_heads = config.num_attention_heads self.head_dim = self.hidden_size // self.num_heads self.num_key_value_heads = config.num_key_value_heads self.num_key_value_groups = self.num_heads // self.num_key_value_heads self.max_position_embeddings = config.max_position_embeddings self.rope_theta = config.rope_theta # assign values after tp tp_size = mpu.get_tensor_model_parallel_world_size() assert self.num_heads % tp_size == 0, f'num_head must be divisible by tp_size. Got num_head={self.num_heads}, tp_size={tp_size}' assert self.num_key_value_heads % tp_size == 0, \\ f'num_key_value_heads must be divisible by tp_size. Got num_key_value_heads={self.num_key_value_heads}, tp_size={tp_size}' self.num_heads_per_tp = self.num_heads // tp_size self.num_key_value_heads_per_tp = self.num_key_value_heads // tp_size self.hidden_size_per_tp = self.hidden_size // tp_size if (self.head_dim * self.num_heads) != self.hidden_size: raise ValueError(f\"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}\" f\" and `num_heads`: {self.num_heads}).\") column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() row_kwargs = tp_utils.get_default_kwargs_for_row_parallel_linear() if megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' assert row_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, megatron_config) tp_utils.update_kwargs_with_config(row_kwargs, megatron_config) # [self.q_size, self.k_size, self.v_size] self.qkv_proj = QKVParallelLinear(input_size=self.hidden_size, num_heads=self.num_heads, num_key_value_heads=self.num_key_value_heads, head_dim=self.head_dim, bias=config.attention_bias, gather_output=False, skip_bias_add=False, **column_kwargs) self.q_size = self.num_heads_per_tp * self.head_dim self.k_size = self.num_key_value_heads_per_tp * self.head_dim self.v_size = self.num_key_value_heads_per_tp * self.head_dim self.o_proj = tensor_parallel.RowParallelLinear(input_size=self.num_heads * self.head_dim, output_size=self.hidden_size, bias=config.attention_bias, input_is_parallel=True, skip_bias_add=False, **row_kwargs) self._init_rope()","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaAttention"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._init_rope","name":"_init_rope","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L204-L229","signature":"def _init_rope(self)","exported":false,"origin":"ast","body_hash":"610b77e59b458abb91fbb270e69669468c9fee9d5df3336a419761e3c8d1e7aa","body_text":"def _init_rope(self): if self.config.rope_scaling is None: self.rotary_emb = LlamaRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, base=self.rope_theta, ) else: scaling_type = self.config.rope_scaling[\"type\"] scaling_factor = self.config.rope_scaling[\"factor\"] if scaling_type == \"linear\": self.rotary_emb = LlamaLinearScalingRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor, base=self.rope_theta, ) elif scaling_type == \"dynamic\": self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding( self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor, base=self.rope_theta, ) else: raise ValueError(f\"Unknown RoPE scaling type {scaling_type}\")","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaAttention"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._shape","name":"_shape","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L231-L232","signature":"def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int)","exported":false,"origin":"ast","body_hash":"72a1970b90b32dd87915b6e371cd23059805a40e653c029d5ae32778af7c14de","body_text":"def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int): return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaAttention"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L234-L280","signature":"def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]","exported":true,"origin":"ast","body_hash":"cd151d0d6aacade4dd84d7cdf1616032c1eb313d40ff206eeca28d959e506461","body_text":"def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: bsz, q_len, _ = hidden_states.size() qkv = self.qkv_proj(hidden_states)[0] query_states, key_states, value_states = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) query_states = query_states.view(bsz, q_len, self.num_heads_per_tp, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads_per_tp, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads_per_tp, self.head_dim).transpose(1, 2) kv_seq_len = key_states.shape[-2] cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len) query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids) key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim) if attn_weights.size() != (bsz, self.num_heads_per_tp, q_len, kv_seq_len): raise ValueError( f\"Attention weights should be of size {(bsz, self.num_heads_per_tp, q_len, kv_seq_len)}, but is\" f\" {attn_weights.size()}\") if attention_mask is not None: if attention_mask.size() != (bsz, 1, q_len, kv_seq_len): raise ValueError( f\"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}\") attn_weights = attn_weights + attention_mask # upcast attention to fp32 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype) attn_output = torch.matmul(attn_weights, value_states) if attn_output.size() != (bsz, self.num_heads_per_tp, q_len, self.head_dim): raise ValueError( f\"`attn_output` should be of size {(bsz, self.num_heads_per_tp, q_len, self.head_dim)}, but is\" f\" {attn_output.size()}\") attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, self.hidden_size_per_tp) attn_output = self.o_proj(attn_output)[0] return attn_output","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaAttention"},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","name":"apply_rotary_pos_emb_rmpad","kind":"function","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L299-L312","signature":"def apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, sequence_length)","exported":true,"origin":"ast","body_hash":"f9358ba866436d9f1ecfbb357800efb08372d8ca48fc658ab660d72895810d8d","body_text":"def apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, sequence_length): batch_size = position_ids.shape[0] q = pad_input(q, indices, batch_size, sequence_length) # (batch_size, seqlen, num_head, head_dim) k = pad_input(k, indices, batch_size, sequence_length) cos = cos[position_ids].unsqueeze(2) # [bs, seq_len, 1, dim] sin = sin[position_ids].unsqueeze(2) # [bs, seq_len, 1, dim] q_embed = (q * cos) + (rotate_half(q) * sin) k_embed = (k * cos) + (rotate_half(k) * sin) q_embed = index_first_axis(rearrange(q_embed, \"b s ... -> (b s) ...\"), indices) k_embed = index_first_axis(rearrange(k_embed, \"b s ... -> (b s) ...\"), indices) return q_embed, k_embed","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad_flash","name":"apply_rotary_pos_emb_rmpad_flash","kind":"function","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L320-L335","signature":"def apply_rotary_pos_emb_rmpad_flash(q, k, cos, sin, cu_seqlens, max_seqlen)","exported":true,"origin":"ast","body_hash":"0cbfefba2d2d99f4de31bbfed26149abd5afe2b5270ccd73e4c371f5fc18354c","body_text":"def apply_rotary_pos_emb_rmpad_flash(q, k, cos, sin, cu_seqlens, max_seqlen): q_embed = apply_rotary_emb(q, cos, sin, interleaved=False, inplace=False, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) k_embed = apply_rotary_emb(k, cos, sin, interleaved=False, inplace=False, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen) return q_embed, k_embed","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad","name":"ParallelLlamaAttentionRmPad","kind":"class","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L338-L418","signature":"class ParallelLlamaAttentionRmPad(ParallelLlamaAttention)","exported":true,"origin":"ast","body_hash":"ee680d58cf2d43d4d8033cfc0eeb065aaee3282834b6c57acfbd53cce0113733","body_text":"class ParallelLlamaAttentionRmPad(ParallelLlamaAttention): def forward(self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: torch.Tensor = None, max_seqlen_in_batch: int = None): total_nnz, _, _ = hidden_states.size() # This is the total_nnz padded after sequence parallel if self.megatron_config.sequence_parallel: total_nnz = total_nnz * mpu.get_tensor_model_parallel_world_size() qkv = self.qkv_proj(hidden_states)[0] query_states, key_states, value_states = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) # (total_nnz, 1, hidden_size) if self.megatron_config.sequence_parallel: sequence_parallel_pad = total_nnz - cu_seqlens[-1] total_nnz = cu_seqlens[-1] # total_nnz before sp padding query_states = query_states[:total_nnz] key_states = key_states[:total_nnz] value_states = value_states[:total_nnz] # Flash attention requires the input to have the shape # batch_size x seq_length x head_dime x hidden_dim # therefore we just need to keep the original shape query_states = query_states.view(total_nnz, self.num_heads_per_tp, self.head_dim) key_states = key_states.view(total_nnz, self.num_key_value_heads_per_tp, self.head_dim) value_states = value_states.view(total_nnz, self.num_key_value_heads_per_tp, self.head_dim) cos, sin = self.rotary_emb(value_states, seq_len=sequence_length) cos, sin = cos[:, :cos.shape[1] // 2], sin[:, :sin.shape[1] // 2] # flash attn only needs half query_states, key_states = apply_rotary_pos_emb_rmpad_flash(query_states, key_states, cos, sin, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen_in_batch) # query_states, key_states = apply_rotary_pos_emb_rmpad(query_states, key_states, cos, sin, position_ids, indices, # TODO: llama does not have dropout in the config?? # It is recommended to use dropout with FA according to the docs # when training. dropout_rate = 0.0 # if not self.training else self.attn_dropout # In PEFT, usually we cast the layer norms in float32 for training stability reasons # therefore the input hidden states gets silently casted in float32. Hence, we need # cast them back in float16 just to be sure everything works as expected. # This might slowdown training & inference so it is recommended to not cast the LayerNorms # in fp32. (LlamaRMSNorm handles it correctly) input_dtype = query_states.dtype if input_dtype == torch.float32: query_states = query_states.to(torch.float16) key_states = key_states.to(torch.float16) value_states = value_states.to(torch.float16) attn_output_unpad = flash_attn_varlen_func( query_states, key_states, value_states, cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens, max_seqlen_q=max_seqlen_in_batch, max_seqlen_k=max_seqlen_in_batch, dropout_p=dropout_rate, softmax_scale=None, causal=True, ) attn_output_unpad = attn_output_unpad.to(input_dtype) attn_output_unpad = attn_output_unpad.reshape(total_nnz, 1, self.hidden_size_per_tp).contiguous() # sequence parallel reduce_scatter is performed inside RowColumnParallel if enabled # Here we need to repad if self.megatron_config.sequence_parallel: attn_output_unpad = F.pad(attn_output_unpad, pad=(0, 0, 0, 0, 0, sequence_parallel_pad)) attn_output_unpad = self.o_proj(attn_output_unpad)[0] return attn_output_unpad","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/layers/parallel_attention.py","span":"L340-L418","signature":"def forward(self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: torch.Tensor = None, max_seqlen_in_batch: int = None)","exported":true,"origin":"ast","body_hash":"0d7c004be6a796db3cc85686de933c0b2d04c7f95009af49f5b297652a8935a8","body_text":"def forward(self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: torch.Tensor = None, max_seqlen_in_batch: int = None): total_nnz, _, _ = hidden_states.size() # This is the total_nnz padded after sequence parallel if self.megatron_config.sequence_parallel: total_nnz = total_nnz * mpu.get_tensor_model_parallel_world_size() qkv = self.qkv_proj(hidden_states)[0] query_states, key_states, value_states = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1) # (total_nnz, 1, hidden_size) if self.megatron_config.sequence_parallel: sequence_parallel_pad = total_nnz - cu_seqlens[-1] total_nnz = cu_seqlens[-1] # total_nnz before sp padding query_states = query_states[:total_nnz] key_states = key_states[:total_nnz] value_states = value_states[:total_nnz] # Flash attention requires the input to have the shape # batch_size x seq_length x head_dime x hidden_dim # therefore we just need to keep the original shape query_states = query_states.view(total_nnz, self.num_heads_per_tp, self.head_dim) key_states = key_states.view(total_nnz, self.num_key_value_heads_per_tp, self.head_dim) value_states = value_states.view(total_nnz, self.num_key_value_heads_per_tp, self.head_dim) cos, sin = self.rotary_emb(value_states, seq_len=sequence_length) cos, sin = cos[:, :cos.shape[1] // 2], sin[:, :sin.shape[1] // 2] # flash attn only needs half query_states, key_states = apply_rotary_pos_emb_rmpad_flash(query_states, key_states, cos, sin, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen_in_batch) # query_states, key_states = apply_rotary_pos_emb_rmpad(query_states, key_states, cos, sin, position_ids, indices, # TODO: llama does not have dropout in the config?? # It is recommended to use dropout with FA according to the docs # when training. dropout_rate = 0.0 # if not self.training else self.attn_dropout # In PEFT, usually we cast the layer norms in float32 for training stability reasons # therefore the input hidden states gets silently casted in float32. Hence, we need # cast them back in float16 just to be sure everything works as expected. # This might slowdown training & inference so it is recommended to not cast the LayerNorms # in fp32. (LlamaRMSNorm handles it correctly) input_dtype = query_states.dtype if input_dtype == torch.float32: query_states = query_states.to(torch.float16) key_states = key_states.to(torch.float16) value_states = value_states.to(torch.float16) attn_output_unpad = flash_attn_varlen_func( query_states, key_states, value_states, cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens, max_seqlen_q=max_seqlen_in_batch, max_seqlen_k=max_seqlen_in_batch, dropout_p=dropout_rate, softmax_scale=None, causal=True, ) attn_output_unpad = attn_output_unpad.to(input_dtype) attn_output_unpad = attn_output_unpad.reshape(total_nnz, 1, self.hidden_size_per_tp).contiguous() # sequence parallel reduce_scatter is performed inside RowColumnParallel if enabled # Here we need to repad if self.megatron_config.sequence_parallel: attn_output_unpad = F.pad(attn_output_unpad, pad=(0, 0, 0, 0, 0, sequence_parallel_pad)) attn_output_unpad = self.o_proj(attn_output_unpad)[0] return attn_output_unpad","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaAttentionRmPad"}],"rawEdges":[{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"math","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"typing","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"transformers","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"verl.models.llama.megatron.layers.parallel_linear","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"verl.utils.megatron","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.__init__","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.__init__","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.__init__","relation":"calls","name":"float","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.__init__","relation":"calls","name":"arange","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.__init__","relation":"calls","name":"register_buffer","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"LlamaRotaryEmbedding"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.__init__","relation":"calls","name":"_set_cos_sin_cache","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"LlamaRotaryEmbedding"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.__init__","relation":"calls","name":"get_default_dtype","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"arange","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"einsum","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"register_buffer","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"LlamaRotaryEmbedding"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"cos","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"register_buffer","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"LlamaRotaryEmbedding"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"sin","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.forward","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.forward","relation":"calls","name":"_set_cos_sin_cache","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"LlamaRotaryEmbedding"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaRotaryEmbedding.forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding","relation":"extends","name":"LlamaRotaryEmbedding","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding.__init__","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"arange","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"einsum","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"register_buffer","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"LlamaLinearScalingRotaryEmbedding"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"cos","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"register_buffer","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"LlamaLinearScalingRotaryEmbedding"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaLinearScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"sin","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding","relation":"extends","name":"LlamaRotaryEmbedding","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding.__init__","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"float","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"arange","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"register_buffer","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"LlamaDynamicNTKScalingRotaryEmbedding"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"arange","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"einsum","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"register_buffer","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"LlamaDynamicNTKScalingRotaryEmbedding"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"cos","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"register_buffer","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"LlamaDynamicNTKScalingRotaryEmbedding"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#LlamaDynamicNTKScalingRotaryEmbedding._set_cos_sin_cache","relation":"calls","name":"sin","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#rotate_half","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#rotate_half","relation":"calls","name":"cat","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb","relation":"calls","name":"rotate_half","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb","relation":"calls","name":"rotate_half","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#repeat_kv","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#repeat_kv","relation":"calls","name":"expand","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#repeat_kv","relation":"calls","name":"reshape","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"get_default_kwargs_for_column_parallel_linear","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"get_default_kwargs_for_row_parallel_linear","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"QKVParallelLinear","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"RowParallelLinear","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.__init__","relation":"calls","name":"_init_rope","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"ParallelLlamaAttention"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._init_rope","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._init_rope","relation":"calls","name":"LlamaRotaryEmbedding","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._init_rope","relation":"calls","name":"LlamaLinearScalingRotaryEmbedding","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._init_rope","relation":"calls","name":"LlamaDynamicNTKScalingRotaryEmbedding","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._init_rope","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._shape","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._shape","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._shape","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention._shape","relation":"calls","name":"view","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"qkv_proj","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"ParallelLlamaAttention"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"split","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"repeat_kv"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"repeat_kv"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"rotary_emb","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"ParallelLlamaAttention"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"apply_rotary_pos_emb","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"repeat_kv","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"repeat_kv","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"matmul","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"repeat_kv"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"sqrt","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"softmax","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"matmul","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"reshape","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttention.forward","relation":"calls","name":"o_proj","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"ParallelLlamaAttention"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"transformers.utils","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"einops","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"calls","name":"is_flash_attn_2_available","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"flash_attn","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"flash_attn.bert_padding","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","relation":"calls","name":"rotate_half","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","relation":"calls","name":"rotate_half","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","relation":"calls","name":"index_first_axis","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","relation":"calls","name":"rearrange","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","relation":"calls","name":"index_first_axis","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad","relation":"calls","name":"rearrange","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"imports","specifier":"flash_attn.layers.rotary","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad_flash","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad_flash","relation":"calls","name":"apply_rotary_emb","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#apply_rotary_pos_emb_rmpad_flash","relation":"calls","name":"apply_rotary_emb","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad","relation":"extends","name":"ParallelLlamaAttention","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"qkv_proj","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"ParallelLlamaAttentionRmPad"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"split","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"rotary_emb","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"ParallelLlamaAttentionRmPad"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"apply_rotary_pos_emb_rmpad_flash","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"flash_attn_varlen_func","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"flash_attn_varlen_func"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"reshape","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"flash_attn_varlen_func"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"pad","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py"},{"source":"verl/models/llama/megatron/layers/parallel_attention.py#ParallelLlamaAttentionRmPad.forward","relation":"calls","name":"o_proj","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_attention.py","recvType":"ParallelLlamaAttentionRmPad"}]},"verl/models/llama/megatron/layers/parallel_decoder.py":{"size":6036,"mtimeMs":1789275830308.2024,"hash":"b9d5db6124b7be26188a159637d65bdab15c50439590745c5736a62ba3c20d5d","nodes":[{"id":"verl/models/llama/megatron/layers/parallel_decoder.py","name":"parallel_decoder.py","kind":"file","path":"verl/models/llama/megatron/layers/parallel_decoder.py","span":"L1-L147","signature":null,"exported":true,"origin":"ast","body_hash":"b9d5db6124b7be26188a159637d65bdab15c50439590745c5736a62ba3c20d5d","chars":6036,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX # and OPT implementations in this library. It has been modified from its # original forms to accommodate minor architectural differences compared # to GPT-NeoX and OPT used by the Meta AI team that trained the model. # # 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 Optional, Tuple import torch from torch import nn from transformers import LlamaConfig from megatron.core import ModelParallelConfig from .parallel_attention import ParallelLlamaAttention, ParallelLlamaAttentionRmPad from .parallel_mlp import ParallelLlamaMLP from .parallel_rmsnorm import ParallelLlamaRMSNorm"},{"id":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer","name":"ParallelLlamaDecoderLayer","kind":"class","path":"verl/models/llama/megatron/layers/parallel_decoder.py","span":"L33-L96","signature":"class ParallelLlamaDecoderLayer(nn.Module)","exported":true,"origin":"ast","body_hash":"e1fc1ee08ae152d9abb2fcc5ef94ef488433b72213d956c42425062fba9b8df3","body_text":"class ParallelLlamaDecoderLayer(nn.Module): def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.hidden_size = config.hidden_size self.self_attn = ParallelLlamaAttention(config=config, megatron_config=megatron_config) self.mlp = ParallelLlamaMLP(config, megatron_config=megatron_config) self.input_layernorm = ParallelLlamaRMSNorm(config, megatron_config) self.post_attention_layernorm = ParallelLlamaRMSNorm(config, megatron_config) def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: \"\"\" Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states \"\"\" residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Note: sequence parallel is hidden inside ColumnParallelLinear # reduce scatter is hidden inside RowParallelLinear # Self Attention hidden_states = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, ) # TODO: add sequence parallel operator reduce_scatter here hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) # TODO: add sequence parallel operator all_gather here hidden_states = self.mlp(hidden_states) # TODO: add sequence parallel operator reduce_scatter here hidden_states = residual + hidden_states outputs = hidden_states return outputs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/layers/parallel_decoder.py","span":"L35-L42","signature":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig)","exported":false,"origin":"ast","body_hash":"ee63a8015114b43fcf640f1c6e236ad939e3dee76254f727e7c3a1e495e8e756","body_text":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.hidden_size = config.hidden_size self.self_attn = ParallelLlamaAttention(config=config, megatron_config=megatron_config) self.mlp = ParallelLlamaMLP(config, megatron_config=megatron_config) self.input_layernorm = ParallelLlamaRMSNorm(config, megatron_config) self.post_attention_layernorm = ParallelLlamaRMSNorm(config, megatron_config)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaDecoderLayer"},{"id":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/layers/parallel_decoder.py","span":"L44-L96","signature":"def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]","exported":true,"origin":"ast","body_hash":"e475fd021252f51c0ea6c211c1ffc1df1da6146075cb8ae26a5e539315e2c715","body_text":"def forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: \"\"\" Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`, *optional*): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states \"\"\" residual = hidden_states hidden_states = self.input_layernorm(hidden_states) # Note: sequence parallel is hidden inside ColumnParallelLinear # reduce scatter is hidden inside RowParallelLinear # Self Attention hidden_states = self.self_attn( hidden_states=hidden_states, attention_mask=attention_mask, position_ids=position_ids, ) # TODO: add sequence parallel operator reduce_scatter here hidden_states = residual + hidden_states # Fully Connected residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) # TODO: add sequence parallel operator all_gather here hidden_states = self.mlp(hidden_states) # TODO: add sequence parallel operator reduce_scatter here hidden_states = residual + hidden_states outputs = hidden_states return outputs","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaDecoderLayer"},{"id":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad","name":"ParallelLlamaDecoderLayerRmPad","kind":"class","path":"verl/models/llama/megatron/layers/parallel_decoder.py","span":"L99-L146","signature":"class ParallelLlamaDecoderLayerRmPad(nn.Module)","exported":true,"origin":"ast","body_hash":"0cd694eb31a72952b4c30dd946ffec0c49e6897f6c3b2eb72a01c26445511e9e","body_text":"class ParallelLlamaDecoderLayerRmPad(nn.Module): def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.config = config self.megatron_config = megatron_config self.hidden_size = config.hidden_size self.self_attn = ParallelLlamaAttentionRmPad(config=config, megatron_config=megatron_config) self.mlp = ParallelLlamaMLP(config, megatron_config=megatron_config) self.input_layernorm = ParallelLlamaRMSNorm(config, megatron_config) self.post_attention_layernorm = ParallelLlamaRMSNorm(config, megatron_config) def forward( self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states # (total_nnz // sp, 1, hidden_size) hidden_states = self.input_layernorm(hidden_states) # Self Attention # (total_nnz // sp, 1, hidden_size) -> all-gather (total_nnz, 1, hidden_size) # -> col + row -> reduce-scatter -> (total_nnz // sp, 1, hidden_size) hidden_states = self.self_attn(hidden_states=hidden_states, position_ids=position_ids, sequence_length=sequence_length, indices=indices, cu_seqlens=cu_seqlens, max_seqlen_in_batch=max_seqlen_in_batch) hidden_states = residual + hidden_states # Fully Connected # shape changes same as attn residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = hidden_states return outputs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/layers/parallel_decoder.py","span":"L101-L110","signature":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig)","exported":false,"origin":"ast","body_hash":"928b34b1a2bfd6d9a715e41e74111973ef5c39af3df6920a37bbce07f85aea51","body_text":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.config = config self.megatron_config = megatron_config self.hidden_size = config.hidden_size self.self_attn = ParallelLlamaAttentionRmPad(config=config, megatron_config=megatron_config) self.mlp = ParallelLlamaMLP(config, megatron_config=megatron_config) self.input_layernorm = ParallelLlamaRMSNorm(config, megatron_config) self.post_attention_layernorm = ParallelLlamaRMSNorm(config, megatron_config)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaDecoderLayerRmPad"},{"id":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/layers/parallel_decoder.py","span":"L112-L146","signature":"def forward( self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]","exported":true,"origin":"ast","body_hash":"9bb467c9944b63c005f7f089cbc9e847227367364b37ae5e35bb7df382d92a90","body_text":"def forward( self, hidden_states: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]: residual = hidden_states # (total_nnz // sp, 1, hidden_size) hidden_states = self.input_layernorm(hidden_states) # Self Attention # (total_nnz // sp, 1, hidden_size) -> all-gather (total_nnz, 1, hidden_size) # -> col + row -> reduce-scatter -> (total_nnz // sp, 1, hidden_size) hidden_states = self.self_attn(hidden_states=hidden_states, position_ids=position_ids, sequence_length=sequence_length, indices=indices, cu_seqlens=cu_seqlens, max_seqlen_in_batch=max_seqlen_in_batch) hidden_states = residual + hidden_states # Fully Connected # shape changes same as attn residual = hidden_states hidden_states = self.post_attention_layernorm(hidden_states) hidden_states = self.mlp(hidden_states) hidden_states = residual + hidden_states outputs = hidden_states return outputs","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaDecoderLayerRmPad"}],"rawEdges":[{"source":"verl/models/llama/megatron/layers/parallel_decoder.py","relation":"imports","specifier":"typing","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py","relation":"imports","specifier":"transformers","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py","relation":"imports","specifier":".parallel_attention","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py","relation":"imports","specifier":".parallel_mlp","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py","relation":"imports","specifier":".parallel_rmsnorm","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.__init__","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.__init__","relation":"calls","name":"ParallelLlamaAttention","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.__init__","relation":"calls","name":"ParallelLlamaMLP","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.__init__","relation":"calls","name":"ParallelLlamaRMSNorm","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.__init__","relation":"calls","name":"ParallelLlamaRMSNorm","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.forward","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.forward","relation":"calls","name":"input_layernorm","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_decoder.py","recvType":"ParallelLlamaDecoderLayer"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.forward","relation":"calls","name":"self_attn","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_decoder.py","recvType":"ParallelLlamaDecoderLayer"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.forward","relation":"calls","name":"post_attention_layernorm","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_decoder.py","recvType":"ParallelLlamaDecoderLayer"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayer.forward","relation":"calls","name":"mlp","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_decoder.py","recvType":"ParallelLlamaDecoderLayer"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.__init__","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.__init__","relation":"calls","name":"ParallelLlamaAttentionRmPad","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.__init__","relation":"calls","name":"ParallelLlamaMLP","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.__init__","relation":"calls","name":"ParallelLlamaRMSNorm","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.__init__","relation":"calls","name":"ParallelLlamaRMSNorm","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.forward","file":"verl/models/llama/megatron/layers/parallel_decoder.py"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.forward","relation":"calls","name":"input_layernorm","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_decoder.py","recvType":"ParallelLlamaDecoderLayerRmPad"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.forward","relation":"calls","name":"self_attn","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_decoder.py","recvType":"ParallelLlamaDecoderLayerRmPad"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.forward","relation":"calls","name":"post_attention_layernorm","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_decoder.py","recvType":"ParallelLlamaDecoderLayerRmPad"},{"source":"verl/models/llama/megatron/layers/parallel_decoder.py#ParallelLlamaDecoderLayerRmPad.forward","relation":"calls","name":"mlp","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_decoder.py","recvType":"ParallelLlamaDecoderLayerRmPad"}]},"verl/models/llama/megatron/layers/parallel_linear.py":{"size":2787,"mtimeMs":1789275830308.2024,"hash":"abea6cdb51b6321ee3bf4ba9f043c32c4ec2dbfad64ebf0981f5beb30293bb11","nodes":[{"id":"verl/models/llama/megatron/layers/parallel_linear.py","name":"parallel_linear.py","kind":"file","path":"verl/models/llama/megatron/layers/parallel_linear.py","span":"L1-L75","signature":null,"exported":true,"origin":"ast","body_hash":"abea6cdb51b6321ee3bf4ba9f043c32c4ec2dbfad64ebf0981f5beb30293bb11","chars":2787,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/linear.py from typing import Optional, Tuple from megatron.core import tensor_parallel"},{"id":"verl/models/llama/megatron/layers/parallel_linear.py#QKVParallelLinear","name":"QKVParallelLinear","kind":"class","path":"verl/models/llama/megatron/layers/parallel_linear.py","span":"L21-L49","signature":"class QKVParallelLinear(tensor_parallel.ColumnParallelLinear)","exported":true,"origin":"ast","body_hash":"aa1cd58c35584fa81ce3cc52e4143f5c1d37d5912274afbf03207e62651f4129","body_text":"class QKVParallelLinear(tensor_parallel.ColumnParallelLinear): def __init__(self, input_size, num_heads, num_key_value_heads, head_dim, *, bias=True, gather_output=True, skip_bias_add=False, **kwargs): # Keep input parameters, and already restrict the head numbers self.input_size = input_size self.q_output_size = num_heads * head_dim self.kv_output_size = num_key_value_heads * head_dim self.head_dim = head_dim self.gather_output = gather_output self.skip_bias_add = skip_bias_add input_size = self.input_size output_size = (num_heads + 2 * num_key_value_heads) * self.head_dim super().__init__(input_size=input_size, output_size=output_size, bias=bias, gather_output=gather_output, skip_bias_add=skip_bias_add, **kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_linear.py#QKVParallelLinear.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/layers/parallel_linear.py","span":"L23-L49","signature":"def __init__(self, input_size, num_heads, num_key_value_heads, head_dim, *, bias=True, gather_output=True, skip_bias_add=False, **kwargs): # Keep input parameters, and already restrict the head numbers","exported":false,"origin":"ast","body_hash":"51a2f7bfe85cfb5a72b58a8430f2ccc68de91afeaca1826cf470d99394a5229e","body_text":"def __init__(self, input_size, num_heads, num_key_value_heads, head_dim, *, bias=True, gather_output=True, skip_bias_add=False, **kwargs): # Keep input parameters, and already restrict the head numbers self.input_size = input_size self.q_output_size = num_heads * head_dim self.kv_output_size = num_key_value_heads * head_dim self.head_dim = head_dim self.gather_output = gather_output self.skip_bias_add = skip_bias_add input_size = self.input_size output_size = (num_heads + 2 * num_key_value_heads) * self.head_dim super().__init__(input_size=input_size, output_size=output_size, bias=bias, gather_output=gather_output, skip_bias_add=skip_bias_add, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"QKVParallelLinear"},{"id":"verl/models/llama/megatron/layers/parallel_linear.py#MergedColumnParallelLinear","name":"MergedColumnParallelLinear","kind":"class","path":"verl/models/llama/megatron/layers/parallel_linear.py","span":"L52-L74","signature":"class MergedColumnParallelLinear(tensor_parallel.ColumnParallelLinear)","exported":true,"origin":"ast","body_hash":"5034e2d99713015e2bf5a34348e0065c59b517261852e2add31ecd518a27f334","body_text":"class MergedColumnParallelLinear(tensor_parallel.ColumnParallelLinear): def __init__(self, input_size, gate_ouput_size, up_output_size, *, bias=True, gather_output=True, skip_bias_add=False, **kwargs): # Keep input parameters, and already restrict the head numbers self.input_size = input_size self.output_size = gate_ouput_size + up_output_size self.gather_output = gather_output self.skip_bias_add = skip_bias_add super().__init__(input_size=self.input_size, output_size=self.output_size, bias=bias, gather_output=gather_output, skip_bias_add=skip_bias_add, **kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_linear.py#MergedColumnParallelLinear.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/layers/parallel_linear.py","span":"L54-L74","signature":"def __init__(self, input_size, gate_ouput_size, up_output_size, *, bias=True, gather_output=True, skip_bias_add=False, **kwargs): # Keep input parameters, and already restrict the head numbers","exported":false,"origin":"ast","body_hash":"84b13f11d261c49e83efc3f8a0e4b689374bd160ce6c8e76ff4a14c0ad36205f","body_text":"def __init__(self, input_size, gate_ouput_size, up_output_size, *, bias=True, gather_output=True, skip_bias_add=False, **kwargs): # Keep input parameters, and already restrict the head numbers self.input_size = input_size self.output_size = gate_ouput_size + up_output_size self.gather_output = gather_output self.skip_bias_add = skip_bias_add super().__init__(input_size=self.input_size, output_size=self.output_size, bias=bias, gather_output=gather_output, skip_bias_add=skip_bias_add, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"MergedColumnParallelLinear"}],"rawEdges":[{"source":"verl/models/llama/megatron/layers/parallel_linear.py","relation":"imports","specifier":"typing","file":"verl/models/llama/megatron/layers/parallel_linear.py"},{"source":"verl/models/llama/megatron/layers/parallel_linear.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/layers/parallel_linear.py"},{"source":"verl/models/llama/megatron/layers/parallel_linear.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_linear.py#QKVParallelLinear","file":"verl/models/llama/megatron/layers/parallel_linear.py"},{"source":"verl/models/llama/megatron/layers/parallel_linear.py#QKVParallelLinear","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_linear.py#QKVParallelLinear.__init__","file":"verl/models/llama/megatron/layers/parallel_linear.py"},{"source":"verl/models/llama/megatron/layers/parallel_linear.py#QKVParallelLinear.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_linear.py"},{"source":"verl/models/llama/megatron/layers/parallel_linear.py#QKVParallelLinear.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_linear.py"},{"source":"verl/models/llama/megatron/layers/parallel_linear.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_linear.py#MergedColumnParallelLinear","file":"verl/models/llama/megatron/layers/parallel_linear.py"},{"source":"verl/models/llama/megatron/layers/parallel_linear.py#MergedColumnParallelLinear","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_linear.py#MergedColumnParallelLinear.__init__","file":"verl/models/llama/megatron/layers/parallel_linear.py"},{"source":"verl/models/llama/megatron/layers/parallel_linear.py#MergedColumnParallelLinear.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_linear.py"},{"source":"verl/models/llama/megatron/layers/parallel_linear.py#MergedColumnParallelLinear.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_linear.py"}]},"verl/models/llama/megatron/layers/parallel_mlp.py":{"size":3394,"mtimeMs":1789275830308.2024,"hash":"752d7674c3746007caa5a111ae15f47d30fb769cf30721ede17fca4c2c5e3bab","nodes":[{"id":"verl/models/llama/megatron/layers/parallel_mlp.py","name":"parallel_mlp.py","kind":"file","path":"verl/models/llama/megatron/layers/parallel_mlp.py","span":"L1-L75","signature":null,"exported":true,"origin":"ast","body_hash":"752d7674c3746007caa5a111ae15f47d30fb769cf30721ede17fca4c2c5e3bab","chars":3394,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX # and OPT implementations in this library. It has been modified from its # original forms to accommodate minor architectural differences compared # to GPT-NeoX and OPT used by the Meta AI team that trained the model. # # 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 megatron.core import parallel_state as mpu from megatron.core import tensor_parallel from megatron.core import ModelParallelConfig from torch import nn from transformers.activations import ACT2FN from verl.models.llama.megatron.layers.parallel_linear import MergedColumnParallelLinear from verl.utils.megatron import tensor_parallel as tp_utils"},{"id":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP","name":"ParallelLlamaMLP","kind":"class","path":"verl/models/llama/megatron/layers/parallel_mlp.py","span":"L31-L74","signature":"class ParallelLlamaMLP(nn.Module)","exported":true,"origin":"ast","body_hash":"2553f44ffc23b8117b73e6462ef0f9220db8613ef2e1e74e97d2b195ab14c4a7","body_text":"class ParallelLlamaMLP(nn.Module): def __init__(self, config, megatron_config: ModelParallelConfig = None) -> None: super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size # The weight is only [hidden_size, intermediate_size // model_parallel_world_size] column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() row_kwargs = tp_utils.get_default_kwargs_for_row_parallel_linear() if megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' assert row_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(row_kwargs, megatron_config) tp_utils.update_kwargs_with_config(column_kwargs, megatron_config) tp_size = mpu.get_tensor_model_parallel_world_size() self.gate_up_proj = MergedColumnParallelLinear( input_size=self.hidden_size, gate_ouput_size=self.intermediate_size, up_output_size=self.intermediate_size, bias=False, gather_output=False, skip_bias_add=False, **column_kwargs, ) self.gate_size = self.intermediate_size // tp_size self.down_proj = tensor_parallel.RowParallelLinear(input_size=self.intermediate_size, output_size=self.hidden_size, bias=False, input_is_parallel=True, skip_bias_add=False, **row_kwargs) self.act_fn = ACT2FN[config.hidden_act] def forward(self, x): gate_up = self.gate_up_proj(x)[0] gate, up = gate_up.split(self.gate_size, dim=-1) return self.down_proj(self.act_fn(gate) * up)[0]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/layers/parallel_mlp.py","span":"L33-L69","signature":"def __init__(self, config, megatron_config: ModelParallelConfig = None) -> None","exported":false,"origin":"ast","body_hash":"d265fb0b8a7504f52d88a04bcde7a4970485f2f9e6b57f0d5f685a3c1c62ea26","body_text":"def __init__(self, config, megatron_config: ModelParallelConfig = None) -> None: super().__init__() self.config = config self.hidden_size = config.hidden_size self.intermediate_size = config.intermediate_size # The weight is only [hidden_size, intermediate_size // model_parallel_world_size] column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() row_kwargs = tp_utils.get_default_kwargs_for_row_parallel_linear() if megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' assert row_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(row_kwargs, megatron_config) tp_utils.update_kwargs_with_config(column_kwargs, megatron_config) tp_size = mpu.get_tensor_model_parallel_world_size() self.gate_up_proj = MergedColumnParallelLinear( input_size=self.hidden_size, gate_ouput_size=self.intermediate_size, up_output_size=self.intermediate_size, bias=False, gather_output=False, skip_bias_add=False, **column_kwargs, ) self.gate_size = self.intermediate_size // tp_size self.down_proj = tensor_parallel.RowParallelLinear(input_size=self.intermediate_size, output_size=self.hidden_size, bias=False, input_is_parallel=True, skip_bias_add=False, **row_kwargs) self.act_fn = ACT2FN[config.hidden_act]","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaMLP"},{"id":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/layers/parallel_mlp.py","span":"L71-L74","signature":"def forward(self, x)","exported":true,"origin":"ast","body_hash":"e5caf04ef0518c1b73102e0d84eaa4538b934bc85d00fe77dc624a1a0b926307","body_text":"def forward(self, x): gate_up = self.gate_up_proj(x)[0] gate, up = gate_up.split(self.gate_size, dim=-1) return self.down_proj(self.act_fn(gate) * up)[0]","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaMLP"}],"rawEdges":[{"source":"verl/models/llama/megatron/layers/parallel_mlp.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py","relation":"imports","specifier":"transformers.activations","file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py","relation":"imports","specifier":"verl.models.llama.megatron.layers.parallel_linear","file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py","relation":"imports","specifier":"verl.utils.megatron","file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP","file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"get_default_kwargs_for_column_parallel_linear","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"get_default_kwargs_for_row_parallel_linear","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"MergedColumnParallelLinear","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.__init__","relation":"calls","name":"RowParallelLinear","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.forward","file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.forward","relation":"calls","name":"gate_up_proj","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py","recvType":"ParallelLlamaMLP"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.forward","relation":"calls","name":"split","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.forward","relation":"calls","name":"down_proj","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py","recvType":"ParallelLlamaMLP"},{"source":"verl/models/llama/megatron/layers/parallel_mlp.py#ParallelLlamaMLP.forward","relation":"calls","name":"act_fn","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_mlp.py","recvType":"ParallelLlamaMLP"}]},"verl/models/llama/megatron/layers/parallel_rmsnorm.py":{"size":1860,"mtimeMs":1789275830308.2024,"hash":"fd9c0692745aae1afab37e9730178f19ea0c2388f08cdb67646e3a371a3c1deb","nodes":[{"id":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","name":"parallel_rmsnorm.py","kind":"file","path":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","span":"L1-L46","signature":null,"exported":true,"origin":"ast","body_hash":"fd9c0692745aae1afab37e9730178f19ea0c2388f08cdb67646e3a371a3c1deb","chars":1860,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 numbers import torch from megatron.core import ModelParallelConfig from torch import nn from transformers import LlamaConfig from apex.normalization.fused_layer_norm import fused_rms_norm_affine from verl.utils.megatron import sequence_parallel as sp_utils"},{"id":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm","name":"ParallelLlamaRMSNorm","kind":"class","path":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","span":"L25-L46","signature":"class ParallelLlamaRMSNorm(nn.Module)","exported":true,"origin":"ast","body_hash":"c00ac3a7234f829bf8a46155175fcde545fd306b6ed36a47394f7de3c1039848","body_text":"class ParallelLlamaRMSNorm(nn.Module): def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): \"\"\" LlamaRMSNorm is equivalent to T5LayerNorm \"\"\" super().__init__() if isinstance(config.hidden_size, numbers.Integral): normalized_shape = (config.hidden_size,) self.normalized_shape = torch.Size(normalized_shape) self.weight = nn.Parameter(torch.ones(self.normalized_shape)) self.variance_epsilon = config.rms_norm_eps if megatron_config.sequence_parallel: sp_utils.mark_parameter_as_sequence_parallel(self.weight) def forward(self, hidden_states): return fused_rms_norm_affine(input=hidden_states, weight=self.weight, normalized_shape=self.normalized_shape, eps=self.variance_epsilon, memory_efficient=True)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","span":"L27-L39","signature":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig)","exported":false,"origin":"ast","body_hash":"c2e1b5438f06d1fb0d2a535f63bf846ba05b14c0704989e9094bb85352538330","body_text":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): \"\"\" LlamaRMSNorm is equivalent to T5LayerNorm \"\"\" super().__init__() if isinstance(config.hidden_size, numbers.Integral): normalized_shape = (config.hidden_size,) self.normalized_shape = torch.Size(normalized_shape) self.weight = nn.Parameter(torch.ones(self.normalized_shape)) self.variance_epsilon = config.rms_norm_eps if megatron_config.sequence_parallel: sp_utils.mark_parameter_as_sequence_parallel(self.weight)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaRMSNorm"},{"id":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","span":"L41-L46","signature":"def forward(self, hidden_states)","exported":true,"origin":"ast","body_hash":"cc76f3c59fbb67cf414ab8ccf45c992bed8fa999b34942af9c50d2d742347879","body_text":"def forward(self, hidden_states): return fused_rms_norm_affine(input=hidden_states, weight=self.weight, normalized_shape=self.normalized_shape, eps=self.variance_epsilon, memory_efficient=True)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaRMSNorm"}],"rawEdges":[{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","relation":"imports","specifier":"numbers","file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","relation":"imports","specifier":"transformers","file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","relation":"imports","specifier":"apex.normalization.fused_layer_norm","file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","relation":"imports","specifier":"verl.utils.megatron","file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm","file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.__init__","file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.__init__","relation":"calls","name":"Size","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.__init__","relation":"calls","name":"Parameter","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.__init__","relation":"calls","name":"ones","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.__init__","relation":"calls","name":"mark_parameter_as_sequence_parallel","viaMember":true,"file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm","relation":"contains","targetId":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.forward","file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"},{"source":"verl/models/llama/megatron/layers/parallel_rmsnorm.py#ParallelLlamaRMSNorm.forward","relation":"calls","name":"fused_rms_norm_affine","viaMember":false,"file":"verl/models/llama/megatron/layers/parallel_rmsnorm.py"}]},"verl/models/llama/megatron/modeling_llama_megatron.py":{"size":29656,"mtimeMs":1789275830309.2026,"hash":"e0c992cb7f3f606e3109758b6ca4114fdc28903e2be716dff1259444492b6b55","nodes":[{"id":"verl/models/llama/megatron/modeling_llama_megatron.py","name":"modeling_llama_megatron.py","kind":"file","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L1-L657","signature":null,"exported":true,"origin":"ast","body_hash":"e0c992cb7f3f606e3109758b6ca4114fdc28903e2be716dff1259444492b6b55","chars":29656,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. # # This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX # and OPT implementations in this library. It has been modified from its # original forms to accommodate minor architectural differences compared # to GPT-NeoX and OPT used by the Meta AI team that trained the model. # # 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. \"\"\"PyTorch LLaMA model with Megatron-style acceleration.\"\"\" from typing import Optional, Tuple, Union import torch import torch.utils.checkpoint from megatron.core import tensor_parallel from megatron.core import ModelParallelConfig from torch import nn from transformers.modeling_outputs import BaseModelOutputWithPast from transformers.models.llama.configuration_llama import LlamaConfig from transformers.models.llama.modeling_llama import CausalLMOutputWithPast from verl.utils.megatron import sequence_parallel as sp_utils from verl.utils.megatron import tensor_parallel as tp_utils from .layers import ParallelLlamaDecoderLayer, ParallelLlamaRMSNorm, ParallelLlamaDecoderLayerRmPad \"\"\" TODO: 1. Add weight initialization. Here we need to be careful on TP weight init. 2. Add sequence parallel 3. Load checkpoint from meta LLama pretrained checkpoint \"\"\" # Copied from transformers.models.bart.modeling_bart._make_causal_mask # Copied from transformers.models.bart.modeling_bart._expand_mask from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa \"\"\" Support pipeline parallelism \"\"\""},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","name":"_make_causal_mask","kind":"function","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L45-L54","signature":"def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device)","exported":false,"origin":"ast","body_hash":"4836799258639bfbfdc83054165da581e95fcb7e6c8522bf4651e42af051f1f7","body_text":"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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#_expand_mask","name":"_expand_mask","kind":"function","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L58-L69","signature":"def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None)","exported":false,"origin":"ast","body_hash":"7eac0af7f47e8bd3c1587db698c5f5583656b24149ec5c2ea1a410e34ae82160","body_text":"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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel","name":"ParallelLlamaModel","kind":"class","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L72-L152","signature":"class ParallelLlamaModel(nn.Module)","exported":true,"origin":"ast","body_hash":"80aca6b60879cb9b90349e21e18267840777bc6e47a5cc7061fa5ed1feee07f4","body_text":"class ParallelLlamaModel(nn.Module): \"\"\" Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] Args: config: LlamaConfig \"\"\" def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding() if megatron_config is not None: assert embedding_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config) self.embed_tokens = tensor_parallel.VocabParallelEmbedding(num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs) self.layers = nn.ModuleList( [ParallelLlamaDecoderLayer(config, megatron_config) for _ in range(config.num_hidden_layers)]) self.norm = ParallelLlamaRMSNorm(config, megatron_config) # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask def _prepare_decoder_attention_mask(self, 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 def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: \"\"\" Args: input_ids: input ids. shape (batch_size, seq_length) attention_mask: attention_mask. shape (batch_size, seq_length) position_ids: position ids. shape (batch_size, seq_length) Returns: \"\"\" batch_size, seq_length = input_ids.shape inputs_embeds = self.embed_tokens(input_ids) # embed positions attention_mask = self._prepare_decoder_attention_mask(attention_mask, (batch_size, seq_length), inputs_embeds) hidden_states = inputs_embeds for idx, decoder_layer in enumerate(self.layers): layer_outputs = decoder_layer( hidden_states, attention_mask=attention_mask, position_ids=position_ids, ) hidden_states = layer_outputs hidden_states = self.norm(hidden_states) return hidden_states","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L80-L94","signature":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig)","exported":false,"origin":"ast","body_hash":"69b5e79ff2bf13e20f6f3f51a440c8462b1ac71ea9ca6a58e4646bc898694d1c","body_text":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding() if megatron_config is not None: assert embedding_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config) self.embed_tokens = tensor_parallel.VocabParallelEmbedding(num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs) self.layers = nn.ModuleList( [ParallelLlamaDecoderLayer(config, megatron_config) for _ in range(config.num_hidden_layers)]) self.norm = ParallelLlamaRMSNorm(config, megatron_config)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaModel"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel._prepare_decoder_attention_mask","name":"_prepare_decoder_attention_mask","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L97-L115","signature":"def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds): # create causal mask # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]","exported":false,"origin":"ast","body_hash":"72baf5f0d6878417af422634e5235f04ff58e2623cb03526d58ba11882033a85","body_text":"def _prepare_decoder_attention_mask(self, 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","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaModel"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L117-L152","signature":"def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]","exported":true,"origin":"ast","body_hash":"41ee69fd163e06bd0cf82eed87986616d5d8ac34f3a8f7ce975082b66d330b9f","body_text":"def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, BaseModelOutputWithPast]: \"\"\" Args: input_ids: input ids. shape (batch_size, seq_length) attention_mask: attention_mask. shape (batch_size, seq_length) position_ids: position ids. shape (batch_size, seq_length) Returns: \"\"\" batch_size, seq_length = input_ids.shape inputs_embeds = self.embed_tokens(input_ids) # embed positions attention_mask = self._prepare_decoder_attention_mask(attention_mask, (batch_size, seq_length), inputs_embeds) hidden_states = inputs_embeds for idx, decoder_layer in enumerate(self.layers): layer_outputs = decoder_layer( hidden_states, attention_mask=attention_mask, position_ids=position_ids, ) hidden_states = layer_outputs hidden_states = self.norm(hidden_states) return hidden_states","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaModel"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM","name":"ParallelLlamaForCausalLM","kind":"class","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L155-L209","signature":"class ParallelLlamaForCausalLM(nn.Module)","exported":true,"origin":"ast","body_hash":"524478e100cf0c0b353d0fe48137b95ef52e684eedfe9679a9d4839d60a8a98f","body_text":"class ParallelLlamaForCausalLM(nn.Module): def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.model = ParallelLlamaModel(config, megatron_config=megatron_config) self.vocab_size = config.vocab_size column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() if megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config) self.lm_head = tensor_parallel.ColumnParallelLinear(input_size=config.hidden_size, output_size=config.vocab_size, bias=False, gather_output=False, skip_bias_add=False, **column_kwargs) def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: r\"\"\" Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Returns: ```\"\"\" # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, ) hidden_states = outputs logits = self.lm_head(hidden_states)[0] logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits) logits = logits.float() return CausalLMOutputWithPast( loss=None, logits=logits, past_key_values=None, hidden_states=None, attentions=None, )","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L157-L172","signature":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig)","exported":false,"origin":"ast","body_hash":"b4c5e761217abf5a12a41d494ce0ca6dac5f0bd8390498e8b8ad6ff78fae82d7","body_text":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.model = ParallelLlamaModel(config, megatron_config=megatron_config) self.vocab_size = config.vocab_size column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() if megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config) self.lm_head = tensor_parallel.ColumnParallelLinear(input_size=config.hidden_size, output_size=config.vocab_size, bias=False, gather_output=False, skip_bias_add=False, **column_kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLM"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L174-L209","signature":"def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]","exported":true,"origin":"ast","body_hash":"974ab318d80cfdc6fefe9ab78806e091db2f21ef841b5da7fcd1736761f3ddd1","body_text":"def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: r\"\"\" Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Returns: ```\"\"\" # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, ) hidden_states = outputs logits = self.lm_head(hidden_states)[0] logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits) logits = logits.float() return CausalLMOutputWithPast( loss=None, logits=logits, past_key_values=None, hidden_states=None, attentions=None, )","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLM"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad","name":"ParallelLlamaModelRmPad","kind":"class","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L215-L276","signature":"class ParallelLlamaModelRmPad(nn.Module)","exported":true,"origin":"ast","body_hash":"4d5cf6c9440402d75e3fb103877842599c93f240250274106902eac1e46706a0","body_text":"class ParallelLlamaModelRmPad(nn.Module): \"\"\" Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] Args: config: LlamaConfig \"\"\" def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding() self.megatron_config = megatron_config if megatron_config is not None: assert embedding_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config) self.embed_tokens = tensor_parallel.VocabParallelEmbedding(num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs) self.layers = nn.ModuleList( [ParallelLlamaDecoderLayerRmPad(config, megatron_config) for _ in range(config.num_hidden_layers)]) self.norm = ParallelLlamaRMSNorm(config, megatron_config) def forward(self, input_ids: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None) -> Union[Tuple, BaseModelOutputWithPast]: \"\"\" Args: input_ids: input ids. shape (1, totol_nnz) position_ids: position ids. shape (batch_size, seq_length) Returns: \"\"\" inputs_embeds = self.embed_tokens(input_ids) # (1, total_nnz) -> (1, total_nnz, hidden_size) # (1, total_nnz, hidden_size) -> (total_nnz, 1, hidden_size) -> (total_nnz // sp, 1, hidden_size) inputs_embeds = inputs_embeds.transpose(0, 1) if self.megatron_config.sequence_parallel: inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region(inputs_embeds) hidden_states = inputs_embeds for idx, decoder_layer in enumerate(self.layers): layer_outputs = decoder_layer(hidden_states, position_ids=position_ids, sequence_length=sequence_length, indices=indices, cu_seqlens=cu_seqlens, max_seqlen_in_batch=max_seqlen_in_batch) hidden_states = layer_outputs hidden_states = self.norm(hidden_states) return hidden_states","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L223-L238","signature":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig)","exported":false,"origin":"ast","body_hash":"ac9130b73bc042d00818540e1476f5333a50360e2d097841864b6e5ecc4a4125","body_text":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding() self.megatron_config = megatron_config if megatron_config is not None: assert embedding_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config) self.embed_tokens = tensor_parallel.VocabParallelEmbedding(num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs) self.layers = nn.ModuleList( [ParallelLlamaDecoderLayerRmPad(config, megatron_config) for _ in range(config.num_hidden_layers)]) self.norm = ParallelLlamaRMSNorm(config, megatron_config)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaModelRmPad"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L240-L276","signature":"def forward(self, input_ids: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None) -> Union[Tuple, BaseModelOutputWithPast]","exported":true,"origin":"ast","body_hash":"1e1407c50b3d976c1848c2866b84eda74366b70adbc1f2950d4a9644575f310c","body_text":"def forward(self, input_ids: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None) -> Union[Tuple, BaseModelOutputWithPast]: \"\"\" Args: input_ids: input ids. shape (1, totol_nnz) position_ids: position ids. shape (batch_size, seq_length) Returns: \"\"\" inputs_embeds = self.embed_tokens(input_ids) # (1, total_nnz) -> (1, total_nnz, hidden_size) # (1, total_nnz, hidden_size) -> (total_nnz, 1, hidden_size) -> (total_nnz // sp, 1, hidden_size) inputs_embeds = inputs_embeds.transpose(0, 1) if self.megatron_config.sequence_parallel: inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region(inputs_embeds) hidden_states = inputs_embeds for idx, decoder_layer in enumerate(self.layers): layer_outputs = decoder_layer(hidden_states, position_ids=position_ids, sequence_length=sequence_length, indices=indices, cu_seqlens=cu_seqlens, max_seqlen_in_batch=max_seqlen_in_batch) hidden_states = layer_outputs hidden_states = self.norm(hidden_states) return hidden_states","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaModelRmPad"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad","name":"ParallelLlamaForCausalLMRmPad","kind":"class","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L279-L363","signature":"class ParallelLlamaForCausalLMRmPad(nn.Module)","exported":true,"origin":"ast","body_hash":"e50847d21c74182a27954b3c4e48997ef87563c1db17b93c37653b93a53a63b2","body_text":"class ParallelLlamaForCausalLMRmPad(nn.Module): def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.config = config self.megatron_config = megatron_config self.model = ParallelLlamaModelRmPad(config, megatron_config=megatron_config) self.vocab_size = config.vocab_size self._init_head() def _init_head(self): column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() if self.megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config) self.lm_head = tensor_parallel.ColumnParallelLinear(input_size=self.config.hidden_size, output_size=self.config.vocab_size, bias=False, gather_output=False, skip_bias_add=False, **column_kwargs) def _forward_head(self, hidden_states): # all_gather from sequence parallel region is performed inside lm_head logits = self.lm_head(hidden_states)[0] logits = logits.float() # (total_nnz_padded, 1, vocab_size // tp) logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits) # (total_nnz_padded, 1, vocab_size) return logits def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: r\"\"\" Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Returns: ```\"\"\" batch_size, sequence_length = input_ids.shape # remove padding here input_ids, indices, cu_seqlens, max_seqlen_in_batch, *_ = unpad_input(input_ids.unsqueeze(dim=-1), attention_mask) # (total_nnz, 1) # pad input_ids to multiple of tp for all tp ranks # TODO: for better performance, the sp padding should be removed at each layer. Not sure the performance gap if self.megatron_config.sequence_parallel: input_ids = sp_utils.pad_to_sequence_parallel(input_ids) input_ids = input_ids.transpose(0, 1) # (1, total_nnz+pad) outputs = self.model(input_ids=input_ids, position_ids=position_ids, sequence_length=sequence_length, indices=indices, cu_seqlens=cu_seqlens, max_seqlen_in_batch=max_seqlen_in_batch) hidden_states = outputs logits = self._forward_head(hidden_states) # remove padding from sequence parallel if self.megatron_config.sequence_parallel: totol_nnz = cu_seqlens[-1] logits = logits[:totol_nnz] # (total_nnz_padded) logits = torch.squeeze(logits, dim=1) # remove the artificial batch dimension # add removed padding back logits = pad_input(logits, indices, batch_size, seqlen=sequence_length) # (batch_size, sequence_length, vocab_size) return CausalLMOutputWithPast( loss=None, logits=logits, past_key_values=None, hidden_states=None, attentions=None, )","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L281-L287","signature":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig)","exported":false,"origin":"ast","body_hash":"4fc8ad94efc85aee77b99f002a3fa237c5f17b5e51983d6fdbe99c514bea792c","body_text":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): super().__init__() self.config = config self.megatron_config = megatron_config self.model = ParallelLlamaModelRmPad(config, megatron_config=megatron_config) self.vocab_size = config.vocab_size self._init_head()","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLMRmPad"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._init_head","name":"_init_head","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L289-L299","signature":"def _init_head(self)","exported":false,"origin":"ast","body_hash":"e26b2ad854017b37b2566b868dc36344aeee40dd59125285722952c20ca6b1cb","body_text":"def _init_head(self): column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() if self.megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config) self.lm_head = tensor_parallel.ColumnParallelLinear(input_size=self.config.hidden_size, output_size=self.config.vocab_size, bias=False, gather_output=False, skip_bias_add=False, **column_kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLMRmPad"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._forward_head","name":"_forward_head","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L301-L306","signature":"def _forward_head(self, hidden_states): # all_gather from sequence parallel region is performed inside lm_head","exported":false,"origin":"ast","body_hash":"86f3116ccac4359e1aed7d51af44b25bd385e5820632ab0c08bc054a5f941d75","body_text":"def _forward_head(self, hidden_states): # all_gather from sequence parallel region is performed inside lm_head logits = self.lm_head(hidden_states)[0] logits = logits.float() # (total_nnz_padded, 1, vocab_size // tp) logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits) # (total_nnz_padded, 1, vocab_size) return logits","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLMRmPad"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L308-L363","signature":"def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]","exported":true,"origin":"ast","body_hash":"f4ee14f475ea62c69ab3566defb5a879d84d3c2daca4833a7e523422e1afef5f","body_text":"def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: r\"\"\" Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Returns: ```\"\"\" batch_size, sequence_length = input_ids.shape # remove padding here input_ids, indices, cu_seqlens, max_seqlen_in_batch, *_ = unpad_input(input_ids.unsqueeze(dim=-1), attention_mask) # (total_nnz, 1) # pad input_ids to multiple of tp for all tp ranks # TODO: for better performance, the sp padding should be removed at each layer. Not sure the performance gap if self.megatron_config.sequence_parallel: input_ids = sp_utils.pad_to_sequence_parallel(input_ids) input_ids = input_ids.transpose(0, 1) # (1, total_nnz+pad) outputs = self.model(input_ids=input_ids, position_ids=position_ids, sequence_length=sequence_length, indices=indices, cu_seqlens=cu_seqlens, max_seqlen_in_batch=max_seqlen_in_batch) hidden_states = outputs logits = self._forward_head(hidden_states) # remove padding from sequence parallel if self.megatron_config.sequence_parallel: totol_nnz = cu_seqlens[-1] logits = logits[:totol_nnz] # (total_nnz_padded) logits = torch.squeeze(logits, dim=1) # remove the artificial batch dimension # add removed padding back logits = pad_input(logits, indices, batch_size, seqlen=sequence_length) # (batch_size, sequence_length, vocab_size) return CausalLMOutputWithPast( loss=None, logits=logits, past_key_values=None, hidden_states=None, attentions=None, )","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLMRmPad"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad","name":"ParallelLlamaForValueRmPad","kind":"class","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L366-L392","signature":"class ParallelLlamaForValueRmPad(ParallelLlamaForCausalLMRmPad)","exported":true,"origin":"ast","body_hash":"18a0674aabfe898215a855adedc54f1c9f3382daf7f1ebb3a4c3b8a187d0a472","body_text":"class ParallelLlamaForValueRmPad(ParallelLlamaForCausalLMRmPad): def _init_head(self): column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() if self.megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config) self.lm_head = nn.Linear(in_features=self.config.hidden_size, out_features=1, bias=False) # lm_head is effectively the same as sequence parallel sp_utils.mark_parameter_as_sequence_parallel(self.lm_head.weight) def _forward_head(self, hidden_states): logits = self.lm_head(hidden_states) # (total_nnz_padded // tp, 1, 1) logits = logits.float() if self.megatron_config.sequence_parallel: logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False) return logits def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: output = super().forward(input_ids, attention_mask, position_ids) output.logits = torch.squeeze(output.logits, dim=-1) return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._init_head","name":"_init_head","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L368-L375","signature":"def _init_head(self)","exported":false,"origin":"ast","body_hash":"51af3f483828793a452f7f59a190039e0c2b511da22b0d80034de041fdeaa753","body_text":"def _init_head(self): column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() if self.megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config) self.lm_head = nn.Linear(in_features=self.config.hidden_size, out_features=1, bias=False) # lm_head is effectively the same as sequence parallel sp_utils.mark_parameter_as_sequence_parallel(self.lm_head.weight)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForValueRmPad"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._forward_head","name":"_forward_head","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L377-L382","signature":"def _forward_head(self, hidden_states)","exported":false,"origin":"ast","body_hash":"2c01174dfc8d7cce4ff60e2ffe09223ba1139598dad6b1f3127d2bc174137186","body_text":"def _forward_head(self, hidden_states): logits = self.lm_head(hidden_states) # (total_nnz_padded // tp, 1, 1) logits = logits.float() if self.megatron_config.sequence_parallel: logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False) return logits","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForValueRmPad"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L384-L392","signature":"def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]","exported":true,"origin":"ast","body_hash":"8890c733ee268b257ff186dd7590d6ff8f5fcde9fbe8afa0a108e97155d305ec","body_text":"def forward( self, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: output = super().forward(input_ids, attention_mask, position_ids) output.logits = torch.squeeze(output.logits, dim=-1) return output","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForValueRmPad"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP","name":"ParallelLlamaModelRmPadPP","kind":"class","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L400-L511","signature":"class ParallelLlamaModelRmPadPP(nn.Module)","exported":true,"origin":"ast","body_hash":"0e171b2885cc440334e021dcb6b84d4354191fa0dbc5147b5d9dd1a4037979ad","body_text":"class ParallelLlamaModelRmPadPP(nn.Module): \"\"\" Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`] This model definition supports pipeline parallelism. To support pp and vpp, - This model only contains layer in this pp stage and vpp chunk - When calling get_model in Megatron, this rank will instantiate all the vpp chunks in this pp. Args: config: LlamaConfig \"\"\" def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pre_process, post_process): super().__init__() self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.pre_process = pre_process self.post_process = post_process self.megatron_config = megatron_config embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding() if megatron_config is not None: assert embedding_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config) if pre_process: self.embed_tokens = tensor_parallel.VocabParallelEmbedding(num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs) else: self.embed_tokens = None # pp_rank = megatron_config.pipeline_model_parallel_rank pp_size = megatron_config.pipeline_model_parallel_size self.num_layer_per_pp = config.num_hidden_layers // pp_size vpp_size = megatron_config.virtual_pipeline_model_parallel_size if vpp_size is not None: self.num_layer_vpp_chunk = self.num_layer_per_pp // vpp_size self.num_layer_this_model = self.num_layer_vpp_chunk # vpp_rank = megatron_config.virtual_pipeline_model_parallel_rank # self.offset = vpp_rank * ( # config.num_hidden_layers // megatron_config.virtual_pipeline_model_parallel_size) + \\ # (megatron_config.pipeline_model_parallel_rank * self.num_layer_vpp_chunk) else: self.num_layer_this_model = self.num_layer_per_pp # self.offset = pp_rank * self.num_layer_per_pp layers = [] for i in range(self.num_layer_this_model): layer = ParallelLlamaDecoderLayerRmPad(config, megatron_config) # setattr(layer, 'hidden_layer_index', self.offset + i) layers.append(layer) self.layers = nn.ModuleList(layers) if post_process: self.norm = ParallelLlamaRMSNorm(config, megatron_config) else: self.norm = None def set_input_tensor(self, input_tensor): \"\"\"Set input tensor to be used instead of forward()'s input. When doing pipeline parallelism the input from the previous stage comes from communication, not from the input, so the model's forward_step_func won't have it. This function is thus used by internal code to bypass the input provided by the forward_step_func\"\"\" self.input_tensor = input_tensor def forward(self, input_ids: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None) -> Union[Tuple, BaseModelOutputWithPast]: \"\"\" Args: input_ids: input ids. shape (1, totol_nnz) position_ids: position ids. shape (batch_size, seq_length) Returns: \"\"\" if self.pre_process: inputs_embeds = self.embed_tokens(input_ids) # (1, total_nnz) -> (1, total_nnz, hidden_size) # vocab parallel embedding will not do sequence parallel reduce-scatter in open source megatron # so need to deal with it by handle here: # (1, total_nnz, hidden_size) -> (total_nnz, 1, hidden_size) -> (total_nnz // sp, 1, hidden_size) inputs_embeds = inputs_embeds.transpose(0, 1) if self.megatron_config.sequence_parallel: inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region(inputs_embeds) hidden_states = inputs_embeds else: # self.hidden_states should be passed by Megatron hidden_states = self.input_tensor for idx, decoder_layer in enumerate(self.layers): layer_outputs = decoder_layer(hidden_states, position_ids=position_ids, sequence_length=sequence_length, indices=indices, cu_seqlens=cu_seqlens, max_seqlen_in_batch=max_seqlen_in_batch) hidden_states = layer_outputs if self.post_process: hidden_states = self.norm(hidden_states) return hidden_states","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L410-L455","signature":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pre_process, post_process)","exported":false,"origin":"ast","body_hash":"2045595e1577ef7a23f739a0c18bd925c64250d1df18e252742fc4a04901db97","body_text":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pre_process, post_process): super().__init__() self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.pre_process = pre_process self.post_process = post_process self.megatron_config = megatron_config embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding() if megatron_config is not None: assert embedding_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config) if pre_process: self.embed_tokens = tensor_parallel.VocabParallelEmbedding(num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs) else: self.embed_tokens = None # pp_rank = megatron_config.pipeline_model_parallel_rank pp_size = megatron_config.pipeline_model_parallel_size self.num_layer_per_pp = config.num_hidden_layers // pp_size vpp_size = megatron_config.virtual_pipeline_model_parallel_size if vpp_size is not None: self.num_layer_vpp_chunk = self.num_layer_per_pp // vpp_size self.num_layer_this_model = self.num_layer_vpp_chunk # vpp_rank = megatron_config.virtual_pipeline_model_parallel_rank # self.offset = vpp_rank * ( # config.num_hidden_layers // megatron_config.virtual_pipeline_model_parallel_size) + \\ # (megatron_config.pipeline_model_parallel_rank * self.num_layer_vpp_chunk) else: self.num_layer_this_model = self.num_layer_per_pp # self.offset = pp_rank * self.num_layer_per_pp layers = [] for i in range(self.num_layer_this_model): layer = ParallelLlamaDecoderLayerRmPad(config, megatron_config) # setattr(layer, 'hidden_layer_index', self.offset + i) layers.append(layer) self.layers = nn.ModuleList(layers) if post_process: self.norm = ParallelLlamaRMSNorm(config, megatron_config) else: self.norm = None","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaModelRmPadPP"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.set_input_tensor","name":"set_input_tensor","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L457-L465","signature":"def set_input_tensor(self, input_tensor)","exported":true,"origin":"ast","body_hash":"d8093744c1b1348ec91bd85ef1b336a7b4bffbaae443f57104a00fdecc092b82","body_text":"def set_input_tensor(self, input_tensor): \"\"\"Set input tensor to be used instead of forward()'s input. When doing pipeline parallelism the input from the previous stage comes from communication, not from the input, so the model's forward_step_func won't have it. This function is thus used by internal code to bypass the input provided by the forward_step_func\"\"\" self.input_tensor = input_tensor","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaModelRmPadPP"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L467-L511","signature":"def forward(self, input_ids: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None) -> Union[Tuple, BaseModelOutputWithPast]","exported":true,"origin":"ast","body_hash":"97877547a17d224bca6506ad138bf8681e4906fcbc62f05c3fe76a40e27b9e88","body_text":"def forward(self, input_ids: torch.Tensor, position_ids: Optional[torch.LongTensor] = None, sequence_length: int = None, indices: torch.Tensor = None, cu_seqlens: int = None, max_seqlen_in_batch: int = None) -> Union[Tuple, BaseModelOutputWithPast]: \"\"\" Args: input_ids: input ids. shape (1, totol_nnz) position_ids: position ids. shape (batch_size, seq_length) Returns: \"\"\" if self.pre_process: inputs_embeds = self.embed_tokens(input_ids) # (1, total_nnz) -> (1, total_nnz, hidden_size) # vocab parallel embedding will not do sequence parallel reduce-scatter in open source megatron # so need to deal with it by handle here: # (1, total_nnz, hidden_size) -> (total_nnz, 1, hidden_size) -> (total_nnz // sp, 1, hidden_size) inputs_embeds = inputs_embeds.transpose(0, 1) if self.megatron_config.sequence_parallel: inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region(inputs_embeds) hidden_states = inputs_embeds else: # self.hidden_states should be passed by Megatron hidden_states = self.input_tensor for idx, decoder_layer in enumerate(self.layers): layer_outputs = decoder_layer(hidden_states, position_ids=position_ids, sequence_length=sequence_length, indices=indices, cu_seqlens=cu_seqlens, max_seqlen_in_batch=max_seqlen_in_batch) hidden_states = layer_outputs if self.post_process: hidden_states = self.norm(hidden_states) return hidden_states","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaModelRmPadPP"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP","name":"ParallelLlamaForCausalLMRmPadPP","kind":"class","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L514-L623","signature":"class ParallelLlamaForCausalLMRmPadPP(nn.Module)","exported":true,"origin":"ast","body_hash":"22f073160388a91588816311ab368575df438b23a98e908de9f98326dc54e47c","body_text":"class ParallelLlamaForCausalLMRmPadPP(nn.Module): def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pre_process, post_process): super().__init__() self.config = config self.megatron_config = megatron_config self.model = ParallelLlamaModelRmPadPP(config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process) self.share_embeddings_and_output_weights = None # workaround, megatron requires this attr self.vocab_size = config.vocab_size self.pre_process = pre_process self.post_process = post_process if post_process: self._init_head() def set_input_tensor(self, input_tensor): \"\"\"Set input tensor to be used instead of forward()'s input. When doing pipeline parallelism the input from the previous stage comes from communication, not from the input, so the model's forward_step_func won't have it. This function is thus used by internal code to bypass the input provided by the forward_step_func\"\"\" assert len(input_tensor) == 1 self.model.set_input_tensor(input_tensor[0]) def _init_head(self): column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() if self.megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config) self.lm_head = tensor_parallel.ColumnParallelLinear(input_size=self.config.hidden_size, output_size=self.config.vocab_size, bias=False, gather_output=False, skip_bias_add=False, **column_kwargs) def _forward_head(self, hidden_states): # all_gather from sequence parallel region is performed inside lm_head # logits shape before forward_head hidden_states.shape: [4, 32, 4096] logits = self.lm_head(hidden_states)[0] # logits shape after forward_head logits.shape: [8, 32, 8] logits = logits.float() # (total_nnz_padded, 1, vocab_size // tp) return logits def forward( self, # original input *, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: r\"\"\" Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Returns: ```\"\"\" # Note that input_ids, attention_mask and position_ids should be passed to every pp layer. # In the first pp, input_ids will be used, in other pp layers hidden_states will be used inside self.model batch_size, sequence_length = input_ids.shape # remove padding here input_ids_rmpad, indices, cu_seqlens, max_seqlen_in_batch, *_ = unpad_input(input_ids.unsqueeze(dim=-1), attention_mask) # (total_nnz, 1) # pad input_ids to multiple of tp for all tp ranks # TODO: for better performance, the sp padding should be removed at each layer. Not sure the performance gap if self.megatron_config.sequence_parallel: input_ids_rmpad = sp_utils.pad_to_sequence_parallel(input_ids_rmpad) input_ids_rmpad = input_ids_rmpad.transpose(0, 1) # (1, total_nnz+pad) outputs = self.model(input_ids=input_ids_rmpad, position_ids=position_ids, sequence_length=sequence_length, indices=indices, cu_seqlens=cu_seqlens, max_seqlen_in_batch=max_seqlen_in_batch) if self.post_process: hidden_states = outputs # print(f'hidden_states.shape = {hidden_states.shape}') # torch.Size([4, 32, 4096]) logits = self._forward_head(hidden_states) logits = torch.squeeze(logits, dim=1) # remove the artificial batch dimension # torch.Size([8, 32, 16]) # remove padding from sequence parallel if self.megatron_config.sequence_parallel: totol_nnz = cu_seqlens[-1] logits = logits[:totol_nnz] # (total_nnz_padded) # add removed padding back. If input is already rmpad, we let the caller pad_input logits = pad_input(logits, indices, batch_size, seqlen=sequence_length) # (batch_size, sequence_length, vocab_size) return CausalLMOutputWithPast( loss=None, logits=logits, past_key_values=None, hidden_states=None, attentions=None, ) else: return outputs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.__init__","name":"__init__","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L516-L529","signature":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pre_process, post_process)","exported":false,"origin":"ast","body_hash":"1fa43fd60df5e786a33c1278f636eec71cf4f3a7c52b8cf116256e7ac63f025d","body_text":"def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pre_process, post_process): super().__init__() self.config = config self.megatron_config = megatron_config self.model = ParallelLlamaModelRmPadPP(config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process) self.share_embeddings_and_output_weights = None # workaround, megatron requires this attr self.vocab_size = config.vocab_size self.pre_process = pre_process self.post_process = post_process if post_process: self._init_head()","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLMRmPadPP"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.set_input_tensor","name":"set_input_tensor","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L531-L540","signature":"def set_input_tensor(self, input_tensor)","exported":true,"origin":"ast","body_hash":"4d42f55d7ad1b74eb5beb556e8bd9962d259a62d91ed98cf3f2c49faded97cdf","body_text":"def set_input_tensor(self, input_tensor): \"\"\"Set input tensor to be used instead of forward()'s input. When doing pipeline parallelism the input from the previous stage comes from communication, not from the input, so the model's forward_step_func won't have it. This function is thus used by internal code to bypass the input provided by the forward_step_func\"\"\" assert len(input_tensor) == 1 self.model.set_input_tensor(input_tensor[0])","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLMRmPadPP"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP._init_head","name":"_init_head","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L542-L552","signature":"def _init_head(self)","exported":false,"origin":"ast","body_hash":"e26b2ad854017b37b2566b868dc36344aeee40dd59125285722952c20ca6b1cb","body_text":"def _init_head(self): column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() if self.megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config) self.lm_head = tensor_parallel.ColumnParallelLinear(input_size=self.config.hidden_size, output_size=self.config.vocab_size, bias=False, gather_output=False, skip_bias_add=False, **column_kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLMRmPadPP"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP._forward_head","name":"_forward_head","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L554-L560","signature":"def _forward_head(self, hidden_states): # all_gather from sequence parallel region is performed inside lm_head # logits shape before forward_head hidden_states.shape: [4, 32, 4096]","exported":false,"origin":"ast","body_hash":"98659534a93f89cdbd042ae177cf95cbdbde79c7b372e6972a6d8f6fc437bd91","body_text":"def _forward_head(self, hidden_states): # all_gather from sequence parallel region is performed inside lm_head # logits shape before forward_head hidden_states.shape: [4, 32, 4096] logits = self.lm_head(hidden_states)[0] # logits shape after forward_head logits.shape: [8, 32, 8] logits = logits.float() # (total_nnz_padded, 1, vocab_size // tp) return logits","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLMRmPadPP"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L562-L623","signature":"def forward( self, # original input *, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]","exported":true,"origin":"ast","body_hash":"3f60af6ea1f373020cc8c0ccae6643f03a105ae946d3f20496d590e9f024fca6","body_text":"def forward( self, # original input *, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: r\"\"\" Args: labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Returns: ```\"\"\" # Note that input_ids, attention_mask and position_ids should be passed to every pp layer. # In the first pp, input_ids will be used, in other pp layers hidden_states will be used inside self.model batch_size, sequence_length = input_ids.shape # remove padding here input_ids_rmpad, indices, cu_seqlens, max_seqlen_in_batch, *_ = unpad_input(input_ids.unsqueeze(dim=-1), attention_mask) # (total_nnz, 1) # pad input_ids to multiple of tp for all tp ranks # TODO: for better performance, the sp padding should be removed at each layer. Not sure the performance gap if self.megatron_config.sequence_parallel: input_ids_rmpad = sp_utils.pad_to_sequence_parallel(input_ids_rmpad) input_ids_rmpad = input_ids_rmpad.transpose(0, 1) # (1, total_nnz+pad) outputs = self.model(input_ids=input_ids_rmpad, position_ids=position_ids, sequence_length=sequence_length, indices=indices, cu_seqlens=cu_seqlens, max_seqlen_in_batch=max_seqlen_in_batch) if self.post_process: hidden_states = outputs # print(f'hidden_states.shape = {hidden_states.shape}') # torch.Size([4, 32, 4096]) logits = self._forward_head(hidden_states) logits = torch.squeeze(logits, dim=1) # remove the artificial batch dimension # torch.Size([8, 32, 16]) # remove padding from sequence parallel if self.megatron_config.sequence_parallel: totol_nnz = cu_seqlens[-1] logits = logits[:totol_nnz] # (total_nnz_padded) # add removed padding back. If input is already rmpad, we let the caller pad_input logits = pad_input(logits, indices, batch_size, seqlen=sequence_length) # (batch_size, sequence_length, vocab_size) return CausalLMOutputWithPast( loss=None, logits=logits, past_key_values=None, hidden_states=None, attentions=None, ) else: return outputs","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForCausalLMRmPadPP"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP","name":"ParallelLlamaForValueRmPadPP","kind":"class","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L626-L656","signature":"class ParallelLlamaForValueRmPadPP(ParallelLlamaForCausalLMRmPadPP)","exported":true,"origin":"ast","body_hash":"5fa71c069e3ad5bcdcae61f07988ceef330e6c90a33e3f96f94be2f2979f506a","body_text":"class ParallelLlamaForValueRmPadPP(ParallelLlamaForCausalLMRmPadPP): def _init_head(self): column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() if self.megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config) self.lm_head = nn.Linear(in_features=self.config.hidden_size, out_features=1, bias=False) # lm_head is effectively the same as sequence parallel sp_utils.mark_parameter_as_sequence_parallel(self.lm_head.weight) def _forward_head(self, hidden_states): logits = self.lm_head(hidden_states) # (total_nnz_padded // tp, 1, 1) logits = logits.float() if self.megatron_config.sequence_parallel: logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False) return logits def forward( self, *, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: output = super().forward(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) if self.post_process: output.logits = torch.squeeze(output.logits, dim=-1) return output else: return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._init_head","name":"_init_head","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L628-L635","signature":"def _init_head(self)","exported":false,"origin":"ast","body_hash":"51af3f483828793a452f7f59a190039e0c2b511da22b0d80034de041fdeaa753","body_text":"def _init_head(self): column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear() if self.megatron_config is not None: assert column_kwargs.get('config', False), 'must have ModelParallelConfig' tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config) self.lm_head = nn.Linear(in_features=self.config.hidden_size, out_features=1, bias=False) # lm_head is effectively the same as sequence parallel sp_utils.mark_parameter_as_sequence_parallel(self.lm_head.weight)","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForValueRmPadPP"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._forward_head","name":"_forward_head","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L637-L642","signature":"def _forward_head(self, hidden_states)","exported":false,"origin":"ast","body_hash":"2c01174dfc8d7cce4ff60e2ffe09223ba1139598dad6b1f3127d2bc174137186","body_text":"def _forward_head(self, hidden_states): logits = self.lm_head(hidden_states) # (total_nnz_padded // tp, 1, 1) logits = logits.float() if self.megatron_config.sequence_parallel: logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False) return logits","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForValueRmPadPP"},{"id":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP.forward","name":"forward","kind":"method","path":"verl/models/llama/megatron/modeling_llama_megatron.py","span":"L644-L656","signature":"def forward( self, *, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]","exported":true,"origin":"ast","body_hash":"9f806adada908311ec51c11e661349389f095116287b7e80bbf5ac06908a0941","body_text":"def forward( self, *, input_ids: torch.LongTensor = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> Union[Tuple, CausalLMOutputWithPast]: output = super().forward(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) if self.post_process: output.logits = torch.squeeze(output.logits, dim=-1) return output else: return output","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelLlamaForValueRmPadPP"}],"rawEdges":[{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"typing","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"torch.utils.checkpoint","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"megatron.core","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"torch","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"transformers.modeling_outputs","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"transformers.models.llama.configuration_llama","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"transformers.models.llama.modeling_llama","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"verl.utils.megatron","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"verl.utils.megatron","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":".layers","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","relation":"calls","name":"full","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","relation":"calls","name":"finfo","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","relation":"calls","name":"arange","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","relation":"calls","name":"masked_fill_","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","relation":"calls","name":"view","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_make_causal_mask","relation":"calls","name":"expand","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#_expand_mask","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_expand_mask","relation":"calls","name":"size","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_expand_mask","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_expand_mask","relation":"calls","name":"expand","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_expand_mask","relation":"calls","name":"masked_fill","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_expand_mask","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#_expand_mask","relation":"calls","name":"finfo","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","relation":"calls","name":"get_default_kwargs_for_parallel_embedding","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","relation":"calls","name":"VocabParallelEmbedding","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","relation":"calls","name":"ModuleList","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","relation":"calls","name":"ParallelLlamaDecoderLayer","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.__init__","relation":"calls","name":"ParallelLlamaRMSNorm","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel._prepare_decoder_attention_mask","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel._prepare_decoder_attention_mask","relation":"calls","name":"_make_causal_mask","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel._prepare_decoder_attention_mask","relation":"calls","name":"to","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel._prepare_decoder_attention_mask","relation":"calls","name":"_expand_mask","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.forward","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.forward","relation":"calls","name":"embed_tokens","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaModel"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.forward","relation":"calls","name":"_prepare_decoder_attention_mask","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaModel"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.forward","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.forward","relation":"calls","name":"decoder_layer","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModel.forward","relation":"calls","name":"norm","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaModel"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.__init__","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.__init__","relation":"calls","name":"ParallelLlamaModel","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.__init__","relation":"calls","name":"get_default_kwargs_for_column_parallel_linear","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.__init__","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.__init__","relation":"calls","name":"ColumnParallelLinear","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.forward","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.forward","relation":"calls","name":"model","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForCausalLM"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.forward","relation":"calls","name":"lm_head","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForCausalLM"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.forward","relation":"calls","name":"gather_from_tensor_model_parallel_region","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.forward","relation":"calls","name":"float","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLM.forward","relation":"calls","name":"CausalLMOutputWithPast","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"imports","specifier":"flash_attn.bert_padding","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","relation":"calls","name":"get_default_kwargs_for_parallel_embedding","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","relation":"calls","name":"VocabParallelEmbedding","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","relation":"calls","name":"ModuleList","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","relation":"calls","name":"ParallelLlamaDecoderLayerRmPad","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.__init__","relation":"calls","name":"ParallelLlamaRMSNorm","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.forward","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.forward","relation":"calls","name":"embed_tokens","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaModelRmPad"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.forward","relation":"calls","name":"scatter_to_sequence_parallel_region","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.forward","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.forward","relation":"calls","name":"decoder_layer","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPad.forward","relation":"calls","name":"norm","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaModelRmPad"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.__init__","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.__init__","relation":"calls","name":"ParallelLlamaModelRmPad","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.__init__","relation":"calls","name":"_init_head","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForCausalLMRmPad"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._init_head","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._init_head","relation":"calls","name":"get_default_kwargs_for_column_parallel_linear","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._init_head","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._init_head","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._init_head","relation":"calls","name":"ColumnParallelLinear","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._forward_head","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._forward_head","relation":"calls","name":"lm_head","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForCausalLMRmPad"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._forward_head","relation":"calls","name":"float","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad._forward_head","relation":"calls","name":"gather_from_tensor_model_parallel_region","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","relation":"calls","name":"unpad_input","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","relation":"calls","name":"pad_to_sequence_parallel","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","relation":"calls","name":"model","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForCausalLMRmPad"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","relation":"calls","name":"_forward_head","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForCausalLMRmPad"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPad.forward","relation":"calls","name":"CausalLMOutputWithPast","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad","relation":"extends","name":"ParallelLlamaForCausalLMRmPad","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._init_head","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._init_head","relation":"calls","name":"get_default_kwargs_for_column_parallel_linear","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._init_head","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._init_head","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._init_head","relation":"calls","name":"Linear","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._init_head","relation":"calls","name":"mark_parameter_as_sequence_parallel","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._forward_head","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._forward_head","relation":"calls","name":"lm_head","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForValueRmPad"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._forward_head","relation":"calls","name":"float","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad._forward_head","relation":"calls","name":"gather_from_sequence_parallel_region","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad.forward","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad.forward","relation":"calls","name":"forward","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad.forward","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPad.forward","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"get_default_kwargs_for_parallel_embedding","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"VocabParallelEmbedding","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"ParallelLlamaDecoderLayerRmPad","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"append","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"ModuleList","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.__init__","relation":"calls","name":"ParallelLlamaRMSNorm","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.set_input_tensor","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.forward","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.forward","relation":"calls","name":"embed_tokens","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaModelRmPadPP"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.forward","relation":"calls","name":"scatter_to_sequence_parallel_region","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.forward","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.forward","relation":"calls","name":"decoder_layer","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaModelRmPadPP.forward","relation":"calls","name":"norm","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaModelRmPadPP"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.__init__","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.__init__","relation":"calls","name":"ParallelLlamaModelRmPadPP","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.__init__","relation":"calls","name":"_init_head","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForCausalLMRmPadPP"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.set_input_tensor","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.set_input_tensor","relation":"calls","name":"len","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.set_input_tensor","relation":"calls","name":"set_input_tensor","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaModelRmPadPP"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP._init_head","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP._init_head","relation":"calls","name":"get_default_kwargs_for_column_parallel_linear","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP._init_head","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP._init_head","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP._init_head","relation":"calls","name":"ColumnParallelLinear","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP._forward_head","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP._forward_head","relation":"calls","name":"lm_head","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForCausalLMRmPadPP"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP._forward_head","relation":"calls","name":"float","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","relation":"calls","name":"unpad_input","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","relation":"calls","name":"pad_to_sequence_parallel","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","relation":"calls","name":"model","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForCausalLMRmPadPP"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","relation":"calls","name":"_forward_head","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForCausalLMRmPadPP"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForCausalLMRmPadPP.forward","relation":"calls","name":"CausalLMOutputWithPast","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP","relation":"extends","name":"ParallelLlamaForCausalLMRmPadPP","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._init_head","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._init_head","relation":"calls","name":"get_default_kwargs_for_column_parallel_linear","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._init_head","relation":"calls","name":"get","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._init_head","relation":"calls","name":"update_kwargs_with_config","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._init_head","relation":"calls","name":"Linear","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._init_head","relation":"calls","name":"mark_parameter_as_sequence_parallel","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._forward_head","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._forward_head","relation":"calls","name":"lm_head","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py","recvType":"ParallelLlamaForValueRmPadPP"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._forward_head","relation":"calls","name":"float","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP._forward_head","relation":"calls","name":"gather_from_sequence_parallel_region","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP","relation":"contains","targetId":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP.forward","file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP.forward","relation":"calls","name":"forward","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP.forward","relation":"calls","name":"super","viaMember":false,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"},{"source":"verl/models/llama/megatron/modeling_llama_megatron.py#ParallelLlamaForValueRmPadPP.forward","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/models/llama/megatron/modeling_llama_megatron.py"}]},"verl/models/registry.py":{"size":2535,"mtimeMs":1789275830309.2026,"hash":"c7a2d330f96d7f793e9b5305bf9f932e92245f23ebb93c572221266391c10494","nodes":[{"id":"verl/models/registry.py","name":"registry.py","kind":"file","path":"verl/models/registry.py","span":"L1-L67","signature":null,"exported":true,"origin":"ast","body_hash":"c7a2d330f96d7f793e9b5305bf9f932e92245f23ebb93c572221266391c10494","chars":2535,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 typing import List, Optional, Type import torch.nn as nn # Supported models using HF Rmpad # TODO(sgm): HF may supported more than listed here, we should add more after testing from transformers import LlamaConfig, MistralConfig, GemmaConfig, Qwen2Config _REOVEPAD_MODELS = {'llama': LlamaConfig, 'mistral': MistralConfig, 'gemma': GemmaConfig, 'qwen2': Qwen2Config} # Supported models in Megatron-LM # Architecture -> (module, class). _MODELS = { \"LlamaForCausalLM\": (\"llama\", (\"ParallelLlamaForCausalLMRmPadPP\", \"ParallelLlamaForValueRmPadPP\", \"ParallelLlamaForCausalLMRmPad\")), \"MistralForCausalLM\": (\"mistral\", (\"ParallelMistralForCausalLMRmPadPP\", \"ParallelMistralForValueRmPadPP\", \"ParallelMistralForCausalLMRmPad\")) } # return model class"},{"id":"verl/models/registry.py#check_model_support_rmpad","name":"check_model_support_rmpad","kind":"function","path":"verl/models/registry.py","span":"L27-L32","signature":"def check_model_support_rmpad(model_type: str)","exported":true,"origin":"ast","body_hash":"b7dae798100438bb893014406ac3edd1094d8c621cb497b12a90dd9a3048ad44","body_text":"def check_model_support_rmpad(model_type: str): assert isinstance(model_type, str) if not model_type in _REOVEPAD_MODELS.keys(): raise ValueError(f\"Model architecture {model_type} is not supported for now. \" f\"RMPad supported architectures: {_REOVEPAD_MODELS.keys()}.\" f\"Please set `use_remove_padding=False` in the model config.\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/registry.py#ModelRegistry","name":"ModelRegistry","kind":"class","path":"verl/models/registry.py","span":"L46-L66","signature":"class ModelRegistry","exported":true,"origin":"ast","body_hash":"7af10eb2e2bbf74c6a76d1d26c5effd68d8e59c679335baa3509e7ac5f44df0d","body_text":"class ModelRegistry: @staticmethod def load_model_cls(model_arch: str, value=False) -> Optional[Type[nn.Module]]: if model_arch not in _MODELS: return None megatron = \"megatron\" module_name, model_cls_name = _MODELS[model_arch] if not value: # actor/ref model_cls_name = model_cls_name[0] elif value: # critic/rm model_cls_name = model_cls_name[1] module = importlib.import_module(f\"verl.models.{module_name}.{megatron}.modeling_{module_name}_megatron\") return getattr(module, model_cls_name, None) @staticmethod def get_supported_archs() -> List[str]: return list(_MODELS.keys())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/registry.py#ModelRegistry.load_model_cls","name":"load_model_cls","kind":"method","path":"verl/models/registry.py","span":"L49-L62","signature":"def load_model_cls(model_arch: str, value=False) -> Optional[Type[nn.Module]]","exported":true,"origin":"ast","body_hash":"8e0a696b8e10323f88a53f2084c5f423a8ce2a7325478f2a45aee3e3875db5d4","body_text":"def load_model_cls(model_arch: str, value=False) -> Optional[Type[nn.Module]]: if model_arch not in _MODELS: return None megatron = \"megatron\" module_name, model_cls_name = _MODELS[model_arch] if not value: # actor/ref model_cls_name = model_cls_name[0] elif value: # critic/rm model_cls_name = model_cls_name[1] module = importlib.import_module(f\"verl.models.{module_name}.{megatron}.modeling_{module_name}_megatron\") return getattr(module, model_cls_name, None)","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRegistry"},{"id":"verl/models/registry.py#ModelRegistry.get_supported_archs","name":"get_supported_archs","kind":"method","path":"verl/models/registry.py","span":"L65-L66","signature":"def get_supported_archs() -> List[str]","exported":true,"origin":"ast","body_hash":"881a969fb9e574cba5880e9894970d05fce39b75e74cfe7d0a03e16beb3d79a6","body_text":"def get_supported_archs() -> List[str]: return list(_MODELS.keys())","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRegistry"}],"rawEdges":[{"source":"verl/models/registry.py","relation":"imports","specifier":"importlib","file":"verl/models/registry.py"},{"source":"verl/models/registry.py","relation":"imports","specifier":"typing","file":"verl/models/registry.py"},{"source":"verl/models/registry.py","relation":"imports","specifier":"transformers","file":"verl/models/registry.py"},{"source":"verl/models/registry.py","relation":"contains","targetId":"verl/models/registry.py#check_model_support_rmpad","file":"verl/models/registry.py"},{"source":"verl/models/registry.py#check_model_support_rmpad","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/models/registry.py"},{"source":"verl/models/registry.py#check_model_support_rmpad","relation":"calls","name":"keys","viaMember":true,"file":"verl/models/registry.py"},{"source":"verl/models/registry.py#check_model_support_rmpad","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/models/registry.py"},{"source":"verl/models/registry.py#check_model_support_rmpad","relation":"calls","name":"keys","viaMember":true,"file":"verl/models/registry.py"},{"source":"verl/models/registry.py","relation":"contains","targetId":"verl/models/registry.py#ModelRegistry","file":"verl/models/registry.py"},{"source":"verl/models/registry.py#ModelRegistry","relation":"contains","targetId":"verl/models/registry.py#ModelRegistry.load_model_cls","file":"verl/models/registry.py"},{"source":"verl/models/registry.py#ModelRegistry.load_model_cls","relation":"calls","name":"import_module","viaMember":true,"file":"verl/models/registry.py"},{"source":"verl/models/registry.py#ModelRegistry.load_model_cls","relation":"calls","name":"getattr","viaMember":false,"file":"verl/models/registry.py"},{"source":"verl/models/registry.py#ModelRegistry","relation":"contains","targetId":"verl/models/registry.py#ModelRegistry.get_supported_archs","file":"verl/models/registry.py"},{"source":"verl/models/registry.py#ModelRegistry.get_supported_archs","relation":"calls","name":"list","viaMember":false,"file":"verl/models/registry.py"},{"source":"verl/models/registry.py#ModelRegistry.get_supported_archs","relation":"calls","name":"keys","viaMember":true,"file":"verl/models/registry.py"}]},"verl/models/transformers/__init__.py":{"size":600,"mtimeMs":1789275830309.2026,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/models/transformers/__init__.py","name":"__init__.py","kind":"file","path":"verl/models/transformers/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/models/transformers/llama.py":{"size":6759,"mtimeMs":1789275830309.2026,"hash":"8ef8564e7b3b45d8aaf2509eff07265b6cd1d6af1da0e588c8976d221d40d323","nodes":[{"id":"verl/models/transformers/llama.py","name":"llama.py","kind":"file","path":"verl/models/transformers/llama.py","span":"L1-L146","signature":null,"exported":true,"origin":"ast","body_hash":"8ef8564e7b3b45d8aaf2509eff07265b6cd1d6af1da0e588c8976d221d40d323","chars":6759,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 torch from typing import Optional, List, Union, Tuple, Unpack, Callable from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv from transformers.cache_utils import Cache from transformers.utils import logging from transformers.modeling_flash_attention_utils import _flash_attention_forward from verl.utils.ulysses import gather_heads_scatter_seq, gather_seq_scatter_heads, get_ulysses_sequence_parallel_world_size logger = logging.get_logger(__name__)"},{"id":"verl/models/transformers/llama.py#llama_flash_attn_forward","name":"llama_flash_attn_forward","kind":"function","path":"verl/models/transformers/llama.py","span":"L27-L145","signature":"def llama_flash_attn_forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]","exported":true,"origin":"ast","body_hash":"ff595c8414e383367423ef7ac48815a1f62911c6e3d74f04e0e9d7d655f544fa","body_text":"def llama_flash_attn_forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.LongTensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 **kwargs, ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]: \"\"\" adapt from transformers 4.47.1 \"\"\" output_attentions = False bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) # Flash attention requires the input to have the shape # batch_size x seq_length x head_dim x hidden_dim # therefore we just need to keep the original shape query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2) # trade off: repeat first and then all to all # key_states = repeat_kv(key_states, self.num_key_value_groups) # value_states = repeat_kv(value_states, self.num_key_value_groups) ########## AlltoAll for Ulysses ########## ulysses_sp_size = get_ulysses_sequence_parallel_world_size() if ulysses_sp_size > 1: # (bsz, n_head, seq_len/n, head_dim) -> (bsz, n_head/n, seq_len, head_dim) query_states = gather_seq_scatter_heads(query_states, seq_dim=2, head_dim=1) key_states = gather_seq_scatter_heads(key_states, seq_dim=2, head_dim=1) value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1) full_q_len = query_states.size(2) # full seq length if position_embeddings is None: logger.warning_once( \"The attention layers in this model are transitioning from computing the RoPE embeddings internally \" \"through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed \" \"`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be \" \"removed and `position_embeddings` will be mandatory.\") cos, sin = self.rotary_emb(value_states, position_ids) else: cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: # sin and cos are specific to RoPE models; cache_position needed for the static cache cache_kwargs = {\"sin\": sin, \"cos\": cos, \"cache_position\": cache_position} key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) # TODO: These transpose are quite inefficient but Flash Attention requires the layout [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache # to be able to avoid many of these transpose/reshape/view. query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) dropout_rate = self.attention_dropout if self.training else 0.0 # In PEFT, usually we cast the layer norms in float32 for training stability reasons # therefore the input hidden states gets silently casted in float32. Hence, we need # cast them back in the correct dtype just to be sure everything works as expected. # This might slowdown training & inference so it is recommended to not cast the LayerNorms # in fp32. (LlamaRMSNorm handles it correctly) input_dtype = query_states.dtype if input_dtype == torch.float32: if torch.is_autocast_enabled(): target_dtype = torch.get_autocast_gpu_dtype() # Handle the case where the model is quantized elif hasattr(self.config, \"_pre_quantization_dtype\"): target_dtype = self.config._pre_quantization_dtype else: target_dtype = self.q_proj.weight.dtype logger.warning_once( f\"The input hidden states seems to be silently casted in float32, this might be related to\" f\" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in\" f\" {target_dtype}.\") query_states = query_states.to(target_dtype) key_states = key_states.to(target_dtype) value_states = value_states.to(target_dtype) attn_output = _flash_attention_forward( query_states, key_states, value_states, attention_mask, full_q_len, position_ids=position_ids, dropout=dropout_rate, sliding_window=getattr(self, \"sliding_window\", None), use_top_left_mask=self._flash_attn_uses_top_left_mask, is_causal=self.is_causal, **kwargs, ) attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous() ########## AlltoAll for Ulysses ########## if ulysses_sp_size > 1: attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2) attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() attn_output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/models/transformers/llama.py","relation":"imports","specifier":"os","file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py","relation":"imports","specifier":"torch","file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py","relation":"imports","specifier":"typing","file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py","relation":"imports","specifier":"transformers.models.llama.modeling_llama","file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py","relation":"imports","specifier":"transformers.cache_utils","file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py","relation":"imports","specifier":"transformers.utils","file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py","relation":"imports","specifier":"transformers.modeling_flash_attention_utils","file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py","relation":"imports","specifier":"verl.utils.ulysses","file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py","relation":"calls","name":"get_logger","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py","relation":"contains","targetId":"verl/models/transformers/llama.py#llama_flash_attn_forward","file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"q_proj","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"k_proj","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"v_proj","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"get_ulysses_sequence_parallel_world_size","viaMember":false,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"gather_seq_scatter_heads","viaMember":false,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"gather_seq_scatter_heads","viaMember":false,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"gather_seq_scatter_heads","viaMember":false,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"warning_once","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"rotary_emb","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"apply_rotary_pos_emb","viaMember":false,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"update","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"is_autocast_enabled","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"get_autocast_gpu_dtype","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"warning_once","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"_flash_attention_forward","viaMember":false,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"getattr","viaMember":false,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"reshape","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_heads_scatter_seq"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"gather_heads_scatter_seq","viaMember":false,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/models/transformers/llama.py"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"reshape","viaMember":true,"file":"verl/models/transformers/llama.py","recvType":"gather_heads_scatter_seq"},{"source":"verl/models/transformers/llama.py#llama_flash_attn_forward","relation":"calls","name":"o_proj","viaMember":true,"file":"verl/models/transformers/llama.py"}]},"verl/models/transformers/monkey_patch.py":{"size":2884,"mtimeMs":1789275830309.2026,"hash":"8a0775fa2927755d1fa16e5cc2291e3ea45ec6f146f331bbca1d30e879b09943","nodes":[{"id":"verl/models/transformers/monkey_patch.py","name":"monkey_patch.py","kind":"file","path":"verl/models/transformers/monkey_patch.py","span":"L1-L75","signature":null,"exported":true,"origin":"ast","body_hash":"8a0775fa2927755d1fa16e5cc2291e3ea45ec6f146f331bbca1d30e879b09943","chars":2884,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Apply monkey-patch function to models \"\"\" #### Open Source Models #### transformers version < 4.48 _PATCH_NAME_TO_FUNC = { 'llama': apply_monkey_patch_to_llama, 'qwen2': apply_monkey_patch_to_qwen2, } from transformers import PretrainedConfig from functools import lru_cache from packaging import version import importlib.metadata @lru_cache()"},{"id":"verl/models/transformers/monkey_patch.py#apply_monkey_patch_to_llama","name":"apply_monkey_patch_to_llama","kind":"function","path":"verl/models/transformers/monkey_patch.py","span":"L22-L25","signature":"def apply_monkey_patch_to_llama()","exported":true,"origin":"ast","body_hash":"a2cb89d90cc9a59e359216ada15c84411dbc941c73411e14f55c7888b17b4de4","body_text":"def apply_monkey_patch_to_llama(): from transformers.models.llama.modeling_llama import LlamaFlashAttention2 from verl.models.transformers.llama import llama_flash_attn_forward LlamaFlashAttention2.forward = llama_flash_attn_forward","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/transformers/monkey_patch.py#apply_monkey_patch_to_qwen2","name":"apply_monkey_patch_to_qwen2","kind":"function","path":"verl/models/transformers/monkey_patch.py","span":"L28-L31","signature":"def apply_monkey_patch_to_qwen2()","exported":true,"origin":"ast","body_hash":"d434ea93c6ff5d31f1a7023ca19ad66aa47e86803e7fec0d1baa0dff172adbd6","body_text":"def apply_monkey_patch_to_qwen2(): from transformers.models.qwen2.modeling_qwen2 import Qwen2FlashAttention2 from verl.models.transformers.qwen2 import qwen2_flash_attn_forward Qwen2FlashAttention2.forward = qwen2_flash_attn_forward","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/transformers/monkey_patch.py#apply_monkey_patch","name":"apply_monkey_patch","kind":"function","path":"verl/models/transformers/monkey_patch.py","span":"L42-L57","signature":"def apply_monkey_patch(config: PretrainedConfig, verbose=True)","exported":true,"origin":"ast","body_hash":"6460d612cce27699cc003a79fa6dde059ed508e3fba5c2679f44266b70156d68","body_text":"def apply_monkey_patch(config: PretrainedConfig, verbose=True): if not is_transformers_version_in_range(\"4.45.0\", \"4.47.1\"): raise AssertionError(\"The installed `transformers` version doesn't support ulysses patch. \" \"Please install a version between 4.45.0 and 4.47.1 to use this ulysses feature.\") success_apply_monkey_patch = False if config.model_type in _PATCH_NAME_TO_FUNC: _PATCH_NAME_TO_FUNC[config.model_type]() success_apply_monkey_patch = True if success_apply_monkey_patch and verbose: print(f'Applying monkey patch to model {config.model_type}') elif not success_apply_monkey_patch: raise NotImplementedError(f'Ulysses for model {config.model_type} is not implemented, \\ please set `ulysses_sequence_parallel_size=1`') return success_apply_monkey_patch","summary_state":"pending","summary":null,"crux":null},{"id":"verl/models/transformers/monkey_patch.py#is_transformers_version_in_range","name":"is_transformers_version_in_range","kind":"function","path":"verl/models/transformers/monkey_patch.py","span":"L66-L74","signature":"def is_transformers_version_in_range(min_version: str, max_version: str) -> bool","exported":true,"origin":"ast","body_hash":"38063fd64d4387629cff8d1c5e74683a8558a524d397cf8af95fa15abe6a252e","body_text":"def is_transformers_version_in_range(min_version: str, max_version: str) -> bool: try: # Get the installed version of the transformers library transformers_version = importlib.metadata.version(\"transformers\") except importlib.metadata.PackageNotFoundError: raise ModuleNotFoundError(\"The `transformers` package is not installed.\") # Check if the version is within the specified range return version.parse(min_version) <= version.parse(transformers_version) <= version.parse(max_version)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/models/transformers/monkey_patch.py","relation":"contains","targetId":"verl/models/transformers/monkey_patch.py#apply_monkey_patch_to_llama","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"imports","specifier":"transformers.models.llama.modeling_llama","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"imports","specifier":"verl.models.transformers.llama","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"contains","targetId":"verl/models/transformers/monkey_patch.py#apply_monkey_patch_to_qwen2","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"imports","specifier":"transformers.models.qwen2.modeling_qwen2","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"imports","specifier":"verl.models.transformers.qwen2","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"imports","specifier":"transformers","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"contains","targetId":"verl/models/transformers/monkey_patch.py#apply_monkey_patch","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py#apply_monkey_patch","relation":"calls","name":"is_transformers_version_in_range","viaMember":false,"file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py#apply_monkey_patch","relation":"calls","name":"AssertionError","viaMember":false,"file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py#apply_monkey_patch","relation":"calls","name":"print","viaMember":false,"file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py#apply_monkey_patch","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"imports","specifier":"functools","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"imports","specifier":"packaging","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"imports","specifier":"importlib.metadata","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"calls","name":"lru_cache","viaMember":false,"file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py","relation":"contains","targetId":"verl/models/transformers/monkey_patch.py#is_transformers_version_in_range","file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py#is_transformers_version_in_range","relation":"calls","name":"version","viaMember":true,"file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py#is_transformers_version_in_range","relation":"calls","name":"ModuleNotFoundError","viaMember":false,"file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py#is_transformers_version_in_range","relation":"calls","name":"parse","viaMember":true,"file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py#is_transformers_version_in_range","relation":"calls","name":"parse","viaMember":true,"file":"verl/models/transformers/monkey_patch.py"},{"source":"verl/models/transformers/monkey_patch.py#is_transformers_version_in_range","relation":"calls","name":"parse","viaMember":true,"file":"verl/models/transformers/monkey_patch.py"}]},"verl/models/transformers/qwen2.py":{"size":6144,"mtimeMs":1789275830309.2026,"hash":"81b67ee203676ce26dbf6d5c1f5436ddf1d4a6070e5f646b2b9b87e7d95dd7d1","nodes":[{"id":"verl/models/transformers/qwen2.py","name":"qwen2.py","kind":"file","path":"verl/models/transformers/qwen2.py","span":"L1-L138","signature":null,"exported":true,"origin":"ast","body_hash":"81b67ee203676ce26dbf6d5c1f5436ddf1d4a6070e5f646b2b9b87e7d95dd7d1","chars":6144,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 torch from typing import Optional, Tuple from transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv from transformers.cache_utils import Cache from transformers.utils import logging from transformers.modeling_flash_attention_utils import _flash_attention_forward from verl.utils.ulysses import gather_heads_scatter_seq, gather_seq_scatter_heads, get_ulysses_sequence_parallel_world_size logger = logging.get_logger(__name__)"},{"id":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","name":"qwen2_flash_attn_forward","kind":"function","path":"verl/models/transformers/qwen2.py","span":"L28-L137","signature":"def qwen2_flash_attn_forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 )","exported":true,"origin":"ast","body_hash":"930831d1187694b0dbc820d3eba79ce924d1cba82339a9d45b13d62ec3abee3b","body_text":"def qwen2_flash_attn_forward( self, hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_value: Optional[Cache] = None, output_attentions: bool = False, use_cache: bool = False, cache_position: Optional[torch.LongTensor] = None, position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 ): bsz, q_len, _ = hidden_states.size() query_states = self.q_proj(hidden_states) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2) ########## AlltoAll for Ulysses ########## ulysses_sp_size = get_ulysses_sequence_parallel_world_size() if ulysses_sp_size > 1: # (bsz, n_head, seq_len/n, head_dim) -> (bsz, n_head/n, seq_len, head_dim) query_states = gather_seq_scatter_heads(query_states, seq_dim=2, head_dim=1) key_states = gather_seq_scatter_heads(key_states, seq_dim=2, head_dim=1) value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1) full_q_len = query_states.size(2) # full seq length if position_embeddings is None: logger.warning_once( \"The attention layers in this model are transitioning from computing the RoPE embeddings internally \" \"through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed \" \"`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be \" \"removed and `position_embeddings` will be mandatory.\") cos, sin = self.rotary_emb(value_states, position_ids) else: cos, sin = position_embeddings query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) if past_key_value is not None: cache_kwargs = {\"sin\": sin, \"cos\": cos, \"cache_position\": cache_position} # Specific to RoPE models key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs) # repeat k/v heads if n_kv_heads < n_heads key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) dropout_rate = 0.0 if not self.training else self.attention_dropout # In PEFT, usually we cast the layer norms in float32 for training stability reasons # therefore the input hidden states gets silently casted in float32. Hence, we need # cast them back in float16 just to be sure everything works as expected. input_dtype = query_states.dtype if input_dtype == torch.float32: if torch.is_autocast_enabled(): target_dtype = torch.get_autocast_gpu_dtype() # Handle the case where the model is quantized elif hasattr(self.config, \"_pre_quantization_dtype\"): target_dtype = self.config._pre_quantization_dtype else: target_dtype = self.q_proj.weight.dtype logger.warning_once( f\"The input hidden states seems to be silently casted in float32, this might be related to\" f\" the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in\" f\" {target_dtype}.\") query_states = query_states.to(target_dtype) key_states = key_states.to(target_dtype) value_states = value_states.to(target_dtype) # Reashape to the expected shape for Flash Attention query_states = query_states.transpose(1, 2) key_states = key_states.transpose(1, 2) value_states = value_states.transpose(1, 2) if (self.config.use_sliding_window and getattr(self.config, \"sliding_window\", None) is not None and self.layer_idx >= self.config.max_window_layers): sliding_window = self.config.sliding_window else: sliding_window = None attn_output = _flash_attention_forward( query_states, key_states, value_states, attention_mask, full_q_len, position_ids=position_ids, dropout=dropout_rate, sliding_window=sliding_window, is_causal=self.is_causal, use_top_left_mask=self._flash_attn_uses_top_left_mask, ) # use full_q_len to reshape attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous() ########## AlltoAll for Ulysses ########## if ulysses_sp_size > 1: attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2) attn_output = attn_output.reshape(bsz, q_len, -1).contiguous() attn_output = self.o_proj(attn_output) if not output_attentions: attn_weights = None return attn_output, attn_weights, past_key_value","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/models/transformers/qwen2.py","relation":"imports","specifier":"os","file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py","relation":"imports","specifier":"torch","file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py","relation":"imports","specifier":"typing","file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py","relation":"imports","specifier":"transformers.models.llama.modeling_llama","file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py","relation":"imports","specifier":"transformers.cache_utils","file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py","relation":"imports","specifier":"transformers.utils","file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py","relation":"imports","specifier":"transformers.modeling_flash_attention_utils","file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py","relation":"imports","specifier":"verl.utils.ulysses","file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py","relation":"calls","name":"get_logger","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py","relation":"contains","targetId":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"q_proj","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"k_proj","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"v_proj","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"repeat_kv"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"view","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"repeat_kv"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"get_ulysses_sequence_parallel_world_size","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"gather_seq_scatter_heads","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"gather_seq_scatter_heads","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"gather_seq_scatter_heads","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"size","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"warning_once","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"rotary_emb","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"apply_rotary_pos_emb","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"update","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"repeat_kv","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"repeat_kv","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"is_autocast_enabled","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"get_autocast_gpu_dtype","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"warning_once","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"repeat_kv"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"to","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"repeat_kv"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"gather_seq_scatter_heads"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"repeat_kv"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"transpose","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"repeat_kv"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"getattr","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"_flash_attention_forward","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"reshape","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"gather_heads_scatter_seq"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"gather_heads_scatter_seq","viaMember":false,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/models/transformers/qwen2.py"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"reshape","viaMember":true,"file":"verl/models/transformers/qwen2.py","recvType":"gather_heads_scatter_seq"},{"source":"verl/models/transformers/qwen2.py#qwen2_flash_attn_forward","relation":"calls","name":"o_proj","viaMember":true,"file":"verl/models/transformers/qwen2.py"}]},"verl/models/weight_loader_registry.py":{"size":1137,"mtimeMs":1789275830309.2026,"hash":"14826b9646576cfe7b157c6bf293bbfe9c76b861c0803c1ad2d4bb93e7820d59","nodes":[{"id":"verl/models/weight_loader_registry.py","name":"weight_loader_registry.py","kind":"file","path":"verl/models/weight_loader_registry.py","span":"L1-L24","signature":null,"exported":true,"origin":"ast","body_hash":"14826b9646576cfe7b157c6bf293bbfe9c76b861c0803c1ad2d4bb93e7820d59","chars":1137,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."},{"id":"verl/models/weight_loader_registry.py#get_weight_loader","name":"get_weight_loader","kind":"function","path":"verl/models/weight_loader_registry.py","span":"L16-L23","signature":"def get_weight_loader(arch: str)","exported":true,"origin":"ast","body_hash":"42ec46896e08824dec411bab38fb25f2e1edef93fb1ef78f6f2d998d2bae80fa","body_text":"def get_weight_loader(arch: str): from verl.models.llama.megatron.checkpoint_utils.llama_loader import load_state_dict_to_megatron_llama _MODEL_WEIGHT_MEGATRON_LOADER_REGISTRY = {'LlamaForCausalLM': load_state_dict_to_megatron_llama} if arch in _MODEL_WEIGHT_MEGATRON_LOADER_REGISTRY: return _MODEL_WEIGHT_MEGATRON_LOADER_REGISTRY[arch] raise ValueError(f\"Model architectures {arch} are not supported for now. \" f\"Supported architectures: {_MODEL_WEIGHT_MEGATRON_LOADER_REGISTRY.keys()}\")","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/models/weight_loader_registry.py","relation":"contains","targetId":"verl/models/weight_loader_registry.py#get_weight_loader","file":"verl/models/weight_loader_registry.py"},{"source":"verl/models/weight_loader_registry.py","relation":"imports","specifier":"verl.models.llama.megatron.checkpoint_utils.llama_loader","file":"verl/models/weight_loader_registry.py"},{"source":"verl/models/weight_loader_registry.py#get_weight_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/models/weight_loader_registry.py"},{"source":"verl/models/weight_loader_registry.py#get_weight_loader","relation":"calls","name":"keys","viaMember":true,"file":"verl/models/weight_loader_registry.py"}]},"verl/protocol.py":{"size":24417,"mtimeMs":1789275830309.2026,"hash":"ceccbe7c9ad3bfe430355f35705054881f32dc254b0862a1f0bf6800e387405d","nodes":[{"id":"verl/protocol.py","name":"protocol.py","kind":"file","path":"verl/protocol.py","span":"L1-L640","signature":null,"exported":true,"origin":"ast","body_hash":"ceccbe7c9ad3bfe430355f35705054881f32dc254b0862a1f0bf6800e387405d","chars":24417,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Implement base data transfer protocol between any two functions, modules. We can subclass Protocol to define more detailed batch info with specific keys \"\"\" import pickle import numpy as np import copy from dataclasses import dataclass, field from typing import Callable, Dict, List, Union import torch import tensordict from tensordict import TensorDict from torch.utils.data import DataLoader, Dataset from verl.utils.py_functional import union_two_dict __all__ = ['DataProto', 'union_tensor_dict'] try: tensordict.set_lazy_legacy(False).set() except: pass @dataclass @dataclass import ray @dataclass"},{"id":"verl/protocol.py#pad_dataproto_to_divisor","name":"pad_dataproto_to_divisor","kind":"function","path":"verl/protocol.py","span":"L40-L57","signature":"def pad_dataproto_to_divisor(data: 'DataProto', size_divisor: int)","exported":true,"origin":"ast","body_hash":"09385d739bd274794c5b64054d84ef2729590add2aa44cfc94f183561fe2274c","body_text":"def pad_dataproto_to_divisor(data: 'DataProto', size_divisor: int): \"\"\"Pad a DataProto to size divisible by size_divisor Args: size_divisor (int): size divisor Returns: data: (DataProto): the padded DataProto pad_size (int) \"\"\" assert isinstance(data, DataProto), 'data must be a DataProto' if len(data) % size_divisor != 0: pad_size = size_divisor - len(data) % size_divisor data_padded = DataProto.concat([data, data[:pad_size]]) else: pad_size = 0 data_padded = data return data_padded, pad_size","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#unpad_dataproto","name":"unpad_dataproto","kind":"function","path":"verl/protocol.py","span":"L60-L63","signature":"def unpad_dataproto(data: 'DataProto', pad_size)","exported":true,"origin":"ast","body_hash":"36f58338b3d685bf67cc6e2692bca04f2817c013a9337bb7806f2cd480ac0765","body_text":"def unpad_dataproto(data: 'DataProto', pad_size): if pad_size != 0: data = data[:-pad_size] return data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#union_tensor_dict","name":"union_tensor_dict","kind":"function","path":"verl/protocol.py","span":"L66-L77","signature":"def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict","exported":true,"origin":"ast","body_hash":"5802fb27fc61adc37d92081b594ddcbe7788685ea0a95abbfb6eb8d9a5d04f7e","body_text":"def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict: \"\"\"Union two tensordicts.\"\"\" 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(): tensor_dict1[key] = tensor_dict2[key] else: assert tensor_dict1[key].equal(tensor_dict2[key]), \\ f'{key} in tensor_dict1 and tensor_dict2 are not the same object' return tensor_dict1","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#union_numpy_dict","name":"union_numpy_dict","kind":"function","path":"verl/protocol.py","span":"L80-L89","signature":"def union_numpy_dict(tensor_dict1: dict[np.ndarray], tensor_dict2: dict[np.ndarray]) -> dict[np.ndarray]","exported":true,"origin":"ast","body_hash":"9190f293c6b2654a7a48b6d0895a1316b66081614fb0e26e361dafc80bc015d3","body_text":"def union_numpy_dict(tensor_dict1: dict[np.ndarray], tensor_dict2: dict[np.ndarray]) -> dict[np.ndarray]: for key, val in tensor_dict2.items(): if key in tensor_dict1: assert isinstance(tensor_dict2[key], np.ndarray) assert isinstance(tensor_dict1[key], np.ndarray) assert np.all(tensor_dict2[key] == tensor_dict1[key]), \\ f'{key} in tensor_dict1 and tensor_dict2 are not the same object' tensor_dict1[key] = val return tensor_dict1","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#list_of_dict_to_dict_of_list","name":"list_of_dict_to_dict_of_list","kind":"function","path":"verl/protocol.py","span":"L92-L101","signature":"def list_of_dict_to_dict_of_list(list_of_dict: list[dict])","exported":true,"origin":"ast","body_hash":"92469650f3a7fc08cfbe452556a7379f469de8a86fd2dfb02137b29fe3acfe4b","body_text":"def list_of_dict_to_dict_of_list(list_of_dict: list[dict]): if len(list_of_dict) == 0: return {} keys = list_of_dict[0].keys() output = {key: [] for key in keys} for data in list_of_dict: for key, item in data.items(): assert key in output output[key].append(item) return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#fold_batch_dim","name":"fold_batch_dim","kind":"function","path":"verl/protocol.py","span":"L104-L121","signature":"def fold_batch_dim(data: 'DataProto', new_batch_size)","exported":true,"origin":"ast","body_hash":"74df5f535e51a3d7fc00627e9898b6f90bffe96e9ebaa5ee39432e64366173a6","body_text":"def fold_batch_dim(data: 'DataProto', new_batch_size): \"\"\" Fold a batch dim from [bsz, xxx] into [new_bsz, bsz // new_bsz, xxx] \"\"\" batch_size = data.batch.batch_size[0] assert batch_size % new_batch_size == 0 tensor: TensorDict = data.batch non_tensor = data.non_tensor_batch tensor = tensor.view(new_batch_size, -1) tensor.auto_batch_size_(batch_dims=1) for key, val in non_tensor.items(): non_tensor[key] = np.reshape(val, newshape=(new_batch_size, -1, *val.shape[1:])) return DataProto(batch=tensor, non_tensor_batch=non_tensor, meta_info=data.meta_info)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#unfold_batch_dim","name":"unfold_batch_dim","kind":"function","path":"verl/protocol.py","span":"L124-L140","signature":"def unfold_batch_dim(data: 'DataProto', batch_dims=2)","exported":true,"origin":"ast","body_hash":"869728770cff0a94c3083559704f1aa89f0544ece88a70f38113d4cfe60a06df","body_text":"def unfold_batch_dim(data: 'DataProto', batch_dims=2): \"\"\" Unfold the first n dims as new batch dim \"\"\" tensor: TensorDict = data.batch non_tensor = data.non_tensor_batch tensor.auto_batch_size_(batch_dims=batch_dims) tensor = tensor.view(-1) batch_size = tensor.batch_size[0] non_tensor_new = {} for key, val in non_tensor.items(): non_tensor_new[key] = np.reshape(val, newshape=(batch_size, *val.shape[batch_dims:])) return DataProto(batch=tensor, non_tensor_batch=non_tensor_new, meta_info=data.meta_info)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#collate_fn","name":"collate_fn","kind":"function","path":"verl/protocol.py","span":"L143-L153","signature":"def collate_fn(x: list['DataProtoItem'])","exported":true,"origin":"ast","body_hash":"e60965207a22fee739c989ca3a257ef3712f2d8197a77fd58bf09cb8acce2e16","body_text":"def collate_fn(x: list['DataProtoItem']): batch = [] non_tensor_batch = [] for data in x: batch.append(data.batch) non_tensor_batch.append(data.non_tensor_batch) batch = torch.stack(batch).contiguous() non_tensor_batch = list_of_dict_to_dict_of_list(non_tensor_batch) for key, val in non_tensor_batch.items(): non_tensor_batch[key] = np.array(val, dtype=object) return DataProto(batch=batch, non_tensor_batch=non_tensor_batch)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#DataProtoItem","name":"DataProtoItem","kind":"class","path":"verl/protocol.py","span":"L157-L161","signature":"class DataProtoItem: # TODO(zhangchi.usc1992) add consistency check","exported":true,"origin":"ast","body_hash":"2db969dfd81f0957e3d34b05a949275630e208dc6cbd0c9c190d706a00c6a849","body_text":"class DataProtoItem: # TODO(zhangchi.usc1992) add consistency check batch: TensorDict = None non_tensor_batch: Dict = field(default_factory=dict) meta_info: Dict = field(default_factory=dict)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#DataProto","name":"DataProto","kind":"class","path":"verl/protocol.py","span":"L165-L589","signature":"class DataProto","exported":true,"origin":"ast","body_hash":"5c388ec1f359b3ade458751836539ba4cb0620384129fe833e83bb672ab72d5a","body_text":"class DataProto: \"\"\" A DataProto is a data structure that aims to provide a standard protocol for data exchange between functions. It contains a batch (TensorDict) and a meta_info (Dict). The batch is a TensorDict https://pytorch.org/tensordict/. TensorDict allows you to manipulate a dictionary of Tensors like a single Tensor. Ideally, the tensors with the same batch size should be put inside batch. \"\"\" batch: TensorDict = None non_tensor_batch: Dict = field(default_factory=dict) meta_info: Dict = field(default_factory=dict) def __post_init__(self): # perform necessary checking self.check_consistency() def __len__(self): if self.batch is not None: return self.batch.batch_size[0] elif self.non_tensor_batch is not None and len(self.non_tensor_batch) > 0: random_key = list(self.non_tensor_batch.keys())[0] return self.non_tensor_batch[random_key].shape[0] else: return 0 def __getitem__(self, item): tensor_data = self.batch[item] non_tensor_data = {key: val[item] for key, val in self.non_tensor_batch.items()} return DataProtoItem(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info) def __getstate__(self): import io buffer = io.BytesIO() if tensordict.__version__ >= '0.5.0' and self.batch is not None: self.batch = self.batch.contiguous() self.batch = self.batch.consolidate() torch.save(self.batch, buffer) buffer_bytes = buffer.getvalue() return buffer_bytes, self.non_tensor_batch, self.meta_info def __setstate__(self, data): import io batch_deserialized_bytes, non_tensor_batch, meta_info = data batch_deserialized = io.BytesIO(initial_bytes=batch_deserialized_bytes) batch = torch.load(batch_deserialized, weights_only=False, map_location='cpu' if not torch.cuda.is_available() else None) self.batch = batch self.non_tensor_batch = non_tensor_batch self.meta_info = meta_info def save_to_disk(self, filepath): with open(filepath, 'wb') as f: pickle.dump(self, f) @staticmethod def load_from_disk(filepath) -> 'DataProto': with open(filepath, 'rb') as f: data = pickle.load(f) return data def print_size(self, prefix=\"\"): size_of_tensordict = 0 for key, tensor in self.batch.items(): size_of_tensordict += tensor.element_size() * tensor.numel() size_of_numpy_array = 0 for key, numpy_array in self.non_tensor_batch.items(): size_of_numpy_array += numpy_array.nbytes size_of_numpy_array /= 1024**3 size_of_tensordict /= 1024**3 message = f'Size of tensordict: {size_of_tensordict} GB, size of non_tensor_batch: {size_of_numpy_array} GB' if prefix: message = f'{prefix}, ' + message print(message) def check_consistency(self): \"\"\"Check the consistency of the DataProto. Mainly for batch and non_tensor_batch We expose this function as a public one so that user can call themselves directly \"\"\" if self.batch is not None: assert len(self.batch.batch_size) == 1, 'only support num_batch_dims=1' if self.non_tensor_batch is not None: for key, val in self.non_tensor_batch.items(): assert isinstance(val, np.ndarray) if self.batch is not None and len(self.non_tensor_batch) != 0: # TODO: we can actually lift this restriction if needed assert len(self.batch.batch_size) == 1, 'only support num_batch_dims=1 when non_tensor_batch is not empty.' batch_size = self.batch.batch_size[0] for key, val in self.non_tensor_batch.items(): assert isinstance( val, np.ndarray ) and val.dtype == object, 'data in the non_tensor_batch must be a numpy.array with dtype=object' assert val.shape[ 0] == batch_size, f'key {key} length {len(val)} is not equal to batch size {batch_size}' @classmethod def from_single_dict(cls, data: Dict[str, Union[torch.Tensor, np.ndarray]], meta_info=None): tensors = {} non_tensors = {} for key, val in data.items(): if isinstance(val, torch.Tensor): tensors[key] = val elif isinstance(val, np.ndarray): non_tensors[key] = val else: raise ValueError(f'Unsupported type in data {type(val)}') return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info) @classmethod def from_dict(cls, tensors: Dict[str, torch.Tensor], non_tensors=None, meta_info=None, num_batch_dims=1): \"\"\"Create a DataProto from a dict of tensors. This assumes that 1. All the tensor in tensors have the same dim0 2. Only dim0 is the batch dim \"\"\" assert len(tensors) > 0, 'tensors must not be empty' assert num_batch_dims > 0, 'num_batch_dims must be greater than zero' if non_tensors is not None: assert num_batch_dims == 1, 'only support num_batch_dims=1 when non_tensors is not None.' if meta_info is None: meta_info = {} if non_tensors is None: non_tensors = {} assert isinstance(non_tensors, dict) # get and check batch size batch_size = None pivot_key = None for key, tensor in tensors.items(): if batch_size is None: batch_size = tensor.shape[:num_batch_dims] pivot_key = key else: current_batch = tensor.shape[:num_batch_dims] assert batch_size == current_batch, \\ f'Not all the tensor in tensors have the same batch size with batch_dims={num_batch_dims}. Got {pivot_key} has {batch_size}, {key} has {current_batch}' for key, val in non_","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#DataProto.__post_init__","name":"__post_init__","kind":"method","path":"verl/protocol.py","span":"L176-L178","signature":"def __post_init__(self): # perform necessary checking","exported":false,"origin":"ast","body_hash":"8aa3118b41795178c1cd04cf63a3eba420d28d7dcf0c296592d1979e7225702a","body_text":"def __post_init__(self): # perform necessary checking self.check_consistency()","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.__len__","name":"__len__","kind":"method","path":"verl/protocol.py","span":"L180-L187","signature":"def __len__(self)","exported":false,"origin":"ast","body_hash":"b6a685265bb0d9e039281cc1064aa6c673ce892741a92b45a0718d207439f444","body_text":"def __len__(self): if self.batch is not None: return self.batch.batch_size[0] elif self.non_tensor_batch is not None and len(self.non_tensor_batch) > 0: random_key = list(self.non_tensor_batch.keys())[0] return self.non_tensor_batch[random_key].shape[0] else: return 0","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.__getitem__","name":"__getitem__","kind":"method","path":"verl/protocol.py","span":"L189-L192","signature":"def __getitem__(self, item)","exported":false,"origin":"ast","body_hash":"88c548dd7cb243c8d588341989d52505a59acf52d333e4bc35cfc7a117c5c63c","body_text":"def __getitem__(self, item): tensor_data = self.batch[item] non_tensor_data = {key: val[item] for key, val in self.non_tensor_batch.items()} return DataProtoItem(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info)","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.__getstate__","name":"__getstate__","kind":"method","path":"verl/protocol.py","span":"L194-L202","signature":"def __getstate__(self)","exported":false,"origin":"ast","body_hash":"fcc79d44099a1c997864fa7888cc3b71e23468ab83c51319c857de584b8c0e37","body_text":"def __getstate__(self): import io buffer = io.BytesIO() if tensordict.__version__ >= '0.5.0' and self.batch is not None: self.batch = self.batch.contiguous() self.batch = self.batch.consolidate() torch.save(self.batch, buffer) buffer_bytes = buffer.getvalue() return buffer_bytes, self.non_tensor_batch, self.meta_info","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.__setstate__","name":"__setstate__","kind":"method","path":"verl/protocol.py","span":"L204-L213","signature":"def __setstate__(self, data)","exported":false,"origin":"ast","body_hash":"a3b08a9eb2930e51522a07c354e49033e249ae96825837e8e6d250a9cc1e165e","body_text":"def __setstate__(self, data): import io batch_deserialized_bytes, non_tensor_batch, meta_info = data batch_deserialized = io.BytesIO(initial_bytes=batch_deserialized_bytes) batch = torch.load(batch_deserialized, weights_only=False, map_location='cpu' if not torch.cuda.is_available() else None) self.batch = batch self.non_tensor_batch = non_tensor_batch self.meta_info = meta_info","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.save_to_disk","name":"save_to_disk","kind":"method","path":"verl/protocol.py","span":"L215-L217","signature":"def save_to_disk(self, filepath)","exported":true,"origin":"ast","body_hash":"7079cdeb7f96c81d89ae1369e6eda369e33ae724d25575ea46e0fc3800001664","body_text":"def save_to_disk(self, filepath): with open(filepath, 'wb') as f: pickle.dump(self, f)","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.load_from_disk","name":"load_from_disk","kind":"method","path":"verl/protocol.py","span":"L220-L223","signature":"def load_from_disk(filepath) -> 'DataProto'","exported":true,"origin":"ast","body_hash":"ea33402d606d4f0f4a1807d76ec335e5c2c9906f04c0d77320011291f07e3f05","body_text":"def load_from_disk(filepath) -> 'DataProto': with open(filepath, 'rb') as f: data = pickle.load(f) return data","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.print_size","name":"print_size","kind":"method","path":"verl/protocol.py","span":"L225-L240","signature":"def print_size(self, prefix=\"\")","exported":true,"origin":"ast","body_hash":"dad52a42b33550642492b207a5215a569109767f02ba7abb5abcf4b3822306ab","body_text":"def print_size(self, prefix=\"\"): size_of_tensordict = 0 for key, tensor in self.batch.items(): size_of_tensordict += tensor.element_size() * tensor.numel() size_of_numpy_array = 0 for key, numpy_array in self.non_tensor_batch.items(): size_of_numpy_array += numpy_array.nbytes size_of_numpy_array /= 1024**3 size_of_tensordict /= 1024**3 message = f'Size of tensordict: {size_of_tensordict} GB, size of non_tensor_batch: {size_of_numpy_array} GB' if prefix: message = f'{prefix}, ' + message print(message)","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.check_consistency","name":"check_consistency","kind":"method","path":"verl/protocol.py","span":"L242-L263","signature":"def check_consistency(self)","exported":true,"origin":"ast","body_hash":"e1da358297569e28bf36fd00bbeb5884ab75f0adc5b56cb184fed68228d8c392","body_text":"def check_consistency(self): \"\"\"Check the consistency of the DataProto. Mainly for batch and non_tensor_batch We expose this function as a public one so that user can call themselves directly \"\"\" if self.batch is not None: assert len(self.batch.batch_size) == 1, 'only support num_batch_dims=1' if self.non_tensor_batch is not None: for key, val in self.non_tensor_batch.items(): assert isinstance(val, np.ndarray) if self.batch is not None and len(self.non_tensor_batch) != 0: # TODO: we can actually lift this restriction if needed assert len(self.batch.batch_size) == 1, 'only support num_batch_dims=1 when non_tensor_batch is not empty.' batch_size = self.batch.batch_size[0] for key, val in self.non_tensor_batch.items(): assert isinstance( val, np.ndarray ) and val.dtype == object, 'data in the non_tensor_batch must be a numpy.array with dtype=object' assert val.shape[ 0] == batch_size, f'key {key} length {len(val)} is not equal to batch size {batch_size}'","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.from_single_dict","name":"from_single_dict","kind":"method","path":"verl/protocol.py","span":"L266-L278","signature":"def from_single_dict(cls, data: Dict[str, Union[torch.Tensor, np.ndarray]], meta_info=None)","exported":true,"origin":"ast","body_hash":"f98728e333b3d9c4d1563f566f6b0aa1947532cba1e54b9330e7017f44411389","body_text":"def from_single_dict(cls, data: Dict[str, Union[torch.Tensor, np.ndarray]], meta_info=None): tensors = {} non_tensors = {} for key, val in data.items(): if isinstance(val, torch.Tensor): tensors[key] = val elif isinstance(val, np.ndarray): non_tensors[key] = val else: raise ValueError(f'Unsupported type in data {type(val)}') return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info)","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.from_dict","name":"from_dict","kind":"method","path":"verl/protocol.py","span":"L281-L314","signature":"def from_dict(cls, tensors: Dict[str, torch.Tensor], non_tensors=None, meta_info=None, num_batch_dims=1)","exported":true,"origin":"ast","body_hash":"1ad4e3a3f2d190d3b07916286973393f0beba76a63c4076fe715386bd7ea88e7","body_text":"def from_dict(cls, tensors: Dict[str, torch.Tensor], non_tensors=None, meta_info=None, num_batch_dims=1): \"\"\"Create a DataProto from a dict of tensors. This assumes that 1. All the tensor in tensors have the same dim0 2. Only dim0 is the batch dim \"\"\" assert len(tensors) > 0, 'tensors must not be empty' assert num_batch_dims > 0, 'num_batch_dims must be greater than zero' if non_tensors is not None: assert num_batch_dims == 1, 'only support num_batch_dims=1 when non_tensors is not None.' if meta_info is None: meta_info = {} if non_tensors is None: non_tensors = {} assert isinstance(non_tensors, dict) # get and check batch size batch_size = None pivot_key = None for key, tensor in tensors.items(): if batch_size is None: batch_size = tensor.shape[:num_batch_dims] pivot_key = key else: current_batch = tensor.shape[:num_batch_dims] assert batch_size == current_batch, \\ f'Not all the tensor in tensors have the same batch size with batch_dims={num_batch_dims}. Got {pivot_key} has {batch_size}, {key} has {current_batch}' for key, val in non_tensors.items(): non_tensors[key] = np.array(val, dtype=object) tensor_dict = TensorDict(source=tensors, batch_size=batch_size) return cls(batch=tensor_dict, non_tensor_batch=non_tensors, meta_info=meta_info)","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.to","name":"to","kind":"method","path":"verl/protocol.py","span":"L316-L328","signature":"def to(self, device) -> 'DataProto'","exported":true,"origin":"ast","body_hash":"8a05eaf0853163433a85c5c430f874c809688e25553e3282884d73a5bff7056e","body_text":"def to(self, device) -> 'DataProto': \"\"\"move the batch to device Args: device (torch.device, str): torch device Returns: DataProto: the current DataProto \"\"\" if self.batch is not None: self.batch = self.batch.to(device) return self","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.select","name":"select","kind":"method","path":"verl/protocol.py","span":"L330-L363","signature":"def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None, deepcopy=False) -> 'DataProto'","exported":true,"origin":"ast","body_hash":"e6d97f84abc94b4295bd4ac54bde5fcb4f104e9d89277c826cc3ff17a32fcdf0","body_text":"def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None, deepcopy=False) -> 'DataProto': \"\"\"Select a subset of the DataProto via batch_keys and meta_info_keys Args: batch_keys (list, optional): a list of strings indicating the keys in batch to select meta_info_keys (list, optional): a list of keys indicating the meta info to select Returns: DataProto: the DataProto with the selected batch_keys and meta_info_keys \"\"\" # TODO (zhangchi.usc1992) whether to copy if batch_keys is not None: batch_keys = tuple(batch_keys) sub_batch = self.batch.select(*batch_keys) else: sub_batch = self.batch if non_tensor_batch_keys is not None: non_tensor_batch = {key: val for key, val in self.non_tensor_batch.items() if key in non_tensor_batch_keys} else: non_tensor_batch = self.non_tensor_batch if deepcopy: non_tensor_batch = copy.deepcopy(non_tensor_batch) if meta_info_keys is not None: sub_meta_info = {key: val for key, val in self.meta_info.items() if key in meta_info_keys} else: sub_meta_info = self.meta_info if deepcopy: sub_meta_info = copy.deepcopy(sub_meta_info) return DataProto(batch=sub_batch, non_tensor_batch=non_tensor_batch, meta_info=sub_meta_info)","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.pop","name":"pop","kind":"method","path":"verl/protocol.py","span":"L365-L395","signature":"def pop(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None) -> 'DataProto'","exported":true,"origin":"ast","body_hash":"dd9aff86a52004710bf6376f22a9f50cb291fcf79a57590a3b226dc8752546f3","body_text":"def pop(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None) -> 'DataProto': \"\"\"Pop a subset of the DataProto via `batch_keys` and `meta_info_keys` Args: batch_keys (list, optional): a list of strings indicating the keys in batch to pop meta_info_keys (list, optional): a list of keys indicating the meta info to pop Returns: DataProto: the DataProto with the poped batch_keys and meta_info_keys \"\"\" assert batch_keys is not None if meta_info_keys is None: meta_info_keys = [] if non_tensor_batch_keys is None: non_tensor_batch_keys = [] tensors = {} # tensor batch for key in batch_keys: assert key in self.batch.keys() tensors[key] = self.batch.pop(key) non_tensors = {} # non tensor batch for key in non_tensor_batch_keys: assert key in self.non_tensor_batch.keys() non_tensors[key] = self.non_tensor_batch.pop(key) meta_info = {} for key in meta_info_keys: assert key in self.meta_info.keys() meta_info[key] = self.meta_info.pop(key) return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info)","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.rename","name":"rename","kind":"method","path":"verl/protocol.py","span":"L397-L421","signature":"def rename(self, old_keys=None, new_keys=None) -> 'DataProto'","exported":true,"origin":"ast","body_hash":"26e31bc4d8452966264be1725ea63090bf2bede58d40be221be9f0b7e5fd7a0c","body_text":"def rename(self, old_keys=None, new_keys=None) -> 'DataProto': \"\"\" Note that this function only rename the key in the batch \"\"\" def validate_input(keys): if keys is not None: if isinstance(keys, str): keys = [keys] elif isinstance(keys, list): pass else: raise TypeError(f'keys must be a list or a string, but got {type(keys)}') return keys old_keys = validate_input(old_keys) new_keys = validate_input(new_keys) if len(new_keys) != len(old_keys): raise ValueError( f'new_keys and old_keys must have the same length, but got {len(new_keys)} and {len(old_keys)}') self.batch.rename_key_(tuple(old_keys), tuple(new_keys)) return self","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.rename.validate_input","name":"validate_input","kind":"function","path":"verl/protocol.py","span":"L402-L410","signature":"def validate_input(keys)","exported":true,"origin":"ast","body_hash":"d8ef0134550c97e0367eedbb8e63bee1f2b7581f2457145e056e6763068b55b8","body_text":"def validate_input(keys): if keys is not None: if isinstance(keys, str): keys = [keys] elif isinstance(keys, list): pass else: raise TypeError(f'keys must be a list or a string, but got {type(keys)}') return keys","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#DataProto.union","name":"union","kind":"method","path":"verl/protocol.py","span":"L423-L439","signature":"def union(self, other: 'DataProto') -> 'DataProto'","exported":true,"origin":"ast","body_hash":"aa2cc50c72a9ab2cf1880baea813b49fd3fe5bb6ff921f49840581491118f9fa","body_text":"def union(self, other: 'DataProto') -> 'DataProto': \"\"\"Union with another DataProto. Union batch and meta_info separately. Throw an error if - there are conflict keys in batch and they are not equal - the batch size of two data batch is not the same - there are conflict keys in meta_info and they are not the same. Args: other (DataProto): another DataProto to union Returns: DataProto: the DataProto after union \"\"\" self.batch = union_tensor_dict(self.batch, other.batch) self.non_tensor_batch = union_numpy_dict(self.non_tensor_batch, other.non_tensor_batch) self.meta_info = union_two_dict(self.meta_info, other.meta_info) return self","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.make_iterator","name":"make_iterator","kind":"method","path":"verl/protocol.py","span":"L441-L480","signature":"def make_iterator(self, mini_batch_size, epochs, seed=None, dataloader_kwargs=None)","exported":true,"origin":"ast","body_hash":"d88c2a95d51c21c71a40fdd31b35d26cd8173104532aadac26c7c1c617581062","body_text":"def make_iterator(self, mini_batch_size, epochs, seed=None, dataloader_kwargs=None): \"\"\"Make an iterator from the DataProto. This is built upon that TensorDict can be used as a normal Pytorch dataset. See https://pytorch.org/tensordict/tutorials/data_fashion for more details. Args: mini_batch_size (int): mini-batch size when iterating the dataset. We require that ``batch.batch_size[0] % mini_batch_size == 0`` epochs (int): number of epochs when iterating the dataset. dataloader_kwargs: internally, it returns a DataLoader over the batch. The dataloader_kwargs is the kwargs passed to the DataLoader Returns: Iterator: an iterator that yields a mini-batch data at a time. The total number of iteration steps is ``self.batch.batch_size * epochs // mini_batch_size`` \"\"\" assert self.batch.batch_size[0] % mini_batch_size == 0, f\"{self.batch.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) train_dataloader = DataLoader(dataset=self, batch_size=mini_batch_size, collate_fn=collate_fn, generator=generator, **dataloader_kwargs) def get_data(): for _ in range(epochs): for d in train_dataloader: d.meta_info = self.meta_info yield d return iter(get_data())","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.make_iterator.get_data","name":"get_data","kind":"function","path":"verl/protocol.py","span":"L474-L478","signature":"def get_data()","exported":true,"origin":"ast","body_hash":"0baca75ddc8cfd88e7e933921df3e85fbca219a5133146e01021fb42a5348292","body_text":"def get_data(): for _ in range(epochs): for d in train_dataloader: d.meta_info = self.meta_info yield d","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#DataProto.chunk","name":"chunk","kind":"method","path":"verl/protocol.py","span":"L482-L512","signature":"def chunk(self, chunks: int) -> List['DataProto']","exported":true,"origin":"ast","body_hash":"da8f79ee59b08d03e11569d3af498e705bc5a7371418cec0b9f57a34234f1ca8","body_text":"def chunk(self, chunks: int) -> List['DataProto']: \"\"\"Split the batch among dim=0 into chunks. The meta_info is passed to each DataProto after split. Args: chunks (int): the number of chunks to split on dim=0 Returns: List[DataProto]: a list of DataProto after splitting \"\"\" assert len( self) % chunks == 0, f'only support equal chunk. Got size of DataProto {len(self)} and chunk {chunks}.' if self.batch is not None: batch_lst = self.batch.chunk(chunks=chunks, dim=0) else: batch_lst = [None for _ in range(chunks)] non_tensor_batch_lst = [{} for _ in range(chunks)] for key, val in self.non_tensor_batch.items(): assert isinstance(val, np.ndarray) non_tensor_lst = np.array_split(val, chunks) assert len(non_tensor_lst) == chunks for i in range(chunks): non_tensor_batch_lst[i][key] = non_tensor_lst[i] output = [] for i in range(chunks): output.append( DataProto(batch=batch_lst[i], non_tensor_batch=non_tensor_batch_lst[i], meta_info=self.meta_info)) return output","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.concat","name":"concat","kind":"method","path":"verl/protocol.py","span":"L515-L537","signature":"def concat(data: List['DataProto']) -> 'DataProto'","exported":true,"origin":"ast","body_hash":"dbe3b9d1f772d45a18f5a783be013db353fbb4a8573f2a6c0a4b45e6887b4925","body_text":"def concat(data: List['DataProto']) -> 'DataProto': \"\"\"Concat a list of DataProto. The batch is concatenated among dim=0. The meta_info is assumed to be identical and will use the first one. Args: data (List[DataProto]): list of DataProto Returns: DataProto: concatenated DataProto \"\"\" batch_lst = [] for batch in data: batch_lst.append(batch.batch) if batch_lst[0] is not None: new_batch = torch.cat(batch_lst, dim=0) else: new_batch = None non_tensor_batch = list_of_dict_to_dict_of_list(list_of_dict=[d.non_tensor_batch for d in data]) for key, val in non_tensor_batch.items(): non_tensor_batch[key] = np.concatenate(val, axis=0) return DataProto(batch=new_batch, non_tensor_batch=non_tensor_batch, meta_info=data[0].meta_info)","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.reorder","name":"reorder","kind":"method","path":"verl/protocol.py","span":"L539-L545","signature":"def reorder(self, indices)","exported":true,"origin":"ast","body_hash":"59ae3d7e628cdfcd7503bfcebb68491da4ddf246f3967391743fe49e24fbde31","body_text":"def reorder(self, indices): \"\"\" Note that this operation is in-place \"\"\" indices_np = indices.detach().numpy() self.batch = self.batch[indices] self.non_tensor_batch = {key: val[indices_np] for key, val in self.non_tensor_batch.items()}","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProto.repeat","name":"repeat","kind":"method","path":"verl/protocol.py","span":"L547-L589","signature":"def repeat(self, repeat_times=2, interleave=True)","exported":true,"origin":"ast","body_hash":"10d780e95c89802d250ba0156332c6ffafc85923da1fd5c496879b4b29f2d42c","body_text":"def repeat(self, repeat_times=2, interleave=True): \"\"\" Repeat the batch data a specified number of times. Args: repeat_times (int): Number of times to repeat the data. interleave (bool): Whether to interleave the repeated data. Returns: DataProto: A new DataProto with repeated data. \"\"\" if self.batch is not None: if interleave: # Interleave the data repeated_tensors = { key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items() } else: # Stack the data repeated_tensors = { key: tensor.unsqueeze(0).expand(repeat_times, *tensor.shape).reshape(-1, *tensor.shape[1:]) for key, tensor in self.batch.items() } repeated_batch = TensorDict( source=repeated_tensors, batch_size=(self.batch.batch_size[0] * repeat_times,), ) else: repeated_batch = None repeated_non_tensor_batch = {} for key, val in self.non_tensor_batch.items(): if interleave: repeated_non_tensor_batch[key] = np.repeat(val, repeat_times, axis=0) else: repeated_non_tensor_batch[key] = np.tile(val, (repeat_times,) + (1,) * (val.ndim - 1)) return DataProto( batch=repeated_batch, non_tensor_batch=repeated_non_tensor_batch, meta_info=self.meta_info, )","summary_state":"pending","summary":null,"crux":null,"owner":"DataProto"},{"id":"verl/protocol.py#DataProtoFuture","name":"DataProtoFuture","kind":"class","path":"verl/protocol.py","span":"L596-L639","signature":"class DataProtoFuture","exported":true,"origin":"ast","body_hash":"55d045d57e3c09b99af3fc25f8004d0f8ca5035c0b9dfb43f0dfbee2e8a53f35","body_text":"class DataProtoFuture: \"\"\" DataProtoFuture aims to eliminate actual data fetching on driver. By doing so, the driver doesn't have to wait for data so that asynchronous execution becomes possible. DataProtoFuture contains a list of futures from another WorkerGroup of size world_size. - collect_fn is a Callable that reduces the list of futures to a DataProto - dispatch_fn is a Callable that partitions the DataProto into a list of DataProto of size world_size and then select Potential issue: we can optimize dispatch_fn(collect_fn) such that only needed data is fetched on destination - DataProtoFuture only supports directly passing from the output of a method to another input. You can't perform any operation on the DataProtoFuture in driver. \"\"\" collect_fn: Callable futures: List[ray.ObjectRef] dispatch_fn: Callable = None @staticmethod def concat(data: List[ray.ObjectRef]) -> 'DataProtoFuture': output = DataProtoFuture(collect_fn=DataProto.concat, futures=data) return output def chunk(self, chunks: int) -> List['DataProtoFuture']: from functools import partial arg_future_lst = [] for i in range(chunks): # note that we can't directly pass i and chunks def dispatch_fn(x, i, chunks): return x.chunk(chunks=chunks)[i] arg_future = DataProtoFuture(collect_fn=self.collect_fn, dispatch_fn=partial(dispatch_fn, i=i, chunks=chunks), futures=self.futures) arg_future_lst.append(arg_future) return arg_future_lst def get(self): output = ray.get(self.futures) # dp_size. for o in output: assert isinstance(o, DataProto) output = self.collect_fn(output) # select dp, concat if self.dispatch_fn is not None: output = self.dispatch_fn(output) # split in batch dim, select using dp return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#DataProtoFuture.concat","name":"concat","kind":"method","path":"verl/protocol.py","span":"L613-L615","signature":"def concat(data: List[ray.ObjectRef]) -> 'DataProtoFuture'","exported":true,"origin":"ast","body_hash":"13acff9ede17597f6dad119dd773b186f60c1379b8e46b92c7f4e261d8acada4","body_text":"def concat(data: List[ray.ObjectRef]) -> 'DataProtoFuture': output = DataProtoFuture(collect_fn=DataProto.concat, futures=data) return output","summary_state":"pending","summary":null,"crux":null,"owner":"DataProtoFuture"},{"id":"verl/protocol.py#DataProtoFuture.chunk","name":"chunk","kind":"method","path":"verl/protocol.py","span":"L617-L630","signature":"def chunk(self, chunks: int) -> List['DataProtoFuture']","exported":true,"origin":"ast","body_hash":"7de6adb92c7eb2fde5a9549b6eaebcb13aeb8d852dcf7188ffaf9af04ae29dfd","body_text":"def chunk(self, chunks: int) -> List['DataProtoFuture']: from functools import partial arg_future_lst = [] for i in range(chunks): # note that we can't directly pass i and chunks def dispatch_fn(x, i, chunks): return x.chunk(chunks=chunks)[i] arg_future = DataProtoFuture(collect_fn=self.collect_fn, dispatch_fn=partial(dispatch_fn, i=i, chunks=chunks), futures=self.futures) arg_future_lst.append(arg_future) return arg_future_lst","summary_state":"pending","summary":null,"crux":null,"owner":"DataProtoFuture"},{"id":"verl/protocol.py#DataProtoFuture.chunk.dispatch_fn","name":"dispatch_fn","kind":"function","path":"verl/protocol.py","span":"L623-L624","signature":"def dispatch_fn(x, i, chunks)","exported":true,"origin":"ast","body_hash":"0a7a2f1caed9c196d4ac408d4a22e83add5f9ce5de29a066a9885d18a2ae0057","body_text":"def dispatch_fn(x, i, chunks): return x.chunk(chunks=chunks)[i]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/protocol.py#DataProtoFuture.get","name":"get","kind":"method","path":"verl/protocol.py","span":"L632-L639","signature":"def get(self)","exported":true,"origin":"ast","body_hash":"85e5077cb1958f776591e674091d6e4b4e2ff748f82cb9198e73ab66c67821dc","body_text":"def get(self): output = ray.get(self.futures) # dp_size. for o in output: assert isinstance(o, DataProto) output = self.collect_fn(output) # select dp, concat if self.dispatch_fn is not None: output = self.dispatch_fn(output) # split in batch dim, select using dp return output","summary_state":"pending","summary":null,"crux":null,"owner":"DataProtoFuture"}],"rawEdges":[{"source":"verl/protocol.py","relation":"imports","specifier":"pickle","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"copy","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"dataclasses","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"typing","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"torch","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"tensordict","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"tensordict","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"torch.utils.data","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"verl.utils.py_functional","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"calls","name":"set","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"calls","name":"set_lazy_legacy","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#pad_dataproto_to_divisor","file":"verl/protocol.py"},{"source":"verl/protocol.py#pad_dataproto_to_divisor","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#pad_dataproto_to_divisor","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#pad_dataproto_to_divisor","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#pad_dataproto_to_divisor","relation":"calls","name":"concat","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#unpad_dataproto","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#union_tensor_dict","file":"verl/protocol.py"},{"source":"verl/protocol.py#union_tensor_dict","relation":"calls","name":"keys","viaMember":true,"file":"verl/protocol.py","recvType":"TensorDict"},{"source":"verl/protocol.py#union_tensor_dict","relation":"calls","name":"keys","viaMember":true,"file":"verl/protocol.py","recvType":"TensorDict"},{"source":"verl/protocol.py#union_tensor_dict","relation":"calls","name":"equal","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#union_numpy_dict","file":"verl/protocol.py"},{"source":"verl/protocol.py#union_numpy_dict","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#union_numpy_dict","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#union_numpy_dict","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#union_numpy_dict","relation":"calls","name":"all","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#list_of_dict_to_dict_of_list","file":"verl/protocol.py"},{"source":"verl/protocol.py#list_of_dict_to_dict_of_list","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#list_of_dict_to_dict_of_list","relation":"calls","name":"keys","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#list_of_dict_to_dict_of_list","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#list_of_dict_to_dict_of_list","relation":"calls","name":"append","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#fold_batch_dim","file":"verl/protocol.py"},{"source":"verl/protocol.py#fold_batch_dim","relation":"calls","name":"view","viaMember":true,"file":"verl/protocol.py","recvType":"TensorDict"},{"source":"verl/protocol.py#fold_batch_dim","relation":"calls","name":"auto_batch_size_","viaMember":true,"file":"verl/protocol.py","recvType":"TensorDict"},{"source":"verl/protocol.py#fold_batch_dim","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#fold_batch_dim","relation":"calls","name":"reshape","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#fold_batch_dim","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#unfold_batch_dim","file":"verl/protocol.py"},{"source":"verl/protocol.py#unfold_batch_dim","relation":"calls","name":"auto_batch_size_","viaMember":true,"file":"verl/protocol.py","recvType":"TensorDict"},{"source":"verl/protocol.py#unfold_batch_dim","relation":"calls","name":"view","viaMember":true,"file":"verl/protocol.py","recvType":"TensorDict"},{"source":"verl/protocol.py#unfold_batch_dim","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#unfold_batch_dim","relation":"calls","name":"reshape","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#unfold_batch_dim","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#collate_fn","file":"verl/protocol.py"},{"source":"verl/protocol.py#collate_fn","relation":"calls","name":"append","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#collate_fn","relation":"calls","name":"append","viaMember":true,"file":"verl/protocol.py","recvType":"list_of_dict_to_dict_of_list"},{"source":"verl/protocol.py#collate_fn","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#collate_fn","relation":"calls","name":"stack","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#collate_fn","relation":"calls","name":"list_of_dict_to_dict_of_list","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#collate_fn","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"list_of_dict_to_dict_of_list"},{"source":"verl/protocol.py#collate_fn","relation":"calls","name":"array","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#collate_fn","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#DataProtoItem","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoItem","relation":"calls","name":"field","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoItem","relation":"calls","name":"field","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#DataProto","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"calls","name":"field","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"calls","name":"field","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.__post_init__","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__post_init__","relation":"calls","name":"check_consistency","viaMember":true,"file":"verl/protocol.py","recvType":"DataProto"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.__len__","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__len__","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__len__","relation":"calls","name":"list","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__len__","relation":"calls","name":"keys","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.__getitem__","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__getitem__","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto.__getitem__","relation":"calls","name":"DataProtoItem","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.__getstate__","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"io","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__getstate__","relation":"calls","name":"BytesIO","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__getstate__","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto.__getstate__","relation":"calls","name":"consolidate","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto.__getstate__","relation":"calls","name":"save","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__getstate__","relation":"calls","name":"getvalue","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.__setstate__","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"io","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__setstate__","relation":"calls","name":"BytesIO","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__setstate__","relation":"calls","name":"load","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.__setstate__","relation":"calls","name":"is_available","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.save_to_disk","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.save_to_disk","relation":"calls","name":"open","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.save_to_disk","relation":"calls","name":"dump","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.load_from_disk","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.load_from_disk","relation":"calls","name":"open","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.load_from_disk","relation":"calls","name":"load","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.print_size","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.print_size","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto.print_size","relation":"calls","name":"element_size","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.print_size","relation":"calls","name":"numel","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.print_size","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto.print_size","relation":"calls","name":"print","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.check_consistency","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.check_consistency","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.check_consistency","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto.check_consistency","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.check_consistency","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.check_consistency","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.check_consistency","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto.check_consistency","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.check_consistency","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.from_single_dict","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_single_dict","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_single_dict","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_single_dict","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_single_dict","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_single_dict","relation":"calls","name":"type","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_single_dict","relation":"calls","name":"from_dict","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.from_dict","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_dict","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_dict","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_dict","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_dict","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_dict","relation":"calls","name":"array","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_dict","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.from_dict","relation":"calls","name":"cls","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.to","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.to","relation":"calls","name":"to","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.select","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.select","relation":"calls","name":"tuple","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.select","relation":"calls","name":"select","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto.select","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto.select","relation":"calls","name":"deepcopy","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.select","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_two_dict"},{"source":"verl/protocol.py#DataProto.select","relation":"calls","name":"deepcopy","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.select","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.pop","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.pop","relation":"calls","name":"keys","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto.pop","relation":"calls","name":"pop","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto.pop","relation":"calls","name":"keys","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto.pop","relation":"calls","name":"pop","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto.pop","relation":"calls","name":"keys","viaMember":true,"file":"verl/protocol.py","recvType":"union_two_dict"},{"source":"verl/protocol.py#DataProto.pop","relation":"calls","name":"pop","viaMember":true,"file":"verl/protocol.py","recvType":"union_two_dict"},{"source":"verl/protocol.py#DataProto.pop","relation":"calls","name":"from_dict","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.rename","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename","relation":"contains","targetId":"verl/protocol.py#DataProto.rename.validate_input","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename.validate_input","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename.validate_input","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename.validate_input","relation":"calls","name":"TypeError","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename.validate_input","relation":"calls","name":"type","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename","relation":"calls","name":"validate_input","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename","relation":"calls","name":"validate_input","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename","relation":"calls","name":"rename_key_","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto.rename","relation":"calls","name":"tuple","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.rename","relation":"calls","name":"tuple","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.union","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.union","relation":"calls","name":"union_tensor_dict","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.union","relation":"calls","name":"union_numpy_dict","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.union","relation":"calls","name":"union_two_dict","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.make_iterator","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.make_iterator","relation":"calls","name":"Generator","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.make_iterator","relation":"calls","name":"manual_seed","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.make_iterator","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.make_iterator","relation":"calls","name":"DataLoader","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.make_iterator","relation":"contains","targetId":"verl/protocol.py#DataProto.make_iterator.get_data","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.make_iterator.get_data","relation":"calls","name":"range","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.make_iterator","relation":"calls","name":"iter","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.make_iterator","relation":"calls","name":"get_data","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.chunk","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"chunk","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"range","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"range","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"array_split","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"len","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"range","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"range","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"append","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.chunk","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.concat","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.concat","relation":"calls","name":"append","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.concat","relation":"calls","name":"cat","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.concat","relation":"calls","name":"list_of_dict_to_dict_of_list","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.concat","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"list_of_dict_to_dict_of_list"},{"source":"verl/protocol.py#DataProto.concat","relation":"calls","name":"concatenate","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.concat","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.reorder","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.reorder","relation":"calls","name":"numpy","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.reorder","relation":"calls","name":"detach","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.reorder","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto","relation":"contains","targetId":"verl/protocol.py#DataProto.repeat","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"repeat_interleave","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"reshape","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"expand","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_tensor_dict"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"items","viaMember":true,"file":"verl/protocol.py","recvType":"union_numpy_dict"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"repeat","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"tile","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProto.repeat","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"ray","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"contains","targetId":"verl/protocol.py#DataProtoFuture","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture","relation":"contains","targetId":"verl/protocol.py#DataProtoFuture.concat","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture.concat","relation":"calls","name":"DataProtoFuture","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture","relation":"contains","targetId":"verl/protocol.py#DataProtoFuture.chunk","file":"verl/protocol.py"},{"source":"verl/protocol.py","relation":"imports","specifier":"functools","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture.chunk","relation":"calls","name":"range","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture.chunk","relation":"contains","targetId":"verl/protocol.py#DataProtoFuture.chunk.dispatch_fn","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture.chunk.dispatch_fn","relation":"calls","name":"chunk","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture.chunk","relation":"calls","name":"DataProtoFuture","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture.chunk","relation":"calls","name":"partial","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture.chunk","relation":"calls","name":"append","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture","relation":"contains","targetId":"verl/protocol.py#DataProtoFuture.get","file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture.get","relation":"calls","name":"get","viaMember":true,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture.get","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/protocol.py"},{"source":"verl/protocol.py#DataProtoFuture.get","relation":"calls","name":"collect_fn","viaMember":true,"file":"verl/protocol.py","recvType":"DataProtoFuture"},{"source":"verl/protocol.py#DataProtoFuture.get","relation":"calls","name":"dispatch_fn","viaMember":true,"file":"verl/protocol.py","recvType":"DataProtoFuture"}]},"verl/single_controller/__init__.py":{"size":787,"mtimeMs":1789275830310.2026,"hash":"504089640d8e047b8b53d032a9278ad3374c0efaae8c2a177468c7b2ceb7a02f","nodes":[{"id":"verl/single_controller/__init__.py","name":"__init__.py","kind":"file","path":"verl/single_controller/__init__.py","span":"L1-L21","signature":null,"exported":true,"origin":"ast","body_hash":"504089640d8e047b8b53d032a9278ad3374c0efaae8c2a177468c7b2ceb7a02f","chars":787,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 version_folder = os.path.dirname(os.path.join(os.path.abspath(__file__))) with open(os.path.join(version_folder, 'version/version')) as f: __version__ = f.read().strip()"}],"rawEdges":[{"source":"verl/single_controller/__init__.py","relation":"imports","specifier":"os","file":"verl/single_controller/__init__.py"},{"source":"verl/single_controller/__init__.py","relation":"calls","name":"dirname","viaMember":true,"file":"verl/single_controller/__init__.py"},{"source":"verl/single_controller/__init__.py","relation":"calls","name":"join","viaMember":true,"file":"verl/single_controller/__init__.py"},{"source":"verl/single_controller/__init__.py","relation":"calls","name":"abspath","viaMember":true,"file":"verl/single_controller/__init__.py"},{"source":"verl/single_controller/__init__.py","relation":"calls","name":"open","viaMember":false,"file":"verl/single_controller/__init__.py"},{"source":"verl/single_controller/__init__.py","relation":"calls","name":"join","viaMember":true,"file":"verl/single_controller/__init__.py"},{"source":"verl/single_controller/__init__.py","relation":"calls","name":"strip","viaMember":true,"file":"verl/single_controller/__init__.py"},{"source":"verl/single_controller/__init__.py","relation":"calls","name":"read","viaMember":true,"file":"verl/single_controller/__init__.py"}]},"verl/single_controller/base/__init__.py":{"size":699,"mtimeMs":1789275830310.2026,"hash":"2ccb73c92abfe5b47ea9af01ae0d9a2877553267d5c9e702f75c12071b7fea6b","nodes":[{"id":"verl/single_controller/base/__init__.py","name":"__init__.py","kind":"file","path":"verl/single_controller/base/__init__.py","span":"L1-L17","signature":null,"exported":true,"origin":"ast","body_hash":"2ccb73c92abfe5b47ea9af01ae0d9a2877553267d5c9e702f75c12071b7fea6b","chars":699,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 .worker import Worker from .worker_group import WorkerGroup, ClassWithInitArgs, ResourcePool"}],"rawEdges":[{"source":"verl/single_controller/base/__init__.py","relation":"imports","specifier":".worker","file":"verl/single_controller/base/__init__.py"},{"source":"verl/single_controller/base/__init__.py","relation":"imports","specifier":".worker_group","file":"verl/single_controller/base/__init__.py"}]},"verl/single_controller/base/decorator.py":{"size":15531,"mtimeMs":1789275830310.2026,"hash":"b6e4cd3c8ad3c4fa2eea07df3ef45344bb5dc2e1e8357c0b94a9d19a6bd4daac","nodes":[{"id":"verl/single_controller/base/decorator.py","name":"decorator.py","kind":"file","path":"verl/single_controller/base/decorator.py","span":"L1-L411","signature":null,"exported":true,"origin":"ast","body_hash":"b6e4cd3c8ad3c4fa2eea07df3ef45344bb5dc2e1e8357c0b94a9d19a6bd4daac","chars":15531,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 enum import Enum from functools import wraps from typing import Dict, List, Tuple from types import FunctionType from verl.protocol import DataProtoFuture # here we add a magic number of avoid user-defined function already have this attribute MAGIC_ATTR = 'attrs_3141562937'"},{"id":"verl/single_controller/base/decorator.py#Dispatch","name":"Dispatch","kind":"class","path":"verl/single_controller/base/decorator.py","span":"L25-L37","signature":"class Dispatch(Enum)","exported":true,"origin":"ast","body_hash":"4697aeb6e34b1615262753c03d6dfc245617e17f85466d040871e1a2bc5754de","body_text":"class Dispatch(Enum): RANK_ZERO = 0 ONE_TO_ALL = 1 ALL_TO_ALL = 2 MEGATRON_COMPUTE = 3 MEGATRON_PP_AS_DP = 4 MEGATRON_PP_ONLY = 5 MEGATRON_COMPUTE_PROTO = 6 MEGATRON_PP_AS_DP_PROTO = 7 DP_COMPUTE = 8 DP_COMPUTE_PROTO = 9 DP_COMPUTE_PROTO_WITH_FUNC = 10 DP_COMPUTE_METRIC = 11","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#Execute","name":"Execute","kind":"class","path":"verl/single_controller/base/decorator.py","span":"L40-L42","signature":"class Execute(Enum)","exported":true,"origin":"ast","body_hash":"5433429b319065c81b65b425ce3a1c0c6e0856dba0638b4a54a0262549cc6c72","body_text":"class Execute(Enum): ALL = 0 RANK_ZERO = 1","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#_split_args_kwargs_data_proto","name":"_split_args_kwargs_data_proto","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L45-L57","signature":"def _split_args_kwargs_data_proto(chunks, *args, **kwargs)","exported":false,"origin":"ast","body_hash":"ac769a90a0e5e18d867f612a9f2e2c49f0259a4ce6383f907300cdff9b662099","body_text":"def _split_args_kwargs_data_proto(chunks, *args, **kwargs): from verl.protocol import DataProto, DataProtoFuture splitted_args = [] for arg in args: assert isinstance(arg, (DataProto, DataProtoFuture)) splitted_args.append(arg.chunk(chunks=chunks)) splitted_kwargs = {} for key, val in kwargs.items(): assert isinstance(val, (DataProto, DataProtoFuture)) splitted_kwargs[key] = val.chunk(chunks=chunks) return splitted_args, splitted_kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#dispatch_one_to_all","name":"dispatch_one_to_all","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L60-L63","signature":"def dispatch_one_to_all(worker_group, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"cf257f3efdc1e55e6cd4d3e2f8de203a36e36f911488e1a62d8f9164f9e0b93b","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#dispatch_all_to_all","name":"dispatch_all_to_all","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L66-L67","signature":"def dispatch_all_to_all(worker_group, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"02102a041713e5cff86786d0adf918db226829becf225c170f30128049551227","body_text":"def dispatch_all_to_all(worker_group, *args, **kwargs): return args, kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#collect_all_to_all","name":"collect_all_to_all","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L70-L71","signature":"def collect_all_to_all(worker_group, output)","exported":true,"origin":"ast","body_hash":"0803a44ac31a4e45e48aea4401f65b32defc3c5e3fb5fcfb79ca16f8cd3fa7d7","body_text":"def collect_all_to_all(worker_group, output): return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","name":"dispatch_megatron_compute","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L74-L100","signature":"def dispatch_megatron_compute(worker_group, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"fd48bf2441850af83f467f5f1eed7699f53f2e90092a7937b36b7338e8180753","body_text":"def dispatch_megatron_compute(worker_group, *args, **kwargs): \"\"\" User passes in dp data. The data is dispatched to all tp/pp ranks with the same dp \"\"\" from verl.single_controller.base.megatron.worker_group import MegatronWorkerGroup assert isinstance(worker_group, MegatronWorkerGroup), f'worker_group must be MegatronWorkerGroup, Got {type(worker_group)}' all_args = [] for arg in args: assert isinstance(arg, (Tuple, List)) and len(arg) == worker_group.dp_size transformed_args = [] for i in range(worker_group.world_size): local_dp_rank = worker_group.get_megatron_rank_info(rank=i).dp_rank transformed_args.append(arg[local_dp_rank]) all_args.append(transformed_args) all_args = tuple(all_args) all_kwargs = {} for k, v in kwargs.items(): assert isinstance(v, (Tuple, List)) and len(v) == worker_group.dp_size transformed_v = [] for i in range(worker_group.world_size): local_dp_rank = worker_group.get_megatron_rank_info(rank=i).dp_rank transformed_v.append(v[local_dp_rank]) all_kwargs[k] = transformed_v return all_args, all_kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#collect_megatron_compute","name":"collect_megatron_compute","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L103-L115","signature":"def collect_megatron_compute(worker_group, output)","exported":true,"origin":"ast","body_hash":"1f0c2a437abc8de8aa713fd722b51c4efde25da78648268ed46f017614ee9810","body_text":"def collect_megatron_compute(worker_group, output): \"\"\" Only collect the data from the tp=0 and pp=last and every dp ranks \"\"\" from verl.single_controller.base.megatron.worker_group import MegatronWorkerGroup assert isinstance(worker_group, MegatronWorkerGroup) output_in_dp = [] pp_size = worker_group.get_megatron_global_info().pp_size for global_rank in range(worker_group.world_size): local_rank_info = worker_group.get_megatron_rank_info(rank=global_rank) if local_rank_info.tp_rank == 0 and local_rank_info.pp_rank == pp_size - 1: output_in_dp.append(output[global_rank]) return output_in_dp","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#dispatch_megatron_compute_data_proto","name":"dispatch_megatron_compute_data_proto","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L118-L126","signature":"def dispatch_megatron_compute_data_proto(worker_group, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"9477a2906f548968a4c7a85143373f3fe4674170280e3f1772739fbfb5d9533a","body_text":"def dispatch_megatron_compute_data_proto(worker_group, *args, **kwargs): \"\"\" All the args and kwargs must be DataProto. The batch will be chunked by dp_size and passed to each rank \"\"\" from verl.single_controller.base.megatron.worker_group import MegatronWorkerGroup assert isinstance(worker_group, MegatronWorkerGroup) splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(worker_group.dp_size, *args, **kwargs) return dispatch_megatron_compute(worker_group, *splitted_args, **splitted_kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#_concat_data_proto_or_future","name":"_concat_data_proto_or_future","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L129-L144","signature":"def _concat_data_proto_or_future(output: List)","exported":false,"origin":"ast","body_hash":"44fb99df7743ebaeb453e1ab3ec67b243b6852aa1fadf21d86e291e4ca3e5741","body_text":"def _concat_data_proto_or_future(output: List): from verl.protocol import DataProto, DataProtoFuture import ray # make sure all the elements in output has the same type for o in output: assert type(o) == type(output[0]) o = output[0] if isinstance(o, DataProto): return DataProto.concat(output) elif isinstance(o, ray.ObjectRef): return DataProtoFuture.concat(output) else: raise NotImplementedError","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#collect_megatron_compute_data_proto","name":"collect_megatron_compute_data_proto","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L147-L158","signature":"def collect_megatron_compute_data_proto(worker_group, output)","exported":true,"origin":"ast","body_hash":"1150d50aff0bb331cbee7f85da7afe696c46b716d245f6fe62ebdf8b957636ca","body_text":"def collect_megatron_compute_data_proto(worker_group, output): \"\"\" Each output must be a DataProto. We concat the dim=0 of output \"\"\" from verl.protocol import DataProto import ray output = collect_megatron_compute(worker_group, output) for o in output: assert isinstance(o, (DataProto, ray.ObjectRef)), f\"expecting {o} to be DataProto, but got {type(o)}\" return _concat_data_proto_or_future(output)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","name":"dispatch_megatron_pp_as_dp","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L161-L206","signature":"def dispatch_megatron_pp_as_dp(worker_group, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"9ac603ed75d6cd8f43f579493e945b2d8e68187b690417f32100782f1b90c2d7","body_text":"def dispatch_megatron_pp_as_dp(worker_group, *args, **kwargs): \"\"\" treat pp as dp. \"\"\" from verl.single_controller.base.megatron.worker_group import MegatronWorkerGroup assert isinstance(worker_group, MegatronWorkerGroup) pp_size = worker_group.pp_size dp_size = worker_group.dp_size pp_dp_size = pp_size * dp_size all_args = [] for arg in args: assert isinstance(arg, (List, Tuple)) and len(arg) == pp_dp_size transformed_args = [] for i in range(worker_group.world_size): local_dp_rank = worker_group.get_megatron_rank_info(rank=i).dp_rank local_pp_rank = worker_group.get_megatron_rank_info(rank=i).pp_rank # compute the rank in arg. Note that the order is dp then pp # Also note that the outputs within a pp group will be firstly allgathered, then only the output of pp0 will be collected. # For pp=2 dp=4, a batch of data \"ABCDEFGH\" should be dispatched and collected in below order: # dispatch: pp_allgther: collect: # dp 0 1 2 3 dp 0 1 2 3 # pp +---------+ pp +-------------+ # 0 | A C E G | 0 | AB CD EF GH | ABCDEFGH # 1 | B D F H | 1 | AB CD EF GH | # +---------+ +-------------+ arg_rank = local_dp_rank * worker_group.pp_size + local_pp_rank transformed_args.append(arg[arg_rank]) all_args.append(transformed_args) all_args = tuple(all_args) all_kwargs = {} for k, v in kwargs.items(): assert isinstance(v, (List, Tuple)) and len(v) == pp_dp_size, f'expect len(v)=={pp_dp_size}, got {len(v)}' transformed_v = [] for i in range(worker_group.world_size): local_dp_rank = worker_group.get_megatron_rank_info(rank=i).dp_rank local_pp_rank = worker_group.get_megatron_rank_info(rank=i).pp_rank # compute the rank in arg. Note that the order is dp then pp arg_rank = local_dp_rank * worker_group.pp_size + local_pp_rank transformed_v.append(v[arg_rank]) all_kwargs[k] = transformed_v return all_args, all_kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp","name":"collect_megatron_pp_as_dp","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L209-L220","signature":"def collect_megatron_pp_as_dp(worker_group, output)","exported":true,"origin":"ast","body_hash":"eac8b2174b7d9fdf9344e1fec5dbdf30ffebb939662282fca6d6e64b7cd80bfc","body_text":"def collect_megatron_pp_as_dp(worker_group, output): \"\"\" treat pp as dp. Only collect data on tp=0 \"\"\" from verl.single_controller.base.megatron.worker_group import MegatronWorkerGroup assert isinstance(worker_group, MegatronWorkerGroup) output_in_dp = [] for global_rank in range(worker_group.world_size): local_rank_info = worker_group.get_megatron_rank_info(rank=global_rank) if local_rank_info.tp_rank == 0 and local_rank_info.pp_rank == 0: output_in_dp.append(output[global_rank]) return output_in_dp","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#collect_megatron_pp_only","name":"collect_megatron_pp_only","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L223-L234","signature":"def collect_megatron_pp_only(worker_group, output)","exported":true,"origin":"ast","body_hash":"10d044338db0ce102ad0cd74f87c4e3dd3e3295be7f748072be35e799f824db1","body_text":"def collect_megatron_pp_only(worker_group, output): \"\"\" Only collect output of megatron pp. This is useful when examine weight names as they are identical in tp/dp \"\"\" from verl.single_controller.base.megatron.worker_group import MegatronWorkerGroup assert isinstance(worker_group, MegatronWorkerGroup) output_in_pp = [] for global_rank in range(worker_group.world_size): local_rank_info = worker_group.get_megatron_rank_info(rank=global_rank) if local_rank_info.tp_rank == 0 and local_rank_info.dp_rank == 0: output_in_pp.append(output[global_rank]) return output_in_pp","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp_data_proto","name":"dispatch_megatron_pp_as_dp_data_proto","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L237-L243","signature":"def dispatch_megatron_pp_as_dp_data_proto(worker_group, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"aa5e0d832bd738d808f4fa4f7fb0e8992d4c184f2c59a55b25521b8c4781a572","body_text":"def dispatch_megatron_pp_as_dp_data_proto(worker_group, *args, **kwargs): from verl.single_controller.base.megatron.worker_group import MegatronWorkerGroup assert isinstance(worker_group, MegatronWorkerGroup) pp_dp_size = worker_group.dp_size * worker_group.pp_size splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(pp_dp_size, *args, **kwargs) return dispatch_megatron_pp_as_dp(worker_group, *splitted_args, **splitted_kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp_data_proto","name":"collect_megatron_pp_as_dp_data_proto","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L246-L252","signature":"def collect_megatron_pp_as_dp_data_proto(worker_group, output)","exported":true,"origin":"ast","body_hash":"65d1b6c7fa98793a9efff756a75c7acc38aed03a2deb791195512fb8268ec40d","body_text":"def collect_megatron_pp_as_dp_data_proto(worker_group, output): from verl.protocol import DataProto from verl.single_controller.base.megatron.worker_group import MegatronWorkerGroup assert isinstance(worker_group, MegatronWorkerGroup) output = collect_megatron_pp_as_dp(worker_group, output) return _concat_data_proto_or_future(output)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#dispatch_dp_compute","name":"dispatch_dp_compute","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L255-L262","signature":"def dispatch_dp_compute(worker_group, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"fd09cc0ff316d40c6c32d15399dceedcb93056f73a4acd4b959a994a4b046d2d","body_text":"def dispatch_dp_compute(worker_group, *args, **kwargs): from verl.single_controller.base.worker_group import WorkerGroup assert isinstance(worker_group, WorkerGroup) for arg in args: assert isinstance(arg, (Tuple, List)) and len(arg) == worker_group.world_size for k, v in kwargs.items(): assert isinstance(v, (Tuple, List)) and len(v) == worker_group.world_size return args, kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#collect_dp_compute","name":"collect_dp_compute","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L265-L269","signature":"def collect_dp_compute(worker_group, output)","exported":true,"origin":"ast","body_hash":"560379a497ec1d60463e8c8ba1155deaac5d0faee2d1160af5e09668a9e3168a","body_text":"def collect_dp_compute(worker_group, output): from verl.single_controller.base.worker_group import WorkerGroup assert isinstance(worker_group, WorkerGroup) assert len(output) == worker_group.world_size return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#dispatch_dp_compute_data_proto","name":"dispatch_dp_compute_data_proto","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L272-L276","signature":"def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"c14a8ff12e1c50860ad31d3f44fd03da5867450853c5bdfbee3faf28ef16819c","body_text":"def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs): from verl.single_controller.base.worker_group import WorkerGroup assert isinstance(worker_group, WorkerGroup) splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(worker_group.world_size, *args, **kwargs) return splitted_args, splitted_kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#dispatch_dp_compute_data_proto_with_func","name":"dispatch_dp_compute_data_proto_with_func","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L279-L286","signature":"def dispatch_dp_compute_data_proto_with_func(worker_group, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"3ea9f0cadd3f8a0bf4936239b7bd3426668b1f281914cd56e207abecb6c45eeb","body_text":"def dispatch_dp_compute_data_proto_with_func(worker_group, *args, **kwargs): from verl.single_controller.base.worker_group import WorkerGroup assert isinstance(worker_group, WorkerGroup) assert type(args[0]) == FunctionType # NOTE: The first one args is a function! splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(worker_group.world_size, *args[1:], **kwargs) splitted_args_with_func = [[args[0]] * worker_group.world_size] + splitted_args return splitted_args_with_func, splitted_kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#collect_dp_compute_data_proto","name":"collect_dp_compute_data_proto","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L289-L297","signature":"def collect_dp_compute_data_proto(worker_group, output)","exported":true,"origin":"ast","body_hash":"fa2783bb5bc337f7f6bc4dbf08763de35e215467ac778a27ece6214f60d0591d","body_text":"def collect_dp_compute_data_proto(worker_group, output): from verl.protocol import DataProto import ray for o in output: assert isinstance(o, (DataProto, ray.ObjectRef)), f\"expecting {o} to be DataProto, but got {type(o)}\" output = collect_dp_compute(worker_group, output) return _concat_data_proto_or_future(output)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#get_predefined_dispatch_fn","name":"get_predefined_dispatch_fn","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L300-L347","signature":"def get_predefined_dispatch_fn(dispatch_mode)","exported":true,"origin":"ast","body_hash":"8db886873baba4eec32dfabb1ca10e0ce5793a287d0617b44c166318c44e6cb5","body_text":"def get_predefined_dispatch_fn(dispatch_mode): predefined_dispatch_mode_fn = { Dispatch.ONE_TO_ALL: { 'dispatch_fn': dispatch_one_to_all, 'collect_fn': collect_all_to_all, }, Dispatch.ALL_TO_ALL: { 'dispatch_fn': dispatch_all_to_all, 'collect_fn': collect_all_to_all, }, Dispatch.MEGATRON_COMPUTE: { 'dispatch_fn': dispatch_megatron_compute, 'collect_fn': collect_megatron_compute, }, Dispatch.MEGATRON_PP_AS_DP: { 'dispatch_fn': dispatch_megatron_pp_as_dp, 'collect_fn': collect_megatron_pp_as_dp, }, Dispatch.MEGATRON_PP_ONLY: { 'dispatch_fn': dispatch_one_to_all, 'collect_fn': collect_megatron_pp_only }, Dispatch.MEGATRON_COMPUTE_PROTO: { 'dispatch_fn': dispatch_megatron_compute_data_proto, 'collect_fn': collect_megatron_compute_data_proto }, Dispatch.MEGATRON_PP_AS_DP_PROTO: { 'dispatch_fn': dispatch_megatron_pp_as_dp_data_proto, 'collect_fn': collect_megatron_pp_as_dp_data_proto }, Dispatch.DP_COMPUTE: { 'dispatch_fn': dispatch_dp_compute, 'collect_fn': collect_dp_compute }, Dispatch.DP_COMPUTE_PROTO: { 'dispatch_fn': dispatch_dp_compute_data_proto, 'collect_fn': collect_dp_compute_data_proto }, Dispatch.DP_COMPUTE_PROTO_WITH_FUNC: { 'dispatch_fn': dispatch_dp_compute_data_proto_with_func, 'collect_fn': collect_dp_compute_data_proto }, Dispatch.DP_COMPUTE_METRIC: { 'dispatch_fn': dispatch_dp_compute_data_proto, 'collect_fn': collect_dp_compute } } return predefined_dispatch_mode_fn[dispatch_mode]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#get_predefined_execute_fn","name":"get_predefined_execute_fn","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L350-L363","signature":"def get_predefined_execute_fn(execute_mode)","exported":true,"origin":"ast","body_hash":"e53992be261d0b3e28cda19a0e53576f3233fdd6e6d414b6a89749b82716a20b","body_text":"def get_predefined_execute_fn(execute_mode): \"\"\" Note that here we only asks execute_all and execute_rank_zero to be implemented Leave the choice of how these two functions handle argument 'blocking' to users \"\"\" predefined_execute_mode_fn = { Execute.ALL: { 'execute_fn_name': 'execute_all' }, Execute.RANK_ZERO: { 'execute_fn_name': 'execute_rank_zero' } } return predefined_execute_mode_fn[execute_mode]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#_check_dispatch_mode","name":"_check_dispatch_mode","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L366-L372","signature":"def _check_dispatch_mode(dispatch_mode)","exported":false,"origin":"ast","body_hash":"36714acabafcc3e7a08509b5b7666b8b421c7061e6ec20837c42d3ce3e9ea279","body_text":"def _check_dispatch_mode(dispatch_mode): assert isinstance(dispatch_mode, (Dispatch, Dict)), f'dispatch_mode must be a Dispatch or a Dict. Got {dispatch_mode}' if isinstance(dispatch_mode, Dict): necessary_keys = ['dispatch_fn', 'collect_fn'] for key in necessary_keys: assert key in dispatch_mode, f'key {key} should be in dispatch_mode if it is a dictionary'","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#_check_execute_mode","name":"_check_execute_mode","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L375-L376","signature":"def _check_execute_mode(execute_mode)","exported":false,"origin":"ast","body_hash":"fb15a7a6a5c7c23a69b5acc7ba3c3d9ef8382c49df043b44ac5f9846f3891610","body_text":"def _check_execute_mode(execute_mode): assert isinstance(execute_mode, Execute), f'execute_mode must be a Execute. Got {execute_mode}'","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#_materialize_futures","name":"_materialize_futures","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L379-L391","signature":"def _materialize_futures(*args, **kwargs)","exported":false,"origin":"ast","body_hash":"656177e62e49238bf76fd81778db687ebb97250e3bdba67331672b435eedd35d","body_text":"def _materialize_futures(*args, **kwargs): new_args = [] for arg in args: if isinstance(arg, DataProtoFuture): arg = arg.get() # add more type to materialize new_args.append(arg) for k, v in kwargs.items(): if isinstance(v, DataProtoFuture): kwargs[k] = v.get() new_args = tuple(new_args) return new_args, kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#register","name":"register","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L394-L410","signature":"def register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, blocking=True, materialize_futures=True)","exported":true,"origin":"ast","body_hash":"2d7d1b14634f64eca51aaa462161e3ba958c97d5ec306854dd8c3e0bbc3116f1","body_text":"def register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, blocking=True, materialize_futures=True): _check_dispatch_mode(dispatch_mode=dispatch_mode) _check_execute_mode(execute_mode=execute_mode) 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#register.decorator","name":"decorator","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L398-L408","signature":"def decorator(func)","exported":true,"origin":"ast","body_hash":"0e1cacf2e74861668e9226c9c61b38c0e81b8375839cf61787fa1d98091d7403","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/decorator.py#register.decorator.inner","name":"inner","kind":"function","path":"verl/single_controller/base/decorator.py","span":"L401-L404","signature":"def inner(*args, **kwargs)","exported":true,"origin":"ast","body_hash":"00e50d2ed2d098dc24a931a9f991e1790c224b29e46dc49cef594396eb17ef41","body_text":"def inner(*args, **kwargs): if materialize_futures: args, kwargs = _materialize_futures(*args, **kwargs) return func(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"enum","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"functools","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"typing","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"types","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.protocol","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#Dispatch","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#Dispatch","relation":"extends","name":"Enum","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#Execute","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#Execute","relation":"extends","name":"Enum","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#_split_args_kwargs_data_proto","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.protocol","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_split_args_kwargs_data_proto","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_split_args_kwargs_data_proto","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_split_args_kwargs_data_proto","relation":"calls","name":"chunk","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_split_args_kwargs_data_proto","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_split_args_kwargs_data_proto","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_split_args_kwargs_data_proto","relation":"calls","name":"chunk","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#dispatch_one_to_all","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_one_to_all","relation":"calls","name":"tuple","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_one_to_all","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#dispatch_all_to_all","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#collect_all_to_all","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"type","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"get_megatron_rank_info","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py","recvType":"tuple"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"tuple","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"get_megatron_rank_info","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#collect_megatron_compute","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_compute","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_compute","relation":"calls","name":"get_megatron_global_info","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_compute","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_compute","relation":"calls","name":"get_megatron_rank_info","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_compute","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#dispatch_megatron_compute_data_proto","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute_data_proto","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute_data_proto","relation":"calls","name":"_split_args_kwargs_data_proto","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_compute_data_proto","relation":"calls","name":"dispatch_megatron_compute","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#_concat_data_proto_or_future","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.protocol","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"ray","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_concat_data_proto_or_future","relation":"calls","name":"type","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_concat_data_proto_or_future","relation":"calls","name":"type","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_concat_data_proto_or_future","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_concat_data_proto_or_future","relation":"calls","name":"concat","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_concat_data_proto_or_future","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_concat_data_proto_or_future","relation":"calls","name":"concat","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#collect_megatron_compute_data_proto","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.protocol","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"ray","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_compute_data_proto","relation":"calls","name":"collect_megatron_compute","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_compute_data_proto","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_compute_data_proto","relation":"calls","name":"type","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_compute_data_proto","relation":"calls","name":"_concat_data_proto_or_future","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"get_megatron_rank_info","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"get_megatron_rank_info","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py","recvType":"tuple"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"tuple","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"get_megatron_rank_info","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"get_megatron_rank_info","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp","relation":"calls","name":"get_megatron_rank_info","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#collect_megatron_pp_only","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_only","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_only","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_only","relation":"calls","name":"get_megatron_rank_info","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_only","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp_data_proto","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp_data_proto","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp_data_proto","relation":"calls","name":"_split_args_kwargs_data_proto","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_megatron_pp_as_dp_data_proto","relation":"calls","name":"dispatch_megatron_pp_as_dp","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp_data_proto","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.protocol","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp_data_proto","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp_data_proto","relation":"calls","name":"collect_megatron_pp_as_dp","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_megatron_pp_as_dp_data_proto","relation":"calls","name":"_concat_data_proto_or_future","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#dispatch_dp_compute","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#collect_dp_compute","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_dp_compute","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_dp_compute","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#dispatch_dp_compute_data_proto","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute_data_proto","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute_data_proto","relation":"calls","name":"_split_args_kwargs_data_proto","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#dispatch_dp_compute_data_proto_with_func","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.single_controller.base.worker_group","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute_data_proto_with_func","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute_data_proto_with_func","relation":"calls","name":"type","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#dispatch_dp_compute_data_proto_with_func","relation":"calls","name":"_split_args_kwargs_data_proto","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#collect_dp_compute_data_proto","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"verl.protocol","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"imports","specifier":"ray","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_dp_compute_data_proto","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_dp_compute_data_proto","relation":"calls","name":"type","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_dp_compute_data_proto","relation":"calls","name":"collect_dp_compute","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#collect_dp_compute_data_proto","relation":"calls","name":"_concat_data_proto_or_future","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#get_predefined_dispatch_fn","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#get_predefined_execute_fn","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#_check_dispatch_mode","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_check_dispatch_mode","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_check_dispatch_mode","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#_check_execute_mode","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_check_execute_mode","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#_materialize_futures","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_materialize_futures","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_materialize_futures","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_materialize_futures","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/decorator.py","recvType":"tuple"},{"source":"verl/single_controller/base/decorator.py#_materialize_futures","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_materialize_futures","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_materialize_futures","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#_materialize_futures","relation":"calls","name":"tuple","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py","relation":"contains","targetId":"verl/single_controller/base/decorator.py#register","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#register","relation":"calls","name":"_check_dispatch_mode","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#register","relation":"calls","name":"_check_execute_mode","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#register","relation":"contains","targetId":"verl/single_controller/base/decorator.py#register.decorator","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#register.decorator","relation":"calls","name":"wraps","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#register.decorator","relation":"contains","targetId":"verl/single_controller/base/decorator.py#register.decorator.inner","file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#register.decorator.inner","relation":"calls","name":"_materialize_futures","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#register.decorator.inner","relation":"calls","name":"func","viaMember":false,"file":"verl/single_controller/base/decorator.py"},{"source":"verl/single_controller/base/decorator.py#register.decorator","relation":"calls","name":"setattr","viaMember":false,"file":"verl/single_controller/base/decorator.py"}]},"verl/single_controller/base/megatron/__init__.py":{"size":600,"mtimeMs":1789275830310.2026,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/single_controller/base/megatron/__init__.py","name":"__init__.py","kind":"file","path":"verl/single_controller/base/megatron/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/single_controller/base/megatron/worker.py":{"size":1595,"mtimeMs":1789275830310.2026,"hash":"37aa599c7e77c72e99f52ebf1019a3e69e6dfeaa67bb852dc1b2e6afdd6e2c08","nodes":[{"id":"verl/single_controller/base/megatron/worker.py","name":"worker.py","kind":"file","path":"verl/single_controller/base/megatron/worker.py","span":"L1-L39","signature":null,"exported":true,"origin":"ast","body_hash":"37aa599c7e77c72e99f52ebf1019a3e69e6dfeaa67bb852dc1b2e6afdd6e2c08","chars":1595,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 dataclasses import dataclass from verl.single_controller.base.worker import Worker, DistRankInfo, DistGlobalInfo"},{"id":"verl/single_controller/base/megatron/worker.py#MegatronWorker","name":"MegatronWorker","kind":"class","path":"verl/single_controller/base/megatron/worker.py","span":"L20-L39","signature":"class MegatronWorker(Worker)","exported":true,"origin":"ast","body_hash":"e61250858323eec210ec5fd0507bd2db4a6fab32266a1359438efa549401d016","body_text":"class MegatronWorker(Worker): def __init__(self, cuda_visible_devices=None) -> None: super().__init__(cuda_visible_devices) def get_megatron_global_info(self): from megatron.core import parallel_state as mpu tp_size = mpu.get_tensor_model_parallel_world_size() dp_size = mpu.get_data_parallel_world_size() pp_size = mpu.get_pipeline_model_parallel_world_size() info = DistGlobalInfo(tp_size=tp_size, dp_size=dp_size, pp_size=pp_size) return info def get_megatron_rank_info(self): from megatron.core import parallel_state as mpu tp_rank = mpu.get_tensor_model_parallel_rank() dp_rank = mpu.get_data_parallel_rank() pp_rank = mpu.get_pipeline_model_parallel_rank() info = DistRankInfo(tp_rank=tp_rank, dp_rank=dp_rank, pp_rank=pp_rank) return info","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/megatron/worker.py#MegatronWorker.__init__","name":"__init__","kind":"method","path":"verl/single_controller/base/megatron/worker.py","span":"L22-L23","signature":"def __init__(self, cuda_visible_devices=None) -> None","exported":false,"origin":"ast","body_hash":"48571abe27cb7bbbc9c3f97b6c37136286421f9e1127f4bde765b2442203fd7e","body_text":"def __init__(self, cuda_visible_devices=None) -> None: super().__init__(cuda_visible_devices)","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronWorker"},{"id":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_global_info","name":"get_megatron_global_info","kind":"method","path":"verl/single_controller/base/megatron/worker.py","span":"L25-L31","signature":"def get_megatron_global_info(self)","exported":true,"origin":"ast","body_hash":"fa729d53ccfea6d803aca0bc4ea8073dbe39f8a7fe637fff2dc35d369c49b73b","body_text":"def get_megatron_global_info(self): from megatron.core import parallel_state as mpu tp_size = mpu.get_tensor_model_parallel_world_size() dp_size = mpu.get_data_parallel_world_size() pp_size = mpu.get_pipeline_model_parallel_world_size() info = DistGlobalInfo(tp_size=tp_size, dp_size=dp_size, pp_size=pp_size) return info","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronWorker"},{"id":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_rank_info","name":"get_megatron_rank_info","kind":"method","path":"verl/single_controller/base/megatron/worker.py","span":"L33-L39","signature":"def get_megatron_rank_info(self)","exported":true,"origin":"ast","body_hash":"6e4afb14ff0321a27c259b610b01be3d4aeb50932c0a019740368e88105edb73","body_text":"def get_megatron_rank_info(self): from megatron.core import parallel_state as mpu tp_rank = mpu.get_tensor_model_parallel_rank() dp_rank = mpu.get_data_parallel_rank() pp_rank = mpu.get_pipeline_model_parallel_rank() info = DistRankInfo(tp_rank=tp_rank, dp_rank=dp_rank, pp_rank=pp_rank) return info","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronWorker"}],"rawEdges":[{"source":"verl/single_controller/base/megatron/worker.py","relation":"imports","specifier":"os","file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py","relation":"imports","specifier":"dataclasses","file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py","relation":"imports","specifier":"verl.single_controller.base.worker","file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py","relation":"contains","targetId":"verl/single_controller/base/megatron/worker.py#MegatronWorker","file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker","relation":"extends","name":"Worker","file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker","relation":"contains","targetId":"verl/single_controller/base/megatron/worker.py#MegatronWorker.__init__","file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker","relation":"contains","targetId":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_global_info","file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py","relation":"imports","specifier":"megatron.core","file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_global_info","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_global_info","relation":"calls","name":"get_data_parallel_world_size","viaMember":true,"file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_global_info","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_global_info","relation":"calls","name":"DistGlobalInfo","viaMember":false,"file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker","relation":"contains","targetId":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_rank_info","file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py","relation":"imports","specifier":"megatron.core","file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_rank_info","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":true,"file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_rank_info","relation":"calls","name":"get_data_parallel_rank","viaMember":true,"file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_rank_info","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/single_controller/base/megatron/worker.py"},{"source":"verl/single_controller/base/megatron/worker.py#MegatronWorker.get_megatron_rank_info","relation":"calls","name":"DistRankInfo","viaMember":false,"file":"verl/single_controller/base/megatron/worker.py"}]},"verl/single_controller/base/megatron/worker_group.py":{"size":2080,"mtimeMs":1789275830310.2026,"hash":"73362d1f9a2c162dcad2470c3ae93cb0867273a5fa236b135357e76807fa987f","nodes":[{"id":"verl/single_controller/base/megatron/worker_group.py","name":"worker_group.py","kind":"file","path":"verl/single_controller/base/megatron/worker_group.py","span":"L1-L52","signature":null,"exported":true,"origin":"ast","body_hash":"73362d1f9a2c162dcad2470c3ae93cb0867273a5fa236b135357e76807fa987f","chars":2080,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 typing import Dict from .worker import DistRankInfo, DistGlobalInfo from verl.single_controller.base import ResourcePool, WorkerGroup"},{"id":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup","name":"MegatronWorkerGroup","kind":"class","path":"verl/single_controller/base/megatron/worker_group.py","span":"L21-L51","signature":"class MegatronWorkerGroup(WorkerGroup)","exported":true,"origin":"ast","body_hash":"02289ff08363e0f170675e1af71e98faab060b6e2784940d545cb57b24cb18df","body_text":"class MegatronWorkerGroup(WorkerGroup): def __init__(self, resource_pool: ResourcePool, **kwargs): super().__init__(resource_pool=resource_pool, **kwargs) self._megatron_rank_info = None self._megatron_global_info: DistGlobalInfo = None def init_megatron(self, default_megatron_kwargs: Dict = None): raise NotImplementedError(f\"MegatronWorkerGroup.init_megatron should be overwritten\") def get_megatron_rank_info(self, rank: int) -> DistRankInfo: assert 0 <= rank < self.world_size, f'rank must be from [0, world_size), Got {rank}' return self._megatron_rank_info[rank] @property def tp_size(self): assert self._megatron_global_info is not None, \"MegatronWorkerGroup._megatron_global_info must be initialized\" return self._megatron_global_info.tp_size @property def dp_size(self): assert self._megatron_global_info is not None, \"MegatronWorkerGroup._megatron_global_info must be initialized\" return self._megatron_global_info.dp_size @property def pp_size(self): assert self._megatron_global_info is not None, \"MegatronWorkerGroup._megatron_global_info must be initialized\" return self._megatron_global_info.pp_size def get_megatron_global_info(self): return self._megatron_global_info","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.__init__","name":"__init__","kind":"method","path":"verl/single_controller/base/megatron/worker_group.py","span":"L23-L26","signature":"def __init__(self, resource_pool: ResourcePool, **kwargs)","exported":false,"origin":"ast","body_hash":"bb2fd1f7d25649ebc0817445cc69ffc1232c30b0b25b78d11f887f179cf6d57d","body_text":"def __init__(self, resource_pool: ResourcePool, **kwargs): super().__init__(resource_pool=resource_pool, **kwargs) self._megatron_rank_info = None self._megatron_global_info: DistGlobalInfo = None","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronWorkerGroup"},{"id":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.init_megatron","name":"init_megatron","kind":"method","path":"verl/single_controller/base/megatron/worker_group.py","span":"L28-L29","signature":"def init_megatron(self, default_megatron_kwargs: Dict = None)","exported":true,"origin":"ast","body_hash":"4a13ea2ef9f1252e3d96fa32d610e4ad7384a9976ba3ced3a014aefdda95596d","body_text":"def init_megatron(self, default_megatron_kwargs: Dict = None): raise NotImplementedError(f\"MegatronWorkerGroup.init_megatron should be overwritten\")","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronWorkerGroup"},{"id":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.get_megatron_rank_info","name":"get_megatron_rank_info","kind":"method","path":"verl/single_controller/base/megatron/worker_group.py","span":"L31-L33","signature":"def get_megatron_rank_info(self, rank: int) -> DistRankInfo","exported":true,"origin":"ast","body_hash":"2e2f70add9898e25fc56effd98f1f038fe15ebb51448b158de57c2cf0323b9fe","body_text":"def get_megatron_rank_info(self, rank: int) -> DistRankInfo: assert 0 <= rank < self.world_size, f'rank must be from [0, world_size), Got {rank}' return self._megatron_rank_info[rank]","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronWorkerGroup"},{"id":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.tp_size","name":"tp_size","kind":"method","path":"verl/single_controller/base/megatron/worker_group.py","span":"L36-L38","signature":"def tp_size(self)","exported":true,"origin":"ast","body_hash":"67478e414ad212e12688bdacf09439aac2704732d909f5f152cb9526a8169bf3","body_text":"def tp_size(self): assert self._megatron_global_info is not None, \"MegatronWorkerGroup._megatron_global_info must be initialized\" return self._megatron_global_info.tp_size","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronWorkerGroup"},{"id":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.dp_size","name":"dp_size","kind":"method","path":"verl/single_controller/base/megatron/worker_group.py","span":"L41-L43","signature":"def dp_size(self)","exported":true,"origin":"ast","body_hash":"83aed3788fea0ab348e2fc56cc2258d7911c0baeef5e8a58c950e74cc17eb7e5","body_text":"def dp_size(self): assert self._megatron_global_info is not None, \"MegatronWorkerGroup._megatron_global_info must be initialized\" return self._megatron_global_info.dp_size","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronWorkerGroup"},{"id":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.pp_size","name":"pp_size","kind":"method","path":"verl/single_controller/base/megatron/worker_group.py","span":"L46-L48","signature":"def pp_size(self)","exported":true,"origin":"ast","body_hash":"668f11b566328b5611d8b177629547ad68d00e40577ebb50a4e4347ab508176e","body_text":"def pp_size(self): assert self._megatron_global_info is not None, \"MegatronWorkerGroup._megatron_global_info must be initialized\" return self._megatron_global_info.pp_size","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronWorkerGroup"},{"id":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.get_megatron_global_info","name":"get_megatron_global_info","kind":"method","path":"verl/single_controller/base/megatron/worker_group.py","span":"L50-L51","signature":"def get_megatron_global_info(self)","exported":true,"origin":"ast","body_hash":"6ae18787579bd69cde739643e5cfeaf50b409f938edd7bf9b2b8281326818556","body_text":"def get_megatron_global_info(self): return self._megatron_global_info","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronWorkerGroup"}],"rawEdges":[{"source":"verl/single_controller/base/megatron/worker_group.py","relation":"imports","specifier":"typing","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py","relation":"imports","specifier":".worker","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py","relation":"imports","specifier":"verl.single_controller.base","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py","relation":"contains","targetId":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup","relation":"extends","name":"WorkerGroup","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup","relation":"contains","targetId":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.__init__","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup","relation":"contains","targetId":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.init_megatron","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.init_megatron","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup","relation":"contains","targetId":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.get_megatron_rank_info","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup","relation":"contains","targetId":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.tp_size","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup","relation":"contains","targetId":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.dp_size","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup","relation":"contains","targetId":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.pp_size","file":"verl/single_controller/base/megatron/worker_group.py"},{"source":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup","relation":"contains","targetId":"verl/single_controller/base/megatron/worker_group.py#MegatronWorkerGroup.get_megatron_global_info","file":"verl/single_controller/base/megatron/worker_group.py"}]},"verl/single_controller/base/register_center/__init__.py":{"size":600,"mtimeMs":1789275830310.2026,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/single_controller/base/register_center/__init__.py","name":"__init__.py","kind":"file","path":"verl/single_controller/base/register_center/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/single_controller/base/register_center/ray.py":{"size":939,"mtimeMs":1789275830310.2026,"hash":"453fb6557bd4dde011730d92d96b8004e280c2dd75bb60f4d3eff9b8947ae702","nodes":[{"id":"verl/single_controller/base/register_center/ray.py","name":"ray.py","kind":"file","path":"verl/single_controller/base/register_center/ray.py","span":"L1-L30","signature":null,"exported":true,"origin":"ast","body_hash":"453fb6557bd4dde011730d92d96b8004e280c2dd75bb60f4d3eff9b8947ae702","chars":939,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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"},{"id":"verl/single_controller/base/register_center/ray.py#WorkerGroupRegisterCenter","name":"WorkerGroupRegisterCenter","kind":"class","path":"verl/single_controller/base/register_center/ray.py","span":"L19-L25","signature":"class WorkerGroupRegisterCenter","exported":true,"origin":"ast","body_hash":"5c75994fa634a57b7ae4467fb69f4e7d10d879c3ed48cceecf2e226bc93ec6f2","body_text":"class WorkerGroupRegisterCenter: def __init__(self, rank_zero_info): self.rank_zero_info = rank_zero_info def get_rank_zero_info(self): return self.rank_zero_info","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/register_center/ray.py#WorkerGroupRegisterCenter.__init__","name":"__init__","kind":"method","path":"verl/single_controller/base/register_center/ray.py","span":"L21-L22","signature":"def __init__(self, rank_zero_info)","exported":false,"origin":"ast","body_hash":"17e4dc06a759d2ce0ca2b8407ac1254c9fcd4418eb1cb5c8ec74b393c93eff1d","body_text":"def __init__(self, rank_zero_info): self.rank_zero_info = rank_zero_info","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerGroupRegisterCenter"},{"id":"verl/single_controller/base/register_center/ray.py#WorkerGroupRegisterCenter.get_rank_zero_info","name":"get_rank_zero_info","kind":"method","path":"verl/single_controller/base/register_center/ray.py","span":"L24-L25","signature":"def get_rank_zero_info(self)","exported":true,"origin":"ast","body_hash":"4f7f9150f6f242e35f2758a13aa25775355b8ec57d13e9c67e8276561fb66bae","body_text":"def get_rank_zero_info(self): return self.rank_zero_info","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerGroupRegisterCenter"},{"id":"verl/single_controller/base/register_center/ray.py#create_worker_group_register_center","name":"create_worker_group_register_center","kind":"function","path":"verl/single_controller/base/register_center/ray.py","span":"L28-L29","signature":"def create_worker_group_register_center(name, info)","exported":true,"origin":"ast","body_hash":"ade25a4898a12e94b6699098f15bee89d50059525f225ed0a0b45e7d47383e03","body_text":"def create_worker_group_register_center(name, info): return WorkerGroupRegisterCenter.options(name=name).remote(info)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/single_controller/base/register_center/ray.py","relation":"imports","specifier":"ray","file":"verl/single_controller/base/register_center/ray.py"},{"source":"verl/single_controller/base/register_center/ray.py","relation":"contains","targetId":"verl/single_controller/base/register_center/ray.py#WorkerGroupRegisterCenter","file":"verl/single_controller/base/register_center/ray.py"},{"source":"verl/single_controller/base/register_center/ray.py#WorkerGroupRegisterCenter","relation":"contains","targetId":"verl/single_controller/base/register_center/ray.py#WorkerGroupRegisterCenter.__init__","file":"verl/single_controller/base/register_center/ray.py"},{"source":"verl/single_controller/base/register_center/ray.py#WorkerGroupRegisterCenter","relation":"contains","targetId":"verl/single_controller/base/register_center/ray.py#WorkerGroupRegisterCenter.get_rank_zero_info","file":"verl/single_controller/base/register_center/ray.py"},{"source":"verl/single_controller/base/register_center/ray.py","relation":"contains","targetId":"verl/single_controller/base/register_center/ray.py#create_worker_group_register_center","file":"verl/single_controller/base/register_center/ray.py"},{"source":"verl/single_controller/base/register_center/ray.py#create_worker_group_register_center","relation":"calls","name":"remote","viaMember":true,"file":"verl/single_controller/base/register_center/ray.py"},{"source":"verl/single_controller/base/register_center/ray.py#create_worker_group_register_center","relation":"calls","name":"options","viaMember":true,"file":"verl/single_controller/base/register_center/ray.py"}]},"verl/single_controller/base/worker.py":{"size":6416,"mtimeMs":1789275830310.2026,"hash":"3a4c40474aab8a2f898cc2fbbcbad7db5a697507065330026f6e7bbc9e11b181","nodes":[{"id":"verl/single_controller/base/worker.py","name":"worker.py","kind":"file","path":"verl/single_controller/base/worker.py","span":"L1-L186","signature":null,"exported":true,"origin":"ast","body_hash":"3a4c40474aab8a2f898cc2fbbcbad7db5a697507065330026f6e7bbc9e11b181","chars":6416,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 class for Worker \"\"\" import os import socket from dataclasses import dataclass from verl.single_controller.base.decorator import register, Dispatch, Execute @dataclass @dataclass # we assume that in each WorkerGroup, there is a Master Worker"},{"id":"verl/single_controller/base/worker.py#DistRankInfo","name":"DistRankInfo","kind":"class","path":"verl/single_controller/base/worker.py","span":"L24-L27","signature":"class DistRankInfo","exported":true,"origin":"ast","body_hash":"46205767bfa1efde3b980dc0949743998af62fc64f7a5e05c7183c801d232609","body_text":"class DistRankInfo: tp_rank: int dp_rank: int pp_rank: int","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/worker.py#DistGlobalInfo","name":"DistGlobalInfo","kind":"class","path":"verl/single_controller/base/worker.py","span":"L31-L34","signature":"class DistGlobalInfo","exported":true,"origin":"ast","body_hash":"cafe25ad598dc623204a97199115dd3e8422efc586569e6aadde03152f1f7a7c","body_text":"class DistGlobalInfo: tp_size: int dp_size: int pp_size: int","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/worker.py#WorkerHelper","name":"WorkerHelper","kind":"class","path":"verl/single_controller/base/worker.py","span":"L37-L67","signature":"class WorkerHelper","exported":true,"origin":"ast","body_hash":"82e62000fc221259f06a87eedb156e3dd19f64e8e6b3067b9d04f88dc1f687a6","body_text":"class WorkerHelper: def _get_node_ip(self): def get_node_ip_by_sdk(): if os.getenv(\"WG_BACKEND\", None) == \"ray\": import ray return ray._private.services.get_node_ip_address() elif os.getenv(\"WG_BACKEND\", None) == \"torch_rpc\": from verl.single_controller.torchrpc.k8s_client import get_ip_addr return get_ip_addr() return None host_ipv4 = os.getenv(\"MY_HOST_IP\", None) host_ipv6 = os.getenv(\"MY_HOST_IPV6\", None) host_ip_by_env = host_ipv4 or host_ipv6 host_ip_by_sdk = get_node_ip_by_sdk() host_ip = host_ip_by_env or host_ip_by_sdk return host_ip def _get_free_port(self): with socket.socket() as sock: sock.bind(('', 0)) return sock.getsockname()[1] def get_availale_master_addr_port(self): return self._get_node_ip(), str(self._get_free_port()) def _get_pid(self): return","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip","name":"_get_node_ip","kind":"method","path":"verl/single_controller/base/worker.py","span":"L39-L56","signature":"def _get_node_ip(self)","exported":false,"origin":"ast","body_hash":"1f05d8562ac5389a3bc068f5566b3662aede2dadaa8006de4febd2562910b043","body_text":"def _get_node_ip(self): def get_node_ip_by_sdk(): if os.getenv(\"WG_BACKEND\", None) == \"ray\": import ray return ray._private.services.get_node_ip_address() elif os.getenv(\"WG_BACKEND\", None) == \"torch_rpc\": from verl.single_controller.torchrpc.k8s_client import get_ip_addr return get_ip_addr() return None host_ipv4 = os.getenv(\"MY_HOST_IP\", None) host_ipv6 = os.getenv(\"MY_HOST_IPV6\", None) host_ip_by_env = host_ipv4 or host_ipv6 host_ip_by_sdk = get_node_ip_by_sdk() host_ip = host_ip_by_env or host_ip_by_sdk return host_ip","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerHelper"},{"id":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip.get_node_ip_by_sdk","name":"get_node_ip_by_sdk","kind":"function","path":"verl/single_controller/base/worker.py","span":"L41-L48","signature":"def get_node_ip_by_sdk()","exported":true,"origin":"ast","body_hash":"fdb0fcff795c16817d66461c60cbf7de1fe31349542326a6387daafb0a9f6e60","body_text":"def get_node_ip_by_sdk(): if os.getenv(\"WG_BACKEND\", None) == \"ray\": import ray return ray._private.services.get_node_ip_address() elif os.getenv(\"WG_BACKEND\", None) == \"torch_rpc\": from verl.single_controller.torchrpc.k8s_client import get_ip_addr return get_ip_addr() return None","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/worker.py#WorkerHelper._get_free_port","name":"_get_free_port","kind":"method","path":"verl/single_controller/base/worker.py","span":"L58-L61","signature":"def _get_free_port(self)","exported":false,"origin":"ast","body_hash":"e7adc0cd353c3ed7360ca5bde4308da79ba207d18b655d9bd4b4961b895d7b23","body_text":"def _get_free_port(self): with socket.socket() as sock: sock.bind(('', 0)) return sock.getsockname()[1]","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerHelper"},{"id":"verl/single_controller/base/worker.py#WorkerHelper.get_availale_master_addr_port","name":"get_availale_master_addr_port","kind":"method","path":"verl/single_controller/base/worker.py","span":"L63-L64","signature":"def get_availale_master_addr_port(self)","exported":true,"origin":"ast","body_hash":"09a95db74930db593d2068e77fb5636dab537d056f6987a5f4a53dd19fa794be","body_text":"def get_availale_master_addr_port(self): return self._get_node_ip(), str(self._get_free_port())","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerHelper"},{"id":"verl/single_controller/base/worker.py#WorkerHelper._get_pid","name":"_get_pid","kind":"method","path":"verl/single_controller/base/worker.py","span":"L66-L67","signature":"def _get_pid(self)","exported":false,"origin":"ast","body_hash":"b8f1746673ca963156d65533e0ca501cd04d6a0530d8c38c4058e5f77618131c","body_text":"def _get_pid(self): return","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerHelper"},{"id":"verl/single_controller/base/worker.py#WorkerMeta","name":"WorkerMeta","kind":"class","path":"verl/single_controller/base/worker.py","span":"L70-L79","signature":"class WorkerMeta","exported":true,"origin":"ast","body_hash":"bedb505206b97658e06333908864d662c57ca9eef49b201155a422dfcc657e7f","body_text":"class WorkerMeta: keys = [ \"WORLD_SIZE\", \"RANK\", \"LOCAL_WORLD_SIZE\", \"LOCAL_RANK\", \"MASTER_ADDR\", \"MASTER_PORT\", \"CUDA_VISIBLE_DEVICES\" ] def __init__(self, store) -> None: self._store = store def to_dict(self): return {f\"_{key.lower()}\": self._store.get(f\"_{key.lower()}\", None) for key in WorkerMeta.keys}","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/worker.py#WorkerMeta.__init__","name":"__init__","kind":"method","path":"verl/single_controller/base/worker.py","span":"L75-L76","signature":"def __init__(self, store) -> None","exported":false,"origin":"ast","body_hash":"53af3f40c2980e8b00f196158ccbae9eefdd8395acc9185a0bbf8062de0b62a6","body_text":"def __init__(self, store) -> None: self._store = store","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerMeta"},{"id":"verl/single_controller/base/worker.py#WorkerMeta.to_dict","name":"to_dict","kind":"method","path":"verl/single_controller/base/worker.py","span":"L78-L79","signature":"def to_dict(self)","exported":true,"origin":"ast","body_hash":"0be649e5967f5ddef8a1ae38a4e131bfe2a56656a2099816189bd7301da8544d","body_text":"def to_dict(self): return {f\"_{key.lower()}\": self._store.get(f\"_{key.lower()}\", None) for key in WorkerMeta.keys}","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerMeta"},{"id":"verl/single_controller/base/worker.py#Worker","name":"Worker","kind":"class","path":"verl/single_controller/base/worker.py","span":"L83-L186","signature":"class Worker(WorkerHelper)","exported":true,"origin":"ast","body_hash":"6e08be07bc59f410bb51a7f9a6d8d4b259b12bddeec43fa35178c5d5a4d3177a","body_text":"class Worker(WorkerHelper): def __new__(cls, *args, **kwargs): instance = super().__new__(cls) # note that here we use int to distinguish disable_worker_init = int(os.environ.get('DISABLE_WORKER_INIT', 0)) if disable_worker_init: return instance rank = os.environ.get(\"RANK\", None) worker_group_prefix = os.environ.get(\"WG_PREFIX\", None) # when decorator @ray.remote applies, __new__ will be called while we don't want to apply _configure_before_init if None not in [rank, worker_group_prefix] and 'ActorClass(' not in cls.__name__: instance._configure_before_init(f\"{worker_group_prefix}_register_center\", int(rank)) return instance def _configure_before_init(self, register_center_name: str, rank: int): assert isinstance(rank, int), f\"rank must be int, instead of {type(rank)}\" if rank == 0: master_addr, master_port = self.get_availale_master_addr_port() rank_zero_info = { \"MASTER_ADDR\": master_addr, \"MASTER_PORT\": master_port, } if os.getenv(\"WG_BACKEND\", None) == \"ray\": from verl.single_controller.base.register_center.ray import create_worker_group_register_center self.register_center = create_worker_group_register_center(name=register_center_name, info=rank_zero_info) os.environ.update(rank_zero_info) def __init__(self, cuda_visible_devices=None) -> None: # construct a meta from envrionment variable. Note that the import must be inside the class because it is executed remotely import os world_size = int(os.environ['WORLD_SIZE']) rank = int(os.environ['RANK']) self._rank = rank self._world_size = world_size master_addr = os.environ[\"MASTER_ADDR\"] master_port = os.environ[\"MASTER_PORT\"] local_world_size = int(os.getenv(\"LOCAL_WORLD_SIZE\", \"1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) store = { '_world_size': world_size, '_rank': rank, '_local_world_size': local_world_size, '_local_rank': local_rank, '_master_addr': master_addr, '_master_port': master_port } if cuda_visible_devices is not None: store['_cuda_visible_devices'] = cuda_visible_devices meta = WorkerMeta(store=store) self._configure_with_meta(meta=meta) def _configure_with_meta(self, meta: WorkerMeta): \"\"\" This function should only be called inside by WorkerGroup \"\"\" assert isinstance(meta, WorkerMeta) self.__dict__.update(meta.to_dict()) # this is hacky # print(f\"__dict__: {self.__dict__}\") for key in WorkerMeta.keys: val = self.__dict__.get(f\"_{key.lower()}\", None) if val is not None: # print(f\"set {key} to {val}\") os.environ[key] = str(val) os.environ[\"REDIS_STORE_SERVER_HOST\"] = str(self._master_addr).replace(\"[\", \"\").replace( \"]\", \"\") if self._master_addr else \"\" def get_master_addr_port(self): return self._master_addr, self._master_port def get_cuda_visible_devices(self): import os cuda_visible_devices = os.environ.get(\"CUDA_VISIBLE_DEVICES\", \"not set\") return cuda_visible_devices @property def world_size(self): return self._world_size @property def rank(self): return self._rank @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO_WITH_FUNC) def execute_with_func_generator(self, func, *args, **kwargs): ret_proto = func(self, *args, **kwargs) return ret_proto @register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.RANK_ZERO) def execute_func_rank_zero(self, func, *args, **kwargs): result = func(*args, **kwargs) return result","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/worker.py#Worker.__new__","name":"__new__","kind":"method","path":"verl/single_controller/base/worker.py","span":"L85-L100","signature":"def __new__(cls, *args, **kwargs)","exported":false,"origin":"ast","body_hash":"eb26d47ad1e770ee5b7cdb4d2b443b0c1719d1cb72e45c52b9e57257900d90e1","body_text":"def __new__(cls, *args, **kwargs): instance = super().__new__(cls) # note that here we use int to distinguish disable_worker_init = int(os.environ.get('DISABLE_WORKER_INIT', 0)) if disable_worker_init: return instance rank = os.environ.get(\"RANK\", None) worker_group_prefix = os.environ.get(\"WG_PREFIX\", None) # when decorator @ray.remote applies, __new__ will be called while we don't want to apply _configure_before_init if None not in [rank, worker_group_prefix] and 'ActorClass(' not in cls.__name__: instance._configure_before_init(f\"{worker_group_prefix}_register_center\", int(rank)) return instance","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/single_controller/base/worker.py#Worker._configure_before_init","name":"_configure_before_init","kind":"method","path":"verl/single_controller/base/worker.py","span":"L102-L117","signature":"def _configure_before_init(self, register_center_name: str, rank: int)","exported":false,"origin":"ast","body_hash":"3ac928c852f120b1dc9df95f47dbafa94d02892ddfe56e54092c82c7f5b10956","body_text":"def _configure_before_init(self, register_center_name: str, rank: int): assert isinstance(rank, int), f\"rank must be int, instead of {type(rank)}\" if rank == 0: master_addr, master_port = self.get_availale_master_addr_port() rank_zero_info = { \"MASTER_ADDR\": master_addr, \"MASTER_PORT\": master_port, } if os.getenv(\"WG_BACKEND\", None) == \"ray\": from verl.single_controller.base.register_center.ray import create_worker_group_register_center self.register_center = create_worker_group_register_center(name=register_center_name, info=rank_zero_info) os.environ.update(rank_zero_info)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/single_controller/base/worker.py#Worker.__init__","name":"__init__","kind":"method","path":"verl/single_controller/base/worker.py","span":"L119-L145","signature":"def __init__(self, cuda_visible_devices=None) -> None: # construct a meta from envrionment variable. Note that the import must be inside the class because it is executed remotely","exported":false,"origin":"ast","body_hash":"6834df0f7784adc70ee5d5dedcae3f1d1f728a6bd40257e7f31f72163e18d81c","body_text":"def __init__(self, cuda_visible_devices=None) -> None: # construct a meta from envrionment variable. Note that the import must be inside the class because it is executed remotely import os world_size = int(os.environ['WORLD_SIZE']) rank = int(os.environ['RANK']) self._rank = rank self._world_size = world_size master_addr = os.environ[\"MASTER_ADDR\"] master_port = os.environ[\"MASTER_PORT\"] local_world_size = int(os.getenv(\"LOCAL_WORLD_SIZE\", \"1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) store = { '_world_size': world_size, '_rank': rank, '_local_world_size': local_world_size, '_local_rank': local_rank, '_master_addr': master_addr, '_master_port': master_port } if cuda_visible_devices is not None: store['_cuda_visible_devices'] = cuda_visible_devices meta = WorkerMeta(store=store) self._configure_with_meta(meta=meta)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/single_controller/base/worker.py#Worker._configure_with_meta","name":"_configure_with_meta","kind":"method","path":"verl/single_controller/base/worker.py","span":"L147-L160","signature":"def _configure_with_meta(self, meta: WorkerMeta)","exported":false,"origin":"ast","body_hash":"d071f3edc18742f5bdc15d319ee770e7953d3bf5993546458095d1580b95d9e8","body_text":"def _configure_with_meta(self, meta: WorkerMeta): \"\"\" This function should only be called inside by WorkerGroup \"\"\" assert isinstance(meta, WorkerMeta) self.__dict__.update(meta.to_dict()) # this is hacky # print(f\"__dict__: {self.__dict__}\") for key in WorkerMeta.keys: val = self.__dict__.get(f\"_{key.lower()}\", None) if val is not None: # print(f\"set {key} to {val}\") os.environ[key] = str(val) os.environ[\"REDIS_STORE_SERVER_HOST\"] = str(self._master_addr).replace(\"[\", \"\").replace( \"]\", \"\") if self._master_addr else \"\"","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/single_controller/base/worker.py#Worker.get_master_addr_port","name":"get_master_addr_port","kind":"method","path":"verl/single_controller/base/worker.py","span":"L162-L163","signature":"def get_master_addr_port(self)","exported":true,"origin":"ast","body_hash":"b7368c8debf440abf7dc60c61d74e9f2f9e15d76a07d0436d37b51cfc969dbb6","body_text":"def get_master_addr_port(self): return self._master_addr, self._master_port","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/single_controller/base/worker.py#Worker.get_cuda_visible_devices","name":"get_cuda_visible_devices","kind":"method","path":"verl/single_controller/base/worker.py","span":"L165-L168","signature":"def get_cuda_visible_devices(self)","exported":true,"origin":"ast","body_hash":"a6f72b1365f288d4a983a82ac99673143bc5aa3ca2ed8e3a58e41ac97406c934","body_text":"def get_cuda_visible_devices(self): import os cuda_visible_devices = os.environ.get(\"CUDA_VISIBLE_DEVICES\", \"not set\") return cuda_visible_devices","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/single_controller/base/worker.py#Worker.world_size","name":"world_size","kind":"method","path":"verl/single_controller/base/worker.py","span":"L171-L172","signature":"def world_size(self)","exported":true,"origin":"ast","body_hash":"417053f3e044003b010d0dc3d1a7213269ba5b6bee8a5bcff37b385f3fcb0427","body_text":"def world_size(self): return self._world_size","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/single_controller/base/worker.py#Worker.rank","name":"rank","kind":"method","path":"verl/single_controller/base/worker.py","span":"L175-L176","signature":"def rank(self)","exported":true,"origin":"ast","body_hash":"831d3b46bb6916182ed19e1bcdca0da7a9c2638ffb618fd72609fa7fc20b47ee","body_text":"def rank(self): return self._rank","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/single_controller/base/worker.py#Worker.execute_with_func_generator","name":"execute_with_func_generator","kind":"method","path":"verl/single_controller/base/worker.py","span":"L179-L181","signature":"def execute_with_func_generator(self, func, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"70050b9762c92ae0d0cb63e7b4858f2acf5516928d3fb5e70d4f5f8e327e15e2","body_text":"def execute_with_func_generator(self, func, *args, **kwargs): ret_proto = func(self, *args, **kwargs) return ret_proto","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/single_controller/base/worker.py#Worker.execute_func_rank_zero","name":"execute_func_rank_zero","kind":"method","path":"verl/single_controller/base/worker.py","span":"L184-L186","signature":"def execute_func_rank_zero(self, func, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"6c8557a3797663b9a6256c85f51594eabe8f72454a4f9d3121c90c033098e379","body_text":"def execute_func_rank_zero(self, func, *args, **kwargs): result = func(*args, **kwargs) return result","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"}],"rawEdges":[{"source":"verl/single_controller/base/worker.py","relation":"imports","specifier":"os","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"imports","specifier":"socket","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"imports","specifier":"dataclasses","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"imports","specifier":"verl.single_controller.base.decorator","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"contains","targetId":"verl/single_controller/base/worker.py#DistRankInfo","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"contains","targetId":"verl/single_controller/base/worker.py#DistGlobalInfo","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"contains","targetId":"verl/single_controller/base/worker.py#WorkerHelper","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper","relation":"contains","targetId":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip","relation":"contains","targetId":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip.get_node_ip_by_sdk","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip.get_node_ip_by_sdk","relation":"calls","name":"getenv","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"imports","specifier":"ray","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip.get_node_ip_by_sdk","relation":"calls","name":"get_node_ip_address","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip.get_node_ip_by_sdk","relation":"calls","name":"getenv","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"imports","specifier":"verl.single_controller.torchrpc.k8s_client","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip.get_node_ip_by_sdk","relation":"calls","name":"get_ip_addr","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip","relation":"calls","name":"getenv","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip","relation":"calls","name":"getenv","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_node_ip","relation":"calls","name":"get_node_ip_by_sdk","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper","relation":"contains","targetId":"verl/single_controller/base/worker.py#WorkerHelper._get_free_port","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_free_port","relation":"calls","name":"socket","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_free_port","relation":"calls","name":"bind","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper._get_free_port","relation":"calls","name":"getsockname","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper","relation":"contains","targetId":"verl/single_controller/base/worker.py#WorkerHelper.get_availale_master_addr_port","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper.get_availale_master_addr_port","relation":"calls","name":"_get_node_ip","viaMember":true,"file":"verl/single_controller/base/worker.py","recvType":"WorkerHelper"},{"source":"verl/single_controller/base/worker.py#WorkerHelper.get_availale_master_addr_port","relation":"calls","name":"str","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerHelper.get_availale_master_addr_port","relation":"calls","name":"_get_free_port","viaMember":true,"file":"verl/single_controller/base/worker.py","recvType":"WorkerHelper"},{"source":"verl/single_controller/base/worker.py#WorkerHelper","relation":"contains","targetId":"verl/single_controller/base/worker.py#WorkerHelper._get_pid","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"contains","targetId":"verl/single_controller/base/worker.py#WorkerMeta","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerMeta","relation":"contains","targetId":"verl/single_controller/base/worker.py#WorkerMeta.__init__","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerMeta","relation":"contains","targetId":"verl/single_controller/base/worker.py#WorkerMeta.to_dict","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerMeta.to_dict","relation":"calls","name":"lower","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerMeta.to_dict","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#WorkerMeta.to_dict","relation":"calls","name":"lower","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"extends","name":"WorkerHelper","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker.__new__","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__new__","relation":"calls","name":"__new__","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__new__","relation":"calls","name":"super","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__new__","relation":"calls","name":"int","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__new__","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__new__","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__new__","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__new__","relation":"calls","name":"_configure_before_init","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__new__","relation":"calls","name":"int","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker._configure_before_init","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_before_init","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_before_init","relation":"calls","name":"type","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_before_init","relation":"calls","name":"get_availale_master_addr_port","viaMember":true,"file":"verl/single_controller/base/worker.py","recvType":"Worker"},{"source":"verl/single_controller/base/worker.py#Worker._configure_before_init","relation":"calls","name":"getenv","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"imports","specifier":"verl.single_controller.base.register_center.ray","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_before_init","relation":"calls","name":"create_worker_group_register_center","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_before_init","relation":"calls","name":"update","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker.__init__","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"imports","specifier":"os","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__init__","relation":"calls","name":"int","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__init__","relation":"calls","name":"int","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__init__","relation":"calls","name":"int","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__init__","relation":"calls","name":"getenv","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__init__","relation":"calls","name":"int","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__init__","relation":"calls","name":"getenv","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__init__","relation":"calls","name":"WorkerMeta","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.__init__","relation":"calls","name":"_configure_with_meta","viaMember":true,"file":"verl/single_controller/base/worker.py","recvType":"Worker"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker._configure_with_meta","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_with_meta","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_with_meta","relation":"calls","name":"update","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_with_meta","relation":"calls","name":"to_dict","viaMember":true,"file":"verl/single_controller/base/worker.py","recvType":"WorkerMeta"},{"source":"verl/single_controller/base/worker.py#Worker._configure_with_meta","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_with_meta","relation":"calls","name":"lower","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_with_meta","relation":"calls","name":"str","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_with_meta","relation":"calls","name":"replace","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_with_meta","relation":"calls","name":"replace","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker._configure_with_meta","relation":"calls","name":"str","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker.get_master_addr_port","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker.get_cuda_visible_devices","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py","relation":"imports","specifier":"os","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.get_cuda_visible_devices","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker.world_size","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker.rank","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"calls","name":"register","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker.execute_with_func_generator","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.execute_with_func_generator","relation":"calls","name":"func","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"calls","name":"register","viaMember":false,"file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker","relation":"contains","targetId":"verl/single_controller/base/worker.py#Worker.execute_func_rank_zero","file":"verl/single_controller/base/worker.py"},{"source":"verl/single_controller/base/worker.py#Worker.execute_func_rank_zero","relation":"calls","name":"func","viaMember":false,"file":"verl/single_controller/base/worker.py"}]},"verl/single_controller/base/worker_group.py":{"size":7494,"mtimeMs":1789275830311.2026,"hash":"aee8ea551404288e27a0d12e29304504a1551468ba388b2f0e1292e8ff4dde79","nodes":[{"id":"verl/single_controller/base/worker_group.py","name":"worker_group.py","kind":"file","path":"verl/single_controller/base/worker_group.py","span":"L1-L197","signature":null,"exported":true,"origin":"ast","body_hash":"aee8ea551404288e27a0d12e29304504a1551468ba388b2f0e1292e8ff4dde79","chars":7494,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 class of WorkerGroup \"\"\" import logging import threading import signal import time from typing import List, Any, Callable, Dict from verl.single_controller.base.decorator import MAGIC_ATTR, Dispatch, get_predefined_dispatch_fn, get_predefined_execute_fn"},{"id":"verl/single_controller/base/worker_group.py#ResourcePool","name":"ResourcePool","kind":"class","path":"verl/single_controller/base/worker_group.py","span":"L26-L57","signature":"class ResourcePool","exported":true,"origin":"ast","body_hash":"71a28b891cd885b29e7970a3fd540bb376df943216ff8f3c3f6270caa256bcaf","body_text":"class ResourcePool: def __init__(self, process_on_nodes=None, max_collocate_count: int = 10, n_gpus_per_node=8) -> None: if process_on_nodes is None: process_on_nodes = [] self._store = process_on_nodes self.max_collocate_count = max_collocate_count self.n_gpus_per_node = n_gpus_per_node # this is left for future huawei GPU that contains 16 GPUs per node def add_node(self, process_count): self._store.append(process_count) @property def world_size(self): return sum(self._store) def __call__(self) -> Any: return self._store @property def store(self): return self._store def local_world_size_list(self) -> List[int]: nested_local_world_size_list = [ [local_world_size for _ in range(local_world_size)] for local_world_size in self._store ] return [item for row in nested_local_world_size_list for item in row] def local_rank_list(self) -> List[int]: nested_local_rank_list = [[i for i in range(local_world_size)] for local_world_size in self._store] return [item for row in nested_local_rank_list for item in row]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/worker_group.py#ResourcePool.__init__","name":"__init__","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L28-L33","signature":"def __init__(self, process_on_nodes=None, max_collocate_count: int = 10, n_gpus_per_node=8) -> None","exported":false,"origin":"ast","body_hash":"92b092738092d7446f5a68e99dda3769225b821d1becc40c33baf3e57aa6f766","body_text":"def __init__(self, process_on_nodes=None, max_collocate_count: int = 10, n_gpus_per_node=8) -> None: if process_on_nodes is None: process_on_nodes = [] self._store = process_on_nodes self.max_collocate_count = max_collocate_count self.n_gpus_per_node = n_gpus_per_node # this is left for future huawei GPU that contains 16 GPUs per node","summary_state":"pending","summary":null,"crux":null,"owner":"ResourcePool"},{"id":"verl/single_controller/base/worker_group.py#ResourcePool.add_node","name":"add_node","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L35-L36","signature":"def add_node(self, process_count)","exported":true,"origin":"ast","body_hash":"c1f57a47bbedbf0edce0dcba80a85cd46147698d14c9835aff1853070365694e","body_text":"def add_node(self, process_count): self._store.append(process_count)","summary_state":"pending","summary":null,"crux":null,"owner":"ResourcePool"},{"id":"verl/single_controller/base/worker_group.py#ResourcePool.world_size","name":"world_size","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L39-L40","signature":"def world_size(self)","exported":true,"origin":"ast","body_hash":"7b83e6dfef46ee4aa03b4a2558440a3f7f269a3ff21288bab38617d0450276d4","body_text":"def world_size(self): return sum(self._store)","summary_state":"pending","summary":null,"crux":null,"owner":"ResourcePool"},{"id":"verl/single_controller/base/worker_group.py#ResourcePool.__call__","name":"__call__","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L42-L43","signature":"def __call__(self) -> Any","exported":false,"origin":"ast","body_hash":"804a62ef0cb7b85a38032165d02b290b3d619ecb12aa1a41f0aea6cbb336b961","body_text":"def __call__(self) -> Any: return self._store","summary_state":"pending","summary":null,"crux":null,"owner":"ResourcePool"},{"id":"verl/single_controller/base/worker_group.py#ResourcePool.store","name":"store","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L46-L47","signature":"def store(self)","exported":true,"origin":"ast","body_hash":"8bac6dc0ed68362b724b3a8e77185adeb774431830fb5a2ee7921b3ac00be447","body_text":"def store(self): return self._store","summary_state":"pending","summary":null,"crux":null,"owner":"ResourcePool"},{"id":"verl/single_controller/base/worker_group.py#ResourcePool.local_world_size_list","name":"local_world_size_list","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L49-L53","signature":"def local_world_size_list(self) -> List[int]","exported":true,"origin":"ast","body_hash":"6afeab9fbb9e112e0151fe736a534ef25e281d3139837321e8d7f8dc139e8454","body_text":"def local_world_size_list(self) -> List[int]: nested_local_world_size_list = [ [local_world_size for _ in range(local_world_size)] for local_world_size in self._store ] return [item for row in nested_local_world_size_list for item in row]","summary_state":"pending","summary":null,"crux":null,"owner":"ResourcePool"},{"id":"verl/single_controller/base/worker_group.py#ResourcePool.local_rank_list","name":"local_rank_list","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L55-L57","signature":"def local_rank_list(self) -> List[int]","exported":true,"origin":"ast","body_hash":"db3fb7533145e0af6bef677d5f2d109db3a5c27d01d5eeb6b79e1135950ec220","body_text":"def local_rank_list(self) -> List[int]: nested_local_rank_list = [[i for i in range(local_world_size)] for local_world_size in self._store] return [item for row in nested_local_rank_list for item in row]","summary_state":"pending","summary":null,"crux":null,"owner":"ResourcePool"},{"id":"verl/single_controller/base/worker_group.py#ClassWithInitArgs","name":"ClassWithInitArgs","kind":"class","path":"verl/single_controller/base/worker_group.py","span":"L60-L78","signature":"class ClassWithInitArgs","exported":true,"origin":"ast","body_hash":"d21ed4e1310ac7918b4c2e5dd19a82a7cea745c2bd4de738d9e3e5aa2ba4c28a","body_text":"class ClassWithInitArgs: \"\"\" This class stores a class constructor and the args/kwargs to construct the class. It is used to instantiate the remote class. \"\"\" def __init__(self, cls, *args, **kwargs) -> None: self.cls = cls self.args = args self.kwargs = kwargs # def add_arg(self, arg): # self.args += (arg,) # def add_kwarg(self, key, value): # self.kwargs[key] = value def __call__(self) -> Any: return self.cls(*self.args, **self.kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/worker_group.py#ClassWithInitArgs.__init__","name":"__init__","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L66-L69","signature":"def __init__(self, cls, *args, **kwargs) -> None","exported":false,"origin":"ast","body_hash":"75cbab9d0a6bc980fb41c1d487c7d0748edd3fd2fdb8b23922260ffad9667999","body_text":"def __init__(self, cls, *args, **kwargs) -> None: self.cls = cls self.args = args self.kwargs = kwargs","summary_state":"pending","summary":null,"crux":null,"owner":"ClassWithInitArgs"},{"id":"verl/single_controller/base/worker_group.py#ClassWithInitArgs.__call__","name":"__call__","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L77-L78","signature":"def __call__(self) -> Any","exported":false,"origin":"ast","body_hash":"9b02a7e107a708606187090a070456b7db0c66acd04f42da05936b19f3d1d045","body_text":"def __call__(self) -> Any: return self.cls(*self.args, **self.kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"ClassWithInitArgs"},{"id":"verl/single_controller/base/worker_group.py#check_workers_alive","name":"check_workers_alive","kind":"function","path":"verl/single_controller/base/worker_group.py","span":"L81-L88","signature":"def check_workers_alive(workers: List, is_alive: Callable, gap_time: float = 1) -> None","exported":true,"origin":"ast","body_hash":"697a16f2c81106e75b3a76d55e95ad3d48d6b8d7dacf4cdbaa38c08b0b45c67b","body_text":"def check_workers_alive(workers: List, is_alive: Callable, gap_time: float = 1) -> None: import time while True: for worker in workers: if not is_alive(worker): logging.warning(f\"worker {worker} is not alive\" + \" sending signal to main thread\") signal.raise_signal(signal.SIGABRT) time.sleep(gap_time)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/worker_group.py#WorkerGroup","name":"WorkerGroup","kind":"class","path":"verl/single_controller/base/worker_group.py","span":"L91-L196","signature":"class WorkerGroup","exported":true,"origin":"ast","body_hash":"d5f99a41bdc56f57c447bcff45a7ad2fdd4fe1a3f1b6e0d8dee6a6250132266e","body_text":"class WorkerGroup: def __init__(self, resource_pool: ResourcePool, **kwargs) -> None: self._is_init_with_detached_workers = True if resource_pool is None else False if resource_pool is not None: # handle the case when WorkGroup is attached to an existing one self._procecss_dispatch_config = resource_pool() else: self._procecss_dispatch_config = None self._workers = [] self._worker_names = [] self._master_addr = None self._master_port = None self._checker_thread: threading.Thread = None def _is_worker_alive(self, worker): raise NotImplementedError(f\"WorkerGroup._is_worker_alive called, should be implemented in derived class.\") def _block_until_all_workers_alive(self) -> None: while True: all_state = [self._is_worker_alive(worker) for worker in self._workers] if False in all_state: time.sleep(1) else: break def start_worker_aliveness_check(self, every_n_seconds=1) -> None: # before starting checking worker aliveness, make sure all workers are already alive self._block_until_all_workers_alive() self._checker_thread = threading.Thread(target=check_workers_alive, args=(self._workers, self._is_worker_alive, every_n_seconds)) self._checker_thread.start() @property def world_size(self): return len(self._workers) # execute_all_async and execute_rank_zero_async should be implemented by RayWorkerGroup, TorchRPCWorkerGroup, # MegatronWorkerGroup, XperfWorkerGroup should skip def _bind_worker_method(self, user_defined_cls, func_generator): \"\"\" Bind the worker method to the WorkerGroup \"\"\" for method_name in dir(user_defined_cls): try: method = getattr(user_defined_cls, method_name) assert callable(method), f\"{method_name} in {user_defined_cls} is not callable\" except Exception as e: # if it is a property, it will fail because Class doesn't have instance property continue if hasattr(method, MAGIC_ATTR): # this method is decorated by register attribute = getattr(method, MAGIC_ATTR) assert isinstance(attribute, Dict), f'attribute must be a dictionary. Got {type(attribute)}' assert 'dispatch_mode' in attribute, f'attribute must contain dispatch_mode in its key' dispatch_mode = attribute['dispatch_mode'] execute_mode = attribute['execute_mode'] blocking = attribute['blocking'] # get dispatch fn if isinstance(dispatch_mode, Dispatch): # get default dispatch fn fn = get_predefined_dispatch_fn(dispatch_mode=dispatch_mode) dispatch_fn = fn['dispatch_fn'] collect_fn = fn['collect_fn'] else: assert isinstance(dispatch_mode, dict) assert 'dispatch_fn' in dispatch_mode assert 'collect_fn' in dispatch_mode dispatch_fn = dispatch_mode['dispatch_fn'] collect_fn = dispatch_mode['collect_fn'] # 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 as e: print(f'execute_fn {wg_execute_fn_name} is invalid') raise # 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) except Exception as e: raise ValueError(f'Fail to set method_name {method_name}')","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/base/worker_group.py#WorkerGroup.__init__","name":"__init__","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L93-L108","signature":"def __init__(self, resource_pool: ResourcePool, **kwargs) -> None","exported":false,"origin":"ast","body_hash":"ef8c49e82dd85901839cd24d733b6193a87919798fdee4f44332604427868864","body_text":"def __init__(self, resource_pool: ResourcePool, **kwargs) -> None: self._is_init_with_detached_workers = True if resource_pool is None else False if resource_pool is not None: # handle the case when WorkGroup is attached to an existing one self._procecss_dispatch_config = resource_pool() else: self._procecss_dispatch_config = None self._workers = [] self._worker_names = [] self._master_addr = None self._master_port = None self._checker_thread: threading.Thread = None","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerGroup"},{"id":"verl/single_controller/base/worker_group.py#WorkerGroup._is_worker_alive","name":"_is_worker_alive","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L110-L111","signature":"def _is_worker_alive(self, worker)","exported":false,"origin":"ast","body_hash":"1be54ab86883dfa0290b60184e5ac8fda598ba43e516fc0c3e6b8face8a2b966","body_text":"def _is_worker_alive(self, worker): raise NotImplementedError(f\"WorkerGroup._is_worker_alive called, should be implemented in derived class.\")","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerGroup"},{"id":"verl/single_controller/base/worker_group.py#WorkerGroup._block_until_all_workers_alive","name":"_block_until_all_workers_alive","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L113-L119","signature":"def _block_until_all_workers_alive(self) -> None","exported":false,"origin":"ast","body_hash":"9906a6236dd47746117e388e245f2a9c57a18c87be027af9a2dca0c3b33ca84a","body_text":"def _block_until_all_workers_alive(self) -> None: while True: all_state = [self._is_worker_alive(worker) for worker in self._workers] if False in all_state: time.sleep(1) else: break","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerGroup"},{"id":"verl/single_controller/base/worker_group.py#WorkerGroup.start_worker_aliveness_check","name":"start_worker_aliveness_check","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L121-L127","signature":"def start_worker_aliveness_check(self, every_n_seconds=1) -> None: # before starting checking worker aliveness, make sure all workers are already alive","exported":true,"origin":"ast","body_hash":"9f938d3e391bd6c8c89ea7cea3f5832b55fa6b8720605140b9a5a6cea31beac2","body_text":"def start_worker_aliveness_check(self, every_n_seconds=1) -> None: # before starting checking worker aliveness, make sure all workers are already alive self._block_until_all_workers_alive() self._checker_thread = threading.Thread(target=check_workers_alive, args=(self._workers, self._is_worker_alive, every_n_seconds)) self._checker_thread.start()","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerGroup"},{"id":"verl/single_controller/base/worker_group.py#WorkerGroup.world_size","name":"world_size","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L130-L131","signature":"def world_size(self)","exported":true,"origin":"ast","body_hash":"f25442dfd2f49f19c9e340c0e549bdaa8952e48737ecc9df229de289d1839978","body_text":"def world_size(self): return len(self._workers)","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerGroup"},{"id":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","name":"_bind_worker_method","kind":"method","path":"verl/single_controller/base/worker_group.py","span":"L136-L196","signature":"def _bind_worker_method(self, user_defined_cls, func_generator)","exported":false,"origin":"ast","body_hash":"e0fda6de434d941712d06bf2cab529319b6c02251ec44a38cfdaaff1924aaae4","body_text":"def _bind_worker_method(self, user_defined_cls, func_generator): \"\"\" Bind the worker method to the WorkerGroup \"\"\" for method_name in dir(user_defined_cls): try: method = getattr(user_defined_cls, method_name) assert callable(method), f\"{method_name} in {user_defined_cls} is not callable\" except Exception as e: # if it is a property, it will fail because Class doesn't have instance property continue if hasattr(method, MAGIC_ATTR): # this method is decorated by register attribute = getattr(method, MAGIC_ATTR) assert isinstance(attribute, Dict), f'attribute must be a dictionary. Got {type(attribute)}' assert 'dispatch_mode' in attribute, f'attribute must contain dispatch_mode in its key' dispatch_mode = attribute['dispatch_mode'] execute_mode = attribute['execute_mode'] blocking = attribute['blocking'] # get dispatch fn if isinstance(dispatch_mode, Dispatch): # get default dispatch fn fn = get_predefined_dispatch_fn(dispatch_mode=dispatch_mode) dispatch_fn = fn['dispatch_fn'] collect_fn = fn['collect_fn'] else: assert isinstance(dispatch_mode, dict) assert 'dispatch_fn' in dispatch_mode assert 'collect_fn' in dispatch_mode dispatch_fn = dispatch_mode['dispatch_fn'] collect_fn = dispatch_mode['collect_fn'] # 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 as e: print(f'execute_fn {wg_execute_fn_name} is invalid') raise # 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) except Exception as e: raise ValueError(f'Fail to set method_name {method_name}')","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerGroup"}],"rawEdges":[{"source":"verl/single_controller/base/worker_group.py","relation":"imports","specifier":"logging","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py","relation":"imports","specifier":"threading","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py","relation":"imports","specifier":"signal","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py","relation":"imports","specifier":"time","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py","relation":"imports","specifier":"typing","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py","relation":"imports","specifier":"verl.single_controller.base.decorator","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ResourcePool","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ResourcePool.__init__","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ResourcePool.add_node","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool.add_node","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ResourcePool.world_size","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool.world_size","relation":"calls","name":"sum","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ResourcePool.__call__","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ResourcePool.store","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ResourcePool.local_world_size_list","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool.local_world_size_list","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ResourcePool.local_rank_list","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ResourcePool.local_rank_list","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ClassWithInitArgs","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ClassWithInitArgs","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ClassWithInitArgs.__init__","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ClassWithInitArgs","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#ClassWithInitArgs.__call__","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#ClassWithInitArgs.__call__","relation":"calls","name":"cls","viaMember":true,"file":"verl/single_controller/base/worker_group.py","recvType":"ClassWithInitArgs"},{"source":"verl/single_controller/base/worker_group.py","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#check_workers_alive","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py","relation":"imports","specifier":"time","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#check_workers_alive","relation":"calls","name":"is_alive","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#check_workers_alive","relation":"calls","name":"warning","viaMember":true,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#check_workers_alive","relation":"calls","name":"raise_signal","viaMember":true,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#check_workers_alive","relation":"calls","name":"sleep","viaMember":true,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#WorkerGroup","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#WorkerGroup.__init__","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup.__init__","relation":"calls","name":"resource_pool","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#WorkerGroup._is_worker_alive","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._is_worker_alive","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#WorkerGroup._block_until_all_workers_alive","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._block_until_all_workers_alive","relation":"calls","name":"_is_worker_alive","viaMember":true,"file":"verl/single_controller/base/worker_group.py","recvType":"WorkerGroup"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._block_until_all_workers_alive","relation":"calls","name":"sleep","viaMember":true,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#WorkerGroup.start_worker_aliveness_check","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup.start_worker_aliveness_check","relation":"calls","name":"_block_until_all_workers_alive","viaMember":true,"file":"verl/single_controller/base/worker_group.py","recvType":"WorkerGroup"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup.start_worker_aliveness_check","relation":"calls","name":"Thread","viaMember":true,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup.start_worker_aliveness_check","relation":"calls","name":"start","viaMember":true,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#WorkerGroup.world_size","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup.world_size","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup","relation":"contains","targetId":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"dir","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"getattr","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"callable","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"getattr","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"type","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"get_predefined_dispatch_fn","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"get_predefined_execute_fn","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"getattr","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"callable","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"print","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"func_generator","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"setattr","viaMember":false,"file":"verl/single_controller/base/worker_group.py"},{"source":"verl/single_controller/base/worker_group.py#WorkerGroup._bind_worker_method","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/single_controller/base/worker_group.py"}]},"verl/single_controller/ray/__init__.py":{"size":778,"mtimeMs":1789275830311.2026,"hash":"fc7a2d9064a035f948c196db973149438f977b41d81e55a5e05a3e599c9101ab","nodes":[{"id":"verl/single_controller/ray/__init__.py","name":"__init__.py","kind":"file","path":"verl/single_controller/ray/__init__.py","span":"L1-L16","signature":null,"exported":true,"origin":"ast","body_hash":"fc7a2d9064a035f948c196db973149438f977b41d81e55a5e05a3e599c9101ab","chars":778,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 RayResourcePool, RayClassWithInitArgs, RayWorkerGroup, create_colocated_worker_cls from .megatron import (MegatronRayWorkerGroup, DistRankInfo, DistGlobalInfo)"}],"rawEdges":[{"source":"verl/single_controller/ray/__init__.py","relation":"imports","specifier":".base","file":"verl/single_controller/ray/__init__.py"},{"source":"verl/single_controller/ray/__init__.py","relation":"imports","specifier":".megatron","file":"verl/single_controller/ray/__init__.py"}]},"verl/single_controller/ray/base.py":{"size":19238,"mtimeMs":1789275830311.2026,"hash":"b0dca8aad381d4e9a60169df108af2035fd0b5f88ab38dfc833689c842d4caa9","nodes":[{"id":"verl/single_controller/ray/base.py","name":"base.py","kind":"file","path":"verl/single_controller/ray/base.py","span":"L1-L460","signature":null,"exported":true,"origin":"ast","body_hash":"b0dca8aad381d4e9a60169df108af2035fd0b5f88ab38dfc833689c842d4caa9","chars":19136,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 time from typing import Dict, List, Any, Tuple import ray from ray.util import list_named_actors from ray.util.placement_group import placement_group, PlacementGroup from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy, NodeAffinitySchedulingStrategy from ray.experimental.state.api import get_actor from verl.single_controller.base import WorkerGroup, ResourcePool, ClassWithInitArgs, Worker __all__ = ['Worker'] \"\"\" Utilities that enables creating workers inside the same ray.Actor, with code written in separate ray.Actors. \"\"\" from unittest.mock import patch from verl.single_controller.base.decorator import MAGIC_ATTR import os"},{"id":"verl/single_controller/ray/base.py#get_random_string","name":"get_random_string","kind":"function","path":"verl/single_controller/ray/base.py","span":"L29-L33","signature":"def get_random_string(length: int) -> str","exported":true,"origin":"ast","body_hash":"3e24edfffba333e2f7f83cf9dd6e63d295e7d51adf299e71dcaf88d25271f8e5","body_text":"def get_random_string(length: int) -> str: import random import string letters_digits = string.ascii_letters + string.digits return ''.join(random.choice(letters_digits) for _ in range(length))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#func_generator","name":"func_generator","kind":"function","path":"verl/single_controller/ray/base.py","span":"L36-L46","signature":"def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking)","exported":true,"origin":"ast","body_hash":"7514162bce755d56f76af616c8cbbfe8ab88fc3e618b9a5e1abf3605a259647f","body_text":"def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking): def func(*args, **kwargs): args, kwargs = dispatch_fn(self, *args, **kwargs) output = execute_fn(method_name, *args, **kwargs) if blocking: output = ray.get(output) output = collect_fn(self, output) return output return func","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#func_generator.func","name":"func","kind":"function","path":"verl/single_controller/ray/base.py","span":"L38-L44","signature":"def func(*args, **kwargs)","exported":true,"origin":"ast","body_hash":"38f209504006c796a82fc34c8dc896b61bd199dc096f677a5bd4618db7a02010","body_text":"def func(*args, **kwargs): args, kwargs = dispatch_fn(self, *args, **kwargs) output = execute_fn(method_name, *args, **kwargs) if blocking: output = ray.get(output) output = collect_fn(self, output) return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#RayResourcePool","name":"RayResourcePool","kind":"class","path":"verl/single_controller/ray/base.py","span":"L49-L88","signature":"class RayResourcePool(ResourcePool)","exported":true,"origin":"ast","body_hash":"aea0756fedd913a2e113435d0b70d2f37c067e112ccdc94514617b5e1ae55128","body_text":"class RayResourcePool(ResourcePool): def __init__(self, process_on_nodes: List[int] = None, use_gpu: bool = True, name_prefix: str = \"\", max_colocate_count: int = 5, detached=False) -> None: super().__init__(process_on_nodes, max_colocate_count) self.use_gpu = use_gpu # print(f\"in RayProcessDispatchConfiguration: name_prefix = {name_prefix}\") self.name_prefix = name_prefix self.pgs = None self.detached = detached def get_placement_groups(self, strategy=\"STRICT_PACK\", name=None): if self.pgs is not None: return self.pgs pg_name_prefix = name if name else \\ f\"{self.name_prefix}verl_group_{'_'.join([str(count) for count in self._store])}:\" # print(f\"pg_name_prefix = {pg_name_prefix}\") pg_scheme = [[{ \"CPU\": self.max_collocate_count, \"GPU\": 1 } if self.use_gpu else { \"CPU\": self.max_collocate_count } for _ in range(process_count)] for process_count in self._store] lifetime = 'detached' if self.detached else None pgs = [ placement_group(bundles=bundles, strategy=strategy, name=pg_name_prefix + str(idx), lifetime=lifetime) for idx, bundles in enumerate(pg_scheme) ] ray.get([pg.ready() for pg in pgs]) self.pgs = pgs return pgs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#RayResourcePool.__init__","name":"__init__","kind":"method","path":"verl/single_controller/ray/base.py","span":"L51-L62","signature":"def __init__(self, process_on_nodes: List[int] = None, use_gpu: bool = True, name_prefix: str = \"\", max_colocate_count: int = 5, detached=False) -> None","exported":false,"origin":"ast","body_hash":"93c0a38e7ec61a7cb0c228c3c243a8bb6127549e54b95f2b587c93eb28800acc","body_text":"def __init__(self, process_on_nodes: List[int] = None, use_gpu: bool = True, name_prefix: str = \"\", max_colocate_count: int = 5, detached=False) -> None: super().__init__(process_on_nodes, max_colocate_count) self.use_gpu = use_gpu # print(f\"in RayProcessDispatchConfiguration: name_prefix = {name_prefix}\") self.name_prefix = name_prefix self.pgs = None self.detached = detached","summary_state":"pending","summary":null,"crux":null,"owner":"RayResourcePool"},{"id":"verl/single_controller/ray/base.py#RayResourcePool.get_placement_groups","name":"get_placement_groups","kind":"method","path":"verl/single_controller/ray/base.py","span":"L64-L88","signature":"def get_placement_groups(self, strategy=\"STRICT_PACK\", name=None)","exported":true,"origin":"ast","body_hash":"18c20a0b8d9530a38bed81637f7a7e7eb420aecb2f47c83f680ebd0790acd56d","body_text":"def get_placement_groups(self, strategy=\"STRICT_PACK\", name=None): if self.pgs is not None: return self.pgs pg_name_prefix = name if name else \\ f\"{self.name_prefix}verl_group_{'_'.join([str(count) for count in self._store])}:\" # print(f\"pg_name_prefix = {pg_name_prefix}\") pg_scheme = [[{ \"CPU\": self.max_collocate_count, \"GPU\": 1 } if self.use_gpu else { \"CPU\": self.max_collocate_count } for _ in range(process_count)] for process_count in self._store] lifetime = 'detached' if self.detached else None pgs = [ placement_group(bundles=bundles, strategy=strategy, name=pg_name_prefix + str(idx), lifetime=lifetime) for idx, bundles in enumerate(pg_scheme) ] ray.get([pg.ready() for pg in pgs]) self.pgs = pgs return pgs","summary_state":"pending","summary":null,"crux":null,"owner":"RayResourcePool"},{"id":"verl/single_controller/ray/base.py#extract_pg_from_exist","name":"extract_pg_from_exist","kind":"function","path":"verl/single_controller/ray/base.py","span":"L91-L111","signature":"def extract_pg_from_exist(resource_pools: Dict[str, RayResourcePool], src_role_names: List[str], resource_pool: RayResourcePool) -> List","exported":true,"origin":"ast","body_hash":"e5d02b00f3194217daeda2d7bb7ea3381df07874c0d00a77a3d4fe0875120b39","body_text":"def extract_pg_from_exist(resource_pools: Dict[str, RayResourcePool], src_role_names: List[str], resource_pool: RayResourcePool) -> List: src_pgs = [ pg for role_name, resource_pool in resource_pools.items() for pg in resource_pool.get_placement_groups() if role_name in src_role_names ] sorted_src_pgs = sorted(src_pgs, key=lambda pg: pg.bundle_count, reverse=True) sorted_process_on_nodes = sorted([(val, idx) for idx, val in enumerate(resource_pool.store)], reverse=True) unsorted_pgs: List[Tuple[int, PlacementGroup]] = [] searching_idx = 0 for request_process, original_idx in sorted_process_on_nodes: assert searching_idx < len(sorted_src_pgs), f\"no enough nodes for request: searching {searching_idx} th node\" assert request_process <= sorted_src_pgs[searching_idx].bundle_count, \\ f\"requesting {request_process} processes, bundle count cannot satisfy\" unsorted_pgs.append((original_idx, sorted_src_pgs[searching_idx])) searching_idx += 1 return [pg for _, pg in sorted(unsorted_pgs)]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#merge_resource_pool","name":"merge_resource_pool","kind":"function","path":"verl/single_controller/ray/base.py","span":"L114-L125","signature":"def merge_resource_pool(rp1: RayResourcePool, rp2: RayResourcePool) -> RayResourcePool","exported":true,"origin":"ast","body_hash":"7fd2fbb11903893a4fc91b0430f4b6192d143da8d0314776f1bfdd683e59387d","body_text":"def merge_resource_pool(rp1: RayResourcePool, rp2: RayResourcePool) -> RayResourcePool: assert rp1.use_gpu == rp2.use_gpu, 'Both RayResourcePool must either use_gpu or not' assert rp1.max_collocate_count == rp2.max_collocate_count, 'Both RayResourcePool must has the same max_collocate_count' assert rp1.n_gpus_per_node == rp2.n_gpus_per_node, 'Both RayResourcePool must has the same n_gpus_per_node' assert rp1.detached == rp2.detached, 'Detached ResourcePool cannot be merged with non-detached ResourcePool' new_store = rp1.store + rp2.store merged = RayResourcePool(new_store, rp1.use_gpu, f\"{rp1.name_prefix}_{rp2.name_prefix}\") merged.pgs = rp1.get_placement_groups() + rp2.get_placement_groups() return merged","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#RayClassWithInitArgs","name":"RayClassWithInitArgs","kind":"class","path":"verl/single_controller/ray/base.py","span":"L128-L173","signature":"class RayClassWithInitArgs(ClassWithInitArgs)","exported":true,"origin":"ast","body_hash":"37df0678dacee83647df7549c8f85569034007293110f309e827ace9e41acfe5","body_text":"class RayClassWithInitArgs(ClassWithInitArgs): def __init__(self, cls, *args, **kwargs) -> None: # self._options = kwargs.pop('options', dict()) super().__init__(cls, *args, **kwargs) self._options = {} self._additional_resource = {} def set_additional_resource(self, additional_resource): self._additional_resource = additional_resource def update_options(self, options: Dict): self._options.update(options) def __call__(self, placement_group, placement_group_bundle_idx, use_gpu: bool = True, num_gpus=1, sharing_with=None) -> Any: if sharing_with is not None: target_node_id = ray.get(sharing_with.get_node_id.remote()) cuda_visible_devices = ray.get(sharing_with.get_cuda_visible_devices.remote()) options = {\"scheduling_strategy\": NodeAffinitySchedulingStrategy(node_id=target_node_id, soft=False)} return self.cls.options(**options).remote(*self.args, cuda_visible_devices=cuda_visible_devices, **self.kwargs) options = { \"scheduling_strategy\": PlacementGroupSchedulingStrategy(placement_group=placement_group, placement_group_bundle_index=placement_group_bundle_idx) } options.update(self._options) if use_gpu: options[\"num_gpus\"] = num_gpus if len(self._additional_resource) > 1: for k, v in self._additional_resource.items(): options[k] = v # print(\"cls:\", self.cls) # print(\"args: \", self.args) # print(\"kwargs: \", self.kwargs) return self.cls.options(**options).remote(*self.args, **self.kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__init__","name":"__init__","kind":"method","path":"verl/single_controller/ray/base.py","span":"L130-L134","signature":"def __init__(self, cls, *args, **kwargs) -> None: # self._options = kwargs.pop('options', dict())","exported":false,"origin":"ast","body_hash":"732c0a12c8f730ae38a59a43c96e859bdda20edb93becb6bede90cfba5b9be70","body_text":"def __init__(self, cls, *args, **kwargs) -> None: # self._options = kwargs.pop('options', dict()) super().__init__(cls, *args, **kwargs) self._options = {} self._additional_resource = {}","summary_state":"pending","summary":null,"crux":null,"owner":"RayClassWithInitArgs"},{"id":"verl/single_controller/ray/base.py#RayClassWithInitArgs.set_additional_resource","name":"set_additional_resource","kind":"method","path":"verl/single_controller/ray/base.py","span":"L136-L137","signature":"def set_additional_resource(self, additional_resource)","exported":true,"origin":"ast","body_hash":"fad07a2c5c3a49a0d5184f9468186d9b8232a0f892b658118c7fd98aa26d5c1b","body_text":"def set_additional_resource(self, additional_resource): self._additional_resource = additional_resource","summary_state":"pending","summary":null,"crux":null,"owner":"RayClassWithInitArgs"},{"id":"verl/single_controller/ray/base.py#RayClassWithInitArgs.update_options","name":"update_options","kind":"method","path":"verl/single_controller/ray/base.py","span":"L139-L140","signature":"def update_options(self, options: Dict)","exported":true,"origin":"ast","body_hash":"f20c0f4f0973e25760eeed7bb9807d69e6a95ab9c782015a444d08209db8028d","body_text":"def update_options(self, options: Dict): self._options.update(options)","summary_state":"pending","summary":null,"crux":null,"owner":"RayClassWithInitArgs"},{"id":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","name":"__call__","kind":"method","path":"verl/single_controller/ray/base.py","span":"L142-L173","signature":"def __call__(self, placement_group, placement_group_bundle_idx, use_gpu: bool = True, num_gpus=1, sharing_with=None) -> Any","exported":false,"origin":"ast","body_hash":"6ac4bc86dd1ed6fc75f33bf496ef7a7fe82613ae44c9e557acced12ca1ddf0ab","body_text":"def __call__(self, placement_group, placement_group_bundle_idx, use_gpu: bool = True, num_gpus=1, sharing_with=None) -> Any: if sharing_with is not None: target_node_id = ray.get(sharing_with.get_node_id.remote()) cuda_visible_devices = ray.get(sharing_with.get_cuda_visible_devices.remote()) options = {\"scheduling_strategy\": NodeAffinitySchedulingStrategy(node_id=target_node_id, soft=False)} return self.cls.options(**options).remote(*self.args, cuda_visible_devices=cuda_visible_devices, **self.kwargs) options = { \"scheduling_strategy\": PlacementGroupSchedulingStrategy(placement_group=placement_group, placement_group_bundle_index=placement_group_bundle_idx) } options.update(self._options) if use_gpu: options[\"num_gpus\"] = num_gpus if len(self._additional_resource) > 1: for k, v in self._additional_resource.items(): options[k] = v # print(\"cls:\", self.cls) # print(\"args: \", self.args) # print(\"kwargs: \", self.kwargs) return self.cls.options(**options).remote(*self.args, **self.kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"RayClassWithInitArgs"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup","name":"RayWorkerGroup","kind":"class","path":"verl/single_controller/ray/base.py","span":"L176-L367","signature":"class RayWorkerGroup(WorkerGroup)","exported":true,"origin":"ast","body_hash":"8b019844908a42d8047067a87ab145cfd145f826f921a613e88b3bd5a7053eed","body_text":"class RayWorkerGroup(WorkerGroup): def __init__(self, resource_pool: RayResourcePool = None, ray_cls_with_init: RayClassWithInitArgs = None, bin_pack: bool = True, name_prefix: str = None, detached=False, worker_names=None, **kwargs) -> None: super().__init__(resource_pool=resource_pool, **kwargs) self.ray_cls_with_init = ray_cls_with_init self.name_prefix = get_random_string(length=6) if name_prefix is None else name_prefix if worker_names is not None: assert self._is_init_with_detached_workers self._worker_names = worker_names if self._is_init_with_detached_workers: self._init_with_detached_workers(worker_names=worker_names) else: self._init_with_resource_pool(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, bin_pack=bin_pack, detached=detached) if ray_cls_with_init is not None: self._bind_worker_method(self.ray_cls_with_init.cls, func_generator) def _is_worker_alive(self, worker: ray.actor.ActorHandle): worker_state_dict = get_actor(worker._actor_id.hex()) return worker_state_dict.get(\"state\", \"undefined\") == \"ALIVE\" if worker_state_dict is not None else False def _init_with_detached_workers(self, worker_names): workers = [ray.get_actor(name=name) for name in worker_names] self._workers = workers self._world_size = len(worker_names) def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, detached): use_gpu = resource_pool.use_gpu strategy = \"PACK\" if bin_pack: strategy = \"STRICT_PACK\" pgs = resource_pool.get_placement_groups(strategy=strategy) world_size = resource_pool.world_size self._world_size = world_size # cia.add_kwarg(\"_world_size\", world_size) num_gpus = 1 / resource_pool.max_collocate_count rank = -1 for pg_idx, local_world_size in enumerate(resource_pool.store): pg = pgs[pg_idx] assert local_world_size <= pg.bundle_count, \\ f\"when generating for {self.name_prefix}, for the \" for local_rank in range(local_world_size): rank += 1 # we pass in environment variable at option so that Worker can use environment variable to set env_vars = { 'WORLD_SIZE': str(world_size), 'RANK': str(rank), 'WG_PREFIX': self.name_prefix, 'WG_BACKEND': 'ray', 'RAY_LOCAL_WORLD_SIZE': str(local_world_size), 'RAY_LOCAL_RANK': str(local_rank), } if rank != 0: env_vars['MASTER_ADDR'] = self._master_addr env_vars['MASTER_PORT'] = self._master_port import re cia_name = type(ray_cls_with_init.cls).__name__ match = re.search(r\"ActorClass\\(([^)]+)\\)\", cia_name) # ray.remote(Obj) -> \"ActorClass(Obj)\" cia_name = match.group(1) if match else cia_name # \"ActorClass(Obj)\" -> \"Obj\" name = f\"{self.name_prefix}{cia_name}_{pg_idx}:{local_rank}\" # e.g. Worker_2:5 ray_cls_with_init.update_options({'runtime_env': {'env_vars': env_vars}, 'name': name}) if detached: ray_cls_with_init.update_options({'lifetime': 'detached'}) # create a worker worker = ray_cls_with_init(placement_group=pg, placement_group_bundle_idx=local_rank, use_gpu=use_gpu, num_gpus=num_gpus) self._workers.append(worker) self._worker_names.append(name) if rank == 0: register_center_actor = None for _ in range(120): if f\"{self.name_prefix}_register_center\" not in list_named_actors(): time.sleep(1) else: register_center_actor = ray.get_actor(f\"{self.name_prefix}_register_center\") break assert register_center_actor is not None, f\"failed to get register_center_actor: {self.name_prefix}_register_center in {list_named_actors(all_namespaces=True)}\" rank_zero_info = ray.get(register_center_actor.get_rank_zero_info.remote()) self._master_addr, self._master_port = rank_zero_info['MASTER_ADDR'], rank_zero_info['MASTER_PORT'] # print(f\"rank_zero_info: {rank_zero_info}\") # print(f\"master_addr: {self._master_addr}, master_port: {self._master_port}\") @property def worker_names(self): return self._worker_names @classmethod def from_detached(cls, worker_names=None, ray_cls_with_init=None): worker_group = cls(resource_pool=None, ray_cls_with_init=ray_cls_with_init, name_prefix=None, worker_names=worker_names) return worker_group def spawn(self, prefix_set): \"\"\" spawn to a dictionary of worker groups, each with a subset of method with prefix. \"\"\" def _rebind_actor_methods(worker_group, actor_name): \"\"\" bind the method with actor_prefix to its original name \"\"\" prefix: str = actor_name + '_' for method_name in dir(worker_group): if method_name.startswith(prefix): # only valid when Python >= 3.9 original_method_name = method_name.removeprefix(prefix) method = getattr(worker_group, method_name) setattr(worker_group, original_method_name, method) new_worker_group_dict = {} for prefix in prefix_set: new_worker_group = self.from_detached(worker_names=self._worker_names, ray_cls_with_init=self.ray_cls_with_init) _rebind_actor_methods(new_worker_group, prefix) new_worker_group_dict[prefix] = new_worker_group return new_worker_group_dict def execute_rank_zero_sync(self, method_name: str, *args, **kwargs): return ray.get(self.execute_all_async(method_name, **args, **kwargs)) def execute_rank_zero_async(self, method_name: str, *args, **kwargs): remote_call = getatt","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.__init__","name":"__init__","kind":"method","path":"verl/single_controller/ray/base.py","span":"L178-L203","signature":"def __init__(self, resource_pool: RayResourcePool = None, ray_cls_with_init: RayClassWithInitArgs = None, bin_pack: bool = True, name_prefix: str = None, detached=False, worker_names=None, **kwargs) -> None","exported":false,"origin":"ast","body_hash":"44c99602174e6c5f30401932370dc1a6788bcf6d1a4936a23e860188c5390347","body_text":"def __init__(self, resource_pool: RayResourcePool = None, ray_cls_with_init: RayClassWithInitArgs = None, bin_pack: bool = True, name_prefix: str = None, detached=False, worker_names=None, **kwargs) -> None: super().__init__(resource_pool=resource_pool, **kwargs) self.ray_cls_with_init = ray_cls_with_init self.name_prefix = get_random_string(length=6) if name_prefix is None else name_prefix if worker_names is not None: assert self._is_init_with_detached_workers self._worker_names = worker_names if self._is_init_with_detached_workers: self._init_with_detached_workers(worker_names=worker_names) else: self._init_with_resource_pool(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, bin_pack=bin_pack, detached=detached) if ray_cls_with_init is not None: self._bind_worker_method(self.ray_cls_with_init.cls, func_generator)","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup._is_worker_alive","name":"_is_worker_alive","kind":"method","path":"verl/single_controller/ray/base.py","span":"L205-L207","signature":"def _is_worker_alive(self, worker: ray.actor.ActorHandle)","exported":false,"origin":"ast","body_hash":"1a42cdb1ee7fe76488de243adc58e3ddb3d1ab65f0fc2cd8f5275d99130caff0","body_text":"def _is_worker_alive(self, worker: ray.actor.ActorHandle): worker_state_dict = get_actor(worker._actor_id.hex()) return worker_state_dict.get(\"state\", \"undefined\") == \"ALIVE\" if worker_state_dict is not None else False","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_detached_workers","name":"_init_with_detached_workers","kind":"method","path":"verl/single_controller/ray/base.py","span":"L209-L212","signature":"def _init_with_detached_workers(self, worker_names)","exported":false,"origin":"ast","body_hash":"52d676a5f8278f1aa40920778bf35a0a1b3c28708deb52612da9231c7c834f6f","body_text":"def _init_with_detached_workers(self, worker_names): workers = [ray.get_actor(name=name) for name in worker_names] self._workers = workers self._world_size = len(worker_names)","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","name":"_init_with_resource_pool","kind":"method","path":"verl/single_controller/ray/base.py","span":"L214-L278","signature":"def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, detached)","exported":false,"origin":"ast","body_hash":"93fd9c9ce2bd17539ec265c6c5cfe5dcea90512436e6f164f525e91ad85faea0","body_text":"def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, detached): use_gpu = resource_pool.use_gpu strategy = \"PACK\" if bin_pack: strategy = \"STRICT_PACK\" pgs = resource_pool.get_placement_groups(strategy=strategy) world_size = resource_pool.world_size self._world_size = world_size # cia.add_kwarg(\"_world_size\", world_size) num_gpus = 1 / resource_pool.max_collocate_count rank = -1 for pg_idx, local_world_size in enumerate(resource_pool.store): pg = pgs[pg_idx] assert local_world_size <= pg.bundle_count, \\ f\"when generating for {self.name_prefix}, for the \" for local_rank in range(local_world_size): rank += 1 # we pass in environment variable at option so that Worker can use environment variable to set env_vars = { 'WORLD_SIZE': str(world_size), 'RANK': str(rank), 'WG_PREFIX': self.name_prefix, 'WG_BACKEND': 'ray', 'RAY_LOCAL_WORLD_SIZE': str(local_world_size), 'RAY_LOCAL_RANK': str(local_rank), } if rank != 0: env_vars['MASTER_ADDR'] = self._master_addr env_vars['MASTER_PORT'] = self._master_port import re cia_name = type(ray_cls_with_init.cls).__name__ match = re.search(r\"ActorClass\\(([^)]+)\\)\", cia_name) # ray.remote(Obj) -> \"ActorClass(Obj)\" cia_name = match.group(1) if match else cia_name # \"ActorClass(Obj)\" -> \"Obj\" name = f\"{self.name_prefix}{cia_name}_{pg_idx}:{local_rank}\" # e.g. Worker_2:5 ray_cls_with_init.update_options({'runtime_env': {'env_vars': env_vars}, 'name': name}) if detached: ray_cls_with_init.update_options({'lifetime': 'detached'}) # create a worker worker = ray_cls_with_init(placement_group=pg, placement_group_bundle_idx=local_rank, use_gpu=use_gpu, num_gpus=num_gpus) self._workers.append(worker) self._worker_names.append(name) if rank == 0: register_center_actor = None for _ in range(120): if f\"{self.name_prefix}_register_center\" not in list_named_actors(): time.sleep(1) else: register_center_actor = ray.get_actor(f\"{self.name_prefix}_register_center\") break assert register_center_actor is not None, f\"failed to get register_center_actor: {self.name_prefix}_register_center in {list_named_actors(all_namespaces=True)}\" rank_zero_info = ray.get(register_center_actor.get_rank_zero_info.remote()) self._master_addr, self._master_port = rank_zero_info['MASTER_ADDR'], rank_zero_info['MASTER_PORT'] # print(f\"rank_zero_info: {rank_zero_info}\") # print(f\"master_addr: {self._master_addr}, master_port: {self._master_port}\")","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.worker_names","name":"worker_names","kind":"method","path":"verl/single_controller/ray/base.py","span":"L281-L282","signature":"def worker_names(self)","exported":true,"origin":"ast","body_hash":"8f8aba5b9d536448b8fec0562593ea43911b7672649f8843c8b36eba7ccb3310","body_text":"def worker_names(self): return self._worker_names","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.from_detached","name":"from_detached","kind":"method","path":"verl/single_controller/ray/base.py","span":"L285-L290","signature":"def from_detached(cls, worker_names=None, ray_cls_with_init=None)","exported":true,"origin":"ast","body_hash":"9bf8230a5b5dc26c87b641cc3b2f52482a246442c6442754ea20fd0698d9ff1b","body_text":"def from_detached(cls, worker_names=None, ray_cls_with_init=None): worker_group = cls(resource_pool=None, ray_cls_with_init=ray_cls_with_init, name_prefix=None, worker_names=worker_names) return worker_group","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn","name":"spawn","kind":"method","path":"verl/single_controller/ray/base.py","span":"L292-L317","signature":"def spawn(self, prefix_set)","exported":true,"origin":"ast","body_hash":"a8f2b4c31d33f96f8ceb9d74bc8125d754dfac848248612ef30bb92897e518b3","body_text":"def spawn(self, prefix_set): \"\"\" spawn to a dictionary of worker groups, each with a subset of method with prefix. \"\"\" def _rebind_actor_methods(worker_group, actor_name): \"\"\" bind the method with actor_prefix to its original name \"\"\" prefix: str = actor_name + '_' for method_name in dir(worker_group): if method_name.startswith(prefix): # only valid when Python >= 3.9 original_method_name = method_name.removeprefix(prefix) method = getattr(worker_group, method_name) setattr(worker_group, original_method_name, method) new_worker_group_dict = {} for prefix in prefix_set: new_worker_group = self.from_detached(worker_names=self._worker_names, ray_cls_with_init=self.ray_cls_with_init) _rebind_actor_methods(new_worker_group, prefix) new_worker_group_dict[prefix] = new_worker_group return new_worker_group_dict","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn._rebind_actor_methods","name":"_rebind_actor_methods","kind":"function","path":"verl/single_controller/ray/base.py","span":"L298-L308","signature":"def _rebind_actor_methods(worker_group, actor_name)","exported":false,"origin":"ast","body_hash":"10ef28021f5967fd3f101ce0125c6f25c1e6e32a9d5d2fc4c22512e4d2d547d6","body_text":"def _rebind_actor_methods(worker_group, actor_name): \"\"\" bind the method with actor_prefix to its original name \"\"\" prefix: str = actor_name + '_' for method_name in dir(worker_group): if method_name.startswith(prefix): # only valid when Python >= 3.9 original_method_name = method_name.removeprefix(prefix) method = getattr(worker_group, method_name) setattr(worker_group, original_method_name, method)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero_sync","name":"execute_rank_zero_sync","kind":"method","path":"verl/single_controller/ray/base.py","span":"L319-L320","signature":"def execute_rank_zero_sync(self, method_name: str, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"f93eb6f37a60191961200706a92ebc3679aaf1c0b2f4cd5ecce3b79adab7aa1e","body_text":"def execute_rank_zero_sync(self, method_name: str, *args, **kwargs): return ray.get(self.execute_all_async(method_name, **args, **kwargs))","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero_async","name":"execute_rank_zero_async","kind":"method","path":"verl/single_controller/ray/base.py","span":"L322-L324","signature":"def execute_rank_zero_async(self, method_name: str, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"7359d9a92b3f1c93fb817037e981d5ef45714fcf49a07167408a764760ae5252","body_text":"def execute_rank_zero_async(self, method_name: str, *args, **kwargs): remote_call = getattr(self._workers[0], method_name) return remote_call.remote(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero","name":"execute_rank_zero","kind":"method","path":"verl/single_controller/ray/base.py","span":"L326-L327","signature":"def execute_rank_zero(self, method_name: str, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"741229d564c8dba11619b935c0a7d280021f31d2d0a1fd5629c21c8673feec20","body_text":"def execute_rank_zero(self, method_name: str, *args, **kwargs): return self.execute_rank_zero_async(method_name, *args, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all","name":"execute_all","kind":"method","path":"verl/single_controller/ray/base.py","span":"L329-L330","signature":"def execute_all(self, method_name: str, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"9f62f32817ddc26f7bc81340c0b1c9632f1798a66e48066678ecb95622816fe6","body_text":"def execute_all(self, method_name: str, *args, **kwargs): return self.execute_all_async(method_name, *args, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_sync","name":"execute_all_sync","kind":"method","path":"verl/single_controller/ray/base.py","span":"L332-L333","signature":"def execute_all_sync(self, method_name: str, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"68232603e13cf83c78286a3d27ea662692db1c997b4f19f52add644f5f145ff0","body_text":"def execute_all_sync(self, method_name: str, *args, **kwargs): return ray.get(self.execute_all_async(method_name, *args, **kwargs))","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","name":"execute_all_async","kind":"method","path":"verl/single_controller/ray/base.py","span":"L335-L351","signature":"def execute_all_async(self, method_name: str, *args, **kwargs): # 这里我们假设，如果 args 和 kwargs 里面所有的参数都是 list，且所有的 list 长度都与 len(self._workers) 一致的话，我们会把 # list 中的每一个分别发到对应的 worker 上去 # print(f\"execute_all_async: method {method_name}({args}, {kwargs})\")","exported":true,"origin":"ast","body_hash":"0774fd9561923510eab36fb34e34831a5af9d2e9e01397e78bf5781d232cd491","body_text":"def execute_all_async(self, method_name: str, *args, **kwargs): # 这里我们假设，如果 args 和 kwargs 里面所有的参数都是 list，且所有的 list 长度都与 len(self._workers) 一致的话，我们会把 # list 中的每一个分别发到对应的 worker 上去 # print(f\"execute_all_async: method {method_name}({args}, {kwargs})\") length = len(self._workers) if all(isinstance(arg, list) for arg in args) and all(isinstance(kwarg, list) for kwarg in kwargs.values()): if all(len(arg) == length for arg in args) and all(len(kwarg) == length for kwarg in kwargs.values()): # print(f\"splitting args and kwargs into {length} shards\") result = [] for i in range(length): sliced_args = tuple(arg[i] for arg in args) sliced_kwargs = {k: v[i] for k, v in kwargs.items()} remote_call = getattr(self._workers[i], method_name) result.append(remote_call.remote(*sliced_args, **sliced_kwargs)) return result return [getattr(worker, method_name).remote(*args, **kwargs) for worker in self._workers]","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.master_address","name":"master_address","kind":"method","path":"verl/single_controller/ray/base.py","span":"L354-L355","signature":"def master_address(self)","exported":true,"origin":"ast","body_hash":"37eab4e50e4a82432da776f741a02fe86627fa0957cebd109d1c955bad6ac64f","body_text":"def master_address(self): return self._master_addr","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.master_port","name":"master_port","kind":"method","path":"verl/single_controller/ray/base.py","span":"L358-L359","signature":"def master_port(self)","exported":true,"origin":"ast","body_hash":"df0e5490698a8d734d28dd6fed144c74e5fa2c9c3c1f3fc8f0e3fb6884ba7021","body_text":"def master_port(self): return self._master_port","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.workers","name":"workers","kind":"method","path":"verl/single_controller/ray/base.py","span":"L362-L363","signature":"def workers(self)","exported":true,"origin":"ast","body_hash":"a657d062b61c60674c4ddc9e10af901e3b911f289bca9ff18f1281eb1e33bfe4","body_text":"def workers(self): return self._workers","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#RayWorkerGroup.world_size","name":"world_size","kind":"method","path":"verl/single_controller/ray/base.py","span":"L366-L367","signature":"def world_size(self)","exported":true,"origin":"ast","body_hash":"417053f3e044003b010d0dc3d1a7213269ba5b6bee8a5bcff37b385f3fcb0427","body_text":"def world_size(self): return self._world_size","summary_state":"pending","summary":null,"crux":null,"owner":"RayWorkerGroup"},{"id":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","name":"_bind_workers_method_to_parent","kind":"function","path":"verl/single_controller/ray/base.py","span":"L380-L411","signature":"def _bind_workers_method_to_parent(cls, key, user_defined_cls)","exported":false,"origin":"ast","body_hash":"14e6f35dc58512a85e8cd7b463daf6b0f31f12815d5d319d64082fd72e97d899","body_text":"def _bind_workers_method_to_parent(cls, key, user_defined_cls): \"\"\" Binds the methods of each worker to the WorkerDict. Note that we only bind public methods that are decorated by register \"\"\" for method_name in dir(user_defined_cls): try: method = getattr(user_defined_cls, method_name) assert callable(method), f\"{method_name} in {user_defined_cls} is not callable\" except Exception as e: # if it is a property, it will fail because Class doesn't have instance property continue if hasattr(method, MAGIC_ATTR): def generate_function(name): def func(self, *args, **kwargs): # dispatch to the actual worker return getattr(self.worker_dict[key], name)(*args, **kwargs) return func func = generate_function(method_name) # pass MAGIC_ATTR for outer worker group setattr(func, MAGIC_ATTR, getattr(method, MAGIC_ATTR)) try: method_name_with_prefix = key + '_' + method_name setattr(cls, method_name_with_prefix, func) # print(f'Binding {method_name_with_prefix}') except Exception as e: raise ValueError(f'Fail to set method_name {method_name}')","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent.generate_function","name":"generate_function","kind":"function","path":"verl/single_controller/ray/base.py","span":"L395-L401","signature":"def generate_function(name)","exported":true,"origin":"ast","body_hash":"f840cd651da776c35677b7d79f9a16dde27fa01cb3c927cacfe74a03c4d9cdf1","body_text":"def generate_function(name): def func(self, *args, **kwargs): # dispatch to the actual worker return getattr(self.worker_dict[key], name)(*args, **kwargs) return func","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent.generate_function.func","name":"func","kind":"function","path":"verl/single_controller/ray/base.py","span":"L397-L399","signature":"def func(self, *args, **kwargs): # dispatch to the actual worker","exported":true,"origin":"ast","body_hash":"5c125b1246b4bac15af3bd95092d9426722add0b1e98dfcbd4303db310b162c9","body_text":"def func(self, *args, **kwargs): # dispatch to the actual worker return getattr(self.worker_dict[key], name)(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#_unwrap_ray_remote","name":"_unwrap_ray_remote","kind":"function","path":"verl/single_controller/ray/base.py","span":"L414-L417","signature":"def _unwrap_ray_remote(cls)","exported":false,"origin":"ast","body_hash":"509b623e9014e61cf5bb06d4b23e1e949796bfaaac60a3184375fe9bac3259b1","body_text":"def _unwrap_ray_remote(cls): if hasattr(cls, '__ray_actor_class__'): cls = cls.__ray_actor_class__ return cls","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#create_colocated_worker_cls","name":"create_colocated_worker_cls","kind":"function","path":"verl/single_controller/ray/base.py","span":"L420-L459","signature":"def create_colocated_worker_cls(class_dict: dict[str, RayClassWithInitArgs])","exported":true,"origin":"ast","body_hash":"a3bc8dce89fdefa639c4b160d2b6c5896b676be0a4cb525fcc067114f726fd48","body_text":"def create_colocated_worker_cls(class_dict: dict[str, RayClassWithInitArgs]): \"\"\" This function should return a class instance that delegates the calls to every cls in cls_dict \"\"\" cls_dict = {} init_args_dict = {} worker_cls = None for key, cls in class_dict.items(): if worker_cls == None: worker_cls = cls.cls.__ray_actor_class__.__base__ else: assert worker_cls == cls.cls.__ray_actor_class__.__base__, \\ 'the worker class should be the same when share the same process' cls_dict[key] = cls.cls init_args_dict[key] = {'args': cls.args, 'kwargs': cls.kwargs} assert cls_dict.keys() == init_args_dict.keys() # TODO: create a class with customizable name class WorkerDict(worker_cls): def __init__(self): super().__init__() self.worker_dict = {} for key, user_defined_cls in cls_dict.items(): user_defined_cls = _unwrap_ray_remote(user_defined_cls) # directly instantiate the class without remote with patch.dict(os.environ, {'DISABLE_WORKER_INIT': '1'}): self.worker_dict[key] = user_defined_cls(*init_args_dict[key].get('args', ()), **init_args_dict[key].get('kwargs', {})) # now monkey-patch the methods from inner class to WorkerDict for key, user_defined_cls in cls_dict.items(): user_defined_cls = _unwrap_ray_remote(user_defined_cls) _bind_workers_method_to_parent(WorkerDict, key, user_defined_cls) remote_cls = ray.remote(WorkerDict) remote_cls = RayClassWithInitArgs(cls=remote_cls) return remote_cls","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict","name":"WorkerDict","kind":"class","path":"verl/single_controller/ray/base.py","span":"L440-L450","signature":"class WorkerDict(worker_cls)","exported":true,"origin":"ast","body_hash":"3904ce1ae46a83e95c110f536eaf0ea6f650aa6f4b0e93ca1cdd3aa6e12792cb","body_text":"class WorkerDict(worker_cls): def __init__(self): super().__init__() self.worker_dict = {} for key, user_defined_cls in cls_dict.items(): user_defined_cls = _unwrap_ray_remote(user_defined_cls) # directly instantiate the class without remote with patch.dict(os.environ, {'DISABLE_WORKER_INIT': '1'}): self.worker_dict[key] = user_defined_cls(*init_args_dict[key].get('args', ()), **init_args_dict[key].get('kwargs', {}))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict.__init__","name":"__init__","kind":"method","path":"verl/single_controller/ray/base.py","span":"L442-L450","signature":"def __init__(self)","exported":false,"origin":"ast","body_hash":"1f0a441cd8c9ec0447969874c51af19b92eea1d5de88a67115436a34a42a3705","body_text":"def __init__(self): super().__init__() self.worker_dict = {} for key, user_defined_cls in cls_dict.items(): user_defined_cls = _unwrap_ray_remote(user_defined_cls) # directly instantiate the class without remote with patch.dict(os.environ, {'DISABLE_WORKER_INIT': '1'}): self.worker_dict[key] = user_defined_cls(*init_args_dict[key].get('args', ()), **init_args_dict[key].get('kwargs', {}))","summary_state":"pending","summary":null,"crux":null,"owner":"WorkerDict"}],"rawEdges":[{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"time","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"typing","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"ray","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"ray.util","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"ray.util.placement_group","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"ray.util.scheduling_strategies","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"ray.experimental.state.api","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"verl.single_controller.base","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"contains","targetId":"verl/single_controller/ray/base.py#get_random_string","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"random","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"string","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#get_random_string","relation":"calls","name":"join","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#get_random_string","relation":"calls","name":"choice","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#get_random_string","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"contains","targetId":"verl/single_controller/ray/base.py#func_generator","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#func_generator","relation":"contains","targetId":"verl/single_controller/ray/base.py#func_generator.func","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#func_generator.func","relation":"calls","name":"dispatch_fn","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#func_generator.func","relation":"calls","name":"execute_fn","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#func_generator.func","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#func_generator.func","relation":"calls","name":"collect_fn","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayResourcePool","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool","relation":"extends","name":"ResourcePool","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayResourcePool.__init__","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayResourcePool.get_placement_groups","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool.get_placement_groups","relation":"calls","name":"join","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool.get_placement_groups","relation":"calls","name":"str","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool.get_placement_groups","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool.get_placement_groups","relation":"calls","name":"placement_group","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool.get_placement_groups","relation":"calls","name":"str","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool.get_placement_groups","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool.get_placement_groups","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayResourcePool.get_placement_groups","relation":"calls","name":"ready","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"contains","targetId":"verl/single_controller/ray/base.py#extract_pg_from_exist","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#extract_pg_from_exist","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#extract_pg_from_exist","relation":"calls","name":"get_placement_groups","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayResourcePool"},{"source":"verl/single_controller/ray/base.py#extract_pg_from_exist","relation":"calls","name":"sorted","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#extract_pg_from_exist","relation":"calls","name":"sorted","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#extract_pg_from_exist","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#extract_pg_from_exist","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#extract_pg_from_exist","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#extract_pg_from_exist","relation":"calls","name":"sorted","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"contains","targetId":"verl/single_controller/ray/base.py#merge_resource_pool","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#merge_resource_pool","relation":"calls","name":"RayResourcePool","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#merge_resource_pool","relation":"calls","name":"get_placement_groups","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayResourcePool"},{"source":"verl/single_controller/ray/base.py#merge_resource_pool","relation":"calls","name":"get_placement_groups","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayResourcePool"},{"source":"verl/single_controller/ray/base.py","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayClassWithInitArgs","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs","relation":"extends","name":"ClassWithInitArgs","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__init__","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayClassWithInitArgs.set_additional_resource","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayClassWithInitArgs.update_options","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.update_options","relation":"calls","name":"update","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"remote","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"remote","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"NodeAffinitySchedulingStrategy","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"remote","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"options","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"PlacementGroupSchedulingStrategy","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"update","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"remote","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayClassWithInitArgs.__call__","relation":"calls","name":"options","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"extends","name":"WorkerGroup","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.__init__","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.__init__","relation":"calls","name":"get_random_string","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.__init__","relation":"calls","name":"_init_with_detached_workers","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayWorkerGroup"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.__init__","relation":"calls","name":"_init_with_resource_pool","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayWorkerGroup"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.__init__","relation":"calls","name":"_bind_worker_method","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayWorkerGroup"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup._is_worker_alive","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._is_worker_alive","relation":"calls","name":"get_actor","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._is_worker_alive","relation":"calls","name":"hex","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._is_worker_alive","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"get_actor"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_detached_workers","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_detached_workers","relation":"calls","name":"get_actor","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_detached_workers","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"get_placement_groups","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"str","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"str","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"str","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"str","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"re","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"type","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"search","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"group","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"update_options","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"update_options","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"ray_cls_with_init","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"list_named_actors","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"sleep","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"get_actor","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"list_named_actors","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup._init_with_resource_pool","relation":"calls","name":"remote","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.worker_names","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.from_detached","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.from_detached","relation":"calls","name":"cls","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn._rebind_actor_methods","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn._rebind_actor_methods","relation":"calls","name":"dir","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn._rebind_actor_methods","relation":"calls","name":"startswith","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn._rebind_actor_methods","relation":"calls","name":"removeprefix","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn._rebind_actor_methods","relation":"calls","name":"getattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn._rebind_actor_methods","relation":"calls","name":"setattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn","relation":"calls","name":"from_detached","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayWorkerGroup"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.spawn","relation":"calls","name":"_rebind_actor_methods","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero_sync","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero_sync","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero_sync","relation":"calls","name":"execute_all_async","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayWorkerGroup"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero_async","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero_async","relation":"calls","name":"getattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero_async","relation":"calls","name":"remote","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"getattr"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_rank_zero","relation":"calls","name":"execute_rank_zero_async","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayWorkerGroup"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all","relation":"calls","name":"execute_all_async","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayWorkerGroup"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_sync","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_sync","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_sync","relation":"calls","name":"execute_all_async","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"RayWorkerGroup"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"all","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"all","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"values","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"all","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"all","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"len","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"values","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"range","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"tuple","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"getattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"append","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"remote","viaMember":true,"file":"verl/single_controller/ray/base.py","recvType":"getattr"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"remote","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup.execute_all_async","relation":"calls","name":"getattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.master_address","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.master_port","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.workers","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#RayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/base.py#RayWorkerGroup.world_size","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"unittest.mock","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"verl.single_controller.base.decorator","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"imports","specifier":"os","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"contains","targetId":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","relation":"calls","name":"dir","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","relation":"calls","name":"getattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","relation":"calls","name":"callable","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","relation":"contains","targetId":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent.generate_function","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent.generate_function","relation":"contains","targetId":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent.generate_function.func","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent.generate_function.func","relation":"calls","name":"getattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","relation":"calls","name":"generate_function","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","relation":"calls","name":"setattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","relation":"calls","name":"getattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","relation":"calls","name":"setattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_bind_workers_method_to_parent","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"contains","targetId":"verl/single_controller/ray/base.py#_unwrap_ray_remote","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#_unwrap_ray_remote","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py","relation":"contains","targetId":"verl/single_controller/ray/base.py#create_colocated_worker_cls","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls","relation":"calls","name":"keys","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls","relation":"calls","name":"keys","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls","relation":"contains","targetId":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict","relation":"extends","name":"worker_cls","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict","relation":"contains","targetId":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict.__init__","file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict.__init__","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict.__init__","relation":"calls","name":"_unwrap_ray_remote","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict.__init__","relation":"calls","name":"dict","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict.__init__","relation":"calls","name":"user_defined_cls","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls.WorkerDict.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls","relation":"calls","name":"items","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls","relation":"calls","name":"_unwrap_ray_remote","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls","relation":"calls","name":"_bind_workers_method_to_parent","viaMember":false,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls","relation":"calls","name":"remote","viaMember":true,"file":"verl/single_controller/ray/base.py"},{"source":"verl/single_controller/ray/base.py#create_colocated_worker_cls","relation":"calls","name":"RayClassWithInitArgs","viaMember":false,"file":"verl/single_controller/ray/base.py"}]},"verl/single_controller/ray/megatron.py":{"size":3035,"mtimeMs":1789275830311.2026,"hash":"f07274afae8e1b14e930b48a7b228fd0afe3e409a5012f790481a953858ccde5","nodes":[{"id":"verl/single_controller/ray/megatron.py","name":"megatron.py","kind":"file","path":"verl/single_controller/ray/megatron.py","span":"L1-L63","signature":null,"exported":true,"origin":"ast","body_hash":"f07274afae8e1b14e930b48a7b228fd0afe3e409a5012f790481a953858ccde5","chars":3035,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 typing import Dict, Optional import ray from .base import RayWorkerGroup, RayResourcePool, RayClassWithInitArgs from verl.single_controller.base.megatron.worker import DistRankInfo, DistGlobalInfo from verl.single_controller.base.megatron.worker_group import MegatronWorkerGroup # NOTE(sgm): for opensource megatron-core"},{"id":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup","name":"NVMegatronRayWorkerGroup","kind":"class","path":"verl/single_controller/ray/megatron.py","span":"L25-L35","signature":"class NVMegatronRayWorkerGroup(RayWorkerGroup, MegatronWorkerGroup)","exported":true,"origin":"ast","body_hash":"1ea5da75ad333445d598f344b0580747ea4494375bf85eb3d1ae07d6322f822d","body_text":"class NVMegatronRayWorkerGroup(RayWorkerGroup, MegatronWorkerGroup): \"\"\" MegatronWorkerGroup will query each worker of its megatron rank info and store it inside the WorkerGroup so that the dispatcher can use it to dispatch data. \"\"\" def __init__(self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWithInitArgs, **kwargs): super().__init__(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, **kwargs) self._megatron_rank_info: DistRankInfo = self.execute_all_sync(method_name='get_megatron_rank_info') self._megatron_global_info: DistGlobalInfo = ray.get( self.execute_rank_zero_async(method_name='get_megatron_global_info'))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup.__init__","name":"__init__","kind":"method","path":"verl/single_controller/ray/megatron.py","span":"L31-L35","signature":"def __init__(self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWithInitArgs, **kwargs)","exported":false,"origin":"ast","body_hash":"a97391613e50834872cfd6f73f95e46a87c8875886c43648c6a0bf1c2b1b928a","body_text":"def __init__(self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWithInitArgs, **kwargs): super().__init__(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, **kwargs) self._megatron_rank_info: DistRankInfo = self.execute_all_sync(method_name='get_megatron_rank_info') self._megatron_global_info: DistGlobalInfo = ray.get( self.execute_rank_zero_async(method_name='get_megatron_global_info'))","summary_state":"pending","summary":null,"crux":null,"owner":"NVMegatronRayWorkerGroup"},{"id":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup","name":"MegatronRayWorkerGroup","kind":"class","path":"verl/single_controller/ray/megatron.py","span":"L38-L62","signature":"class MegatronRayWorkerGroup(RayWorkerGroup, MegatronWorkerGroup)","exported":true,"origin":"ast","body_hash":"1419287f95a495c465228ca4fd9ac3004f9df40807d15e80660d2ef608cbd530","body_text":"class MegatronRayWorkerGroup(RayWorkerGroup, MegatronWorkerGroup): \"\"\" MegatronWorkerGroup will query each worker of its megatron rank info and store it inside the WorkerGroup so that the dispatcher can use it to dispatch data. \"\"\" def __init__(self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWithInitArgs, default_megatron_kwargs: Dict = None, **kwargs): super().__init__(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, default_megatron_kwargs=default_megatron_kwargs, **kwargs) self.init_megatron(default_megatron_kwargs=default_megatron_kwargs) self._megatron_rank_info: DistRankInfo = self.execute_all_sync(method_name='get_megatron_rank_info') self._megatron_global_info: DistGlobalInfo = ray.get( self.execute_rank_zero_async(method_name='get_megatron_global_info')) def init_megatron(self, default_megatron_kwargs: Optional[Dict] = None): # after super, we will call init of each worker if not self._is_init_with_detached_workers: # only init_megatron if the WorkerGroup is created from scratch self.execute_all_sync(method_name='init_megatron', default_megatron_kwargs=default_megatron_kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.__init__","name":"__init__","kind":"method","path":"verl/single_controller/ray/megatron.py","span":"L44-L56","signature":"def __init__(self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWithInitArgs, default_megatron_kwargs: Dict = None, **kwargs)","exported":false,"origin":"ast","body_hash":"8603076b5b866a0c5932955312c312264af13b53f28d8b59642b9eefdfce33e0","body_text":"def __init__(self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWithInitArgs, default_megatron_kwargs: Dict = None, **kwargs): super().__init__(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, default_megatron_kwargs=default_megatron_kwargs, **kwargs) self.init_megatron(default_megatron_kwargs=default_megatron_kwargs) self._megatron_rank_info: DistRankInfo = self.execute_all_sync(method_name='get_megatron_rank_info') self._megatron_global_info: DistGlobalInfo = ray.get( self.execute_rank_zero_async(method_name='get_megatron_global_info'))","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronRayWorkerGroup"},{"id":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.init_megatron","name":"init_megatron","kind":"method","path":"verl/single_controller/ray/megatron.py","span":"L58-L62","signature":"def init_megatron(self, default_megatron_kwargs: Optional[Dict] = None): # after super, we will call init of each worker","exported":true,"origin":"ast","body_hash":"a6f29739a8ac354f8270ec67fd3d0636d64d3a7b73783cf89354cedbc9061519","body_text":"def init_megatron(self, default_megatron_kwargs: Optional[Dict] = None): # after super, we will call init of each worker if not self._is_init_with_detached_workers: # only init_megatron if the WorkerGroup is created from scratch self.execute_all_sync(method_name='init_megatron', default_megatron_kwargs=default_megatron_kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronRayWorkerGroup"}],"rawEdges":[{"source":"verl/single_controller/ray/megatron.py","relation":"imports","specifier":"typing","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py","relation":"imports","specifier":"ray","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py","relation":"imports","specifier":".base","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker_group","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py","relation":"contains","targetId":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup","relation":"extends","name":"RayWorkerGroup","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup","relation":"extends","name":"MegatronWorkerGroup","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup.__init__","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup.__init__","relation":"calls","name":"execute_all_sync","viaMember":true,"file":"verl/single_controller/ray/megatron.py","recvType":"NVMegatronRayWorkerGroup"},{"source":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#NVMegatronRayWorkerGroup.__init__","relation":"calls","name":"execute_rank_zero_async","viaMember":true,"file":"verl/single_controller/ray/megatron.py","recvType":"NVMegatronRayWorkerGroup"},{"source":"verl/single_controller/ray/megatron.py","relation":"contains","targetId":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup","relation":"extends","name":"RayWorkerGroup","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup","relation":"extends","name":"MegatronWorkerGroup","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.__init__","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.__init__","relation":"calls","name":"init_megatron","viaMember":true,"file":"verl/single_controller/ray/megatron.py","recvType":"MegatronRayWorkerGroup"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.__init__","relation":"calls","name":"execute_all_sync","viaMember":true,"file":"verl/single_controller/ray/megatron.py","recvType":"MegatronRayWorkerGroup"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.__init__","relation":"calls","name":"execute_rank_zero_async","viaMember":true,"file":"verl/single_controller/ray/megatron.py","recvType":"MegatronRayWorkerGroup"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup","relation":"contains","targetId":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.init_megatron","file":"verl/single_controller/ray/megatron.py"},{"source":"verl/single_controller/ray/megatron.py#MegatronRayWorkerGroup.init_megatron","relation":"calls","name":"execute_all_sync","viaMember":true,"file":"verl/single_controller/ray/megatron.py","recvType":"MegatronRayWorkerGroup"}]},"verl/third_party/__init__.py":{"size":600,"mtimeMs":1789275830311.2026,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/third_party/__init__.py","name":"__init__.py","kind":"file","path":"verl/third_party/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/third_party/vllm/__init__.py":{"size":1746,"mtimeMs":1789275830311.2026,"hash":"9a1b9a8a579f4de9cc7e9d0419631a3d006e45ed5dc2d8504a96aab0954cac4a","nodes":[{"id":"verl/third_party/vllm/__init__.py","name":"__init__.py","kind":"file","path":"verl/third_party/vllm/__init__.py","span":"L1-L52","signature":null,"exported":true,"origin":"ast","body_hash":"9a1b9a8a579f4de9cc7e9d0419631a3d006e45ed5dc2d8504a96aab0954cac4a","chars":1746,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 importlib.metadata import version, PackageNotFoundError package_name = 'vllm' package_version = get_version(package_name) if package_version == '0.3.1': vllm_version = '0.3.1' from .vllm_v_0_3_1.llm import LLM from .vllm_v_0_3_1.llm import LLMEngine from .vllm_v_0_3_1 import parallel_state elif package_version == '0.4.2': vllm_version = '0.4.2' from .vllm_v_0_4_2.llm import LLM from .vllm_v_0_4_2.llm import LLMEngine from .vllm_v_0_4_2 import parallel_state elif package_version == '0.5.4': vllm_version = '0.5.4' from .vllm_v_0_5_4.llm import LLM from .vllm_v_0_5_4.llm import LLMEngine from .vllm_v_0_5_4 import parallel_state elif package_version == '0.6.3': vllm_version = '0.6.3' from .vllm_v_0_6_3.llm import LLM from .vllm_v_0_6_3.llm import LLMEngine from .vllm_v_0_6_3 import parallel_state else: raise ValueError( f'vllm version {package_version} not supported. Currently supported versions are 0.3.1, 0.4.2, 0.5.4 and 0.6.3.' )"},{"id":"verl/third_party/vllm/__init__.py#get_version","name":"get_version","kind":"function","path":"verl/third_party/vllm/__init__.py","span":"L18-L22","signature":"def get_version(pkg)","exported":true,"origin":"ast","body_hash":"691d9af2743cd95a2e41da2dc27059e090848819d78c6b450eea8b19cb6ebec6","body_text":"def get_version(pkg): try: return version(pkg) except PackageNotFoundError: return None","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":"importlib.metadata","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"contains","targetId":"verl/third_party/vllm/__init__.py#get_version","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py#get_version","relation":"calls","name":"version","viaMember":false,"file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"calls","name":"get_version","viaMember":false,"file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_3_1.llm","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_3_1.llm","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_3_1","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_4_2.llm","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_4_2.llm","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_4_2","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_5_4.llm","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_5_4.llm","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_5_4","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_6_3.llm","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_6_3.llm","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"imports","specifier":".vllm_v_0_6_3","file":"verl/third_party/vllm/__init__.py"},{"source":"verl/third_party/vllm/__init__.py","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/__init__.py"}]},"verl/third_party/vllm/vllm_v_0_3_1/__init__.py":{"size":600,"mtimeMs":1789275830312.2026,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/__init__.py","name":"__init__.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py":{"size":11957,"mtimeMs":1789275830312.2026,"hash":"241d3e47e1bf283647237be71099df7ff857e11ce24c1f6082f9225929a7d4be","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","name":"arg_utils.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","span":"L1-L229","signature":null,"exported":true,"origin":"ast","body_hash":"241d3e47e1bf283647237be71099df7ff857e11ce24c1f6082f9225929a7d4be","chars":11957,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/engine/arg_utils.py import argparse import dataclasses from dataclasses import dataclass from typing import Dict, Optional, Tuple import torch.nn as nn from vllm.config import (CacheConfig, DeviceConfig, ModelConfig, ParallelConfig, SchedulerConfig, LoRAConfig) from transformers import PretrainedConfig from .config import ModelConfig @dataclass @dataclass"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs","name":"EngineArgs","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","span":"L27-L204","signature":"class EngineArgs","exported":true,"origin":"ast","body_hash":"ed241c0f5181f17b004cdd8918940a783ffda283aba5684586513b66b674a680","body_text":"class EngineArgs: \"\"\"Arguments for vLLM engine.\"\"\" model_hf_config: PretrainedConfig = None dtype: str = 'auto' kv_cache_dtype: str = 'auto' seed: int = 0 max_model_len: Optional[int] = None worker_use_ray: bool = False pipeline_parallel_size: int = 1 tensor_parallel_size: int = 1 max_parallel_loading_workers: Optional[int] = None block_size: int = 16 swap_space: int = 4 # GiB gpu_memory_utilization: float = 0.90 max_num_batched_tokens: Optional[int] = None max_num_seqs: int = 256 max_paddings: int = 256 disable_log_stats: bool = False revision: Optional[str] = None tokenizer_revision: Optional[str] = None quantization: Optional[str] = None load_format: str = 'model' enforce_eager: bool = False max_context_len_to_capture: int = 8192 disable_custom_all_reduce: bool = False enable_lora: bool = False max_loras: int = 1 max_lora_rank: int = 16 lora_extra_vocab_size: int = 256 lora_dtype = 'auto' max_cpu_loras: Optional[int] = None device: str = 'cuda' @staticmethod def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: \"\"\"Shared CLI arguments for vLLM engine.\"\"\" # Model arguments # TODO(shengguangming): delete the unused args parser.add_argument('--model', type=str, default='facebook/opt-125m', help='name or path of the huggingface model to use') parser.add_argument('--tokenizer', type=str, default=EngineArgs.tokenizer, help='name or path of the huggingface tokenizer to use') parser.add_argument('--revision', type=str, default=None, help='the specific model version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-revision', type=str, default=None, help='the specific tokenizer version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-mode', type=str, default=EngineArgs.tokenizer_mode, choices=['auto', 'slow'], help='tokenizer mode. \"auto\" will use the fast ' 'tokenizer if available, and \"slow\" will ' 'always use the slow tokenizer.') parser.add_argument('--trust-remote-code', action='store_true', help='trust remote code from huggingface') parser.add_argument('--download-dir', type=str, default=EngineArgs.download_dir, help='directory to download and load the weights, ' 'default to the default cache dir of ' 'huggingface') parser.add_argument('--load-format', type=str, default=EngineArgs.load_format, choices=['auto', 'pt', 'safetensors', 'npcache', 'dummy'], help='The format of the model weights to load. ' '\"auto\" will try to load the weights in the safetensors format ' 'and fall back to the pytorch bin format if safetensors format ' 'is not available. ' '\"pt\" will load the weights in the pytorch bin format. ' '\"safetensors\" will load the weights in the safetensors format. ' '\"npcache\" will load the weights in pytorch format and store ' 'a numpy cache to speed up the loading. ' '\"dummy\" will initialize the weights with random values, ' 'which is mainly for profiling.') parser.add_argument('--dtype', type=str, default=EngineArgs.dtype, choices=['auto', 'half', 'float16', 'bfloat16', 'float', 'float32'], help='data type for model weights and activations. ' 'The \"auto\" option will use FP16 precision ' 'for FP32 and FP16 models, and BF16 precision ' 'for BF16 models.') parser.add_argument('--max-model-len', type=int, default=None, help='model context length. If unspecified, ' 'will be automatically derived from the model.') # Parallel arguments parser.add_argument('--worker-use-ray', action='store_true', help='use Ray for distributed serving, will be ' 'automatically set when using more than 1 GPU') parser.add_argument('--pipeline-parallel-size', '-pp', type=int, default=EngineArgs.pipeline_parallel_size, help='number of pipeline stages') parser.add_argument('--tensor-parallel-size', '-tp', type=int, default=EngineArgs.tensor_parallel_size, help='number of tensor parallel replicas') # KV cache arguments parser.add_argument('--block-size', type=int, default=EngineArgs.block_size, choices=[8, 16, 32], help='token block size') # TODO(woosuk): Support fine-grained seeds (e.g., seed per request). parser.add_argument('--seed', type=int, default=EngineArgs.seed, help='random seed') parser.add_argument('--swap-space', type=int, default=EngineArgs.swap_space, help='CPU swap space size (GiB) per GPU') parser.add_argument('--gpu-memory-utilization', type=float, default=EngineArgs.gpu_memory_utilization, help='the percentage of GPU memory to be used for' 'the model executor') parser.add_argument('--max-num-batched-tokens', type=int, default=EngineArgs.max_num_batched_tokens, help='maximum number of batched tokens per ' 'iteration') parser.add_argument('--max-num-seqs', type=int, default=EngineArgs.max_num_seqs, help='maximum number of sequences per iteration') parser.add_argument('--disable-log-stats', action='store_true', help='disable logging statistics') # Quantization settings. parser.add_argument('--quantiza","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","name":"add_cli_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","span":"L61-L175","signature":"def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser","exported":true,"origin":"ast","body_hash":"cf1908f71aee396b15d02e5aad05a2803adf4e9fd0905ff5fcb9c8727ba16fb5","body_text":"def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: \"\"\"Shared CLI arguments for vLLM engine.\"\"\" # Model arguments # TODO(shengguangming): delete the unused args parser.add_argument('--model', type=str, default='facebook/opt-125m', help='name or path of the huggingface model to use') parser.add_argument('--tokenizer', type=str, default=EngineArgs.tokenizer, help='name or path of the huggingface tokenizer to use') parser.add_argument('--revision', type=str, default=None, help='the specific model version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-revision', type=str, default=None, help='the specific tokenizer version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-mode', type=str, default=EngineArgs.tokenizer_mode, choices=['auto', 'slow'], help='tokenizer mode. \"auto\" will use the fast ' 'tokenizer if available, and \"slow\" will ' 'always use the slow tokenizer.') parser.add_argument('--trust-remote-code', action='store_true', help='trust remote code from huggingface') parser.add_argument('--download-dir', type=str, default=EngineArgs.download_dir, help='directory to download and load the weights, ' 'default to the default cache dir of ' 'huggingface') parser.add_argument('--load-format', type=str, default=EngineArgs.load_format, choices=['auto', 'pt', 'safetensors', 'npcache', 'dummy'], help='The format of the model weights to load. ' '\"auto\" will try to load the weights in the safetensors format ' 'and fall back to the pytorch bin format if safetensors format ' 'is not available. ' '\"pt\" will load the weights in the pytorch bin format. ' '\"safetensors\" will load the weights in the safetensors format. ' '\"npcache\" will load the weights in pytorch format and store ' 'a numpy cache to speed up the loading. ' '\"dummy\" will initialize the weights with random values, ' 'which is mainly for profiling.') parser.add_argument('--dtype', type=str, default=EngineArgs.dtype, choices=['auto', 'half', 'float16', 'bfloat16', 'float', 'float32'], help='data type for model weights and activations. ' 'The \"auto\" option will use FP16 precision ' 'for FP32 and FP16 models, and BF16 precision ' 'for BF16 models.') parser.add_argument('--max-model-len', type=int, default=None, help='model context length. If unspecified, ' 'will be automatically derived from the model.') # Parallel arguments parser.add_argument('--worker-use-ray', action='store_true', help='use Ray for distributed serving, will be ' 'automatically set when using more than 1 GPU') parser.add_argument('--pipeline-parallel-size', '-pp', type=int, default=EngineArgs.pipeline_parallel_size, help='number of pipeline stages') parser.add_argument('--tensor-parallel-size', '-tp', type=int, default=EngineArgs.tensor_parallel_size, help='number of tensor parallel replicas') # KV cache arguments parser.add_argument('--block-size', type=int, default=EngineArgs.block_size, choices=[8, 16, 32], help='token block size') # TODO(woosuk): Support fine-grained seeds (e.g., seed per request). parser.add_argument('--seed', type=int, default=EngineArgs.seed, help='random seed') parser.add_argument('--swap-space', type=int, default=EngineArgs.swap_space, help='CPU swap space size (GiB) per GPU') parser.add_argument('--gpu-memory-utilization', type=float, default=EngineArgs.gpu_memory_utilization, help='the percentage of GPU memory to be used for' 'the model executor') parser.add_argument('--max-num-batched-tokens', type=int, default=EngineArgs.max_num_batched_tokens, help='maximum number of batched tokens per ' 'iteration') parser.add_argument('--max-num-seqs', type=int, default=EngineArgs.max_num_seqs, help='maximum number of sequences per iteration') parser.add_argument('--disable-log-stats', action='store_true', help='disable logging statistics') # Quantization settings. parser.add_argument('--quantization', '-q', type=str, choices=['awq', None], default=None, help='Method used to quantize the weights') return parser","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.from_cli_args","name":"from_cli_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","span":"L178-L183","signature":"def from_cli_args(cls, args: argparse.Namespace) -> 'EngineArgs': # Get the list of attributes of this dataclass.","exported":true,"origin":"ast","body_hash":"21089f327dd4125f76a223a109fdee39de77f84d10cd1fad1f24de8837c1b6f1","body_text":"def from_cli_args(cls, args: argparse.Namespace) -> 'EngineArgs': # Get the list of attributes of this dataclass. attrs = [attr.name for attr in dataclasses.fields(cls)] # Set the attributes from the parsed arguments. engine_args = cls(**{attr: getattr(args, attr) for attr in attrs}) return engine_args","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.create_engine_configs","name":"create_engine_configs","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","span":"L185-L204","signature":"def create_engine_configs( self, ) -> Tuple[ModelConfig, CacheConfig, ParallelConfig, SchedulerConfig]","exported":true,"origin":"ast","body_hash":"5c9617a83dc302ccdb2ceed92bbef03f5229b03ffa9e53bcfa019e9eac5a6741","body_text":"def create_engine_configs( self, ) -> Tuple[ModelConfig, CacheConfig, ParallelConfig, SchedulerConfig]: device_config = DeviceConfig(self.device) model_config = ModelConfig(self.model_hf_config, self.dtype, self.seed, self.load_format, self.revision, self.tokenizer_revision, self.max_model_len, self.quantization, self.enforce_eager, self.max_context_len_to_capture) cache_config = CacheConfig(self.block_size, self.gpu_memory_utilization, self.swap_space, self.kv_cache_dtype, model_config.get_sliding_window()) parallel_config = ParallelConfig(self.pipeline_parallel_size, self.tensor_parallel_size, self.worker_use_ray, self.max_parallel_loading_workers, self.disable_custom_all_reduce) scheduler_config = SchedulerConfig(self.max_num_batched_tokens, self.max_num_seqs, model_config.max_model_len, self.max_paddings) lora_config = LoRAConfig(max_lora_rank=self.max_lora_rank, max_loras=self.max_loras, lora_extra_vocab_size=self.lora_extra_vocab_size, lora_dtype=self.lora_dtype, max_cpu_loras=self.max_cpu_loras if self.max_cpu_loras and self.max_cpu_loras > 0 else None) if self.enable_lora else None return (model_config, cache_config, parallel_config, scheduler_config, device_config, lora_config)","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#AsyncEngineArgs","name":"AsyncEngineArgs","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","span":"L208-L228","signature":"class AsyncEngineArgs(EngineArgs)","exported":true,"origin":"ast","body_hash":"8c7317a410aad5330961f47069be6623ae2c016c0e59cca4934260e1d18038bf","body_text":"class AsyncEngineArgs(EngineArgs): \"\"\"Arguments for asynchronous vLLM engine.\"\"\" engine_use_ray: bool = False disable_log_requests: bool = False max_log_len: Optional[int] = None @staticmethod def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser = EngineArgs.add_cli_args(parser) parser.add_argument('--engine-use-ray', action='store_true', help='use Ray to start the LLM engine in a ' 'separate process as the server process.') parser.add_argument('--disable-log-requests', action='store_true', help='disable logging requests') parser.add_argument('--max-log-len', type=int, default=None, help='max number of prompt characters or prompt ' 'ID numbers being printed in log. ' 'Default: unlimited.') return parser","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#AsyncEngineArgs.add_cli_args","name":"add_cli_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","span":"L215-L228","signature":"def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser","exported":true,"origin":"ast","body_hash":"ee7f8b07c4a2eb8b4183456b8d92e916f4d9c31d59dd82f78e69a61e716ccaa3","body_text":"def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: parser = EngineArgs.add_cli_args(parser) parser.add_argument('--engine-use-ray', action='store_true', help='use Ray to start the LLM engine in a ' 'separate process as the server process.') parser.add_argument('--disable-log-requests', action='store_true', help='disable logging requests') parser.add_argument('--max-log-len', type=int, default=None, help='max number of prompt characters or prompt ' 'ID numbers being printed in log. ' 'Default: unlimited.') return parser","summary_state":"pending","summary":null,"crux":null,"owner":"AsyncEngineArgs"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","relation":"imports","specifier":"argparse","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.from_cli_args","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.from_cli_args","relation":"calls","name":"fields","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.from_cli_args","relation":"calls","name":"cls","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.from_cli_args","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.create_engine_configs","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.create_engine_configs","relation":"calls","name":"DeviceConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.create_engine_configs","relation":"calls","name":"ModelConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.create_engine_configs","relation":"calls","name":"CacheConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.create_engine_configs","relation":"calls","name":"get_sliding_window","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.create_engine_configs","relation":"calls","name":"ParallelConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.create_engine_configs","relation":"calls","name":"SchedulerConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#EngineArgs.create_engine_configs","relation":"calls","name":"LoRAConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#AsyncEngineArgs","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#AsyncEngineArgs","relation":"extends","name":"EngineArgs","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#AsyncEngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#AsyncEngineArgs.add_cli_args","file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#AsyncEngineArgs.add_cli_args","relation":"calls","name":"add_cli_args","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#AsyncEngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#AsyncEngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py#AsyncEngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/arg_utils.py"}]},"verl/third_party/vllm/vllm_v_0_3_1/config.py":{"size":25666,"mtimeMs":1789275830312.2026,"hash":"868aa707443c67986ef8b2c5437b617d640f4d47a53026f7e378ddde22ce7965","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py","name":"config.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L1-L578","signature":null,"exported":true,"origin":"ast","body_hash":"868aa707443c67986ef8b2c5437b617d640f4d47a53026f7e378ddde22ce7965","chars":25666,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py from typing import Optional, Union, ClassVar from dataclasses import dataclass import torch from transformers import PretrainedConfig from packaging.version import Version from vllm.logger import init_logger from vllm.transformers_utils.config import get_config from vllm.utils import get_cpu_memory, is_hip, get_nvcc_cuda_version logger = init_logger(__name__) _GB = 1 << 30 @dataclass _STR_DTYPE_TO_TORCH_DTYPE = { \"half\": torch.float16, \"float16\": torch.float16, \"float\": torch.float32, \"float32\": torch.float32, \"bfloat16\": torch.bfloat16, } _ROCM_NOT_SUPPORTED_DTYPE = [\"float\", \"float32\"]"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","name":"ModelConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L31-L237","signature":"class ModelConfig","exported":true,"origin":"ast","body_hash":"b0829ccaac02eb5de5f75279e1a154e404237642e84d92d147d140c3754c6b8a","body_text":"class ModelConfig: \"\"\"Configuration for the model. Args: model: Name or path of the huggingface model to use. tokenizer: Name or path of the huggingface tokenizer to use. tokenizer_mode: Tokenizer mode. \"auto\" will use the fast tokenizer if available, and \"slow\" will always use the slow tokenizer. trust_remote_code: Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer. download_dir: Directory to download and load the weights, default to the default cache directory of huggingface. load_format: The format of the model weights to load: \"auto\" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. \"pt\" will load the weights in the pytorch bin format. \"safetensors\" will load the weights in the safetensors format. \"npcache\" will load the weights in pytorch format and store a numpy cache to speed up the loading. \"dummy\" will initialize the weights with random values, which is mainly for profiling. dtype: Data type for model weights and activations. The \"auto\" option will use FP16 precision for FP32 and FP16 models, and BF16 precision for BF16 models. seed: Random seed for reproducibility. revision: The specific model version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. tokenizer_revision: The specific tokenizer version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. max_model_len: Maximum length of a sequence (including prompt and output). If None, will be derived from the model. quantization: Quantization method that was used to quantize the model weights. If None, we assume the model weights are not quantized. enforce_eager: Whether to enforce eager execution. If True, we will disable CUDA graph and always execute the model in eager mode. If False, we will use CUDA graph and eager execution in hybrid. max_context_len_to_capture: Maximum context len covered by CUDA graphs. When a sequence has context length larger than this, we fall back to eager mode. \"\"\" def __init__( self, hf_config: PretrainedConfig, dtype: str, seed: int, load_format: str = 'model', revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, max_model_len: Optional[int] = None, quantization: Optional[str] = None, trust_remote_code: Optional[bool] = True, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, ) -> None: self.model = hf_config._name_or_path self.tokenizer = hf_config._name_or_path self.load_format = load_format self.seed = seed self.revision = revision self.tokenizer_revision = tokenizer_revision self.quantization = quantization self.trust_remote_code = trust_remote_code self.enforce_eager = enforce_eager self.max_context_len_to_capture = max_context_len_to_capture # self.hf_config = get_config(model, trust_remote_code, revision) self.hf_config = hf_config self.dtype = _get_and_verify_dtype(self.hf_config, dtype) self.max_model_len = _get_and_verify_max_len(self.hf_config, max_model_len) # self._verify_load_format() # self._verify_tokenizer_mode() self._verify_quantization() self._verify_cuda_graph() def _verify_load_format(self) -> None: load_format = self.load_format.lower() if load_format not in [\"auto\", \"pt\", \"safetensors\", \"npcache\", \"dummy\", \"model\"]: raise ValueError(f\"Unknown load format: {self.load_format}. Must be one of \" \"'auto', 'pt', 'safetensors', 'npcache', 'dummy' or 'model'.\") self.load_format = load_format # def _verify_tokenizer_mode(self) -> None: # tokenizer_mode = self.tokenizer_mode.lower() # if tokenizer_mode not in [\"auto\", \"slow\"]: # raise ValueError( # f\"Unknown tokenizer mode: {self.tokenizer_mode}. Must be \" # \"either 'auto' or 'slow'.\") # self.tokenizer_mode = tokenizer_mode def _verify_quantization(self) -> None: supported_quantization = [\"awq\", \"gptq\", \"squeezellm\"] rocm_not_supported_quantization = [\"awq\", \"gptq\"] if self.quantization is not None: self.quantization = self.quantization.lower() # Parse quantization method from the HF model config, if available. hf_quant_config = getattr(self.hf_config, \"quantization_config\", None) if hf_quant_config is not None: hf_quant_method = str(hf_quant_config[\"quant_method\"]).lower() if self.quantization is None: self.quantization = hf_quant_method elif self.quantization != hf_quant_method: raise ValueError(\"Quantization method specified in the model config \" f\"({hf_quant_method}) does not match the quantization \" f\"method specified in the `quantization` argument \" f\"({self.quantization}).\") if self.quantization is not None: if self.quantization not in supported_quantization: raise ValueError(f\"Unknown quantization method: {self.quantization}. Must \" f\"be one of {supported_quantization}.\") if is_hip() and self.quantization in rocm_not_supported_quantization: raise ValueError(f\"{self.quantization} quantization is currently not supported \" f\"in ROCm.\") logger.warning(f\"{self.quantization} qua","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L75-L107","signature":"def __init__( self, hf_config: PretrainedConfig, dtype: str, seed: int, load_format: str = 'model', revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, max_model_len: Optional[int] = None, quantization: Optional[str] = None, trust_remote_code: Optional[bool] = True, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, ) -> None","exported":false,"origin":"ast","body_hash":"00f60fdc0e2531dd6bba00181007d1c2a30302957f43a2f87112757082a1746f","body_text":"def __init__( self, hf_config: PretrainedConfig, dtype: str, seed: int, load_format: str = 'model', revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, max_model_len: Optional[int] = None, quantization: Optional[str] = None, trust_remote_code: Optional[bool] = True, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, ) -> None: self.model = hf_config._name_or_path self.tokenizer = hf_config._name_or_path self.load_format = load_format self.seed = seed self.revision = revision self.tokenizer_revision = tokenizer_revision self.quantization = quantization self.trust_remote_code = trust_remote_code self.enforce_eager = enforce_eager self.max_context_len_to_capture = max_context_len_to_capture # self.hf_config = get_config(model, trust_remote_code, revision) self.hf_config = hf_config self.dtype = _get_and_verify_dtype(self.hf_config, dtype) self.max_model_len = _get_and_verify_max_len(self.hf_config, max_model_len) # self._verify_load_format() # self._verify_tokenizer_mode() self._verify_quantization() self._verify_cuda_graph()","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_load_format","name":"_verify_load_format","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L109-L114","signature":"def _verify_load_format(self) -> None","exported":false,"origin":"ast","body_hash":"3c1b201453a0857c0bf891aee1791b345941532a86ce3847b138b35256c28ec2","body_text":"def _verify_load_format(self) -> None: load_format = self.load_format.lower() if load_format not in [\"auto\", \"pt\", \"safetensors\", \"npcache\", \"dummy\", \"model\"]: raise ValueError(f\"Unknown load format: {self.load_format}. Must be one of \" \"'auto', 'pt', 'safetensors', 'npcache', 'dummy' or 'model'.\") self.load_format = load_format","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","name":"_verify_quantization","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L124-L151","signature":"def _verify_quantization(self) -> None","exported":false,"origin":"ast","body_hash":"e2e54ce7f92f3c043e953ff5ee0f32b4f6e5bc3ee39806089a18d8591828e121","body_text":"def _verify_quantization(self) -> None: supported_quantization = [\"awq\", \"gptq\", \"squeezellm\"] rocm_not_supported_quantization = [\"awq\", \"gptq\"] if self.quantization is not None: self.quantization = self.quantization.lower() # Parse quantization method from the HF model config, if available. hf_quant_config = getattr(self.hf_config, \"quantization_config\", None) if hf_quant_config is not None: hf_quant_method = str(hf_quant_config[\"quant_method\"]).lower() if self.quantization is None: self.quantization = hf_quant_method elif self.quantization != hf_quant_method: raise ValueError(\"Quantization method specified in the model config \" f\"({hf_quant_method}) does not match the quantization \" f\"method specified in the `quantization` argument \" f\"({self.quantization}).\") if self.quantization is not None: if self.quantization not in supported_quantization: raise ValueError(f\"Unknown quantization method: {self.quantization}. Must \" f\"be one of {supported_quantization}.\") if is_hip() and self.quantization in rocm_not_supported_quantization: raise ValueError(f\"{self.quantization} quantization is currently not supported \" f\"in ROCm.\") logger.warning(f\"{self.quantization} quantization is not fully \" \"optimized yet. The speed can be slower than \" \"non-quantized models.\")","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_cuda_graph","name":"_verify_cuda_graph","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L153-L161","signature":"def _verify_cuda_graph(self) -> None","exported":false,"origin":"ast","body_hash":"52ae4d1079f69b7c6c3ca340f3af6744184b37cb7e6bb2a1f5f392ad9144ad03","body_text":"def _verify_cuda_graph(self) -> None: if self.max_context_len_to_capture is None: self.max_context_len_to_capture = self.max_model_len self.max_context_len_to_capture = min(self.max_context_len_to_capture, self.max_model_len) if (self.quantization in [\"gptq\", \"squeezellm\"] and not self.enforce_eager): # Related issue: https://github.com/vllm-project/vllm/issues/2147 logger.warning(f\"{self.quantization} does not support CUDA graph \" \"yet. Disabling CUDA graph.\") self.enforce_eager = True","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.verify_with_parallel_config","name":"verify_with_parallel_config","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L163-L179","signature":"def verify_with_parallel_config( self, parallel_config: \"ParallelConfig\", ) -> None","exported":true,"origin":"ast","body_hash":"50a72aa8a7247c533a60ab8bee34ce9562f17519bedda2ac4c2f8a6969e7ca73","body_text":"def verify_with_parallel_config( self, parallel_config: \"ParallelConfig\", ) -> None: total_num_attention_heads = self.hf_config.num_attention_heads tensor_parallel_size = parallel_config.tensor_parallel_size if total_num_attention_heads % tensor_parallel_size != 0: raise ValueError(f\"Total number of attention heads ({total_num_attention_heads})\" \" must be divisible by tensor parallel size \" f\"({tensor_parallel_size}).\") total_num_hidden_layers = self.hf_config.num_hidden_layers pipeline_parallel_size = parallel_config.pipeline_parallel_size if total_num_hidden_layers % pipeline_parallel_size != 0: raise ValueError(f\"Total number of hidden layers ({total_num_hidden_layers}) \" \"must be divisible by pipeline parallel size \" f\"({pipeline_parallel_size}).\")","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_sliding_window","name":"get_sliding_window","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L181-L182","signature":"def get_sliding_window(self) -> Optional[int]","exported":true,"origin":"ast","body_hash":"65b5f7f717fb0d12bc698c7922a5787b8d66eaae7dbbf6915ec264e2c53dde7f","body_text":"def get_sliding_window(self) -> Optional[int]: return getattr(self.hf_config, \"sliding_window\", None)","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_vocab_size","name":"get_vocab_size","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L184-L185","signature":"def get_vocab_size(self) -> int","exported":true,"origin":"ast","body_hash":"5866e1009e3f9883d88e21eedb9e22502b439ef9d2c2e26fc51e37eefed7bcec","body_text":"def get_vocab_size(self) -> int: return self.hf_config.vocab_size","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_hidden_size","name":"get_hidden_size","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L187-L188","signature":"def get_hidden_size(self) -> int","exported":true,"origin":"ast","body_hash":"ccf67b8bfb779d99633031ad4043460463c4943b97fc1680926c61bc38b22950","body_text":"def get_hidden_size(self) -> int: return self.hf_config.hidden_size","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_head_size","name":"get_head_size","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L190-L192","signature":"def get_head_size(self) -> int: # FIXME(woosuk): This may not be true for all models.","exported":true,"origin":"ast","body_hash":"2df45d8fd6aad6d737d893165d5bbbf978f57ffb1ed07871191da4c92a5cdaf8","body_text":"def get_head_size(self) -> int: # FIXME(woosuk): This may not be true for all models. return self.hf_config.hidden_size // self.hf_config.num_attention_heads","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_total_num_kv_heads","name":"get_total_num_kv_heads","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L194-L224","signature":"def get_total_num_kv_heads(self) -> int","exported":true,"origin":"ast","body_hash":"ff654b49dde1db484022e103e600ccfa35c769d3cac8bff38f63891532694a20","body_text":"def get_total_num_kv_heads(self) -> int: \"\"\"Returns the total number of KV heads.\"\"\" # For GPTBigCode & Falcon: # NOTE: for falcon, when new_decoder_architecture is True, the # multi_query flag is ignored and we use n_head_kv for the number of # KV heads. falcon_model_types = [\"falcon\", \"RefinedWeb\", \"RefinedWebModel\"] new_decoder_arch_falcon = (self.hf_config.model_type in falcon_model_types and getattr(self.hf_config, \"new_decoder_architecture\", False)) if not new_decoder_arch_falcon and getattr(self.hf_config, \"multi_query\", False): # Multi-query attention, only one KV head. # Currently, tensor parallelism is not supported in this case. return 1 attributes = [ # For Falcon: \"n_head_kv\", \"num_kv_heads\", # For LLaMA-2: \"num_key_value_heads\", # For ChatGLM: \"multi_query_group_num\", ] for attr in attributes: num_kv_heads = getattr(self.hf_config, attr, None) if num_kv_heads is not None: return num_kv_heads # For non-grouped-query attention models, the number of KV heads is # equal to the number of attention heads. return self.hf_config.num_attention_heads","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_num_kv_heads","name":"get_num_kv_heads","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L226-L233","signature":"def get_num_kv_heads(self, parallel_config: \"ParallelConfig\") -> int","exported":true,"origin":"ast","body_hash":"728297fd368aa8412d13b81234869a4484f831c2636bc76bcfd687aa5ad7df83","body_text":"def get_num_kv_heads(self, parallel_config: \"ParallelConfig\") -> int: \"\"\"Returns the number of KV heads per GPU.\"\"\" total_num_kv_heads = self.get_total_num_kv_heads() # If tensor parallelism is used, we divide the number of KV heads by # the tensor parallel size. We will replicate the KV heads in the # case where the number of KV heads is smaller than the tensor # parallel size so each GPU has at least one KV head. return max(1, total_num_kv_heads // parallel_config.tensor_parallel_size)","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_num_layers","name":"get_num_layers","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L235-L237","signature":"def get_num_layers(self, parallel_config: \"ParallelConfig\") -> int","exported":true,"origin":"ast","body_hash":"86a3bc5f0bbae80c7d59376952c94538a7c5fb8ef0e602bdfd9b9d79cbaae254","body_text":"def get_num_layers(self, parallel_config: \"ParallelConfig\") -> int: total_num_hidden_layers = self.hf_config.num_hidden_layers return total_num_hidden_layers // parallel_config.pipeline_parallel_size","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig","name":"CacheConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L240-L310","signature":"class CacheConfig","exported":true,"origin":"ast","body_hash":"8ffd3516b171c2cff20298fa5f26333c1a68c65874e9ee87b16b613d3c9b63fb","body_text":"class CacheConfig: \"\"\"Configuration for the KV cache. Args: block_size: Size of a cache block in number of tokens. gpu_memory_utilization: Fraction of GPU memory to use for the vLLM execution. swap_space: Size of the CPU swap space per GPU (in GiB). cache_dtype: Data type for kv cache storage. \"\"\" def __init__( self, block_size: int, gpu_memory_utilization: float, swap_space: int, cache_dtype: str, sliding_window: Optional[int] = None, ) -> None: self.block_size = block_size self.gpu_memory_utilization = gpu_memory_utilization self.swap_space_bytes = swap_space * _GB self.cache_dtype = cache_dtype self.sliding_window = sliding_window self._verify_args() self._verify_cache_dtype() # Will be set after profiling. self.num_gpu_blocks = None self.num_cpu_blocks = None def _verify_args(self) -> None: if self.gpu_memory_utilization > 1.0: raise ValueError(\"GPU memory utilization must be less than 1.0. Got \" f\"{self.gpu_memory_utilization}.\") def _verify_cache_dtype(self) -> None: if self.cache_dtype == \"auto\": pass elif self.cache_dtype == \"fp8_e5m2\": nvcc_cuda_version = get_nvcc_cuda_version() if nvcc_cuda_version < Version(\"11.8\"): raise ValueError(\"FP8 is not supported when cuda version is lower than 11.8.\") device_name = torch.cuda.get_device_name() if \"AMD\" in device_name: raise NotImplementedError(\"FP8_E5M2 KV Cache on AMD GPU has not been supported yet.\") logger.info(\"Using fp8_e5m2 data type to store kv cache. It reduces \" \"the GPU memory footprint and boosts the performance. \" \"But it may cause slight accuracy drop. \" \"Currently we only support fp8 without scaling factors and \" \"make e5m2 as a default format.\") else: raise ValueError(f\"Unknown kv cache dtype: {self.cache_dtype}\") def verify_with_parallel_config( self, parallel_config: \"ParallelConfig\", ) -> None: total_cpu_memory = get_cpu_memory() # FIXME(woosuk): Here, it is assumed that the GPUs in a tensor parallel # group are in the same node. However, the GPUs may span multiple nodes. num_gpus_per_node = parallel_config.tensor_parallel_size cpu_memory_usage = self.swap_space_bytes * num_gpus_per_node msg = (f\"{cpu_memory_usage / _GB:.2f} GiB out of \" f\"the {total_cpu_memory / _GB:.2f} GiB total CPU memory is \" \"allocated for the swap space.\") if cpu_memory_usage > 0.7 * total_cpu_memory: raise ValueError(\"Too large swap space. \" + msg) elif cpu_memory_usage > 0.4 * total_cpu_memory: logger.warning(\"Possibly too large swap space. \" + msg)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L251-L269","signature":"def __init__( self, block_size: int, gpu_memory_utilization: float, swap_space: int, cache_dtype: str, sliding_window: Optional[int] = None, ) -> None","exported":false,"origin":"ast","body_hash":"ebb5f593ab2edac24275e6e0d6d76b0be7c6a071e255259205f96e24ea8d0a4b","body_text":"def __init__( self, block_size: int, gpu_memory_utilization: float, swap_space: int, cache_dtype: str, sliding_window: Optional[int] = None, ) -> None: self.block_size = block_size self.gpu_memory_utilization = gpu_memory_utilization self.swap_space_bytes = swap_space * _GB self.cache_dtype = cache_dtype self.sliding_window = sliding_window self._verify_args() self._verify_cache_dtype() # Will be set after profiling. self.num_gpu_blocks = None self.num_cpu_blocks = None","summary_state":"pending","summary":null,"crux":null,"owner":"CacheConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_args","name":"_verify_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L271-L274","signature":"def _verify_args(self) -> None","exported":false,"origin":"ast","body_hash":"4885b811e514f8a3060c8565e3f735fabcb0b2efcd65f123ddeabf51fb3d1e33","body_text":"def _verify_args(self) -> None: if self.gpu_memory_utilization > 1.0: raise ValueError(\"GPU memory utilization must be less than 1.0. Got \" f\"{self.gpu_memory_utilization}.\")","summary_state":"pending","summary":null,"crux":null,"owner":"CacheConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_cache_dtype","name":"_verify_cache_dtype","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L276-L292","signature":"def _verify_cache_dtype(self) -> None","exported":false,"origin":"ast","body_hash":"c1bfd23959d67f6cc5993c149dbfd80203a8cf01d21453f39402ae4cd9bf0c87","body_text":"def _verify_cache_dtype(self) -> None: if self.cache_dtype == \"auto\": pass elif self.cache_dtype == \"fp8_e5m2\": nvcc_cuda_version = get_nvcc_cuda_version() if nvcc_cuda_version < Version(\"11.8\"): raise ValueError(\"FP8 is not supported when cuda version is lower than 11.8.\") device_name = torch.cuda.get_device_name() if \"AMD\" in device_name: raise NotImplementedError(\"FP8_E5M2 KV Cache on AMD GPU has not been supported yet.\") logger.info(\"Using fp8_e5m2 data type to store kv cache. It reduces \" \"the GPU memory footprint and boosts the performance. \" \"But it may cause slight accuracy drop. \" \"Currently we only support fp8 without scaling factors and \" \"make e5m2 as a default format.\") else: raise ValueError(f\"Unknown kv cache dtype: {self.cache_dtype}\")","summary_state":"pending","summary":null,"crux":null,"owner":"CacheConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig.verify_with_parallel_config","name":"verify_with_parallel_config","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L294-L310","signature":"def verify_with_parallel_config( self, parallel_config: \"ParallelConfig\", ) -> None","exported":true,"origin":"ast","body_hash":"941b39cfd15e60839411a6b42b4cb00a12a843a20e04ba179de471e4f6485285","body_text":"def verify_with_parallel_config( self, parallel_config: \"ParallelConfig\", ) -> None: total_cpu_memory = get_cpu_memory() # FIXME(woosuk): Here, it is assumed that the GPUs in a tensor parallel # group are in the same node. However, the GPUs may span multiple nodes. num_gpus_per_node = parallel_config.tensor_parallel_size cpu_memory_usage = self.swap_space_bytes * num_gpus_per_node msg = (f\"{cpu_memory_usage / _GB:.2f} GiB out of \" f\"the {total_cpu_memory / _GB:.2f} GiB total CPU memory is \" \"allocated for the swap space.\") if cpu_memory_usage > 0.7 * total_cpu_memory: raise ValueError(\"Too large swap space. \" + msg) elif cpu_memory_usage > 0.4 * total_cpu_memory: logger.warning(\"Possibly too large swap space. \" + msg)","summary_state":"pending","summary":null,"crux":null,"owner":"CacheConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig","name":"ParallelConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L313-L367","signature":"class ParallelConfig","exported":true,"origin":"ast","body_hash":"5bc42bf14ef04bfd1f0883c7e7c02669a042446909ee914863ab93f469b26e01","body_text":"class ParallelConfig: \"\"\"Configuration for the distributed execution. Args: pipeline_parallel_size: Number of pipeline parallel groups. tensor_parallel_size: Number of tensor parallel groups. worker_use_ray: Whether to use Ray for model workers. Will be set to True if either pipeline_parallel_size or tensor_parallel_size is greater than 1. max_parallel_loading_workers: Maximum number of multiple batches when load model sequentially. To avoid RAM OOM when using tensor parallel and large models. disable_custom_all_reduce: Disable the custom all-reduce kernel and fall back to NCCL. \"\"\" def __init__( self, pipeline_parallel_size: int, tensor_parallel_size: int, worker_use_ray: bool, max_parallel_loading_workers: Optional[int] = None, disable_custom_all_reduce: bool = False, ) -> None: self.pipeline_parallel_size = pipeline_parallel_size self.tensor_parallel_size = tensor_parallel_size self.worker_use_ray = worker_use_ray self.max_parallel_loading_workers = max_parallel_loading_workers self.disable_custom_all_reduce = disable_custom_all_reduce self.world_size = pipeline_parallel_size * tensor_parallel_size if self.world_size > 1: self.worker_use_ray = True self._verify_args() def _verify_args(self) -> None: if self.pipeline_parallel_size > 1: raise NotImplementedError(\"Pipeline parallelism is not supported yet.\") if not self.disable_custom_all_reduce and self.world_size > 1: if is_hip(): self.disable_custom_all_reduce = True logger.info(\"Disabled the custom all-reduce kernel because it is not \" \"supported on AMD GPUs.\") elif self.pipeline_parallel_size > 1: self.disable_custom_all_reduce = True logger.info(\"Disabled the custom all-reduce kernel because it is not \" \"supported with pipeline parallelism.\") # FIXME(woosuk): Fix the stability issues and re-enable the custom # all-reduce kernel. if not self.disable_custom_all_reduce and self.world_size > 1: self.disable_custom_all_reduce = True logger.info(\"Custom all-reduce kernels are temporarily disabled due to \" \"stability issues. We will re-enable them once the issues are \" \"resolved.\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L329-L346","signature":"def __init__( self, pipeline_parallel_size: int, tensor_parallel_size: int, worker_use_ray: bool, max_parallel_loading_workers: Optional[int] = None, disable_custom_all_reduce: bool = False, ) -> None","exported":false,"origin":"ast","body_hash":"223dbab83ab2ffef7ef09fcd81ab6316cab91c7dd418a559f23d446028de2af9","body_text":"def __init__( self, pipeline_parallel_size: int, tensor_parallel_size: int, worker_use_ray: bool, max_parallel_loading_workers: Optional[int] = None, disable_custom_all_reduce: bool = False, ) -> None: self.pipeline_parallel_size = pipeline_parallel_size self.tensor_parallel_size = tensor_parallel_size self.worker_use_ray = worker_use_ray self.max_parallel_loading_workers = max_parallel_loading_workers self.disable_custom_all_reduce = disable_custom_all_reduce self.world_size = pipeline_parallel_size * tensor_parallel_size if self.world_size > 1: self.worker_use_ray = True self._verify_args()","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig._verify_args","name":"_verify_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L348-L367","signature":"def _verify_args(self) -> None","exported":false,"origin":"ast","body_hash":"c94f15f6d712eccd0c3dcaca68b51f7a7a35053769cbd1c6f9dd99933392f91b","body_text":"def _verify_args(self) -> None: if self.pipeline_parallel_size > 1: raise NotImplementedError(\"Pipeline parallelism is not supported yet.\") if not self.disable_custom_all_reduce and self.world_size > 1: if is_hip(): self.disable_custom_all_reduce = True logger.info(\"Disabled the custom all-reduce kernel because it is not \" \"supported on AMD GPUs.\") elif self.pipeline_parallel_size > 1: self.disable_custom_all_reduce = True logger.info(\"Disabled the custom all-reduce kernel because it is not \" \"supported with pipeline parallelism.\") # FIXME(woosuk): Fix the stability issues and re-enable the custom # all-reduce kernel. if not self.disable_custom_all_reduce and self.world_size > 1: self.disable_custom_all_reduce = True logger.info(\"Custom all-reduce kernels are temporarily disabled due to \" \"stability issues. We will re-enable them once the issues are \" \"resolved.\")","summary_state":"pending","summary":null,"crux":null,"owner":"ParallelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig","name":"SchedulerConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L370-L412","signature":"class SchedulerConfig","exported":true,"origin":"ast","body_hash":"a68fce0d2ee9ab0ab24f28d7629b178f74093ce60a0ce57452cc808c61e7a305","body_text":"class SchedulerConfig: \"\"\"Scheduler configuration. Args: max_num_batched_tokens: Maximum number of tokens to be processed in a single iteration. max_num_seqs: Maximum number of sequences to be processed in a single iteration. max_model_len: Maximum length of a sequence (including prompt and generated text). max_paddings: Maximum number of paddings to be added to a batch. \"\"\" def __init__( self, max_num_batched_tokens: Optional[int], max_num_seqs: int, max_model_len: int, max_paddings: int, ) -> None: if max_num_batched_tokens is not None: self.max_num_batched_tokens = max_num_batched_tokens else: # If max_model_len is too short, use 2048 as the default value for # higher throughput. self.max_num_batched_tokens = max(max_model_len, 2048) self.max_num_seqs = max_num_seqs self.max_model_len = max_model_len self.max_paddings = max_paddings self._verify_args() def _verify_args(self) -> None: if self.max_num_batched_tokens < self.max_model_len: raise ValueError(f\"max_num_batched_tokens ({self.max_num_batched_tokens}) is \" f\"smaller than max_model_len ({self.max_model_len}). \" \"This effectively limits the maximum sequence length to \" \"max_num_batched_tokens and makes vLLM reject longer \" \"sequences. Please increase max_num_batched_tokens or \" \"decrease max_model_len.\") if self.max_num_batched_tokens < self.max_num_seqs: raise ValueError(f\"max_num_batched_tokens ({self.max_num_batched_tokens}) must \" \"be greater than or equal to max_num_seqs \" f\"({self.max_num_seqs}).\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L383-L399","signature":"def __init__( self, max_num_batched_tokens: Optional[int], max_num_seqs: int, max_model_len: int, max_paddings: int, ) -> None","exported":false,"origin":"ast","body_hash":"cbda8f0ebaab2f580086185bf0730e2c208302f8ce667c4cf43325a53e7fc945","body_text":"def __init__( self, max_num_batched_tokens: Optional[int], max_num_seqs: int, max_model_len: int, max_paddings: int, ) -> None: if max_num_batched_tokens is not None: self.max_num_batched_tokens = max_num_batched_tokens else: # If max_model_len is too short, use 2048 as the default value for # higher throughput. self.max_num_batched_tokens = max(max_model_len, 2048) self.max_num_seqs = max_num_seqs self.max_model_len = max_model_len self.max_paddings = max_paddings self._verify_args()","summary_state":"pending","summary":null,"crux":null,"owner":"SchedulerConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig._verify_args","name":"_verify_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L401-L412","signature":"def _verify_args(self) -> None","exported":false,"origin":"ast","body_hash":"1ab89a8f8c97edc525fecdbd0bdcaac213d19cfdc7525087f029a595bc54ecab","body_text":"def _verify_args(self) -> None: if self.max_num_batched_tokens < self.max_model_len: raise ValueError(f\"max_num_batched_tokens ({self.max_num_batched_tokens}) is \" f\"smaller than max_model_len ({self.max_model_len}). \" \"This effectively limits the maximum sequence length to \" \"max_num_batched_tokens and makes vLLM reject longer \" \"sequences. Please increase max_num_batched_tokens or \" \"decrease max_model_len.\") if self.max_num_batched_tokens < self.max_num_seqs: raise ValueError(f\"max_num_batched_tokens ({self.max_num_batched_tokens}) must \" \"be greater than or equal to max_num_seqs \" f\"({self.max_num_seqs}).\")","summary_state":"pending","summary":null,"crux":null,"owner":"SchedulerConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#DeviceConfig","name":"DeviceConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L415-L418","signature":"class DeviceConfig","exported":true,"origin":"ast","body_hash":"60273a7ecdc1fd496506a19bb0aed3216b3f8c39dec67d50f1698f0b938e4dd7","body_text":"class DeviceConfig: def __init__(self, device: str = \"cuda\") -> None: self.device = torch.device(device)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#DeviceConfig.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L417-L418","signature":"def __init__(self, device: str = \"cuda\") -> None","exported":false,"origin":"ast","body_hash":"2d49d5f1c846b32d743263046334ef386121ac6721b0a08a36e483fac1d1bb18","body_text":"def __init__(self, device: str = \"cuda\") -> None: self.device = torch.device(device)","summary_state":"pending","summary":null,"crux":null,"owner":"DeviceConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig","name":"LoRAConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L422-L461","signature":"class LoRAConfig","exported":true,"origin":"ast","body_hash":"f0cb02af36b0ec84bbd79bedeec497dabbe9c8b5eea3e630de03b1ffba427fba","body_text":"class LoRAConfig: max_lora_rank: int max_loras: int max_cpu_loras: Optional[int] = None lora_dtype: Optional[torch.dtype] = None lora_extra_vocab_size: int = 256 # This is a constant. lora_vocab_padding_size: ClassVar[int] = 256 def __post_init__(self): # Keep this in sync with csrc/punica/bgmv/bgmv_config.h possible_max_ranks = (8, 16, 32, 64) possible_lora_extra_vocab_size = (0, 256, 512) if self.max_lora_rank not in possible_max_ranks: raise ValueError(f\"max_lora_rank ({self.max_lora_rank}) must be one of \" f\"{possible_max_ranks}.\") if self.lora_extra_vocab_size not in possible_lora_extra_vocab_size: raise ValueError(f\"lora_extra_vocab_size ({self.lora_extra_vocab_size}) \" f\"must be one of {possible_lora_extra_vocab_size}.\") if self.max_loras < 1: raise ValueError(f\"max_loras ({self.max_loras}) must be >= 1.\") if self.max_cpu_loras is None: self.max_cpu_loras = self.max_loras elif self.max_cpu_loras < self.max_loras: raise ValueError(f\"max_cpu_loras ({self.max_cpu_loras}) must be >= \" f\"max_loras ({self.max_loras})\") def verify_with_model_config(self, model_config: ModelConfig): if self.lora_dtype in (None, \"auto\"): self.lora_dtype = model_config.dtype elif isinstance(self.lora_dtype, str): self.lora_dtype = getattr(torch, self.lora_dtype) if model_config.quantization is not None: raise ValueError(\"LoRA is not supported with quantized models yet.\") def verify_with_scheduler_config(self, scheduler_config: SchedulerConfig): if scheduler_config.max_num_batched_tokens > 65528: raise ValueError(\"Due to limitations of the custom LoRA CUDA kernel, \" \"max_num_batched_tokens must be <= 65528 when \" \"LoRA is enabled.\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.__post_init__","name":"__post_init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L431-L447","signature":"def __post_init__(self): # Keep this in sync with csrc/punica/bgmv/bgmv_config.h","exported":false,"origin":"ast","body_hash":"60c0ee4ec17f4422e59f2bae58b9bc9c83b93ddab654bdfd306dc6b5f76fc038","body_text":"def __post_init__(self): # Keep this in sync with csrc/punica/bgmv/bgmv_config.h possible_max_ranks = (8, 16, 32, 64) possible_lora_extra_vocab_size = (0, 256, 512) if self.max_lora_rank not in possible_max_ranks: raise ValueError(f\"max_lora_rank ({self.max_lora_rank}) must be one of \" f\"{possible_max_ranks}.\") if self.lora_extra_vocab_size not in possible_lora_extra_vocab_size: raise ValueError(f\"lora_extra_vocab_size ({self.lora_extra_vocab_size}) \" f\"must be one of {possible_lora_extra_vocab_size}.\") if self.max_loras < 1: raise ValueError(f\"max_loras ({self.max_loras}) must be >= 1.\") if self.max_cpu_loras is None: self.max_cpu_loras = self.max_loras elif self.max_cpu_loras < self.max_loras: raise ValueError(f\"max_cpu_loras ({self.max_cpu_loras}) must be >= \" f\"max_loras ({self.max_loras})\")","summary_state":"pending","summary":null,"crux":null,"owner":"LoRAConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.verify_with_model_config","name":"verify_with_model_config","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L449-L455","signature":"def verify_with_model_config(self, model_config: ModelConfig)","exported":true,"origin":"ast","body_hash":"4d6293703a2e7bceee5ed54997894d351a93d5125a1f3ca1d537763e14727529","body_text":"def verify_with_model_config(self, model_config: ModelConfig): if self.lora_dtype in (None, \"auto\"): self.lora_dtype = model_config.dtype elif isinstance(self.lora_dtype, str): self.lora_dtype = getattr(torch, self.lora_dtype) if model_config.quantization is not None: raise ValueError(\"LoRA is not supported with quantized models yet.\")","summary_state":"pending","summary":null,"crux":null,"owner":"LoRAConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.verify_with_scheduler_config","name":"verify_with_scheduler_config","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L457-L461","signature":"def verify_with_scheduler_config(self, scheduler_config: SchedulerConfig)","exported":true,"origin":"ast","body_hash":"3952c44a54de8786ef895a5ddad8e5e31dac5ee2488771dea45dff8d1a5c1cd2","body_text":"def verify_with_scheduler_config(self, scheduler_config: SchedulerConfig): if scheduler_config.max_num_batched_tokens > 65528: raise ValueError(\"Due to limitations of the custom LoRA CUDA kernel, \" \"max_num_batched_tokens must be <= 65528 when \" \"LoRA is enabled.\")","summary_state":"pending","summary":null,"crux":null,"owner":"LoRAConfig"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","name":"_get_and_verify_dtype","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L475-L522","signature":"def _get_and_verify_dtype( config: PretrainedConfig, dtype: Union[str, torch.dtype], ) -> torch.dtype: # NOTE: getattr(config, \"torch_dtype\", torch.float32) is not correct # because config.torch_dtype can be None.","exported":false,"origin":"ast","body_hash":"dd6102b21807bd60a7050824861fba84b542f7e00ec780d2f2d1188a9d3fb835","body_text":"def _get_and_verify_dtype( config: PretrainedConfig, dtype: Union[str, torch.dtype], ) -> torch.dtype: # NOTE: getattr(config, \"torch_dtype\", torch.float32) is not correct # because config.torch_dtype can be None. config_dtype = getattr(config, \"torch_dtype\", None) if config_dtype is None: config_dtype = torch.float32 if isinstance(dtype, str): dtype = dtype.lower() if dtype == \"auto\": if config_dtype == torch.float32: # Following the common practice, we use float16 for float32 # models. torch_dtype = torch.float16 else: torch_dtype = config_dtype else: if dtype not in _STR_DTYPE_TO_TORCH_DTYPE: raise ValueError(f\"Unknown dtype: {dtype}\") torch_dtype = _STR_DTYPE_TO_TORCH_DTYPE[dtype] elif isinstance(dtype, torch.dtype): torch_dtype = dtype else: raise ValueError(f\"Unknown dtype: {dtype}\") if is_hip() and torch_dtype == torch.float32: rocm_supported_dtypes = [ k for k, v in _STR_DTYPE_TO_TORCH_DTYPE.items() if (k not in _ROCM_NOT_SUPPORTED_DTYPE) ] raise ValueError(f\"dtype \\'{dtype}\\' is not supported in ROCm. \" f\"Supported dtypes are {rocm_supported_dtypes}\") # Verify the dtype. if torch_dtype != config_dtype: if torch_dtype == torch.float32: # Upcasting to float32 is allowed. pass elif config_dtype == torch.float32: # Downcasting from float32 to float16 or bfloat16 is allowed. pass else: # Casting between float16 and bfloat16 is allowed with a warning. logger.warning(f\"Casting {config_dtype} to {torch_dtype}.\") return torch_dtype","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_max_len","name":"_get_and_verify_max_len","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/config.py","span":"L525-L577","signature":"def _get_and_verify_max_len( hf_config: PretrainedConfig, max_model_len: Optional[int], ) -> int","exported":false,"origin":"ast","body_hash":"b11edf512bd0ebe4dd07b980217e254ff14bcf0c749cc3bed2b1a0e89bfdc84b","body_text":"def _get_and_verify_max_len( hf_config: PretrainedConfig, max_model_len: Optional[int], ) -> int: \"\"\"Get and verify the model's maximum length.\"\"\" derived_max_model_len = float(\"inf\") possible_keys = [ # OPT \"max_position_embeddings\", # GPT-2 \"n_positions\", # MPT \"max_seq_len\", # ChatGLM2 \"seq_length\", # Others \"max_sequence_length\", \"max_seq_length\", \"seq_len\", ] for key in possible_keys: max_len_key = getattr(hf_config, key, None) if max_len_key is not None: derived_max_model_len = min(derived_max_model_len, max_len_key) if derived_max_model_len == float(\"inf\"): if max_model_len is not None: # If max_model_len is specified, we use it. return max_model_len default_max_len = 2048 logger.warning(\"The model's config.json does not contain any of the following \" \"keys to determine the original maximum length of the model: \" f\"{possible_keys}. Assuming the model's maximum length is \" f\"{default_max_len}.\") derived_max_model_len = default_max_len rope_scaling = getattr(hf_config, \"rope_scaling\", None) if rope_scaling is not None: assert \"factor\" in rope_scaling scaling_factor = rope_scaling[\"factor\"] if rope_scaling[\"type\"] == \"yarn\": derived_max_model_len = rope_scaling[\"original_max_position_embeddings\"] derived_max_model_len *= scaling_factor if max_model_len is None: max_model_len = derived_max_model_len elif max_model_len > derived_max_model_len: raise ValueError(f\"User-specified max_model_len ({max_model_len}) is greater than \" f\"the derived max_model_len ({max_len_key}={derived_max_model_len}\" \" in model's config.json). This may lead to incorrect model \" \"outputs or CUDA errors. Make sure the value is correct and \" \"within the model context size.\") return int(max_model_len)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"imports","specifier":"packaging.version","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"imports","specifier":"vllm.transformers_utils.config","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.__init__","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.__init__","relation":"calls","name":"_get_and_verify_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.__init__","relation":"calls","name":"_get_and_verify_max_len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.__init__","relation":"calls","name":"_verify_quantization","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.__init__","relation":"calls","name":"_verify_cuda_graph","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_load_format","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_load_format","relation":"calls","name":"lower","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_load_format","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","relation":"calls","name":"lower","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","relation":"calls","name":"lower","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","relation":"calls","name":"str","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","relation":"calls","name":"is_hip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_quantization","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_cuda_graph","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_cuda_graph","relation":"calls","name":"min","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig._verify_cuda_graph","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.verify_with_parallel_config","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.verify_with_parallel_config","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.verify_with_parallel_config","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_sliding_window","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_sliding_window","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_vocab_size","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_hidden_size","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_head_size","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_total_num_kv_heads","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_total_num_kv_heads","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_total_num_kv_heads","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_total_num_kv_heads","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_num_kv_heads","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_num_kv_heads","relation":"calls","name":"get_total_num_kv_heads","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_num_kv_heads","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ModelConfig.get_num_layers","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig.__init__","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig.__init__","relation":"calls","name":"_verify_args","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"CacheConfig"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig.__init__","relation":"calls","name":"_verify_cache_dtype","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"CacheConfig"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_args","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_args","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_cache_dtype","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_cache_dtype","relation":"calls","name":"get_nvcc_cuda_version","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_cache_dtype","relation":"calls","name":"Version","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_cache_dtype","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_cache_dtype","relation":"calls","name":"get_device_name","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_cache_dtype","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_cache_dtype","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig._verify_cache_dtype","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig.verify_with_parallel_config","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig.verify_with_parallel_config","relation":"calls","name":"get_cpu_memory","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig.verify_with_parallel_config","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#CacheConfig.verify_with_parallel_config","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig.__init__","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig.__init__","relation":"calls","name":"_verify_args","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"ParallelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig._verify_args","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig._verify_args","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig._verify_args","relation":"calls","name":"is_hip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig._verify_args","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig._verify_args","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#ParallelConfig._verify_args","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig.__init__","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig.__init__","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig.__init__","relation":"calls","name":"_verify_args","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"SchedulerConfig"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig._verify_args","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig._verify_args","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#SchedulerConfig._verify_args","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#DeviceConfig","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#DeviceConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#DeviceConfig.__init__","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#DeviceConfig.__init__","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.__post_init__","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.__post_init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.__post_init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.__post_init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.__post_init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.verify_with_model_config","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.verify_with_model_config","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.verify_with_model_config","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.verify_with_model_config","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.verify_with_scheduler_config","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#LoRAConfig.verify_with_scheduler_config","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","relation":"calls","name":"lower","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","relation":"calls","name":"is_hip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_dtype","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_max_len","file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_max_len","relation":"calls","name":"float","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_max_len","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_max_len","relation":"calls","name":"min","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_max_len","relation":"calls","name":"float","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_max_len","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_max_len","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_max_len","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/config.py#_get_and_verify_max_len","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/config.py"}]},"verl/third_party/vllm/vllm_v_0_3_1/llm.py":{"size":13604,"mtimeMs":1789275830312.2026,"hash":"bb5dd8a628d2d58b9daf33fe97fd471778b11e07212f5a7ac75c47a48273dca7","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","name":"llm.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L1-L276","signature":null,"exported":true,"origin":"ast","body_hash":"bb5dd8a628d2d58b9daf33fe97fd471778b11e07212f5a7ac75c47a48273dca7","chars":13604,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/llm.py from typing import Dict, List, Optional, Tuple, Union from tqdm import tqdm from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from transformers import PretrainedConfig import torch.nn as nn from .arg_utils import EngineArgs from .llm_engine_sp import LLMEngine from vllm.lora.request import LoRARequest from vllm.outputs import RequestOutput from vllm.sampling_params import SamplingParams from vllm.utils import Counter import torch from torch.nn.utils.rnn import pad_sequence from verl.workers.rollout.tokenizer import HybridEngineBaseTokenizer"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","name":"LLM","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L33-L275","signature":"class LLM","exported":true,"origin":"ast","body_hash":"c999f578590769057ca38c4d7650d7e356fd8ff00b594026de989aba0a981f3a","body_text":"class LLM: \"\"\"An LLM for generating texts from given prompts and sampling parameters. This class includes a tokenizer, a language model (possibly distributed across multiple GPUs), and GPU memory space allocated for intermediate states (aka KV cache). Given a batch of prompts and sampling parameters, this class generates texts from the model, using an intelligent batching mechanism and efficient memory management. NOTE: This class is intended to be used for offline inference. For online serving, use the `AsyncLLMEngine` class instead. NOTE: For the comprehensive list of arguments, see `EngineArgs`. Args: model: A HuggingFace Transformers model instance. tokenizer: A HuggingFace Transformers tokenizer instance. tokenizer_mode: The tokenizer mode. \"auto\" will use the fast tokenizer if available, and \"slow\" will always use the slow tokenizer. trust_remote_code: Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer. tensor_parallel_size: The number of GPUs to use for distributed execution with tensor parallelism. dtype: The data type for the model weights and activations. Currently, we support `float32`, `float16`, and `bfloat16`. If `auto`, we use the `torch_dtype` attribute specified in the model config file. However, if the `torch_dtype` in the config is `float32`, we will use `float16` instead. quantization: The method used to quantize the model weights. Currently, we support \"awq\". If None, we assume the model weights are not quantized and use `dtype` to determine the data type of the weights. revision: The specific model version to use. It can be a branch name, a tag name, or a commit id. tokenizer_revision: The specific tokenizer version to use. It can be a branch name, a tag name, or a commit id. seed: The seed to initialize the random number generator for sampling. gpu_memory_utilization: The ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV cache. Higher values will increase the KV cache size and thus improve the model's throughput. However, if the value is too high, it may cause out-of- memory (OOM) errors. swap_space: The size (GiB) of CPU memory per GPU to use as swap space. This can be used for temporarily storing the states of the requests when their `best_of` sampling parameters are larger than 1. If all requests will have `best_of=1`, you can safely set this to 0. Otherwise, too small values may cause out-of-memory (OOM) errors. enforce_eager: Whether to enforce eager execution. If True, we will disable CUDA graph and always execute the model in eager mode. If False, we will use CUDA graph and eager execution in hybrid. max_context_len_to_capture: Maximum context len covered by CUDA graphs. When a sequence has context length larger than this, we fall back to eager mode. disable_custom_all_reduce: See ParallelConfig \"\"\" def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, enforce_eager: bool = False, max_context_len_to_capture: int = 8192, disable_custom_all_reduce: bool = False, **kwargs, ) -> None: if \"disable_log_stats\" not in kwargs: kwargs[\"disable_log_stats\"] = True engine_args = EngineArgs( model_hf_config=model_hf_config, tensor_parallel_size=tensor_parallel_size, dtype=dtype, quantization=quantization, revision=revision, tokenizer_revision=tokenizer_revision, seed=seed, gpu_memory_utilization=gpu_memory_utilization, swap_space=swap_space, enforce_eager=enforce_eager, max_context_len_to_capture=max_context_len_to_capture, disable_custom_all_reduce=disable_custom_all_reduce, **kwargs, ) tokenizer_cls = (PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer) if not isinstance(tokenizer, tokenizer_cls): raise ValueError( f\"Unexpected tokenizer type: {type(tokenizer)}. Must be\" \"one of the following: PreTrainedTokenizer, PreTrainedTokenizerFast, verl.workers.rollout.HybridEngineBaseTokenizer\" ) self.llm_engine = LLMEngine.from_engine_args(model, tokenizer, engine_args) self.request_counter = Counter() def init_cache_engine(self): self.llm_engine.init_cache_engine() def free_cache_engine(self): self.llm_engine.free_cache_engine() def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: return self.llm_engine.tokenizer def set_tokenizer( self, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], ) -> None: self.llm_engine.tokenizer = tokenizer def generate( self, prompts: Optional[Union[str, List[str]]] = None, sampling_params: Optional[SamplingParams] = None, prompt_token_ids: Optional[List[List[int]]] = None, prefix_pos","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L87-L131","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, enforce_eager: bool = False, max_context_len_to_capture: int = 8192, disable_custom_all_reduce: bool = False, **kwargs, ) -> None","exported":false,"origin":"ast","body_hash":"61dc731ea33fbd719712e8097f991f8cf957a6aa9a99e043e2e166a9d7c6ad6a","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, enforce_eager: bool = False, max_context_len_to_capture: int = 8192, disable_custom_all_reduce: bool = False, **kwargs, ) -> None: if \"disable_log_stats\" not in kwargs: kwargs[\"disable_log_stats\"] = True engine_args = EngineArgs( model_hf_config=model_hf_config, tensor_parallel_size=tensor_parallel_size, dtype=dtype, quantization=quantization, revision=revision, tokenizer_revision=tokenizer_revision, seed=seed, gpu_memory_utilization=gpu_memory_utilization, swap_space=swap_space, enforce_eager=enforce_eager, max_context_len_to_capture=max_context_len_to_capture, disable_custom_all_reduce=disable_custom_all_reduce, **kwargs, ) tokenizer_cls = (PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer) if not isinstance(tokenizer, tokenizer_cls): raise ValueError( f\"Unexpected tokenizer type: {type(tokenizer)}. Must be\" \"one of the following: PreTrainedTokenizer, PreTrainedTokenizerFast, verl.workers.rollout.HybridEngineBaseTokenizer\" ) self.llm_engine = LLMEngine.from_engine_args(model, tokenizer, engine_args) self.request_counter = Counter()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L133-L134","signature":"def init_cache_engine(self)","exported":true,"origin":"ast","body_hash":"04d68c22194d8053c926485d043932fb2526f0d73a880157acee875717c5937f","body_text":"def init_cache_engine(self): self.llm_engine.init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L136-L137","signature":"def free_cache_engine(self)","exported":true,"origin":"ast","body_hash":"179e8502fa237bb23acb9f06952cf4ae62f3f9010408f63734f8adf6a5163378","body_text":"def free_cache_engine(self): self.llm_engine.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.get_tokenizer","name":"get_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L139-L140","signature":"def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]","exported":true,"origin":"ast","body_hash":"cebad3b0992af1915c8c8a1543b493069d6507f491421f924191ee4b98e60a2f","body_text":"def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: return self.llm_engine.tokenizer","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.set_tokenizer","name":"set_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L142-L146","signature":"def set_tokenizer( self, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], ) -> None","exported":true,"origin":"ast","body_hash":"a9d3cbc026686b5d94c85093062fb18f30a9d4182a5bf7eb9679c456cbc101fa","body_text":"def set_tokenizer( self, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], ) -> None: self.llm_engine.tokenizer = tokenizer","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","name":"generate","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L148-L199","signature":"def generate( self, prompts: Optional[Union[str, List[str]]] = None, sampling_params: Optional[SamplingParams] = None, prompt_token_ids: Optional[List[List[int]]] = None, prefix_pos: Optional[Union[int, List[int]]] = None, use_tqdm: bool = True, lora_request: Optional[LoRARequest] = None, ) -> List[RequestOutput]","exported":true,"origin":"ast","body_hash":"ba58cfbf0008fdb315ae0f2abbd2147231626d275e1e11a0bf3ea688f69adaca","body_text":"def generate( self, prompts: Optional[Union[str, List[str]]] = None, sampling_params: Optional[SamplingParams] = None, prompt_token_ids: Optional[List[List[int]]] = None, prefix_pos: Optional[Union[int, List[int]]] = None, use_tqdm: bool = True, lora_request: Optional[LoRARequest] = None, ) -> List[RequestOutput]: \"\"\"Generates the completions for the input prompts. NOTE: This class automatically batches the given prompts, considering the memory constraint. For the best performance, put all of your prompts into a single list and pass it to this method. Args: prompts: A list of prompts to generate completions for. sampling_params: The sampling parameters for text generation. If None, we use the default sampling parameters. prompt_token_ids: A list of token IDs for the prompts. If None, we use the tokenizer to convert the prompts to token IDs. use_tqdm: Whether to use tqdm to display the progress bar. Returns: A list of `RequestOutput` objects containing the generated completions in the same order as the input prompts. \"\"\" if prompts is None and prompt_token_ids is None: raise ValueError(\"Either prompts or prompt_token_ids must be \" \"provided.\") if isinstance(prompts, str): # Convert a single prompt to a list. prompts = [prompts] if prompts is not None and prompt_token_ids is not None: if len(prompts) != len(prompt_token_ids): raise ValueError(\"The lengths of prompts and prompt_token_ids \" \"must be the same.\") if sampling_params is None: # Use default sampling params. sampling_params = SamplingParams() # Add requests to the engine. num_requests = len(prompts) if prompts is not None else len(prompt_token_ids) for i in range(num_requests): prompt = prompts[i] if prompts is not None else None prefix_pos_i = prefix_pos[i] if prefix_pos is not None else None token_ids = None if prompt_token_ids is None else prompt_token_ids[i] if not isinstance(token_ids, list): # NOTE(shengguangming): convert the rollout input into List[str] token_ids = self._pre_process_inputs(token_ids) self._add_request(prompt, sampling_params, token_ids, lora_request=lora_request, prefix_pos=prefix_pos_i) return self._run_engine(use_tqdm)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._add_request","name":"_add_request","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L201-L215","signature":"def _add_request( self, prompt: Optional[str], sampling_params: SamplingParams, prompt_token_ids: Optional[List[int]], lora_request: Optional[LoRARequest] = None, prefix_pos: Optional[int] = None, ) -> None","exported":false,"origin":"ast","body_hash":"ef97c1429fe284c8141a06feafe15e442e1919fffd9ad3d402842e8a34c42b1d","body_text":"def _add_request( self, prompt: Optional[str], sampling_params: SamplingParams, prompt_token_ids: Optional[List[int]], lora_request: Optional[LoRARequest] = None, prefix_pos: Optional[int] = None, ) -> None: request_id = str(next(self.request_counter)) self.llm_engine.add_request(request_id, prompt, sampling_params, prompt_token_ids, lora_request=lora_request, prefix_pos=prefix_pos)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","name":"_run_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L217-L238","signature":"def _run_engine(self, use_tqdm: bool) -> List[RequestOutput]: # Initialize tqdm.","exported":false,"origin":"ast","body_hash":"7b389b73a1eac511edb8d203eb0dd3a9987bc7f898d2792e16bd63e3b5b1a040","body_text":"def _run_engine(self, use_tqdm: bool) -> List[RequestOutput]: # Initialize tqdm. if use_tqdm: num_requests = self.llm_engine.get_num_unfinished_requests() pbar = tqdm(total=num_requests, desc=\"Processed prompts\") # Run the engine. outputs: List[RequestOutput] = [] while self.llm_engine.has_unfinished_requests(): step_outputs = self.llm_engine.step() for output in step_outputs: if output.finished: outputs.append(output) if use_tqdm: pbar.update(1) if use_tqdm: pbar.close() # Sort the outputs by request ID. # This is necessary because some requests may be finished earlier than # its previous requests. outputs = sorted(outputs, key=lambda x: int(x.request_id)) # TODO(shengguangming): maybe we can hack the autoregressive logics without only apply post process for better performance return self._post_process_outputs(outputs)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._pre_process_inputs","name":"_pre_process_inputs","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L242-L247","signature":"def _pre_process_inputs(self, prompt_token_ids: torch.Tensor) -> List[int]: # remove the left padding in the prompt token_id","exported":false,"origin":"ast","body_hash":"fdf721950a6eec9276740f600d40ca1c9190d1b345d858eefb8f2b9f83c0c799","body_text":"def _pre_process_inputs(self, prompt_token_ids: torch.Tensor) -> List[int]: # remove the left padding in the prompt token_id pad_token_id = self.llm_engine.tokenizer.pad_token_id if self.llm_engine.tokenizer.pad_token_id is not None else self.llm_engine.tokenizer.eos_token_id non_pad_index = torch.nonzero(prompt_token_ids != pad_token_id, as_tuple=False)[0][0] token_ids = prompt_token_ids[non_pad_index:].tolist() return token_ids","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","name":"_post_process_outputs","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L250-L269","signature":"def _post_process_outputs(self, outputs: List[RequestOutput]) -> Tuple[torch.Tensor, torch.Tensor]","exported":false,"origin":"ast","body_hash":"79f98445ec654840b37bcfe5eec45725e8180fa78967e34987d5b6f41cdff140","body_text":"def _post_process_outputs(self, outputs: List[RequestOutput]) -> Tuple[torch.Tensor, torch.Tensor]: output_token_ids = [] logprobs = [] for output in outputs: # List[RequestOutput] output = output.outputs for output in output: # List[CompletionOutput], usually len == 1 output_token_ids.append(torch.tensor(output.token_ids)) # TODO(shengguangming): can be optimzied by rewrite the Sampler._get_logprobs() logits logprobs_dicts = output.logprobs if logprobs_dicts is not None: logprob = [] for logprobs_dict, id in zip(logprobs_dicts, output.token_ids): logprob.append(logprobs_dict[id]) logprobs.append(torch.tensor(logprob)) pad_token_id = self.llm_engine.tokenizer.pad_token_id if self.llm_engine.tokenizer.pad_token_id is not None else self.llm_engine.tokenizer.eos_token_id output_token_ids = pad_sequence(output_token_ids, batch_first=True, padding_value=pad_token_id) if len(logprobs) > 0: logprobs = pad_sequence(logprobs, batch_first=True, padding_value=pad_token_id) return output_token_ids, logprobs","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L271-L272","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor]) -> None","exported":true,"origin":"ast","body_hash":"835f8c2306c34dac4e7e8813254e954963b210496f18b5dce7e9fcaec7504a2f","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor]) -> None: self.llm_engine.sync_model_weights(actor_weights=actor_weights)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","span":"L274-L275","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"27a73de1e582c115374d023e58aeaf54824c75ca2740c19f48f07d7d67773621","body_text":"def offload_model_weights(self) -> None: self.llm_engine.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"tqdm","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":".arg_utils","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":".llm_engine_sp","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"vllm.outputs","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"vllm.sampling_params","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"torch.nn.utils.rnn","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"imports","specifier":"verl.workers.rollout.tokenizer","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.__init__","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.__init__","relation":"calls","name":"EngineArgs","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.__init__","relation":"calls","name":"type","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.__init__","relation":"calls","name":"from_engine_args","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.__init__","relation":"calls","name":"Counter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.init_cache_engine","relation":"calls","name":"init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.get_tokenizer","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.set_tokenizer","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"SamplingParams","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"_pre_process_inputs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","recvType":"LLM"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"_add_request","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","recvType":"LLM"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.generate","relation":"calls","name":"_run_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","recvType":"LLM"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._add_request","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._add_request","relation":"calls","name":"str","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._add_request","relation":"calls","name":"next","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._add_request","relation":"calls","name":"add_request","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","relation":"calls","name":"get_num_unfinished_requests","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","relation":"calls","name":"tqdm","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","relation":"calls","name":"has_unfinished_requests","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","relation":"calls","name":"step","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","recvType":"sorted"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","relation":"calls","name":"update","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","recvType":"tqdm"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","relation":"calls","name":"close","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","recvType":"tqdm"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","relation":"calls","name":"sorted","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._run_engine","relation":"calls","name":"_post_process_outputs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","recvType":"LLM"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._pre_process_inputs","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._pre_process_inputs","relation":"calls","name":"nonzero","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._pre_process_inputs","relation":"calls","name":"tolist","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","recvType":"pad_sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","relation":"calls","name":"zip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py","recvType":"pad_sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","relation":"calls","name":"pad_sequence","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM._post_process_outputs","relation":"calls","name":"pad_sequence","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm.py#LLM.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm.py"}]},"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py":{"size":35525,"mtimeMs":1789275830313.2026,"hash":"ae76e88278e330c7a82b9ed149e76022e1e9522b5199a48f272ea514ed80b6cd","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","name":"llm_engine_sp.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L1-L766","signature":null,"exported":true,"origin":"ast","body_hash":"ae76e88278e330c7a82b9ed149e76022e1e9522b5199a48f272ea514ed80b6cd","chars":35525,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/engine/llm_engine.py import os import socket import time import torch from typing import TYPE_CHECKING, Any, Dict, Iterable, List, Optional, Tuple, Union from vllm.lora.request import LoRARequest from vllm.config import (CacheConfig, DeviceConfig, ModelConfig, ParallelConfig, SchedulerConfig, LoRAConfig) from vllm.core.scheduler import Scheduler, SchedulerOutputs from vllm.logger import init_logger from vllm.outputs import RequestOutput from vllm.sampling_params import SamplingParams from vllm.sequence import (SamplerOutput, Sequence, SequenceGroup, SequenceGroupMetadata, SequenceGroupOutput, SequenceOutput, SequenceStatus) from vllm.transformers_utils.tokenizer import detokenize_incrementally from vllm.engine.metrics import StatLogger, Stats from vllm.utils import Counter import torch.nn as nn from .arg_utils import EngineArgs from .tokenizer import TokenizerGroup logger = init_logger(__name__) _LOCAL_LOGGING_INTERVAL_SEC = 5"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","name":"LLMEngine","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L41-L731","signature":"class LLMEngine","exported":true,"origin":"ast","body_hash":"f9f97eab1518b67b62516f08971ee5501a27322a67df174ad3ce6fb6d3aabffd","body_text":"class LLMEngine: \"\"\"An LLM engine that receives requests and generates texts. This is the main class for the vLLM engine. It receives requests from clients and generates texts from the LLM. It includes a tokenizer, a language model (possibly distributed across multiple GPUs), and GPU memory space allocated for intermediate states (aka KV cache). This class utilizes iteration-level scheduling and efficient memory management to maximize the serving throughput. The `LLM` class wraps this class for offline batched inference and the `AsyncLLMEngine` class wraps this class for online serving. NOTE: The config arguments are derived from the `EngineArgs` class. For the comprehensive list of arguments, see `EngineArgs`. Args: model_config: The configuration related to the LLM model. cache_config: The configuration related to the KV cache memory management. parallel_config: The configuration related to distributed execution. scheduler_config: The configuration related to the request scheduler. distributed_init_method: The initialization method for distributed execution. See `torch.distributed.init_process_group` for details. placement_group: Ray placement group for distributed execution. Required for distributed execution. log_stats: Whether to log statistics. \"\"\" def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], distributed_init_method: str, placement_group: Optional[None], log_stats: bool, ) -> None: logger.info(\"Initializing an LLM engine with config: \" f\"model={model_config.model!r}, \" f\"tokenizer={model_config.tokenizer!r}, \" # f\"tokenizer_mode={model_config.tokenizer_mode}, \" f\"revision={model_config.revision}, \" f\"tokenizer_revision={model_config.tokenizer_revision}, \" # f\"trust_remote_code={model_config.trust_remote_code}, \" f\"dtype={model_config.dtype}, \" f\"max_seq_len={model_config.max_model_len}, \" # f\"download_dir={model_config.download_dir!r}, \" # f\"load_format={model_config.load_format}, \" f\"disable_custom_all_reduce={parallel_config.disable_custom_all_reduce}, \" f\"tensor_parallel_size={parallel_config.tensor_parallel_size}, \" f\"quantization={model_config.quantization}, \" f\"seed={model_config.seed})\") # TODO(woosuk): Print more configs in debug mode. self.model_config = model_config # TODO: currently is hfconfig self.cache_config = cache_config self.lora_config = lora_config assert self.cache_config.sliding_window == getattr(self.model_config.hf_config, \"sliding_window\", None) self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.log_stats = log_stats self._verify_args() # self.model = model # should not store the model, it should be deleted # TODO(shengguangming): maybe we can choose init here or from arguments self._init_tokenizer(tokenizer) self.seq_counter = Counter() # Create the parallel GPU workers. self._init_workers_sp(model, distributed_init_method) # Profile the memory usage and initialize the cache. self._init_cache_sp() # Create the scheduler. # NOTE(shengguangming): each process will have independent scheduler self.scheduler = Scheduler(scheduler_config, cache_config, lora_config) # Metric Logging. if self.log_stats: self.stat_logger = StatLogger(local_interval=_LOCAL_LOGGING_INTERVAL_SEC) # Logging. self.last_logging_time = 0.0 # List of (timestamp, num_tokens) self.num_prompt_tokens: List[Tuple[float, int]] = [] # List of (timestamp, num_tokens) self.num_generation_tokens: List[Tuple[float, int]] = [] def _init_tokenizer(self, tokenizer, **tokenizer_init_kwargs): init_kwargs = dict(enable_lora=bool(self.lora_config), max_num_seqs=self.scheduler_config.max_num_seqs, max_input_length=None) init_kwargs.update(tokenizer_init_kwargs) self.tokenizer: TokenizerGroup = TokenizerGroup(tokenizer, **init_kwargs) # TODO: check get_lora_tokenizer func def get_tokenizer_for_seq(self, sequence: Sequence): return self.tokenizer.get_lora_tokenizer(sequence.lora_request) def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker from .worker import Worker # pylint: disable=import-outside-toplevel rank = int(os.getenv(\"RANK\")) self.worker = Worker( model, self.model_config, self.parallel_config, self.scheduler_config, self.device_config, rank, distributed_init_method, lora_config=self.lora_config, kv_cache_dtype=self.cache_config.cache_dtype, ) # NOTE(shengguangming): torch.distributed.init_process_group will be called inside the init_model() self.worker.init_model() self.worker.load_model() def _verify_args(self) -> None: self.model_config.verify_with_parallel_config(self.parallel_config) self.cache_config.verify_with_parallel_config(self.parallel_config) def _init_cache_sp(self) ","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L70-L136","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], distributed_init_method: str, placement_group: Optional[None], log_stats: bool, ) -> None","exported":false,"origin":"ast","body_hash":"d1657530e7c7086e43a2bfe583a3db3d067875a69d8f184ac65987236b951b9b","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], distributed_init_method: str, placement_group: Optional[None], log_stats: bool, ) -> None: logger.info(\"Initializing an LLM engine with config: \" f\"model={model_config.model!r}, \" f\"tokenizer={model_config.tokenizer!r}, \" # f\"tokenizer_mode={model_config.tokenizer_mode}, \" f\"revision={model_config.revision}, \" f\"tokenizer_revision={model_config.tokenizer_revision}, \" # f\"trust_remote_code={model_config.trust_remote_code}, \" f\"dtype={model_config.dtype}, \" f\"max_seq_len={model_config.max_model_len}, \" # f\"download_dir={model_config.download_dir!r}, \" # f\"load_format={model_config.load_format}, \" f\"disable_custom_all_reduce={parallel_config.disable_custom_all_reduce}, \" f\"tensor_parallel_size={parallel_config.tensor_parallel_size}, \" f\"quantization={model_config.quantization}, \" f\"seed={model_config.seed})\") # TODO(woosuk): Print more configs in debug mode. self.model_config = model_config # TODO: currently is hfconfig self.cache_config = cache_config self.lora_config = lora_config assert self.cache_config.sliding_window == getattr(self.model_config.hf_config, \"sliding_window\", None) self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.log_stats = log_stats self._verify_args() # self.model = model # should not store the model, it should be deleted # TODO(shengguangming): maybe we can choose init here or from arguments self._init_tokenizer(tokenizer) self.seq_counter = Counter() # Create the parallel GPU workers. self._init_workers_sp(model, distributed_init_method) # Profile the memory usage and initialize the cache. self._init_cache_sp() # Create the scheduler. # NOTE(shengguangming): each process will have independent scheduler self.scheduler = Scheduler(scheduler_config, cache_config, lora_config) # Metric Logging. if self.log_stats: self.stat_logger = StatLogger(local_interval=_LOCAL_LOGGING_INTERVAL_SEC) # Logging. self.last_logging_time = 0.0 # List of (timestamp, num_tokens) self.num_prompt_tokens: List[Tuple[float, int]] = [] # List of (timestamp, num_tokens) self.num_generation_tokens: List[Tuple[float, int]] = []","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_tokenizer","name":"_init_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L138-L143","signature":"def _init_tokenizer(self, tokenizer, **tokenizer_init_kwargs)","exported":false,"origin":"ast","body_hash":"012ad286206dcdfe6339e376b29596d7da1ef7041db76fb84ccc7865d8c919fa","body_text":"def _init_tokenizer(self, tokenizer, **tokenizer_init_kwargs): init_kwargs = dict(enable_lora=bool(self.lora_config), max_num_seqs=self.scheduler_config.max_num_seqs, max_input_length=None) init_kwargs.update(tokenizer_init_kwargs) self.tokenizer: TokenizerGroup = TokenizerGroup(tokenizer, **init_kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.get_tokenizer_for_seq","name":"get_tokenizer_for_seq","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L146-L147","signature":"def get_tokenizer_for_seq(self, sequence: Sequence)","exported":true,"origin":"ast","body_hash":"c6ada676f406dfef6ace501efefd446c012ae7ea6f966609dad11508ef66916b","body_text":"def get_tokenizer_for_seq(self, sequence: Sequence): return self.tokenizer.get_lora_tokenizer(sequence.lora_request)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_workers_sp","name":"_init_workers_sp","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L149-L170","signature":"def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker","exported":false,"origin":"ast","body_hash":"d926fd86b48786fd090672a0fdbcbc497bf4a27807ce156b82ee7627d66d8588","body_text":"def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker from .worker import Worker # pylint: disable=import-outside-toplevel rank = int(os.getenv(\"RANK\")) self.worker = Worker( model, self.model_config, self.parallel_config, self.scheduler_config, self.device_config, rank, distributed_init_method, lora_config=self.lora_config, kv_cache_dtype=self.cache_config.cache_dtype, ) # NOTE(shengguangming): torch.distributed.init_process_group will be called inside the init_model() self.worker.init_model() self.worker.load_model()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._verify_args","name":"_verify_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L172-L174","signature":"def _verify_args(self) -> None","exported":false,"origin":"ast","body_hash":"c8614c546959a77402ba8c34ede50da2f3d9baf7b239fdae7e2aaee550ac6871","body_text":"def _verify_args(self) -> None: self.model_config.verify_with_parallel_config(self.parallel_config) self.cache_config.verify_with_parallel_config(self.parallel_config)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_cache_sp","name":"_init_cache_sp","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L176-L213","signature":"def _init_cache_sp(self) -> None","exported":false,"origin":"ast","body_hash":"bc430c2db3ea93a28ec5d1a7093462c652e3c2bc0d9cb5e4171dfcb1012785cc","body_text":"def _init_cache_sp(self) -> None: \"\"\"Profiles the memory usage and initializes the KV cache.\"\"\" # Get the maximum number of blocks that can be allocated on GPU and CPU. num_blocks = self.worker.profile_num_available_blocks( block_size=self.cache_config.block_size, gpu_memory_utilization=self.cache_config.gpu_memory_utilization, cpu_swap_space=self.cache_config.swap_space_bytes, cache_dtype=self.cache_config.cache_dtype, ) # NOTE(shengguangming): Now we don't use a shared centralized controler but each process will # have its own scheduler num_gpu_blocks = num_blocks[0] num_cpu_blocks = num_blocks[1] # FIXME(woosuk): Change to debug log. logger.info(f\"# GPU blocks: {num_gpu_blocks}, \" f\"# CPU blocks: {num_cpu_blocks}\") if num_gpu_blocks <= 0: raise ValueError(\"No available memory for the cache blocks. \" \"Try increasing `gpu_memory_utilization` when \" \"initializing the engine.\") max_seq_len = self.cache_config.block_size * num_gpu_blocks if self.model_config.max_model_len > max_seq_len: raise ValueError(f\"The model's max seq len ({self.model_config.max_model_len}) \" \"is larger than the maximum number of tokens that can be \" f\"stored in KV cache ({max_seq_len}). Try increasing \" \"`gpu_memory_utilization` or decreasing `max_model_len` when \" \"initializing the engine.\") self.cache_config.num_gpu_blocks = num_gpu_blocks self.cache_config.num_cpu_blocks = num_cpu_blocks # Initialize the cache. self.worker.init_cache_engine(cache_config=self.cache_config) self.worker.warm_up_model()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L215-L216","signature":"def init_cache_engine(self)","exported":true,"origin":"ast","body_hash":"1d1217aff16ce0bdfd02b6b7f314acf30b40cc7d6bec7a8595f256e550a2b422","body_text":"def init_cache_engine(self): self.worker.init_cache_engine(cache_config=self.cache_config)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L218-L219","signature":"def free_cache_engine(self)","exported":true,"origin":"ast","body_hash":"b8175956306389b8d63aed00eb271d08fec8433fffbad1748c4738339b5e3121","body_text":"def free_cache_engine(self): self.worker.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.from_engine_args","name":"from_engine_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L222-L236","signature":"def from_engine_args(cls, model, tokenizer, engine_args: EngineArgs) -> \"LLMEngine\"","exported":true,"origin":"ast","body_hash":"ea9aea34f0e4b836a959327d2002507b3a8ad21168b40b1500d8046205865fa0","body_text":"def from_engine_args(cls, model, tokenizer, engine_args: EngineArgs) -> \"LLMEngine\": \"\"\"Creates an LLM engine from the engine arguments.\"\"\" # Create the engine configs. engine_configs = engine_args.create_engine_configs() parallel_config = engine_configs[2] # Initialize the cluster. distributed_init_method, placement_group = initialize_cluster(parallel_config) # Create the LLM engine. engine = cls(model, tokenizer, *engine_configs, distributed_init_method, placement_group, log_stats=not engine_args.disable_log_stats) return engine","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_request","name":"add_request","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L238-L315","signature":"def add_request( self, request_id: str, prompt: Optional[str], sampling_params: SamplingParams, prompt_token_ids: Optional[List[int]] = None, arrival_time: Optional[float] = None, lora_request: Optional[LoRARequest] = None, prefix_pos: Optional[int] = None, ) -> None","exported":true,"origin":"ast","body_hash":"c974da4872f24778f3ca90fac278f079a6aa4fe720539470091c08ed1034b52b","body_text":"def add_request( self, request_id: str, prompt: Optional[str], sampling_params: SamplingParams, prompt_token_ids: Optional[List[int]] = None, arrival_time: Optional[float] = None, lora_request: Optional[LoRARequest] = None, prefix_pos: Optional[int] = None, ) -> None: \"\"\"Add a request to the engine's request pool. The request is added to the request pool and will be processed by the scheduler as `engine.step()` is called. The exact scheduling policy is determined by the scheduler. Args: request_id: The unique ID of the request. prompt: The prompt string. Can be None if prompt_token_ids is provided. sampling_params: The sampling parameters for text generation. prompt_token_ids: The token IDs of the prompt. If None, we use the tokenizer to convert the prompts to token IDs. arrival_time: The arrival time of the request. If None, we use the current monotonic time. prefix_pos: If not None, we use the given position as the prefix position for each prompt. We will cache the prefix's KV cache and reuse it for the next request with the same prefix. This is an experimental feature, and may be replaced with automatic prefix caching in the future. Details: - Set arrival_time to the current time if it is None. - Set prompt_token_ids to the encoded prompt if it is None. - Create `best_of` number of :class:`~vllm.Sequence` objects. - Create a :class:`~vllm.SequenceGroup` object from the list of :class:`~vllm.Sequence`. - Add the :class:`~vllm.SequenceGroup` object to the scheduler. Example: >>> # initialize engine >>> engine = LLMEngine.from_engine_args(engine_args) >>> # set request arguments >>> example_prompt = \"Who is the president of the United States?\" >>> sampling_params = SamplingParams(temperature=0.0) >>> request_id = 0 >>> >>> # add the request to the engine >>> engine.add_request( >>> str(request_id), >>> example_prompt, >>> SamplingParams(temperature=0.0)) >>> # continue the request processing >>> ... \"\"\" if lora_request is not None and not self.lora_config: raise ValueError(f\"Got lora_request {lora_request} but LoRA is \" \"not enabled!\") if arrival_time is None: arrival_time = time.monotonic() if prompt_token_ids is None: assert prompt is not None prompt_token_ids = self.tokenizer.encode(prompt) # Create the sequences. block_size = self.cache_config.block_size seq_id = next(self.seq_counter) seq = Sequence(seq_id, prompt, prompt_token_ids, block_size, lora_request) # Check whether the input specifies prefix prefix = self.scheduler.prefix_pool.add_or_get_prefix(prompt_token_ids[:prefix_pos], lora_request.lora_int_id if lora_request else 0) if prefix_pos is not None else None # Create the sequence group. seq_group = SequenceGroup(request_id, [seq], sampling_params, arrival_time, lora_request, prefix) # Add the sequence group to the scheduler. self.scheduler.add_seq_group(seq_group)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.abort_request","name":"abort_request","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L317-L334","signature":"def abort_request(self, request_id: Union[str, Iterable[str]]) -> None","exported":true,"origin":"ast","body_hash":"28ac714169c023051b34d4870621893a5635bbbcb5a6918fffbca2826f26d794","body_text":"def abort_request(self, request_id: Union[str, Iterable[str]]) -> None: \"\"\"Aborts a request(s) with the given ID. Args: request_id: The ID(s) of the request to abort. Details: - Refer to the :meth:`~vllm.core.scheduler.Scheduler.abort_seq_group` from class :class:`~vllm.core.scheduler.Scheduler`. Example: >>> # initialize engine and add a request with request_id >>> request_id = str(0) >>> # abort the request >>> engine.abort_request(request_id) \"\"\" self.scheduler.abort_seq_group(request_id)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.get_model_config","name":"get_model_config","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L336-L338","signature":"def get_model_config(self) -> ModelConfig","exported":true,"origin":"ast","body_hash":"32b82a855d9acb5d442955e4adeabc7b3cdbeda338e897593d6ae9a925ff3fd5","body_text":"def get_model_config(self) -> ModelConfig: \"\"\"Gets the model configuration.\"\"\" return self.model_config","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.get_num_unfinished_requests","name":"get_num_unfinished_requests","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L340-L342","signature":"def get_num_unfinished_requests(self) -> int","exported":true,"origin":"ast","body_hash":"6f6ed8b9e83dec13a5a67979b0cf9eb6855610b0ac2e228176baaebb5ee3af00","body_text":"def get_num_unfinished_requests(self) -> int: \"\"\"Gets the number of unfinished requests.\"\"\" return self.scheduler.get_num_unfinished_seq_groups()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.has_unfinished_requests","name":"has_unfinished_requests","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L344-L346","signature":"def has_unfinished_requests(self) -> bool","exported":true,"origin":"ast","body_hash":"48371289db35ee6d8d6c16a65f0e60a73d6f176b45c6eba517c69d60e150861b","body_text":"def has_unfinished_requests(self) -> bool: \"\"\"Returns True if there are unfinished requests.\"\"\" return self.scheduler.has_unfinished_seqs()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","name":"_check_beam_search_early_stopping","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L348-L383","signature":"def _check_beam_search_early_stopping( self, early_stopping: Union[bool, str], sampling_params: SamplingParams, best_running_seq: Sequence, current_worst_seq: Sequence, ) -> bool","exported":false,"origin":"ast","body_hash":"5a0246efd15b10e3d23379cf6fd2d43b535f4db8a49339e8cc8d62b7275dcb7a","body_text":"def _check_beam_search_early_stopping( self, early_stopping: Union[bool, str], sampling_params: SamplingParams, best_running_seq: Sequence, current_worst_seq: Sequence, ) -> bool: assert sampling_params.use_beam_search length_penalty = sampling_params.length_penalty if early_stopping is True: return True current_worst_score = (current_worst_seq.get_beam_search_score( length_penalty=length_penalty, eos_token_id=self.get_tokenizer_for_seq(current_worst_seq).eos_token_id)) if early_stopping is False: highest_attainable_score = (best_running_seq.get_beam_search_score( length_penalty=length_penalty, eos_token_id=self.get_tokenizer_for_seq(best_running_seq).eos_token_id)) else: assert early_stopping == \"never\" if length_penalty > 0.0: # If length_penalty > 0.0, beam search will prefer longer # sequences. The highest attainable score calculation is # based on the longest possible sequence length in this case. max_possible_length = max(best_running_seq.get_prompt_len() + sampling_params.max_tokens, self.scheduler_config.max_model_len) highest_attainable_score = (best_running_seq.get_beam_search_score( length_penalty=length_penalty, eos_token_id=self.get_tokenizer_for_seq(best_running_seq).eos_token_id, seq_len=max_possible_length)) else: # Otherwise, beam search will prefer shorter sequences. The # highest attainable score calculation is based on the current # sequence length. highest_attainable_score = (best_running_seq.get_beam_search_score( length_penalty=length_penalty, eos_token_id=self.get_tokenizer_for_seq(best_running_seq).eos_token_id))","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","name":"_process_sequence_group_outputs","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L385-L543","signature":"def _process_sequence_group_outputs(self, seq_group: SequenceGroup, outputs: SequenceGroupOutput) -> None: # Process prompt logprobs","exported":false,"origin":"ast","body_hash":"c33b1dfd348555fe47f457b21405775a9d35070729771aa599f838428ec8c38f","body_text":"def _process_sequence_group_outputs(self, seq_group: SequenceGroup, outputs: SequenceGroupOutput) -> None: # Process prompt logprobs prompt_logprobs = outputs.prompt_logprobs if prompt_logprobs is not None: seq_group.prompt_logprobs = prompt_logprobs # Process samples samples = outputs.samples parent_seqs = seq_group.get_seqs(status=SequenceStatus.RUNNING) existing_finished_seqs = seq_group.get_finished_seqs() parent_child_dict = {parent_seq.seq_id: [] for parent_seq in parent_seqs} for sample in samples: parent_child_dict[sample.parent_seq_id].append(sample) # List of (child, parent) child_seqs: List[Tuple[Sequence, Sequence]] = [] # Process the child samples for each parent sequence for parent in parent_seqs: child_samples: List[SequenceOutput] = parent_child_dict[parent.seq_id] if len(child_samples) == 0: # This parent sequence has no children samples. Remove # the parent sequence from the sequence group since it will # not be used in the future iterations. parent.status = SequenceStatus.FINISHED_ABORTED seq_group.remove(parent.seq_id) self.scheduler.free_seq(parent) continue # Fork the parent sequence if there are multiple child samples. for child_sample in child_samples[:-1]: new_child_seq_id = next(self.seq_counter) child = parent.fork(new_child_seq_id) child.append_token_id(child_sample.output_token, child_sample.logprobs) child_seqs.append((child, parent)) # Continue the parent sequence for the last child sample. # We reuse the parent sequence here to reduce redundant memory # copies, especially when using non-beam search sampling methods. last_child_sample = child_samples[-1] parent.append_token_id(last_child_sample.output_token, last_child_sample.logprobs) child_seqs.append((parent, parent)) for seq, _ in child_seqs: # self._decode_sequence(seq, seq_group.sampling_params) self._check_stop(seq, seq_group.sampling_params) # Non-beam search case if not seq_group.sampling_params.use_beam_search: # For newly created child sequences, add them to the sequence group # and fork them in block manager if they are not finished. for seq, parent in child_seqs: if seq is not parent: seq_group.add(seq) if not seq.is_finished(): self.scheduler.fork_seq(parent, seq) # Free the finished and selected parent sequences' memory in block # manager. Keep them in the sequence group as candidate output. # NOTE: we need to fork the new sequences before freeing the # old sequences. for seq, parent in child_seqs: if seq is parent and seq.is_finished(): self.scheduler.free_seq(seq) return # Beam search case # Select the child sequences to keep in the sequence group. selected_child_seqs = [] unselected_child_seqs = [] beam_width = seq_group.sampling_params.best_of length_penalty = seq_group.sampling_params.length_penalty # Select the newly finished sequences with the highest scores # to replace existing finished sequences. # Tuple of (seq, parent, is_new) existing_finished_seqs = [(seq, None, False) for seq in existing_finished_seqs] new_finished_seqs = [(seq, parent, True) for seq, parent in child_seqs if seq.is_finished()] all_finished_seqs = existing_finished_seqs + new_finished_seqs # Sort the finished sequences by their scores. all_finished_seqs.sort(key=lambda x: x[0].get_beam_search_score( length_penalty=length_penalty, eos_token_id=self.get_tokenizer_for_seq(x[0]).eos_token_id), reverse=True) for seq, parent, is_new in all_finished_seqs[:beam_width]: if is_new: # A newly generated child sequence finishes and has a high # score, so we will add it into the sequence group. selected_child_seqs.append((seq, parent)) for seq, parent, is_new in all_finished_seqs[beam_width:]: if is_new: # A newly generated child sequence finishes but has a low # score, so we will not add it into the sequence group. # Additionally, if this sequence is a continuation of a # parent sequence, we will need remove the parent sequence # from the sequence group. unselected_child_seqs.append((seq, parent)) else: # An existing finished sequence has a low score, so we will # remove it from the sequence group. seq_group.remove(seq.seq_id) # select the top beam_width sequences from the running # sequences for the next iteration to continue the beam # search. running_child_seqs = [(seq, parent) for seq, parent in child_seqs if not seq.is_finished()] # Sort the running sequences by their scores. running_child_seqs.sort(key=lambda x: x[0].get_beam_search_score( length_penalty=length_penalty, eos_token_id=self.get_tokenizer_for_seq(x[0]).eos_token_id), reverse=True) # Check if we can stop the beam search. if len(running_child_seqs) == 0: # No running sequences, stop the beam search. stop_beam_search = True elif len(all_finished_seqs) < beam_width: # Not enough finished sequences, continue the beam search. stop_beam_search = False else: # Check the early stopping criteria best_running_seq = running_child_seqs[0][0] current_worst_seq = all_finished_seqs[beam_width - 1][0] stop_beam_search = self._check_beam_search_early_stopping(seq_group.sampling_params.early_st","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","name":"_process_model_outputs","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L545-L572","signature":"def _process_model_outputs(self, output: SamplerOutput, scheduler_outputs: SchedulerOutputs) -> List[RequestOutput]: # Update the scheduled sequence groups with the model outputs.","exported":false,"origin":"ast","body_hash":"75e9aae8ccd61de44ca6e0c1f960f9cbd3ffbd522552edb9d25b179edf6c2236","body_text":"def _process_model_outputs(self, output: SamplerOutput, scheduler_outputs: SchedulerOutputs) -> List[RequestOutput]: # Update the scheduled sequence groups with the model outputs. scheduled_seq_groups = scheduler_outputs.scheduled_seq_groups for seq_group, outputs in zip(scheduled_seq_groups, output): self._process_sequence_group_outputs(seq_group, outputs) # Free the finished sequence groups. self.scheduler.free_finished_seq_groups() # Create the outputs. request_outputs: List[RequestOutput] = [] for seq_group in scheduled_seq_groups: request_output = RequestOutput.from_seq_group(seq_group) request_outputs.append(request_output) for seq_group in scheduler_outputs.ignored_seq_groups: request_output = RequestOutput.from_seq_group(seq_group) request_outputs.append(request_output) # Update prefix state, now all the uncomputed prefixes are computed. for seq_group in scheduled_seq_groups: if (seq_group.prefix is not None and seq_group.prefix.allocated and not seq_group.prefix.computed): seq_group.prefix.computed = True # Log stats. if self.log_stats: self.stat_logger.log(self._get_stats(scheduler_outputs)) return request_outputs","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.step","name":"step","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L574-L593","signature":"def step(self) -> List[RequestOutput]","exported":true,"origin":"ast","body_hash":"df724a8f2cffee228fb5fca7b1b3c842c77f3aa0698295dcc511cd4ed8731ee9","body_text":"def step(self) -> List[RequestOutput]: \"\"\"Performs one decoding iteration and returns newly generated results. This function performs one decoding iteration of the engine. It first schedules the sequences to be executed in the next iteration and the token blocks to be swapped in/out/copy. Then, it executes the model and updates the scheduler with the model outputs. Finally, it decodes the sequences and returns the newly generated results. \"\"\" seq_group_metadata_list, scheduler_outputs = self.scheduler.schedule() if not scheduler_outputs.is_empty(): output = self.worker.execute_model( seq_group_metadata_list=seq_group_metadata_list, # TODO: check this input blocks_to_swap_in=scheduler_outputs.blocks_to_swap_in, blocks_to_swap_out=scheduler_outputs.blocks_to_swap_out, blocks_to_copy=scheduler_outputs.blocks_to_copy,) else: return [RequestOutput.from_seq_group(seq_group) for seq_group in scheduler_outputs.ignored_seq_groups] return self._process_model_outputs(output, scheduler_outputs)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.do_log_stats","name":"do_log_stats","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L595-L598","signature":"def do_log_stats(self) -> None","exported":true,"origin":"ast","body_hash":"5aa75138203dad2abe279edeb2dddb6fdc6340bff7e0f6d300330001eb4172aa","body_text":"def do_log_stats(self) -> None: \"\"\"Forced log when no requests active.\"\"\" if self.log_stats: self.stat_logger.log(self._get_stats(scheduler_outputs=None))","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","name":"_get_stats","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L600-L659","signature":"def _get_stats(self, scheduler_outputs: Optional[SchedulerOutputs]) -> Stats","exported":false,"origin":"ast","body_hash":"ec13c47194f678079f52e2e1c8e00afab87e0193236539c181998833258b934d","body_text":"def _get_stats(self, scheduler_outputs: Optional[SchedulerOutputs]) -> Stats: \"\"\"Get Stats to be Logged to Prometheus.\"\"\" now = time.monotonic() # KV Cache Usage in %. num_total_gpu = self.cache_config.num_gpu_blocks num_free_gpu = self.scheduler.block_manager.get_num_free_gpu_blocks() gpu_cache_usage = 1.0 - (num_free_gpu / num_total_gpu) num_total_cpu = self.cache_config.num_cpu_blocks cpu_cache_usage = 0. if num_total_cpu > 0: num_free_cpu = self.scheduler.block_manager.get_num_free_cpu_blocks() cpu_cache_usage = 1.0 - (num_free_cpu / num_total_cpu) # Scheduler State num_running = len(self.scheduler.running) num_swapped = len(self.scheduler.swapped) num_waiting = len(self.scheduler.waiting) # Iteration stats if we have scheduler output. num_prompt_tokens = 0 num_generation_tokens = 0 time_to_first_tokens = [] time_per_output_tokens = [] time_e2e_requests = [] if scheduler_outputs is not None: prompt_run = scheduler_outputs.prompt_run # Number of Tokens. if prompt_run: num_prompt_tokens = scheduler_outputs.num_batched_tokens else: num_generation_tokens = scheduler_outputs.num_batched_tokens # Latency Timings. time_last_iters = [] for seq_group in scheduler_outputs.scheduled_seq_groups: # Time since last token. (n.b. updates seq_group.last_token_time) time_last_iters.append(seq_group.get_last_latency(now)) # Time since arrival for all finished requests. if seq_group.is_finished(): time_e2e_requests.append(now - seq_group.arrival_time) time_to_first_tokens = time_last_iters if prompt_run else [] time_per_output_tokens = [] if prompt_run else time_last_iters return Stats( now=now, num_running=num_running, num_swapped=num_swapped, num_waiting=num_waiting, gpu_cache_usage=gpu_cache_usage, cpu_cache_usage=cpu_cache_usage, num_prompt_tokens=num_prompt_tokens, num_generation_tokens=num_generation_tokens, time_to_first_tokens=time_to_first_tokens, time_per_output_tokens=time_per_output_tokens, time_e2e_requests=time_e2e_requests, )","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._decode_sequence","name":"_decode_sequence","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L662-L679","signature":"def _decode_sequence(self, seq: Sequence, prms: SamplingParams) -> None","exported":false,"origin":"ast","body_hash":"1d81941f2f30da38e9f8c821805f9fd1a964adaec6a05312db57506de5f62578","body_text":"def _decode_sequence(self, seq: Sequence, prms: SamplingParams) -> None: \"\"\"Decodes the new token for a sequence.\"\"\" (new_tokens, new_output_text, prefix_offset, read_offset) = detokenize_incrementally( self.get_tokenizer_for_seq(seq), all_input_ids=seq.get_token_ids(), prev_tokens=seq.tokens, prefix_offset=seq.prefix_offset, read_offset=seq.read_offset, skip_special_tokens=prms.skip_special_tokens, spaces_between_special_tokens=prms.spaces_between_special_tokens, ) if seq.tokens is None: seq.tokens = new_tokens else: seq.tokens.extend(new_tokens) seq.prefix_offset = prefix_offset seq.read_offset = read_offset seq.output_text += new_output_text","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_stop","name":"_check_stop","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L681-L708","signature":"def _check_stop(self, seq: Sequence, sampling_params: SamplingParams) -> None","exported":false,"origin":"ast","body_hash":"dc9db787c3124e5b3a0bfd222e772a05a2ffcf8032538afd493dc6bb14bba9cd","body_text":"def _check_stop(self, seq: Sequence, sampling_params: SamplingParams) -> None: \"\"\"Stop the finished sequences.\"\"\" # for stop_str in sampling_params.stop: # if seq.output_text.endswith(stop_str): # self._finalize_sequence(seq, sampling_params, stop_str) # seq.status = SequenceStatus.FINISHED_STOPPED # return # if seq.get_last_token_id() in sampling_params.stop_token_ids: # stop_str = self.get_tokenizer_for_seq(seq).convert_ids_to_tokens(seq.get_last_token_id()) # self._finalize_sequence(seq, sampling_params, stop_str) # seq.status = SequenceStatus.FINISHED_STOPPED # return # Check if the sequence has reached max_model_len. if seq.get_len() > self.scheduler_config.max_model_len: seq.status = SequenceStatus.FINISHED_LENGTH_CAPPED return # Check if the sequence has reached max_tokens. if seq.get_output_len() == sampling_params.max_tokens: seq.status = SequenceStatus.FINISHED_LENGTH_CAPPED return # Check if the sequence has generated the EOS token. if ((not sampling_params.ignore_eos) and seq.get_last_token_id() == self.get_tokenizer_for_seq(seq).eos_token_id): seq.status = SequenceStatus.FINISHED_STOPPED return","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._finalize_sequence","name":"_finalize_sequence","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L710-L714","signature":"def _finalize_sequence(self, seq: Sequence, sampling_params: SamplingParams, stop_string: str) -> None","exported":false,"origin":"ast","body_hash":"3264012aa3c3d90def47d7311703f7459b3fdc100259586f4c8b8a8f68cdf6ff","body_text":"def _finalize_sequence(self, seq: Sequence, sampling_params: SamplingParams, stop_string: str) -> None: if not sampling_params.include_stop_str_in_output and stop_string: # Truncate the output text so that the stop string is # not included in the output. seq.output_text = seq.output_text[:-len(stop_string)]","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_lora","name":"add_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L716-L718","signature":"def add_lora(self, lora_request: LoRARequest) -> bool","exported":true,"origin":"ast","body_hash":"e9f7ffc405171ec302d578962a651d53756db775c9e591e6d077841345e61f6d","body_text":"def add_lora(self, lora_request: LoRARequest) -> bool: assert lora_request.lora_int_id > 0, \"lora_id must be greater than 0.\" return self.worker.add_lora(lora_request)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.remove_lora","name":"remove_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L720-L722","signature":"def remove_lora(self, lora_id: int) -> bool","exported":true,"origin":"ast","body_hash":"cbe1cdb44759c7aabdb73749db619e9e5ee7c34edfc86f0be28a4261075173d1","body_text":"def remove_lora(self, lora_id: int) -> bool: assert lora_id > 0, \"lora_id must be greater than 0.\" return self.worker.remove_lora(lora_id)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.list_loras","name":"list_loras","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L724-L725","signature":"def list_loras(self) -> List[int]","exported":true,"origin":"ast","body_hash":"8fafde07f6d693c4cc72fd0586ad4441efc9fef84432839fd46f9747a20070b5","body_text":"def list_loras(self) -> List[int]: return self.worker.list_loras()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L727-L728","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor]) -> None","exported":true,"origin":"ast","body_hash":"b6e06ff63f4cbb4b57c22841ec40b0708e2550fb9a96e4a8d5f147e2fa3e0c99","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor]) -> None: self.worker.sync_model_weights(actor_weights=actor_weights)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L730-L731","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"c346b8cbee5a5fe63905573385e113fb7b59c46e58b26ecc132e331ba728ea49","body_text":"def offload_model_weights(self) -> None: self.worker.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#initialize_cluster","name":"initialize_cluster","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L734-L759","signature":"def initialize_cluster( parallel_config: ParallelConfig, engine_use_ray: bool = False, ray_address: Optional[str] = None, ) -> Tuple[str, Optional[None]]","exported":true,"origin":"ast","body_hash":"8477b50c95bb8a43b90a29db8bcaa647ca553e5f2226a1db7ed4996e9ca5a180","body_text":"def initialize_cluster( parallel_config: ParallelConfig, engine_use_ray: bool = False, ray_address: Optional[str] = None, ) -> Tuple[str, Optional[None]]: \"\"\"Initialize the distributed cluster probably with Ray. Args: parallel_config: The configurations for parallel execution. engine_use_ray: Whether to use Ray for async engine. ray_address: The address of the Ray cluster. If None, uses the default Ray cluster address. Returns: A tuple of (`distributed_init_method`, `placement_group`). The `distributed_init_method` is the address for initializing the distributed backend. `placement_group` includes the specification of the resources for each distributed worker. \"\"\" # Initialize cluster locally. port = get_open_port() # We need to setup the distributed init method to make sure # the distributed megatron code (e.g., get world size) works correctly. distributed_init_method = f\"tcp://localhost:{port}\" return distributed_init_method, None","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#get_open_port","name":"get_open_port","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","span":"L762-L765","signature":"def get_open_port()","exported":true,"origin":"ast","body_hash":"d88039216a61e6ec2a0e1328331a7b83f0af92fe7306e2625e8d4f34b557c26e","body_text":"def get_open_port(): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind((\"\", 0)) return s.getsockname()[1]","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"socket","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"time","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"vllm.core.scheduler","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"vllm.outputs","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"vllm.sampling_params","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"vllm.transformers_utils.tokenizer","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.metrics","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":".arg_utils","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":".tokenizer","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_verify_args","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_init_tokenizer","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Counter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_init_workers_sp","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_init_cache_sp","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Scheduler","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"StatLogger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_tokenizer","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"bool","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"update","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"dict"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"TokenizerGroup","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.get_tokenizer_for_seq","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.get_tokenizer_for_seq","relation":"calls","name":"get_lora_tokenizer","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"TokenizerGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_workers_sp","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"imports","specifier":".worker","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_workers_sp","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_workers_sp","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_workers_sp","relation":"calls","name":"Worker","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_workers_sp","relation":"calls","name":"init_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_workers_sp","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._verify_args","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._verify_args","relation":"calls","name":"verify_with_parallel_config","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._verify_args","relation":"calls","name":"verify_with_parallel_config","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_cache_sp","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_cache_sp","relation":"calls","name":"profile_num_available_blocks","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_cache_sp","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_cache_sp","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_cache_sp","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_cache_sp","relation":"calls","name":"init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._init_cache_sp","relation":"calls","name":"warm_up_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.init_cache_engine","relation":"calls","name":"init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.from_engine_args","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"create_engine_configs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"EngineArgs"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"initialize_cluster","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"cls","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_request","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_request","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_request","relation":"calls","name":"monotonic","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_request","relation":"calls","name":"encode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"TokenizerGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_request","relation":"calls","name":"next","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_request","relation":"calls","name":"Sequence","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_request","relation":"calls","name":"add_or_get_prefix","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_request","relation":"calls","name":"SequenceGroup","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_request","relation":"calls","name":"add_seq_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.abort_request","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.abort_request","relation":"calls","name":"abort_seq_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.get_model_config","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.get_num_unfinished_requests","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.get_num_unfinished_requests","relation":"calls","name":"get_num_unfinished_seq_groups","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.has_unfinished_requests","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.has_unfinished_requests","relation":"calls","name":"has_unfinished_seqs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","relation":"calls","name":"get_beam_search_score","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","relation":"calls","name":"get_tokenizer_for_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","relation":"calls","name":"get_beam_search_score","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","relation":"calls","name":"get_tokenizer_for_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","relation":"calls","name":"get_prompt_len","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","relation":"calls","name":"get_beam_search_score","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","relation":"calls","name":"get_tokenizer_for_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","relation":"calls","name":"get_beam_search_score","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_beam_search_early_stopping","relation":"calls","name":"get_tokenizer_for_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"get_seqs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"SequenceGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"get_finished_seqs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"SequenceGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"remove","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"SequenceGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"free_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"next","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"fork","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"append_token_id","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"append_token_id","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"_check_stop","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"add","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"SequenceGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"is_finished","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"fork_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"is_finished","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"free_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"is_finished","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"sort","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"get_beam_search_score","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"get_tokenizer_for_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"remove","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"SequenceGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"is_finished","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"sort","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"get_beam_search_score","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"get_tokenizer_for_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"_check_beam_search_early_stopping","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"add","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"SequenceGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"is_finished","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"fork_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"is_finished","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"free_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"remove","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"SequenceGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_sequence_group_outputs","relation":"calls","name":"free_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","relation":"calls","name":"zip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","relation":"calls","name":"_process_sequence_group_outputs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","relation":"calls","name":"free_finished_seq_groups","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","relation":"calls","name":"from_seq_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","relation":"calls","name":"from_seq_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","relation":"calls","name":"log","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"StatLogger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._process_model_outputs","relation":"calls","name":"_get_stats","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.step","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.step","relation":"calls","name":"schedule","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Scheduler"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.step","relation":"calls","name":"is_empty","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.step","relation":"calls","name":"execute_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.step","relation":"calls","name":"from_seq_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.step","relation":"calls","name":"_process_model_outputs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.do_log_stats","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.do_log_stats","relation":"calls","name":"log","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"StatLogger"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.do_log_stats","relation":"calls","name":"_get_stats","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"monotonic","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"get_num_free_gpu_blocks","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"get_num_free_cpu_blocks","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"get_last_latency","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"is_finished","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._get_stats","relation":"calls","name":"Stats","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._decode_sequence","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._decode_sequence","relation":"calls","name":"detokenize_incrementally","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._decode_sequence","relation":"calls","name":"get_tokenizer_for_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._decode_sequence","relation":"calls","name":"get_token_ids","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._decode_sequence","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_stop","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_stop","relation":"calls","name":"get_len","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_stop","relation":"calls","name":"get_output_len","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_stop","relation":"calls","name":"get_last_token_id","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Sequence"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._check_stop","relation":"calls","name":"get_tokenizer_for_seq","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._finalize_sequence","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine._finalize_sequence","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_lora","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.add_lora","relation":"calls","name":"add_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.remove_lora","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.remove_lora","relation":"calls","name":"remove_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.list_loras","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.list_loras","relation":"calls","name":"list_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#LLMEngine.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#initialize_cluster","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#initialize_cluster","relation":"calls","name":"get_open_port","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#get_open_port","file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#get_open_port","relation":"calls","name":"socket","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#get_open_port","relation":"calls","name":"bind","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py#get_open_port","relation":"calls","name":"getsockname","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/llm_engine_sp.py"}]},"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py":{"size":12265,"mtimeMs":1789275830313.2026,"hash":"9178cfe4c30b03a55856a0a59bdd2cf14856181185050814502fe9012cfcde10","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","name":"model_loader.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","span":"L1-L276","signature":null,"exported":true,"origin":"ast","body_hash":"9178cfe4c30b03a55856a0a59bdd2cf14856181185050814502fe9012cfcde10","chars":12265,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/model_loader \"\"\"Utilities for selecting and loading models.\"\"\" import contextlib from typing import Dict, Type, Union import torch import torch.nn as nn from transformers import PretrainedConfig, PreTrainedModel from megatron.core.tensor_parallel.utils import VocabUtility from vllm.model_executor.models import ModelRegistry from vllm.model_executor.weight_utils import (get_quant_config, initialize_dummy_weights) from .config import ModelConfig from vllm.config import DeviceConfig, LoRAConfig from .weight_loaders import * from vllm.model_executor.sampling_metadata import SamplingMetadata, SamplingTensors from vllm.sequence import SamplerOutput from typing import Optional from vllm.model_executor.layers.sampler import Sampler from vllm.model_executor.layers.sampler import _prune_hidden_states, _apply_logits_processors, _apply_penalties, _apply_top_k_top_p, _apply_min_p, _apply_penalties, _sample, _get_logprobs, _build_sampler_output @contextlib.contextmanager from vllm.model_executor.layers.linear import * from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding, ParallelLMHead from vllm.model_executor.layers.activation import ScaledActivation __LAYER_WEIGHT_LOADER_REGISTRY__ = { ColumnParallelLinear: parallel_weight_loader, MergedColumnParallelLinear: parallel_weight_loader, QKVParallelLinear: parallel_weight_loader, RowParallelLinear: parallel_weight_loader, VocabParallelEmbedding: parallel_weight_loader, ParallelLMHead: parallel_weight_loader # \"ScaledActivation.weight_loader\": ScaledActivation, # TODO(shengguangming): latest commit in vllm fix awq for this function and add load_weights # \"default_weight_loader\": default_weight_loader } # NOTE(gmsheng): change the weight_loader function in runtime for layer_class, weight_loader in __LAYER_WEIGHT_LOADER_REGISTRY__.items(): layer_class.weight_loader = weight_loader __MODEL_WEIGHT_LOADER_REGISTRY__ = { 'GPT2LMHeadModel': gpt2_weight_loader, 'LlamaForCausalLM': llama_weight_loader, 'LLaMAForCausalLM': llama_weight_loader, 'MistralForCausalLM': mistral_weight_loader, } # FIXME(shengguangming): the vLLM vocab will pad to 64, which may incur out of bounds # so we need to rewrite the init function of vocab DEFAULT_VOCAB_PADDING_SIZE = 64 VocabParallelEmbedding.__init__ = vocab_init # the actor model is .state_dict() # FIXME(sgm): hack the Sampler function in vllm v0.3.1 # as they use ray, the sampler result will only need to return to the driver node, # therefore gather is enough. However, we use SPMD instead of a central scheduler, # all_gather is required (aligned with v0.2.6) from vllm.model_executor.layers.sampler import Sampler Sampler._get_logits = _get_logits Sampler.forward = forward"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_set_default_torch_dtype","name":"_set_default_torch_dtype","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","span":"L38-L43","signature":"def _set_default_torch_dtype(dtype: torch.dtype)","exported":false,"origin":"ast","body_hash":"48b953f6c3f3e0663c5ac537a307c4d32c61124be77f2483fb19be7a100809f4","body_text":"def _set_default_torch_dtype(dtype: torch.dtype): \"\"\"Sets the default torch dtype to the given dtype.\"\"\" old_dtype = torch.get_default_dtype() torch.set_default_dtype(dtype) yield torch.set_default_dtype(old_dtype)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_model_architecture","name":"_get_model_architecture","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","span":"L46-L53","signature":"def _get_model_architecture(config: PretrainedConfig) -> Type[nn.Module]","exported":false,"origin":"ast","body_hash":"d1123e2dd93cc66a5f39b2e2aafc5bcc16f364911ece8403fe5dc6f8ff9c8280","body_text":"def _get_model_architecture(config: PretrainedConfig) -> Type[nn.Module]: architectures = getattr(config, \"architectures\", []) for arch in architectures: model_cls = ModelRegistry.load_model_cls(arch) if model_cls is not None: return model_cls raise ValueError(f\"Model architectures {architectures} are not supported for now. \" f\"Supported architectures: {ModelRegistry.get_supported_archs()}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","name":"vocab_init","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","span":"L87-L117","signature":"def vocab_init(self, num_embeddings: int, embedding_dim: int, params_dtype: Optional[torch.dtype] = None, org_num_embeddings: Optional[int] = None, padding_size: int = DEFAULT_VOCAB_PADDING_SIZE)","exported":true,"origin":"ast","body_hash":"52161e7f299170a24d7859e33e31079b406e7ac904ed11d4f34aac2b2a4eb591","body_text":"def vocab_init(self, num_embeddings: int, embedding_dim: int, params_dtype: Optional[torch.dtype] = None, org_num_embeddings: Optional[int] = None, padding_size: int = DEFAULT_VOCAB_PADDING_SIZE): super(VocabParallelEmbedding, self).__init__() # Keep the input dimensions. # TODO (pad to be divided by 4) self.num_embeddings = num_embeddings self.org_vocab_size = org_num_embeddings or num_embeddings # self.num_embeddings_padded = pad_vocab_size(num_embeddings, # padding_size) self.embedding_dim = embedding_dim if params_dtype is None: params_dtype = torch.get_default_dtype() self.tp_size = get_tensor_model_parallel_world_size() # Divide the weight matrix along the vocaburaly dimension. self.vocab_start_index, self.vocab_end_index = (VocabUtility.vocab_range_from_global_vocab_size( self.num_embeddings, get_tensor_model_parallel_rank(), self.tp_size)) self.num_embeddings_per_partition = (self.vocab_end_index - self.vocab_start_index) self.weight = Parameter( torch.empty( self.num_embeddings_per_partition, self.embedding_dim, # device=torch.cuda.current_device(), dtype=params_dtype)) set_weight_attrs(self.weight, {\"parallel_dim\": 0, \"weight_loader\": self.weight_loader})","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_model_weight_loader","name":"_get_model_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","span":"L123-L127","signature":"def _get_model_weight_loader(arch: str)","exported":false,"origin":"ast","body_hash":"ddc10a750c10f76e9a2ee47256e4d6b0d7498b307e1cd2153be49f1057d65b2b","body_text":"def _get_model_weight_loader(arch: str): if arch in __MODEL_WEIGHT_LOADER_REGISTRY__: return __MODEL_WEIGHT_LOADER_REGISTRY__[arch] raise ValueError(f\"Model architectures {arch} are not supported for now. \" f\"Supported architectures: {ModelRegistry.get_supported_archs()}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","name":"get_model","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","span":"L130-L177","signature":"def get_model(actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig] = None) -> nn.Module","exported":true,"origin":"ast","body_hash":"e5111832759e10434b7ba3580b196628855892190547dfedbaabc8a8ad2e25a6","body_text":"def get_model(actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig] = None) -> nn.Module: model_class = _get_model_architecture(model_config.hf_config) # Get the quantization config. linear_method = None quant_config = None if model_config.quantization is not None: quant_config = get_quant_config(model_config.quantization, model_config.model, model_config.hf_config, model_config.download_dir) capability = torch.cuda.get_device_capability() capability = capability[0] * 10 + capability[1] if capability < quant_config.get_min_capability(): raise ValueError(f\"The quantization method {model_config.quantization} is not \" \"supported for the current GPU. \" f\"Minimum capability: {quant_config.get_min_capability()}. \" f\"Current capability: {capability}.\") supported_dtypes = quant_config.get_supported_act_dtypes() if model_config.dtype not in supported_dtypes: raise ValueError(f\"{model_config.dtype} is not supported for quantization \" f\"method {model_config.quantization}. Supported dtypes: \" f\"{supported_dtypes}\") linear_method = quant_config.get_linear_method() with _set_default_torch_dtype(model_config.dtype): # Create a model instance. # The weights will be initialized as empty tensors. # with torch.device(device_config.device): # NOTE(sgm): init the model in cpu model = model_class(model_config.hf_config, linear_method) if model_config.load_format == \"dummy\": model = model.cuda() # NOTE(woosuk): For accurate performance evaluation, we assign # random values to the weights. initialize_dummy_weights(model) elif model_config.load_format == 'model' or model_config.load_format == 'auto': # NOTE(shengguangming) Load the weights from the actor model if isinstance(actor_model, nn.Module): load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_weights(actor_weights=actor_model, vllm_model=model) # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#load_weights","name":"load_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","span":"L181-L186","signature":"def load_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"fd9e069a8b1eb95a72aa61709e47b7213e050562a6c1db2af328a530671ef03f","body_text":"def load_weights(actor_weights: Dict, vllm_model: nn.Module): weight_loader = _get_model_weight_loader(vllm_model.__class__.__name__) weight_loader(actor_weights, vllm_model) # NOTE(sgm) to reduce peak memory usage, we offload vllm model to cpu # after init, and we need this after sync model weights for in first iter. vllm_model = vllm_model.cuda()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_logits","name":"_get_logits","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","span":"L193-L203","signature":"def _get_logits(self, hidden_states: torch.Tensor, embedding: torch.Tensor, embedding_bias: Optional[torch.Tensor]) -> torch.Tensor: # Get the logits for the next tokens.","exported":false,"origin":"ast","body_hash":"d294e846526a09f100321505cbacf5666d45ef972200a51c4e4c25557ec585e0","body_text":"def _get_logits(self, hidden_states: torch.Tensor, embedding: torch.Tensor, embedding_bias: Optional[torch.Tensor]) -> torch.Tensor: # Get the logits for the next tokens. logits = torch.matmul(hidden_states, embedding.t()) if embedding_bias is not None: logits += embedding_bias logits = tensor_model_parallel_all_gather(logits) # Remove paddings in vocab (if any). if logits is not None: logits = logits[:, :self.org_vocab_size] return logits","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","name":"forward","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","span":"L206-L269","signature":"def forward( self, embedding: torch.Tensor, hidden_states: torch.Tensor, sampling_metadata: SamplingMetadata, embedding_bias: Optional[torch.Tensor] = None, ) -> Optional[SamplerOutput]: # Get the hidden states that we use for sampling.","exported":true,"origin":"ast","body_hash":"e3cb0c9e173774e28f940721d89fe753a35f19f187f213e9d753f66a76323359","body_text":"def forward( self, embedding: torch.Tensor, hidden_states: torch.Tensor, sampling_metadata: SamplingMetadata, embedding_bias: Optional[torch.Tensor] = None, ) -> Optional[SamplerOutput]: # Get the hidden states that we use for sampling. hidden_states = _prune_hidden_states(hidden_states, sampling_metadata) # Get the logits for the next tokens. logits = self._get_logits(hidden_states, embedding, embedding_bias) # save origin logprobs for sampler_output origin_logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float) # Only perform sampling in the driver worker. # Note: `_get_logits` is still distributed across TP workers because # the `embedding` weight is distributed across TP workers. # TODO(zhuohan): Change the get_logits part to a separate stage. if not sampling_metadata.perform_sampling: return None assert logits is not None _, vocab_size = logits.shape # Apply logits processors (if any). logits = _apply_logits_processors(logits, sampling_metadata) # Prepare sampling tensors with pinned memory to avoid blocking. (sampling_tensors, do_penalties, do_top_p_top_k, do_min_p) = SamplingTensors.from_sampling_metadata(sampling_metadata, vocab_size, logits.device, logits.dtype) # Apply presence and frequency penalties. if do_penalties: logits = _apply_penalties(logits, sampling_tensors.prompt_tokens, sampling_tensors.output_tokens, sampling_tensors.presence_penalties, sampling_tensors.frequency_penalties, sampling_tensors.repetition_penalties) # Apply temperature scaling. # Use in-place division to avoid creating a new tensor. logits.div_(sampling_tensors.temperatures.unsqueeze_(dim=1)) if do_top_p_top_k: logits = _apply_top_k_top_p(logits, sampling_tensors.top_ps, sampling_tensors.top_ks) if do_min_p: logits = _apply_min_p(logits, sampling_tensors.min_ps) # We use float32 for probabilities and log probabilities. # Compute the probabilities. probs = torch.softmax(logits, dim=-1, dtype=torch.float) # Compute the log probabilities. # Use log_softmax to ensure numerical stability. logprobs = torch.log_softmax(logits, dim=-1, dtype=torch.float) # Sample the next tokens. sample_results = _sample(probs, logprobs, sampling_metadata) # Get the logprobs query results. # prompt_logprobs, sample_logprobs = _get_logprobs( # logprobs, sampling_metadata, sample_results) prompt_logprobs, sample_logprobs = _get_logprobs(origin_logprobs, sampling_metadata, sample_results) return _build_sampler_output(sample_results, sampling_metadata, prompt_logprobs, sample_logprobs)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"contextlib","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"megatron.core.tensor_parallel.utils","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.model_executor.models","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.model_executor.weight_utils","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":".weight_loaders","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.model_executor.sampling_metadata","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.model_executor.layers.sampler","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.model_executor.layers.sampler","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_set_default_torch_dtype","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_set_default_torch_dtype","relation":"calls","name":"get_default_dtype","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_set_default_torch_dtype","relation":"calls","name":"set_default_dtype","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_set_default_torch_dtype","relation":"calls","name":"set_default_dtype","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_model_architecture","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_model_architecture","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_model_architecture","relation":"calls","name":"load_model_cls","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_model_architecture","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_model_architecture","relation":"calls","name":"get_supported_archs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.model_executor.layers.linear","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.model_executor.layers.vocab_parallel_embedding","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.model_executor.layers.activation","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","relation":"calls","name":"get_default_dtype","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","relation":"calls","name":"vocab_range_from_global_vocab_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","relation":"calls","name":"Parameter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","relation":"calls","name":"empty","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#vocab_init","relation":"calls","name":"set_weight_attrs","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_model_weight_loader","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_model_weight_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_model_weight_loader","relation":"calls","name":"get_supported_archs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"_get_model_architecture","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"get_quant_config","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"get_device_capability","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"get_min_capability","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","recvType":"get_quant_config"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"get_min_capability","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","recvType":"get_quant_config"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"get_supported_act_dtypes","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","recvType":"get_quant_config"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"get_linear_method","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","recvType":"get_quant_config"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"_set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"model_class","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","recvType":"model_class"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"initialize_dummy_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"load_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"load_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","recvType":"model_class"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#get_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","recvType":"model_class"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#load_weights","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#load_weights","relation":"calls","name":"_get_model_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#load_weights","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_logits","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_logits","relation":"calls","name":"matmul","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_logits","relation":"calls","name":"t","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#_get_logits","relation":"calls","name":"tensor_model_parallel_all_gather","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"_prune_hidden_states","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"_get_logits","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"log_softmax","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"_apply_logits_processors","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"from_sampling_metadata","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"_apply_penalties","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"div_","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","recvType":"_apply_min_p"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"unsqueeze_","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"_apply_top_k_top_p","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"_apply_min_p","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"softmax","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"log_softmax","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"_sample","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"_get_logprobs","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py#forward","relation":"calls","name":"_build_sampler_output","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py","relation":"imports","specifier":"vllm.model_executor.layers.sampler","file":"verl/third_party/vllm/vllm_v_0_3_1/model_loader.py"}]},"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py":{"size":13381,"mtimeMs":1789275830313.2026,"hash":"3ea4154e940f51a1d1f4770240c905b46ed9b923e32f9aa399a715b9ac9c5ff7","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","name":"model_runner.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","span":"L1-L286","signature":null,"exported":true,"origin":"ast","body_hash":"3ea4154e940f51a1d1f4770240c905b46ed9b923e32f9aa399a715b9ac9c5ff7","chars":13381,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/worker/model_runner.py from typing import Dict, List, Optional, Tuple, Set, Union import contextlib import time import numpy as np import torch import torch.nn as nn from vllm.config import (DeviceConfig, ModelConfig, LoRAConfig, ParallelConfig, SchedulerConfig) from vllm.logger import init_logger from vllm.model_executor import InputMetadata, SamplingMetadata from vllm.sampling_params import SamplingParams, SamplingType from vllm.sequence import SamplerOutput, SequenceData, SequenceGroupMetadata from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager from vllm.lora.layers import LoRAMapping from vllm.lora.request import LoRARequest from vllm.utils import in_wsl from vllm.worker.model_runner import ModelRunner, CUDAGraphRunner, _async_h2d from .model_loader import get_model logger = init_logger(__name__) KVCache = Tuple[torch.Tensor, torch.Tensor] _PAD_SLOT_ID = -1 LORA_WARMUP_RANK = 8 # Capture graphs for batch size 1, 2, 4, 8, 16, 24, 32, 40, ..., 256. # NOTE: _get_graph_batch_size needs to be updated if this list is changed. _BATCH_SIZES_TO_CAPTURE = [1, 2, 4] + [8 * i for i in range(1, 33)]"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","name":"ModelRunner","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","span":"L46-L285","signature":"class ModelRunner(ModelRunner)","exported":true,"origin":"ast","body_hash":"ee4b24ff52e99a590178a4901341654c8bde2acbd145711b6134a99a72c03706","body_text":"class ModelRunner(ModelRunner): def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", ): self.model_config = model_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.lora_config = lora_config # model_config can be None in tests/samplers/test_sampler.py. # FIXME(woosuk): This is a hack to make the tests work. Refactor this. self.sliding_window = (model_config.get_sliding_window() if model_config is not None else None) self.device_config = (device_config if device_config is not None else DeviceConfig()) self.device = self.device_config.device self.model = model # this will be replaced by get_model() self.block_size = None # Set after initial profiling. self.lora_manager = None self.graph_runners: Dict[int, CUDAGraphRunner] = {} self.graph_memory_pool = None # Set during graph capture. self.max_context_len_to_capture = (self.model_config.max_context_len_to_capture if self.model_config is not None else 0) # When using CUDA graph, the input block tables must be padded to # max_context_len_to_capture. However, creating the block table in # Python can be expensive. To optimize this, we cache the block table # in numpy and only copy the actual input content at every iteration. # The shape of the cached block table will be # (max batch size to capture, max context len to capture / block size). self.graph_block_tables = None # Set after initial profiling. # cache in_wsl result self.in_wsl = in_wsl() self.kv_cache_dtype = kv_cache_dtype def load_model(self) -> None: self.model = get_model(actor_model=self.model, model_config=self.model_config, device_config=self.device_config, lora_config=self.lora_config) vocab_size = self.model.config.vocab_size if self.lora_config: assert hasattr( self.model, \"supported_lora_modules\") and self.model.supported_lora_modules, \"Model does not support LoRA\" assert hasattr(self.model, \"embedding_modules\"), \"Model does not have embedding_modules\" assert hasattr(self.model, \"embedding_padding_modules\"), \"Model does not have embedding_padding_modules\" self.lora_manager = LRUCacheWorkerLoRAManager( self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens + self.scheduler_config.max_paddings, vocab_size, self.lora_config, self.device, self.model.embedding_modules, self.model.embedding_padding_modules) self.model = self.lora_manager.create_lora_manager(self.model) def _prepare_sample( self, seq_group_metadata_list: List[SequenceGroupMetadata], prompt_lens: List[int], subquery_lens: Optional[List[int]], ) -> SamplingMetadata: seq_groups: List[Tuple[List[int], SamplingParams]] = [] selected_token_indices: List[int] = [] selected_token_start_idx = 0 categorized_sample_indices = {t: [] for t in SamplingType} categorized_sample_indices_start_idx = 0 max_subquery_len = max(subquery_lens) if subquery_lens else 1 for i, seq_group_metadata in enumerate(seq_group_metadata_list): seq_ids = list(seq_group_metadata.seq_data.keys()) sampling_params = seq_group_metadata.sampling_params seq_groups.append((seq_ids, sampling_params)) if seq_group_metadata.is_prompt: assert len(seq_ids) == 1 assert subquery_lens is not None subquery_len = subquery_lens[i] if sampling_params.prompt_logprobs is not None: # NOTE: prompt token positions do not need sample, skip categorized_sample_indices_start_idx += subquery_len - 1 categorized_sample_indices[sampling_params.sampling_type].append(categorized_sample_indices_start_idx) categorized_sample_indices_start_idx += 1 if sampling_params.prompt_logprobs is not None: selected_token_indices.extend( range(selected_token_start_idx, selected_token_start_idx + subquery_len - 1)) selected_token_indices.append(selected_token_start_idx + subquery_len - 1) selected_token_start_idx += max_subquery_len else: num_seqs = len(seq_ids) selected_token_indices.extend(range(selected_token_start_idx, selected_token_start_idx + num_seqs)) selected_token_start_idx += num_seqs categorized_sample_indices[sampling_params.sampling_type].extend( range(categorized_sample_indices_start_idx, categorized_sample_indices_start_idx + num_seqs)) categorized_sample_indices_start_idx += num_seqs selected_token_indices = _async_h2d(selected_token_indices, dtype=torch.long, target_device=self.device, pin_memory=not self.in_wsl) categorized_sample_indices = { t: _async_h2d(seq_ids, dtype=torch.int, target_device=self.device, pin_memory=not self.in_wsl) for t, seq_ids in categorized_sample_indices.items() } seq_data: Dict[int, SequenceData] = {} for seq_group_metadata in seq_group_metadata_list: seq_data.update(seq_group_metadata.seq_data) sampling_metadata = SamplingMetadata( seq_groups=seq_groups, seq_data=seq_data, prompt_lens=prompt_lens, selected_token_indices=selected_token_indices, categorized_sample_indices=categori","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","span":"L48-L88","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", )","exported":false,"origin":"ast","body_hash":"74c9c53a8f7ed01a4e10d3cb0f4d4fd25aa3729f770b92dd75c772bd46c87be5","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", ): self.model_config = model_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.lora_config = lora_config # model_config can be None in tests/samplers/test_sampler.py. # FIXME(woosuk): This is a hack to make the tests work. Refactor this. self.sliding_window = (model_config.get_sliding_window() if model_config is not None else None) self.device_config = (device_config if device_config is not None else DeviceConfig()) self.device = self.device_config.device self.model = model # this will be replaced by get_model() self.block_size = None # Set after initial profiling. self.lora_manager = None self.graph_runners: Dict[int, CUDAGraphRunner] = {} self.graph_memory_pool = None # Set during graph capture. self.max_context_len_to_capture = (self.model_config.max_context_len_to_capture if self.model_config is not None else 0) # When using CUDA graph, the input block tables must be padded to # max_context_len_to_capture. However, creating the block table in # Python can be expensive. To optimize this, we cache the block table # in numpy and only copy the actual input content at every iteration. # The shape of the cached block table will be # (max batch size to capture, max context len to capture / block size). self.graph_block_tables = None # Set after initial profiling. # cache in_wsl result self.in_wsl = in_wsl() self.kv_cache_dtype = kv_cache_dtype","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","span":"L90-L107","signature":"def load_model(self) -> None","exported":true,"origin":"ast","body_hash":"b76750e0d9f83140e555d97fd4595d5a466eecebc4201c09bcbb77a52882066a","body_text":"def load_model(self) -> None: self.model = get_model(actor_model=self.model, model_config=self.model_config, device_config=self.device_config, lora_config=self.lora_config) vocab_size = self.model.config.vocab_size if self.lora_config: assert hasattr( self.model, \"supported_lora_modules\") and self.model.supported_lora_modules, \"Model does not support LoRA\" assert hasattr(self.model, \"embedding_modules\"), \"Model does not have embedding_modules\" assert hasattr(self.model, \"embedding_padding_modules\"), \"Model does not have embedding_padding_modules\" self.lora_manager = LRUCacheWorkerLoRAManager( self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens + self.scheduler_config.max_paddings, vocab_size, self.lora_config, self.device, self.model.embedding_modules, self.model.embedding_padding_modules) self.model = self.lora_manager.create_lora_manager(self.model)","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","name":"_prepare_sample","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","span":"L109-L172","signature":"def _prepare_sample( self, seq_group_metadata_list: List[SequenceGroupMetadata], prompt_lens: List[int], subquery_lens: Optional[List[int]], ) -> SamplingMetadata","exported":false,"origin":"ast","body_hash":"d5b0e97199beb65828fed42e780fe7b36f1ef49058b80c888c1fb3d29cdbcbaf","body_text":"def _prepare_sample( self, seq_group_metadata_list: List[SequenceGroupMetadata], prompt_lens: List[int], subquery_lens: Optional[List[int]], ) -> SamplingMetadata: seq_groups: List[Tuple[List[int], SamplingParams]] = [] selected_token_indices: List[int] = [] selected_token_start_idx = 0 categorized_sample_indices = {t: [] for t in SamplingType} categorized_sample_indices_start_idx = 0 max_subquery_len = max(subquery_lens) if subquery_lens else 1 for i, seq_group_metadata in enumerate(seq_group_metadata_list): seq_ids = list(seq_group_metadata.seq_data.keys()) sampling_params = seq_group_metadata.sampling_params seq_groups.append((seq_ids, sampling_params)) if seq_group_metadata.is_prompt: assert len(seq_ids) == 1 assert subquery_lens is not None subquery_len = subquery_lens[i] if sampling_params.prompt_logprobs is not None: # NOTE: prompt token positions do not need sample, skip categorized_sample_indices_start_idx += subquery_len - 1 categorized_sample_indices[sampling_params.sampling_type].append(categorized_sample_indices_start_idx) categorized_sample_indices_start_idx += 1 if sampling_params.prompt_logprobs is not None: selected_token_indices.extend( range(selected_token_start_idx, selected_token_start_idx + subquery_len - 1)) selected_token_indices.append(selected_token_start_idx + subquery_len - 1) selected_token_start_idx += max_subquery_len else: num_seqs = len(seq_ids) selected_token_indices.extend(range(selected_token_start_idx, selected_token_start_idx + num_seqs)) selected_token_start_idx += num_seqs categorized_sample_indices[sampling_params.sampling_type].extend( range(categorized_sample_indices_start_idx, categorized_sample_indices_start_idx + num_seqs)) categorized_sample_indices_start_idx += num_seqs selected_token_indices = _async_h2d(selected_token_indices, dtype=torch.long, target_device=self.device, pin_memory=not self.in_wsl) categorized_sample_indices = { t: _async_h2d(seq_ids, dtype=torch.int, target_device=self.device, pin_memory=not self.in_wsl) for t, seq_ids in categorized_sample_indices.items() } seq_data: Dict[int, SequenceData] = {} for seq_group_metadata in seq_group_metadata_list: seq_data.update(seq_group_metadata.seq_data) sampling_metadata = SamplingMetadata( seq_groups=seq_groups, seq_data=seq_data, prompt_lens=prompt_lens, selected_token_indices=selected_token_indices, categorized_sample_indices=categorized_sample_indices, ) return sampling_metadata","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.prepare_input_tensors","name":"prepare_input_tensors","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","span":"L174-L200","signature":"def prepare_input_tensors( self, seq_group_metadata_list: Optional[List[SequenceGroupMetadata]], ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata, SamplingMetadata, Set[int], LoRAMapping]: # NOTE: We assume that all sequences in the group are all prompts or # all decodes.","exported":true,"origin":"ast","body_hash":"18539d850f729af5cbafae7677003c29f22e3145786cbe0e2625f29093db4c8c","body_text":"def prepare_input_tensors( self, seq_group_metadata_list: Optional[List[SequenceGroupMetadata]], ) -> Tuple[torch.Tensor, torch.Tensor, InputMetadata, SamplingMetadata, Set[int], LoRAMapping]: # NOTE: We assume that all sequences in the group are all prompts or # all decodes. is_prompt = seq_group_metadata_list[0].is_prompt # Prepare input tensors. if is_prompt: (input_tokens, input_positions, input_metadata, prompt_lens, subquery_lens, lora_index_mapping, lora_prompt_mapping, lora_requests) = self._prepare_prompt(seq_group_metadata_list) else: (input_tokens, input_positions, input_metadata, lora_index_mapping, lora_prompt_mapping, lora_requests) = self._prepare_decode(seq_group_metadata_list) prompt_lens = [] subquery_lens = None sampling_metadata = self._prepare_sample(seq_group_metadata_list, prompt_lens, subquery_lens) if self.lora_config: flat_lora_index_mapping = [item for sublist in lora_index_mapping for item in sublist] lora_mapping = LoRAMapping( flat_lora_index_mapping, lora_prompt_mapping, ) else: lora_mapping = None return (input_tokens, input_positions, input_metadata, sampling_metadata, lora_requests, lora_mapping)","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.execute_model","name":"execute_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","span":"L203-L232","signature":"def execute_model( self, seq_group_metadata_list: Optional[List[SequenceGroupMetadata]], kv_caches: List[Tuple[torch.Tensor, torch.Tensor]], ) -> Optional[SamplerOutput]","exported":true,"origin":"ast","body_hash":"3c98596b4149c7a0fa7a7e1b5723bebc63464fcee92155a8268aa62959e87fe4","body_text":"def execute_model( self, seq_group_metadata_list: Optional[List[SequenceGroupMetadata]], kv_caches: List[Tuple[torch.Tensor, torch.Tensor]], ) -> Optional[SamplerOutput]: (input_tokens, input_positions, input_metadata, sampling_metadata, lora_requests, lora_mapping) = self.prepare_input_tensors(seq_group_metadata_list) if self.lora_config: self.set_active_loras(lora_requests, lora_mapping) # Execute the model. if input_metadata.use_cuda_graph: graph_batch_size = input_tokens.shape[0] model_executable = self.graph_runners[graph_batch_size] else: model_executable = self.model hidden_states = model_executable( input_ids=input_tokens, positions=input_positions, kv_caches=kv_caches, input_metadata=input_metadata, ) # Sample the next token. output = self.model.sample( hidden_states=hidden_states, sampling_metadata=sampling_metadata, ) return output","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","name":"profile_run","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","span":"L235-L285","signature":"def profile_run(self) -> None: # Enable top-k sampling to reflect the accurate memory usage.","exported":true,"origin":"ast","body_hash":"741b4e9e6d5eada5c06999f461253385d3319156d2e4caadb4db3865864c97d9","body_text":"def profile_run(self) -> None: # Enable top-k sampling to reflect the accurate memory usage. vocab_size = self.model_config.get_vocab_size() # FIXME(sgm): this sampling params will call cumsum(), causing the # deterministic cumsum throw error sampling_params = SamplingParams(top_p=0.99, top_k=vocab_size - 1) max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens max_num_seqs = self.scheduler_config.max_num_seqs # This represents the maximum number of different requests # that will have unique loras, an therefore the max amount of memory # consumption create dummy lora request copies from the lora request # passed in, which contains a lora from the lora warmup path. dummy_lora_requests = [] dummy_lora_requests_per_seq = [] if self.lora_config: for idx in range(self.lora_config.max_loras): lora_id = idx + 1 dummy_lora_request = LoRARequest( lora_name=f\"warmup_{lora_id}\", lora_int_id=lora_id, lora_local_path=\"/not/a/real/path\", ) self.lora_manager.add_dummy_lora(dummy_lora_request, rank=LORA_WARMUP_RANK) dummy_lora_requests.append(dummy_lora_request) dummy_lora_requests_per_seq = [ dummy_lora_requests[idx % len(dummy_lora_requests)] for idx in range(max_num_seqs) ] # Profile memory usage with max_num_sequences sequences and the total # number of tokens equal to max_num_batched_tokens. seqs: List[SequenceGroupMetadata] = [] for group_id in range(max_num_seqs): seq_len = (max_num_batched_tokens // max_num_seqs + (group_id < max_num_batched_tokens % max_num_seqs)) seq_data = SequenceData([0] * seq_len) seq = SequenceGroupMetadata( request_id=str(group_id), is_prompt=True, seq_data={group_id: seq_data}, sampling_params=sampling_params, block_tables=None, lora_request=dummy_lora_requests_per_seq[group_id] if dummy_lora_requests_per_seq else None, ) seqs.append(seq) # Run the model with the dummy inputs. num_layers = self.model_config.get_num_layers(self.parallel_config) kv_caches = [(None, None)] * num_layers self.execute_model(seqs, kv_caches) torch.cuda.synchronize() return","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"contextlib","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"time","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"vllm.model_executor","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"vllm.sampling_params","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"vllm.lora.worker_manager","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"vllm.lora.layers","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":"vllm.worker.model_runner","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"imports","specifier":".model_loader","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","relation":"extends","name":"ModelRunner","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.__init__","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.__init__","relation":"calls","name":"get_sliding_window","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.__init__","relation":"calls","name":"DeviceConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.__init__","relation":"calls","name":"in_wsl","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.load_model","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.load_model","relation":"calls","name":"get_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.load_model","relation":"calls","name":"LRUCacheWorkerLoRAManager","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.load_model","relation":"calls","name":"create_lora_manager","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"LRUCacheWorkerLoRAManager"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"_async_h2d"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"_async_h2d"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"_async_h2d"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"_async_h2d","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"_async_h2d","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"update","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner._prepare_sample","relation":"calls","name":"SamplingMetadata","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.prepare_input_tensors","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"_prepare_prompt","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"_prepare_decode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"_prepare_sample","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"LoRAMapping","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","relation":"calls","name":"inference_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.execute_model","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.execute_model","relation":"calls","name":"prepare_input_tensors","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.execute_model","relation":"calls","name":"set_active_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.execute_model","relation":"calls","name":"model_executable","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.execute_model","relation":"calls","name":"sample","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"get_model"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","relation":"calls","name":"inference_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"get_vocab_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"SamplingParams","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"LoRARequest","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"add_dummy_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"LRUCacheWorkerLoRAManager"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"SequenceData","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"SequenceGroupMetadata","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"str","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"get_num_layers","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"execute_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py#ModelRunner.profile_run","relation":"calls","name":"synchronize","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/model_runner.py"}]},"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py":{"size":6075,"mtimeMs":1789275830313.2026,"hash":"7129fc091140197f245ef083b829dafb7b68341f5f7e83959a31e0cbc063b76d","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","name":"parallel_state.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","span":"L1-L148","signature":null,"exported":true,"origin":"ast","body_hash":"7129fc091140197f245ef083b829dafb7b68341f5f7e83959a31e0cbc063b76d","chars":6075,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM team. # Adapted from # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. \"\"\"Model and data parallel groups.\"\"\" import torch import torch.distributed import vllm.model_executor.parallel_utils.parallel_state as ps \"\"\" This version is strongly tied with Megatron to implement HybridEngine and weight sharing between vllm and Megatron. - We assume the Megatron tp+dp+pp world is already established before calling this function. \"\"\" # Tensor model parallel group that the current rank belongs to. _TENSOR_MODEL_PARALLEL_GROUP = None # Micro Data parallel group. Micro data parallel group is additional dp group that origins from splitting training tp # into infer_tp and micro_tp. By default, we use order micro_dp - tp _MICRO_DATA_PARALLEL_GROUP = None \"\"\" Tensor model parallel utilities \"\"\" \"\"\" Micro Data parallel group \"\"\""},{"id":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","name":"initialize_model_parallel_from_megatron","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","span":"L26-L100","signature":"def initialize_model_parallel_from_megatron( tensor_model_parallel_size=None # we set None for backward compatibility to set infer_tp = train_tp ) -> None","exported":true,"origin":"ast","body_hash":"0f01650b61cc21c27a7bbadad130f3c2b43e2cb5b37f410b2f7ba8ef3780e253","body_text":"def initialize_model_parallel_from_megatron( tensor_model_parallel_size=None # we set None for backward compatibility to set infer_tp = train_tp ) -> None: from megatron.core import parallel_state as mpu from megatron.distributed import new_group # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() if tensor_model_parallel_size is None: tensor_model_parallel_size = mpu.get_tensor_model_parallel_world_size() else: assert isinstance(tensor_model_parallel_size, int) # Build the tensor model-parallel groups. assert ps._TENSOR_MODEL_PARALLEL_GROUP is None, (\"tensor model parallel group is already initialized\") assert tensor_model_parallel_size <= mpu.get_tensor_model_parallel_world_size( ), 'Not implemented for infer_tp > train_tp' global _TENSOR_MODEL_PARALLEL_GROUP global _MICRO_DATA_PARALLEL_GROUP assert mpu.get_tensor_model_parallel_world_size() % tensor_model_parallel_size == 0 micro_dp_size = mpu.get_tensor_model_parallel_world_size() // tensor_model_parallel_size world_size: int = torch.distributed.get_world_size() num_micro_dp_groups = world_size // micro_dp_size rank = torch.distributed.get_rank() # Build the micro dp groups. assert _MICRO_DATA_PARALLEL_GROUP is None, (\"micro data parallel group is already initialized\") for i in range(num_micro_dp_groups): ranks = range(i * micro_dp_size, (i + 1) * micro_dp_size) group = new_group(rank=rank, ranks=ranks, group_type='micro_dp') if rank in ranks: _MICRO_DATA_PARALLEL_GROUP = group if tensor_model_parallel_size == mpu.get_tensor_model_parallel_world_size(): # using the same tp group as Megatron ps._TENSOR_MODEL_PARALLEL_GROUP = mpu.get_tensor_model_parallel_group() _TENSOR_MODEL_PARALLEL_GROUP = mpu.get_tensor_model_parallel_group() # no _MICRO_DATA_PARALLEL_GROUP else: # initialize a micro_dp group and a tp group # assume training tp=4, infer tp=2, then, weight is partitioned as # [1], [2], [3], [4] for training and [1,2], [1,2], [3,4], [3,4] for inference # Build the inference tp groups train_tp = mpu.get_tensor_model_parallel_world_size() num_tensor_model_parallel_groups_per_train_tp = train_tp // tensor_model_parallel_size num_tensor_model_parallel_groups = world_size // tensor_model_parallel_size assert _TENSOR_MODEL_PARALLEL_GROUP is None, (\"tensor model parallel group is already initialized\") for i in range(num_tensor_model_parallel_groups // num_tensor_model_parallel_groups_per_train_tp): start = train_tp * i end = train_tp * (i + 1) for j in range(num_tensor_model_parallel_groups_per_train_tp): ranks = list(range(start, end, num_tensor_model_parallel_groups_per_train_tp)) for i in range(len(ranks)): ranks[i] += j # group = torch.distributed.new_group(ranks) group = new_group(rank=rank, ranks=ranks, group_type='infer_tp') if rank in ranks: _TENSOR_MODEL_PARALLEL_GROUP = group ps._TENSOR_MODEL_PARALLEL_GROUP = _TENSOR_MODEL_PARALLEL_GROUP # Build the pipeline model-parallel groups. # global _PIPELINE_MODEL_PARALLEL_GROUP # global _PIPELINE_GLOBAL_RANKS # assert ps._PIPELINE_MODEL_PARALLEL_GROUP is None, (\"pipeline model parallel group is already initialized\") # ps._PIPELINE_MODEL_PARALLEL_GROUP = mpu.get_pipeline_model_parallel_group() # ps._PIPELINE_GLOBAL_RANKS = mpu.get_pipeline_model_parallel_ranks()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_group","name":"get_tensor_model_parallel_group","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","span":"L108-L111","signature":"def get_tensor_model_parallel_group()","exported":true,"origin":"ast","body_hash":"56cf5f4c8b168a0dbe1ae009e714684970069cd42df6d06e4e7e8dd43b1bd278","body_text":"def get_tensor_model_parallel_group(): \"\"\"Get the tensor model parallel group the caller rank belongs to.\"\"\" assert _TENSOR_MODEL_PARALLEL_GROUP is not None, (\"tensor model parallel group is not initialized\") return _TENSOR_MODEL_PARALLEL_GROUP","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_world_size","name":"get_tensor_model_parallel_world_size","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","span":"L114-L116","signature":"def get_tensor_model_parallel_world_size()","exported":true,"origin":"ast","body_hash":"c1af0d2aaf0c6f73d3dc5df62d1ec18f42133f1075e5262abc2daf509ff9dc98","body_text":"def get_tensor_model_parallel_world_size(): \"\"\"Return world size for the tensor model parallel group.\"\"\" return torch.distributed.get_world_size(group=get_tensor_model_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_rank","name":"get_tensor_model_parallel_rank","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","span":"L119-L121","signature":"def get_tensor_model_parallel_rank()","exported":true,"origin":"ast","body_hash":"7d7e3527eb0fbc1e8f6ca37c5449d0a44e76ded85a0b824cf4f92bda58558364","body_text":"def get_tensor_model_parallel_rank(): \"\"\"Return my rank for the tensor model parallel group.\"\"\" return torch.distributed.get_rank(group=get_tensor_model_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_src_rank","name":"get_tensor_model_parallel_src_rank","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","span":"L124-L129","signature":"def get_tensor_model_parallel_src_rank()","exported":true,"origin":"ast","body_hash":"cfa50af80e4b8151b634ff31e0fe23018462a6f55018e05a8cc9c4918ef4e48e","body_text":"def get_tensor_model_parallel_src_rank(): \"\"\"Calculate the global rank corresponding to the first local rank in the tensor model parallel group.\"\"\" global_rank = torch.distributed.get_rank() local_world_size = get_tensor_model_parallel_world_size() return (global_rank // local_world_size) * local_world_size","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_micro_data_parallel_group","name":"get_micro_data_parallel_group","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","span":"L137-L139","signature":"def get_micro_data_parallel_group()","exported":true,"origin":"ast","body_hash":"2f7b380a27daa844595b0cfc44c14ac7b99261bb2450d06387b395f1dba78cf5","body_text":"def get_micro_data_parallel_group(): assert _MICRO_DATA_PARALLEL_GROUP is not None return _MICRO_DATA_PARALLEL_GROUP","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_micro_data_parallel_world_size","name":"get_micro_data_parallel_world_size","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","span":"L142-L143","signature":"def get_micro_data_parallel_world_size()","exported":true,"origin":"ast","body_hash":"a1bfe7b553e575378267954148b439163e9ad2696ed12153bf9be879575bb9b8","body_text":"def get_micro_data_parallel_world_size(): return torch.distributed.get_world_size(group=get_micro_data_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_micro_data_parallel_rank","name":"get_micro_data_parallel_rank","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","span":"L146-L147","signature":"def get_micro_data_parallel_rank()","exported":true,"origin":"ast","body_hash":"3557d8db20db94e469e45ac4f980e16bed562a5b5699954aa2b53229cc537cb7","body_text":"def get_micro_data_parallel_rank(): return torch.distributed.get_rank(group=get_micro_data_parallel_group())","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"imports","specifier":"torch.distributed","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"imports","specifier":"megatron.core","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"imports","specifier":"megatron.distributed","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"new_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#initialize_model_parallel_from_megatron","relation":"calls","name":"new_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_group","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_world_size","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_world_size","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_world_size","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_rank","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_rank","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_src_rank","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_src_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_tensor_model_parallel_src_rank","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_micro_data_parallel_group","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_micro_data_parallel_world_size","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_micro_data_parallel_world_size","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_micro_data_parallel_world_size","relation":"calls","name":"get_micro_data_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_micro_data_parallel_rank","file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_micro_data_parallel_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py#get_micro_data_parallel_rank","relation":"calls","name":"get_micro_data_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/parallel_state.py"}]},"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py":{"size":3047,"mtimeMs":1789275830313.2026,"hash":"557315a3db0e6e7d6cf5bd1dae035e15d69c5d0fb707ee395b629f2f9443e7f1","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","name":"tokenizer.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","span":"L1-L73","signature":null,"exported":true,"origin":"ast","body_hash":"557315a3db0e6e7d6cf5bd1dae035e15d69c5d0fb707ee395b629f2f9443e7f1","chars":3047,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/tokenizer_group/tokenizer_group.py from typing import List, Optional, Tuple, Union from transformers import (AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast) from vllm.lora.request import LoRARequest from vllm.utils import make_async, LRUCache from vllm.transformers_utils.tokenizers import *"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup","name":"TokenizerGroup","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","span":"L25-L72","signature":"class TokenizerGroup","exported":true,"origin":"ast","body_hash":"fe23adaded3d2074f7e219ed963ba791eecb3ef983fe5f91192353ce88b61616","body_text":"class TokenizerGroup: \"\"\"A group of tokenizers that can be used for LoRA adapters.\"\"\" def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int]): self.enable_lora = enable_lora self.max_input_length = max_input_length self.tokenizer = tokenizer if enable_lora: self.lora_tokenizers = LRUCache(capacity=max_num_seqs) else: self.lora_tokenizers = None def encode(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = self.get_lora_tokenizer(lora_request) return tokenizer.encode(prompt) async def encode_async(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = await self.get_lora_tokenizer_async(lora_request) return tokenizer.encode(prompt) def get_lora_tokenizer(self, lora_request: Optional[LoRARequest]) -> \"PreTrainedTokenizer\": if not lora_request or not self.enable_lora: return self.tokenizer if lora_request.lora_int_id not in self.lora_tokenizers: # TODO(sgm): the lora tokenizer is also passed, but may be different tokenizer = self.tokenizer # tokenizer = (get_lora_tokenizer( # lora_request, **self.tokenizer_config) or self.tokenizer) self.lora_tokenizers.put(lora_request.lora_int_id, tokenizer) return tokenizer else: return self.lora_tokenizers.get(lora_request.lora_int_id) # FIXME(sgm): for simplicity, we assign the special token here @property def pad_token_id(self): return self.tokenizer.pad_token_id @property def eos_token_id(self): return self.tokenizer.eos_token_id","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","span":"L28-L36","signature":"def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int])","exported":false,"origin":"ast","body_hash":"e454782207e5143e5c7ca637b19ca66e6c14442c6c03acd587edce106ef857aa","body_text":"def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int]): self.enable_lora = enable_lora self.max_input_length = max_input_length self.tokenizer = tokenizer if enable_lora: self.lora_tokenizers = LRUCache(capacity=max_num_seqs) else: self.lora_tokenizers = None","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.encode","name":"encode","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","span":"L38-L43","signature":"def encode(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]","exported":true,"origin":"ast","body_hash":"c7cd151da6e5d1fb514dcc275e13f4c289111810f56df69c36b64c0d3f2cbc9c","body_text":"def encode(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = self.get_lora_tokenizer(lora_request) return tokenizer.encode(prompt)","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.encode_async","name":"encode_async","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","span":"L45-L50","signature":"async def encode_async(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]","exported":true,"origin":"ast","body_hash":"e0ccf679cf9c6ab3aeebd6ff19c82fdfdb9328b3357f16319f00fadca9561589","body_text":"async def encode_async(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = await self.get_lora_tokenizer_async(lora_request) return tokenizer.encode(prompt)","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.get_lora_tokenizer","name":"get_lora_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","span":"L52-L63","signature":"def get_lora_tokenizer(self, lora_request: Optional[LoRARequest]) -> \"PreTrainedTokenizer\"","exported":true,"origin":"ast","body_hash":"66dbb721b6874d42fc9c39c04fd34dd635d82d307c22dcbff76c9464a47c01ca","body_text":"def get_lora_tokenizer(self, lora_request: Optional[LoRARequest]) -> \"PreTrainedTokenizer\": if not lora_request or not self.enable_lora: return self.tokenizer if lora_request.lora_int_id not in self.lora_tokenizers: # TODO(sgm): the lora tokenizer is also passed, but may be different tokenizer = self.tokenizer # tokenizer = (get_lora_tokenizer( # lora_request, **self.tokenizer_config) or self.tokenizer) self.lora_tokenizers.put(lora_request.lora_int_id, tokenizer) return tokenizer else: return self.lora_tokenizers.get(lora_request.lora_int_id)","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.pad_token_id","name":"pad_token_id","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","span":"L67-L68","signature":"def pad_token_id(self)","exported":true,"origin":"ast","body_hash":"09c75a3902b0bb0ca3d7cec97618ef037c84e09c3213f62cd5ee9aa99c2b6669","body_text":"def pad_token_id(self): return self.tokenizer.pad_token_id","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.eos_token_id","name":"eos_token_id","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","span":"L71-L72","signature":"def eos_token_id(self)","exported":true,"origin":"ast","body_hash":"57455769f69da68b62a323b65045e6baf5ecb3dde6a703266578aeeb1bf096a2","body_text":"def eos_token_id(self): return self.tokenizer.eos_token_id","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","relation":"imports","specifier":"vllm.transformers_utils.tokenizers","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.__init__","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.__init__","relation":"calls","name":"LRUCache","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.encode","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.encode","relation":"calls","name":"get_lora_tokenizer","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","recvType":"TokenizerGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.encode","relation":"calls","name":"encode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.encode_async","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.encode_async","relation":"calls","name":"get_lora_tokenizer_async","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","recvType":"TokenizerGroup"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.encode_async","relation":"calls","name":"encode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.get_lora_tokenizer","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.get_lora_tokenizer","relation":"calls","name":"put","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","recvType":"LRUCache"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.get_lora_tokenizer","relation":"calls","name":"get","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py","recvType":"LRUCache"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.pad_token_id","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py#TokenizerGroup.eos_token_id","file":"verl/third_party/vllm/vllm_v_0_3_1/tokenizer.py"}]},"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py":{"size":4215,"mtimeMs":1789275830313.2026,"hash":"5f5e0fe0298bb98c6b8c453f5d8b403c14786b95ccfaa2e02c83db5f91a2f6df","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","name":"weight_loaders.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","span":"L1-L96","signature":null,"exported":true,"origin":"ast","body_hash":"5f5e0fe0298bb98c6b8c453f5d8b403c14786b95ccfaa2e02c83db5f91a2f6df","chars":4215,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models from typing import Dict import torch import torch.nn as nn # NOTE(shengguangming): replace the origin weight loader function in the class"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#parallel_weight_loader","name":"parallel_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","span":"L22-L29","signature":"def parallel_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None","exported":true,"origin":"ast","body_hash":"0f25a6818304178a092b3b28b97800a19b34e8bb12872e7a788083897c25df2a","body_text":"def parallel_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: \"\"\"Parallel Linear weight loader.\"\"\" assert param.size() == loaded_weight.size( ), 'the parameter size is not align with the loaded weight size, param size: {}, loaded_weight size: {}'.format( param.size(), loaded_weight.size()) assert param.data.dtype == loaded_weight.data.dtype, \"if we want to shared weights, the data type should also be the same\" param.data = loaded_weight.data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#default_weight_loader","name":"default_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","span":"L32-L37","signature":"def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None","exported":true,"origin":"ast","body_hash":"27bc592da2cd05ea4807cdcdd1b8feb48351a0a66ebfda38aaa38edd7f0f770b","body_text":"def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: \"\"\"Default weight loader.\"\"\" assert param.size() == loaded_weight.size() assert param.data.dtype == loaded_weight.data.dtype, \"if we want to shared weights, the data type should also be the same\" param.data = loaded_weight.data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#gpt2_weight_loader","name":"gpt2_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","span":"L40-L65","signature":"def gpt2_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"d7af0ddffa9a06f7c8b64a06c96f5ce1cbbbd5f2c03114b137268dd5d69048eb","body_text":"def gpt2_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"lm_head.weight\" in name: # GPT-2 ties the weights of the embedding layer and the final # linear layer. continue if \".attn.bias\" in name or \".attn.masked_bias\" in name: # Skip attention mask. # NOTE: \"c_attn.bias\" should not be skipped. continue if not name.startswith(\"transformer.\"): name = \"transformer.\" + name param = params_dict[name] # The HF's GPT-2 implementation uses Conv1D instead of Linear. # Because of this, we need to transpose the weights. # Note(zhuohan): the logic below might break quantized models. for conv1d_weight_name in [\"c_attn\", \"c_proj\", \"c_fc\"]: if conv1d_weight_name not in name: continue if not name.endswith(\".weight\"): continue # TODO: check megatron loaded_weight = loaded_weight.t() weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#llama_weight_loader","name":"llama_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","span":"L68-L80","signature":"def llama_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # NOTE(shengguangming): the megatron llama may have this prefix","exported":true,"origin":"ast","body_hash":"e0ee22c45f1d5bc99ffe6251b9c2620fc20d91ca9ad3994c72e5df42355f2257","body_text":"def llama_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # NOTE(shengguangming): the megatron llama may have this prefix prefix = '0.module.module.' params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if name[:len(prefix)] == prefix: name = name[len(prefix):] if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#mistral_weight_loader","name":"mistral_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","span":"L83-L95","signature":"def mistral_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # TODO: need to implement a general way to deal with prefix","exported":true,"origin":"ast","body_hash":"7eb07f84e5663210d1b0d0e97f76e1339d35bd6d45d313d59255c15da1f24d98","body_text":"def mistral_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # TODO: need to implement a general way to deal with prefix prefix = '0.module.module.' params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if name[:len(prefix)] == prefix: name = name[len(prefix):] if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#parallel_weight_loader","file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#parallel_weight_loader","relation":"calls","name":"format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#default_weight_loader","file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#default_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#default_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#gpt2_weight_loader","file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"startswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"t","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#llama_weight_loader","file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#llama_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#llama_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#llama_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#llama_weight_loader","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#llama_weight_loader","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#llama_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#llama_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#mistral_weight_loader","file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#mistral_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#mistral_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#mistral_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#mistral_weight_loader","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#mistral_weight_loader","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#mistral_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py#mistral_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/weight_loaders.py"}]},"verl/third_party/vllm/vllm_v_0_3_1/worker.py":{"size":13214,"mtimeMs":1789275830314.2026,"hash":"20cdda355891836b9c8e80e1ba6c8af30634ec5a8f345259d4cc2463b3c04d13","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","name":"worker.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L1-L315","signature":null,"exported":true,"origin":"ast","body_hash":"20cdda355891836b9c8e80e1ba6c8af30634ec5a8f345259d4cc2463b3c04d13","chars":13214,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/worker/worker.py \"\"\"A GPU worker class.\"\"\" import os import gc from typing import Dict, List, Tuple, Optional, Union, Set import torch import torch.distributed import torch.nn as nn from vllm.config import (CacheConfig, DeviceConfig, ModelConfig, ParallelConfig, SchedulerConfig, LoRAConfig) from vllm.model_executor import InputMetadata, set_random_seed from vllm.model_executor.parallel_utils.parallel_state import (initialize_model_parallel) from vllm.sampling_params import SamplingParams, SamplingType from vllm.sequence import SamplerOutput, SequenceData, SequenceGroupMetadata from vllm.worker.cache_engine import CacheEngine from vllm.model_executor.parallel_utils.custom_all_reduce import init_custom_ar from vllm.model_executor.parallel_utils.parallel_state import get_tensor_model_parallel_group from .model_runner import ModelRunner from .model_loader import load_weights from .parallel_state import initialize_model_parallel_from_megatron from vllm.lora.request import LoRARequest"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","name":"Worker","kind":"class","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L39-L267","signature":"class Worker","exported":true,"origin":"ast","body_hash":"fe733f3805e4436bd7067f3e9fa2eb44416e4f76aa506b00625a22979ec3c494","body_text":"class Worker: \"\"\"A worker class that executes (a partition of) the model on a GPU. Each worker is associated with a single GPU. The worker is responsible for maintaining the KV cache and executing the model on the GPU. In case of distributed inference, each worker is assigned a partition of the model. \"\"\" def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, rank: Optional[int] = None, distributed_init_method: Optional[str] = None, lora_config: Optional[LoRAConfig] = None, kv_cache_dtype: Optional[str] = \"auto\", ) -> None: # self.model = model # will be replaced in the init_model self.model_config = model_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.rank = rank self.distributed_init_method = distributed_init_method self.lora_config = lora_config self.model_runner = ModelRunner( model, model_config, parallel_config, scheduler_config, device_config, lora_config=self.lora_config, kv_cache_dtype=kv_cache_dtype, ) # Uninitialized cache engine. Will be initialized by # self.init_cache_engine(). self.cache_config = None self.block_size = None self.sliding_window = None self.cache_engine = None self.cache_events = None self.gpu_cache = None # For offloading inference engine params self.cpu_model = None def init_model(self, cupy_port: Optional[int] = None): # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # Env vars will be set by TORCHRUN. self.rank = self.rank if self.rank is not None else int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) self.device = torch.device(f\"cuda:{local_rank}\") if self.rank < 0: raise ValueError(\"Invalid or unspecified rank.\") torch.cuda.set_device(self.device) _check_if_gpu_supports_dtype(self.model_config.dtype) # Initialize the distributed environment. # TODO: do not use cupy _init_distributed_environment(self.parallel_config, self.rank, self.distributed_init_method) if not self.parallel_config.disable_custom_all_reduce: init_custom_ar() # Initialize the model. set_random_seed(self.model_config.seed) # self.model = get_model(actor_model=self.model, model_config=self.model_config) def load_model(self): self.model_runner.load_model() @torch.inference_mode() def profile_num_available_blocks( self, block_size: int, gpu_memory_utilization: float, cpu_swap_space: int, cache_dtype: str, ) -> Tuple[int, int]: # Profile the memory usage of the model and get the maximum number of # cache blocks that can be allocated with the remaining free memory. torch.cuda.empty_cache() # torch.cuda.reset_peak_memory_stats() # Execute a forward pass with dummy inputs to profile the memory usage # of the model. self.model_runner.profile_run() # Calculate the number of blocks that can be allocated with the # profiled peak memory. torch.cuda.synchronize() free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info() peak_memory = total_gpu_memory - free_gpu_memory cache_block_size = CacheEngine.get_cache_block_size(block_size, cache_dtype, self.model_config, self.parallel_config) # NOTE(sgm) use the remaining memory num_gpu_blocks = int((free_gpu_memory * gpu_memory_utilization) // cache_block_size) # num_gpu_blocks = int((total_gpu_memory * gpu_memory_utilization - peak_memory) // cache_block_size) num_cpu_blocks = int(cpu_swap_space // cache_block_size) num_gpu_blocks = max(num_gpu_blocks, 0) num_cpu_blocks = max(num_cpu_blocks, 0) if self.model_runner.lora_manager: self.model_runner.remove_all_loras() gc.collect() torch.cuda.empty_cache() # Synchronize number of blocks with all the rank num_gpu_blocks = torch.tensor([num_gpu_blocks], device='cuda') num_cpu_blocks = torch.tensor([num_cpu_blocks], device='cuda') torch.distributed.all_reduce(num_gpu_blocks, op=torch.distributed.ReduceOp.MIN, group=get_tensor_model_parallel_group()) torch.distributed.all_reduce(num_cpu_blocks, op=torch.distributed.ReduceOp.MIN, group=get_tensor_model_parallel_group()) num_gpu_blocks = num_gpu_blocks.item() num_cpu_blocks = num_cpu_blocks.item() return num_gpu_blocks, num_cpu_blocks def init_cache_engine(self, cache_config: CacheConfig) -> None: if self.cache_engine is None and self.gpu_cache is None: self.cache_config = cache_config self.cache_engine = CacheEngine(self.cache_config, self.model_config, self.parallel_config) self.cache_events = self.cache_engine.events self.gpu_cache = self.cache_engine.gpu_cache self.model_runner.set_block_size(self.cache_engine.block_size) def free_cache_engine(self): # ensure `enforce_eager=True` self.cache_engine = None self.gpu_cache","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L47-L87","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, rank: Optional[int] = None, distributed_init_method: Optional[str] = None, lora_config: Optional[LoRAConfig] = None, kv_cache_dtype: Optional[str] = \"auto\", ) -> None: # self.model = model # will be replaced in the init_model","exported":false,"origin":"ast","body_hash":"5cfd051917773f968a933ff2a4e25b7be3b4297f2069a9a5ee3c65bac4725a8f","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, rank: Optional[int] = None, distributed_init_method: Optional[str] = None, lora_config: Optional[LoRAConfig] = None, kv_cache_dtype: Optional[str] = \"auto\", ) -> None: # self.model = model # will be replaced in the init_model self.model_config = model_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.rank = rank self.distributed_init_method = distributed_init_method self.lora_config = lora_config self.model_runner = ModelRunner( model, model_config, parallel_config, scheduler_config, device_config, lora_config=self.lora_config, kv_cache_dtype=kv_cache_dtype, ) # Uninitialized cache engine. Will be initialized by # self.init_cache_engine(). self.cache_config = None self.block_size = None self.sliding_window = None self.cache_engine = None self.cache_events = None self.gpu_cache = None # For offloading inference engine params self.cpu_model = None","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","name":"init_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L89-L115","signature":"def init_model(self, cupy_port: Optional[int] = None): # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573","exported":true,"origin":"ast","body_hash":"0894f18903d9cdc7f6573c5ae81835ccf893b6c0aef403aef2ab3b26f30d5b1e","body_text":"def init_model(self, cupy_port: Optional[int] = None): # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # Env vars will be set by TORCHRUN. self.rank = self.rank if self.rank is not None else int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) self.device = torch.device(f\"cuda:{local_rank}\") if self.rank < 0: raise ValueError(\"Invalid or unspecified rank.\") torch.cuda.set_device(self.device) _check_if_gpu_supports_dtype(self.model_config.dtype) # Initialize the distributed environment. # TODO: do not use cupy _init_distributed_environment(self.parallel_config, self.rank, self.distributed_init_method) if not self.parallel_config.disable_custom_all_reduce: init_custom_ar() # Initialize the model. set_random_seed(self.model_config.seed) # self.model = get_model(actor_model=self.model, model_config=self.model_config)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L117-L118","signature":"def load_model(self)","exported":true,"origin":"ast","body_hash":"b053f08da81ebbfb4694b883ab12cc5ba5cbfc957774f3660deabc1843bc7afa","body_text":"def load_model(self): self.model_runner.load_model()","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","name":"profile_num_available_blocks","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L121-L166","signature":"def profile_num_available_blocks( self, block_size: int, gpu_memory_utilization: float, cpu_swap_space: int, cache_dtype: str, ) -> Tuple[int, int]: # Profile the memory usage of the model and get the maximum number of # cache blocks that can be allocated with the remaining free memory.","exported":true,"origin":"ast","body_hash":"cdfde4dcb71e28d03a0b2625143270bb35f80b5def215e9a76a57804cf02b3b4","body_text":"def profile_num_available_blocks( self, block_size: int, gpu_memory_utilization: float, cpu_swap_space: int, cache_dtype: str, ) -> Tuple[int, int]: # Profile the memory usage of the model and get the maximum number of # cache blocks that can be allocated with the remaining free memory. torch.cuda.empty_cache() # torch.cuda.reset_peak_memory_stats() # Execute a forward pass with dummy inputs to profile the memory usage # of the model. self.model_runner.profile_run() # Calculate the number of blocks that can be allocated with the # profiled peak memory. torch.cuda.synchronize() free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info() peak_memory = total_gpu_memory - free_gpu_memory cache_block_size = CacheEngine.get_cache_block_size(block_size, cache_dtype, self.model_config, self.parallel_config) # NOTE(sgm) use the remaining memory num_gpu_blocks = int((free_gpu_memory * gpu_memory_utilization) // cache_block_size) # num_gpu_blocks = int((total_gpu_memory * gpu_memory_utilization - peak_memory) // cache_block_size) num_cpu_blocks = int(cpu_swap_space // cache_block_size) num_gpu_blocks = max(num_gpu_blocks, 0) num_cpu_blocks = max(num_cpu_blocks, 0) if self.model_runner.lora_manager: self.model_runner.remove_all_loras() gc.collect() torch.cuda.empty_cache() # Synchronize number of blocks with all the rank num_gpu_blocks = torch.tensor([num_gpu_blocks], device='cuda') num_cpu_blocks = torch.tensor([num_cpu_blocks], device='cuda') torch.distributed.all_reduce(num_gpu_blocks, op=torch.distributed.ReduceOp.MIN, group=get_tensor_model_parallel_group()) torch.distributed.all_reduce(num_cpu_blocks, op=torch.distributed.ReduceOp.MIN, group=get_tensor_model_parallel_group()) num_gpu_blocks = num_gpu_blocks.item() num_cpu_blocks = num_cpu_blocks.item() return num_gpu_blocks, num_cpu_blocks","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L168-L174","signature":"def init_cache_engine(self, cache_config: CacheConfig) -> None","exported":true,"origin":"ast","body_hash":"f7d2764319d8c9975a7b89db2d9825fb4aeeee8f7f568571357d1d1bec6d2017","body_text":"def init_cache_engine(self, cache_config: CacheConfig) -> None: if self.cache_engine is None and self.gpu_cache is None: self.cache_config = cache_config self.cache_engine = CacheEngine(self.cache_config, self.model_config, self.parallel_config) self.cache_events = self.cache_engine.events self.gpu_cache = self.cache_engine.gpu_cache self.model_runner.set_block_size(self.cache_engine.block_size)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L176-L179","signature":"def free_cache_engine(self): # ensure `enforce_eager=True`","exported":true,"origin":"ast","body_hash":"d384948b80b2c04ee44319bcd88c9b941e3967b68e2050a10073ade04af52223","body_text":"def free_cache_engine(self): # ensure `enforce_eager=True` self.cache_engine = None self.gpu_cache = None","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.warm_up_model","name":"warm_up_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L181-L186","signature":"def warm_up_model(self) -> None","exported":true,"origin":"ast","body_hash":"eb320b283f005ad95d749d466dfb34c7c566316e5fef7db0d24bd33b7adfab8b","body_text":"def warm_up_model(self) -> None: if not self.model_config.enforce_eager: self.model_runner.capture_model(self.gpu_cache) # Reset the seed to ensure that the random state is not affected by # the model initialization and profiling. set_random_seed(self.model_config.seed)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.cache_swap","name":"cache_swap","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L188-L212","signature":"def cache_swap( self, blocks_to_swap_in: Dict[int, int], blocks_to_swap_out: Dict[int, int], blocks_to_copy: Dict[int, List[int]], ) -> None: # Issue cache operations.","exported":true,"origin":"ast","body_hash":"c9f462500147afa5b4120561a68ac67c878aa89c837e038eadca093ce886c4e9","body_text":"def cache_swap( self, blocks_to_swap_in: Dict[int, int], blocks_to_swap_out: Dict[int, int], blocks_to_copy: Dict[int, List[int]], ) -> None: # Issue cache operations. issued_cache_op = False if blocks_to_swap_in: self.cache_engine.swap_in(blocks_to_swap_in) issued_cache_op = True if blocks_to_swap_out: self.cache_engine.swap_out(blocks_to_swap_out) issued_cache_op = True if blocks_to_copy: self.cache_engine.copy(blocks_to_copy) issued_cache_op = True cache_events = self.cache_events if issued_cache_op else None # Wait for cache operations to finish. # TODO(woosuk): Profile swapping overhead and optimize if needed. if cache_events is not None: for event in cache_events: event.wait()","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.execute_model","name":"execute_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L215-L244","signature":"def execute_model( self, seq_group_metadata_list: List[SequenceGroupMetadata], blocks_to_swap_in: Dict[int, int], blocks_to_swap_out: Dict[int, int], blocks_to_copy: Dict[int, List[int]], ) -> SamplerOutput","exported":true,"origin":"ast","body_hash":"029ac34860bb7b545d0012529a805db910dbb36beacf0a7972d40b7ac800c68e","body_text":"def execute_model( self, seq_group_metadata_list: List[SequenceGroupMetadata], blocks_to_swap_in: Dict[int, int], blocks_to_swap_out: Dict[int, int], blocks_to_copy: Dict[int, List[int]], ) -> SamplerOutput: num_seq_groups = len(seq_group_metadata_list) self.cache_swap(blocks_to_swap_in, blocks_to_swap_out, blocks_to_copy) # If there is no input, we don't need to execute the model. if num_seq_groups == 0: return {} output = self.model_runner.execute_model(seq_group_metadata_list, self.gpu_cache) return output # # Prepare input tensors. # # NOTE(shengguangming): currently we pad in our dataloader and unpad it in pre_process_input, j # # we can just input un-padded sequence for better performance # input_tokens, input_positions, input_metadata = self._prepare_inputs(seq_group_metadata_list) # # Execute the model. # output = self.model( # input_ids=input_tokens, # positions=input_positions, # kv_caches=self.gpu_cache, # input_metadata=input_metadata, # cache_events=cache_events, # ) # return output","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L247-L248","signature":"def sync_model_weights(self, actor_weights: Dict)","exported":true,"origin":"ast","body_hash":"534fea98ce5f439a414993fa7071258b7a03f93d5aea35e9024371ca6d332b37","body_text":"def sync_model_weights(self, actor_weights: Dict): load_weights(actor_weights, self.model_runner.model)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L250-L258","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"58d8d5e64a0765436d968d3a9641e3432be0a43a9ca286f99b1f827937e19e31","body_text":"def offload_model_weights(self) -> None: if self.cpu_model == None: self.cpu_model = {} for name, params in self.model_runner.model.named_parameters(): self.cpu_model[name] = torch.empty_like(params, device='cpu') params.data = self.cpu_model[name] else: for name, params in self.model_runner.model.named_parameters(): params.data = self.cpu_model[name]","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.add_lora","name":"add_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L260-L261","signature":"def add_lora(self, lora_request: LoRARequest) -> bool","exported":true,"origin":"ast","body_hash":"36e26f28d65abc9f1720adf1d88f8801d8d1d6efb2ae59a91b43df38b144ddff","body_text":"def add_lora(self, lora_request: LoRARequest) -> bool: return self.model_runner.add_lora(lora_request)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.remove_lora","name":"remove_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L263-L264","signature":"def remove_lora(self, lora_id: int) -> bool","exported":true,"origin":"ast","body_hash":"ce1bea5fb16f7443da73b665b008a4e81300daf50030a6b963fcfa8f8781dde6","body_text":"def remove_lora(self, lora_id: int) -> bool: return self.model_runner.remove_lora(lora_id)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.list_loras","name":"list_loras","kind":"method","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L266-L267","signature":"def list_loras(self) -> Set[int]","exported":true,"origin":"ast","body_hash":"4d91fe3674c2f89266165cff6f10528b63decdd1849b54c33f151b71166ffe44","body_text":"def list_loras(self) -> Set[int]: return self.model_runner.list_loras()","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","name":"_init_distributed_environment","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L270-L295","signature":"def _init_distributed_environment( parallel_config: ParallelConfig, rank: int, distributed_init_method: Optional[str] = None, ) -> None","exported":false,"origin":"ast","body_hash":"b70429de640547fbf0293c5a7e9b60b0694da7701b7c356cdd58789f03c27fdc","body_text":"def _init_distributed_environment( parallel_config: ParallelConfig, rank: int, distributed_init_method: Optional[str] = None, ) -> None: \"\"\"Initialize the distributed environment.\"\"\" if torch.distributed.is_initialized(): print('The distributed environment has been initialized before vLLM') elif not distributed_init_method: raise ValueError(\"distributed_init_method must be set if torch.distributed \" \"is not already initialized\") else: torch.distributed.init_process_group( backend=\"nccl\", world_size=parallel_config.world_size, rank=rank, # init_method=distributed_init_method, ) # A small all_reduce for warmup. torch.distributed.all_reduce(torch.zeros(1).cuda()) # TODO (shengguangming): maybe we should also flag the megatron is initialized if torch.distributed.get_world_size() > 1: initialize_model_parallel_from_megatron(tensor_model_parallel_size=parallel_config.tensor_parallel_size) else: initialize_model_parallel()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_pad_to_alignment","name":"_pad_to_alignment","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L298-L299","signature":"def _pad_to_alignment(x: List[int], multiple_of: int, pad: int) -> List[int]","exported":false,"origin":"ast","body_hash":"5cb00cafd8f85cd9bfd95babc1dec2d2bfb10c98507130f4bfc6309ab0cccebb","body_text":"def _pad_to_alignment(x: List[int], multiple_of: int, pad: int) -> List[int]: return x + [pad] * ((-len(x)) % multiple_of)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_pad_to_max","name":"_pad_to_max","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L302-L303","signature":"def _pad_to_max(x: List[int], max_len: int, pad: int) -> List[int]","exported":false,"origin":"ast","body_hash":"255569351242d5a0b1d1f5637d925534181da37cc4fbcd63343576bb6ea3ac69","body_text":"def _pad_to_max(x: List[int], max_len: int, pad: int) -> List[int]: return x + [pad] * (max_len - len(x))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_check_if_gpu_supports_dtype","name":"_check_if_gpu_supports_dtype","kind":"function","path":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","span":"L306-L314","signature":"def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype): # Check if the GPU supports the dtype.","exported":false,"origin":"ast","body_hash":"de15eb8a0647e4dbe5b722d3cbbfe2827bedcbb1b863198256e11be2545f3ab6","body_text":"def _check_if_gpu_supports_dtype(torch_dtype: torch.dtype): # Check if the GPU supports the dtype. if torch_dtype == torch.bfloat16: compute_capability = torch.cuda.get_device_capability() if compute_capability[0] < 8: gpu_name = torch.cuda.get_device_name() raise ValueError(\"Bfloat16 is only supported on GPUs with compute capability \" f\"of at least 8.0. Your {gpu_name} GPU has compute capability \" f\"{compute_capability[0]}.{compute_capability[1]}.\")","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"gc","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"torch.distributed","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"vllm.model_executor","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"vllm.model_executor.parallel_utils.parallel_state","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"vllm.sampling_params","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"vllm.worker.cache_engine","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"vllm.model_executor.parallel_utils.custom_all_reduce","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"vllm.model_executor.parallel_utils.parallel_state","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":".model_runner","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":".model_loader","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":".parallel_state","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.__init__","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.__init__","relation":"calls","name":"ModelRunner","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"set_device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"_check_if_gpu_supports_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"_init_distributed_environment","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"init_custom_ar","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_model","relation":"calls","name":"set_random_seed","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.load_model","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.load_model","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"calls","name":"inference_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"profile_run","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"synchronize","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"mem_get_info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"get_cache_block_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"remove_all_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"collect","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"item","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"max"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.profile_num_available_blocks","relation":"calls","name":"item","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"max"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_cache_engine","relation":"calls","name":"CacheEngine","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.init_cache_engine","relation":"calls","name":"set_block_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.warm_up_model","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.warm_up_model","relation":"calls","name":"capture_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.warm_up_model","relation":"calls","name":"set_random_seed","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.cache_swap","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.cache_swap","relation":"calls","name":"swap_in","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"CacheEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.cache_swap","relation":"calls","name":"swap_out","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"CacheEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.cache_swap","relation":"calls","name":"copy","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"CacheEngine"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.cache_swap","relation":"calls","name":"wait","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"calls","name":"inference_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.execute_model","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.execute_model","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.execute_model","relation":"calls","name":"cache_swap","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.execute_model","relation":"calls","name":"execute_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.sync_model_weights","relation":"calls","name":"load_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.offload_model_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.offload_model_weights","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.offload_model_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.add_lora","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.add_lora","relation":"calls","name":"add_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.remove_lora","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.remove_lora","relation":"calls","name":"remove_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.list_loras","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#Worker.list_loras","relation":"calls","name":"list_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","relation":"calls","name":"init_process_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","relation":"calls","name":"zeros","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","relation":"calls","name":"initialize_model_parallel_from_megatron","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_init_distributed_environment","relation":"calls","name":"initialize_model_parallel","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_pad_to_alignment","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_pad_to_alignment","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_pad_to_max","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_pad_to_max","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_check_if_gpu_supports_dtype","file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_check_if_gpu_supports_dtype","relation":"calls","name":"get_device_capability","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_check_if_gpu_supports_dtype","relation":"calls","name":"get_device_name","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_3_1/worker.py#_check_if_gpu_supports_dtype","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_3_1/worker.py"}]},"verl/third_party/vllm/vllm_v_0_4_2/__init__.py":{"size":600,"mtimeMs":1789275830314.2026,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/__init__.py","name":"__init__.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py":{"size":16002,"mtimeMs":1789275830314.2026,"hash":"072ff25d639aac9716f353a7f9573db9296e6f2cdbde8c905b0966f193c5570c","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","name":"arg_utils.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","span":"L1-L321","signature":null,"exported":true,"origin":"ast","body_hash":"072ff25d639aac9716f353a7f9573db9296e6f2cdbde8c905b0966f193c5570c","chars":16002,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/engine/arg_utils.py import os import argparse import dataclasses from dataclasses import dataclass from typing import List, Optional, Union import torch.nn as nn from transformers import PretrainedConfig from .config import ModelConfig, LoadConfig from vllm.config import (CacheConfig, DecodingConfig, DeviceConfig, EngineConfig, LoRAConfig, ParallelConfig, SchedulerConfig, SpeculativeConfig, TokenizerPoolConfig, VisionLanguageConfig) from vllm.model_executor.layers.quantization import QUANTIZATION_METHODS from vllm.utils import str_to_int_tuple @dataclass"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#nullable_str","name":"nullable_str","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","span":"L33-L36","signature":"def nullable_str(val: str)","exported":true,"origin":"ast","body_hash":"900c08104995d0c284480fe8066923076c935ca69209e3a86885fbcc330eac6b","body_text":"def nullable_str(val: str): if not val or val == \"None\": return None return val","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs","name":"EngineArgs","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","span":"L40-L320","signature":"class EngineArgs","exported":true,"origin":"ast","body_hash":"9d799ad555c4f1c7e7d22ae002d5a4051f80a06baf22ccc7541accb014166497","body_text":"class EngineArgs: \"\"\"Arguments for vLLM engine.\"\"\" model_hf_config: PretrainedConfig = None skip_tokenizer_init: bool = False served_model_name: Optional[Union[str, List[str]]] = None # TODO download_dir: Optional[str] = None load_format: str = 'auto' dtype: str = 'auto' kv_cache_dtype: str = 'auto' quantization_param_path: Optional[str] = None seed: int = 0 max_model_len: Optional[int] = None worker_use_ray: bool = False pipeline_parallel_size: int = 1 tensor_parallel_size: int = 1 max_parallel_loading_workers: Optional[int] = None block_size: int = 16 enable_prefix_caching: bool = False use_v2_block_manager: bool = False swap_space: int = 4 # GiB gpu_memory_utilization: float = 0.90 max_num_batched_tokens: Optional[int] = None max_num_seqs: int = 256 max_logprobs: int = 5 # OpenAI default value disable_log_stats: bool = False revision: Optional[str] = None code_revision: Optional[str] = None tokenizer_revision: Optional[str] = None quantization: Optional[str] = None enforce_eager: bool = False max_context_len_to_capture: Optional[int] = None max_seq_len_to_capture: int = 8192 disable_custom_all_reduce: bool = False tokenizer_pool_size: int = 0 tokenizer_pool_type: str = \"ray\" tokenizer_pool_extra_config: Optional[dict] = None enable_lora: bool = False max_loras: int = 1 max_lora_rank: int = 16 fully_sharded_loras: bool = False lora_extra_vocab_size: int = 256 lora_dtype = 'auto' max_cpu_loras: Optional[int] = None device: str = 'auto' ray_workers_use_nsight: bool = False num_gpu_blocks_override: Optional[int] = None num_lookahead_slots: int = 0 model_loader_extra_config: Optional[dict] = None # Related to Vision-language models such as llava image_input_type: Optional[str] = None image_token_id: Optional[int] = None image_input_shape: Optional[str] = None image_feature_size: Optional[int] = None scheduler_delay_factor: float = 0.0 enable_chunked_prefill: bool = False guided_decoding_backend: str = 'outlines' # Speculative decoding configuration. speculative_model: Optional[str] = None num_speculative_tokens: Optional[int] = None speculative_max_model_len: Optional[int] = None ngram_prompt_lookup_max: Optional[int] = None ngram_prompt_lookup_min: Optional[int] = None @staticmethod def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: \"\"\"Shared CLI arguments for vLLM engine.\"\"\" # Model arguments # TODO(shengguangming): delete the unused args parser.add_argument('--model', type=str, default='facebook/opt-125m', help='name or path of the huggingface model to use') parser.add_argument('--tokenizer', type=str, default=EngineArgs.tokenizer, help='name or path of the huggingface tokenizer to use') parser.add_argument('--revision', type=str, default=None, help='the specific model version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-revision', type=str, default=None, help='the specific tokenizer version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-mode', type=str, default=EngineArgs.tokenizer_mode, choices=['auto', 'slow'], help='tokenizer mode. \"auto\" will use the fast ' 'tokenizer if available, and \"slow\" will ' 'always use the slow tokenizer.') parser.add_argument('--trust-remote-code', action='store_true', help='trust remote code from huggingface') parser.add_argument('--download-dir', type=str, default=EngineArgs.download_dir, help='directory to download and load the weights, ' 'default to the default cache dir of ' 'huggingface') parser.add_argument('--load-format', type=str, default=EngineArgs.load_format, choices=['auto', 'pt', 'safetensors', 'npcache', 'dummy'], help='The format of the model weights to load. ' '\"auto\" will try to load the weights in the safetensors format ' 'and fall back to the pytorch bin format if safetensors format ' 'is not available. ' '\"pt\" will load the weights in the pytorch bin format. ' '\"safetensors\" will load the weights in the safetensors format. ' '\"npcache\" will load the weights in pytorch format and store ' 'a numpy cache to speed up the loading. ' '\"dummy\" will initialize the weights with random values, ' 'which is mainly for profiling.') parser.add_argument('--dtype', type=str, default=EngineArgs.dtype, choices=['auto', 'half', 'float16', 'bfloat16', 'float', 'float32'], help='data type for model weights and activations. ' 'The \"auto\" option will use FP16 precision ' 'for FP32 and FP16 models, and BF16 precision ' 'for BF16 models.') parser.add_argument('--max-model-len', type=int, default=None, help='model context length. If unspecified, ' 'will be automatically derived from the model.') # Parallel arguments parser.add_argument('--worker-use-ray', action='store_true', help='use Ray for distributed serving, will be ' 'automatically set when using more than 1 GPU') parser.add_argument('--pipeline-parallel-size', '-pp', type=int, default=EngineArg","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","name":"add_cli_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","span":"L106-L220","signature":"def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser","exported":true,"origin":"ast","body_hash":"cf1908f71aee396b15d02e5aad05a2803adf4e9fd0905ff5fcb9c8727ba16fb5","body_text":"def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: \"\"\"Shared CLI arguments for vLLM engine.\"\"\" # Model arguments # TODO(shengguangming): delete the unused args parser.add_argument('--model', type=str, default='facebook/opt-125m', help='name or path of the huggingface model to use') parser.add_argument('--tokenizer', type=str, default=EngineArgs.tokenizer, help='name or path of the huggingface tokenizer to use') parser.add_argument('--revision', type=str, default=None, help='the specific model version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-revision', type=str, default=None, help='the specific tokenizer version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-mode', type=str, default=EngineArgs.tokenizer_mode, choices=['auto', 'slow'], help='tokenizer mode. \"auto\" will use the fast ' 'tokenizer if available, and \"slow\" will ' 'always use the slow tokenizer.') parser.add_argument('--trust-remote-code', action='store_true', help='trust remote code from huggingface') parser.add_argument('--download-dir', type=str, default=EngineArgs.download_dir, help='directory to download and load the weights, ' 'default to the default cache dir of ' 'huggingface') parser.add_argument('--load-format', type=str, default=EngineArgs.load_format, choices=['auto', 'pt', 'safetensors', 'npcache', 'dummy'], help='The format of the model weights to load. ' '\"auto\" will try to load the weights in the safetensors format ' 'and fall back to the pytorch bin format if safetensors format ' 'is not available. ' '\"pt\" will load the weights in the pytorch bin format. ' '\"safetensors\" will load the weights in the safetensors format. ' '\"npcache\" will load the weights in pytorch format and store ' 'a numpy cache to speed up the loading. ' '\"dummy\" will initialize the weights with random values, ' 'which is mainly for profiling.') parser.add_argument('--dtype', type=str, default=EngineArgs.dtype, choices=['auto', 'half', 'float16', 'bfloat16', 'float', 'float32'], help='data type for model weights and activations. ' 'The \"auto\" option will use FP16 precision ' 'for FP32 and FP16 models, and BF16 precision ' 'for BF16 models.') parser.add_argument('--max-model-len', type=int, default=None, help='model context length. If unspecified, ' 'will be automatically derived from the model.') # Parallel arguments parser.add_argument('--worker-use-ray', action='store_true', help='use Ray for distributed serving, will be ' 'automatically set when using more than 1 GPU') parser.add_argument('--pipeline-parallel-size', '-pp', type=int, default=EngineArgs.pipeline_parallel_size, help='number of pipeline stages') parser.add_argument('--tensor-parallel-size', '-tp', type=int, default=EngineArgs.tensor_parallel_size, help='number of tensor parallel replicas') # KV cache arguments parser.add_argument('--block-size', type=int, default=EngineArgs.block_size, choices=[8, 16, 32], help='token block size') # TODO(woosuk): Support fine-grained seeds (e.g., seed per request). parser.add_argument('--seed', type=int, default=EngineArgs.seed, help='random seed') parser.add_argument('--swap-space', type=int, default=EngineArgs.swap_space, help='CPU swap space size (GiB) per GPU') parser.add_argument('--gpu-memory-utilization', type=float, default=EngineArgs.gpu_memory_utilization, help='the percentage of GPU memory to be used for' 'the model executor') parser.add_argument('--max-num-batched-tokens', type=int, default=EngineArgs.max_num_batched_tokens, help='maximum number of batched tokens per ' 'iteration') parser.add_argument('--max-num-seqs', type=int, default=EngineArgs.max_num_seqs, help='maximum number of sequences per iteration') parser.add_argument('--disable-log-stats', action='store_true', help='disable logging statistics') # Quantization settings. parser.add_argument('--quantization', '-q', type=str, choices=['awq', None], default=None, help='Method used to quantize the weights') return parser","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.from_cli_args","name":"from_cli_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","span":"L223-L228","signature":"def from_cli_args(cls, args: argparse.Namespace) -> 'EngineArgs': # Get the list of attributes of this dataclass.","exported":true,"origin":"ast","body_hash":"21089f327dd4125f76a223a109fdee39de77f84d10cd1fad1f24de8837c1b6f1","body_text":"def from_cli_args(cls, args: argparse.Namespace) -> 'EngineArgs': # Get the list of attributes of this dataclass. attrs = [attr.name for attr in dataclasses.fields(cls)] # Set the attributes from the parsed arguments. engine_args = cls(**{attr: getattr(args, attr) for attr in attrs}) return engine_args","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","name":"create_engine_config","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","span":"L230-L320","signature":"def create_engine_config( self, ) -> EngineConfig","exported":true,"origin":"ast","body_hash":"ff7931b96ebf007c0e7482f6b20bdc1caf8ed654266bf78325dea65b331cf6a4","body_text":"def create_engine_config( self, ) -> EngineConfig: device_config = DeviceConfig(self.device) # NOTE(sgm): we only modify ModelConfig, other configs are import from vllm model_config = ModelConfig(self.model_hf_config, self.dtype, self.seed, self.revision, self.code_revision, self.tokenizer_revision, self.max_model_len, self.quantization, self.quantization_param_path, self.enforce_eager, self.max_context_len_to_capture, self.max_seq_len_to_capture, self.max_logprobs, self.skip_tokenizer_init, self.served_model_name) cache_config = CacheConfig(self.block_size, self.gpu_memory_utilization, self.swap_space, self.kv_cache_dtype, self.num_gpu_blocks_override, model_config.get_sliding_window(), self.enable_prefix_caching) parallel_config = ParallelConfig( self.pipeline_parallel_size, self.tensor_parallel_size, self.worker_use_ray, self.max_parallel_loading_workers, self.disable_custom_all_reduce, TokenizerPoolConfig.create_config( self.tokenizer_pool_size, self.tokenizer_pool_type, self.tokenizer_pool_extra_config, ), self.ray_workers_use_nsight) # Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" parallel_config.world_size = world_size # TODO: spec config speculative_config = SpeculativeConfig.maybe_create_spec_config( target_model_config=model_config, target_parallel_config=parallel_config, target_dtype=self.dtype, speculative_model=self.speculative_model, num_speculative_tokens=self.num_speculative_tokens, speculative_max_model_len=self.speculative_max_model_len, enable_chunked_prefill=self.enable_chunked_prefill, use_v2_block_manager=self.use_v2_block_manager, ngram_prompt_lookup_max=self.ngram_prompt_lookup_max, ngram_prompt_lookup_min=self.ngram_prompt_lookup_min, ) scheduler_config = SchedulerConfig( self.max_num_batched_tokens, self.max_num_seqs, model_config.max_model_len, self.use_v2_block_manager, num_lookahead_slots=(self.num_lookahead_slots if speculative_config is None else speculative_config.num_lookahead_slots), delay_factor=self.scheduler_delay_factor, enable_chunked_prefill=self.enable_chunked_prefill, ) lora_config = LoRAConfig(max_lora_rank=self.max_lora_rank, max_loras=self.max_loras, fully_sharded_loras=self.fully_sharded_loras, lora_extra_vocab_size=self.lora_extra_vocab_size, lora_dtype=self.lora_dtype, max_cpu_loras=self.max_cpu_loras if self.max_cpu_loras and self.max_cpu_loras > 0 else None) if self.enable_lora else None load_config = LoadConfig( load_format=self.load_format, download_dir=self.download_dir, model_loader_extra_config=self.model_loader_extra_config, ) if self.image_input_type: if (not self.image_token_id or not self.image_input_shape or not self.image_feature_size): raise ValueError('Specify `image_token_id`, `image_input_shape` and ' '`image_feature_size` together with `image_input_type`.') vision_language_config = VisionLanguageConfig( image_input_type=VisionLanguageConfig.get_image_input_enum_type(self.image_input_type), image_token_id=self.image_token_id, image_input_shape=str_to_int_tuple(self.image_input_shape), image_feature_size=self.image_feature_size, ) else: vision_language_config = None decoding_config = DecodingConfig(guided_decoding_backend=self.guided_decoding_backend) return EngineConfig(model_config=model_config, cache_config=cache_config, parallel_config=parallel_config, scheduler_config=scheduler_config, device_config=device_config, lora_config=lora_config, vision_language_config=vision_language_config, speculative_config=speculative_config, load_config=load_config, decoding_config=decoding_config)","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"imports","specifier":"argparse","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"imports","specifier":"vllm.model_executor.layers.quantization","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#nullable_str","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.from_cli_args","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.from_cli_args","relation":"calls","name":"fields","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.from_cli_args","relation":"calls","name":"cls","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.from_cli_args","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"DeviceConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"ModelConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"CacheConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"get_sliding_window","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"ParallelConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"create_config","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"maybe_create_spec_config","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"SchedulerConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"LoRAConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"LoadConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"VisionLanguageConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"get_image_input_enum_type","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"str_to_int_tuple","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"DecodingConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"EngineConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/arg_utils.py"}]},"verl/third_party/vllm/vllm_v_0_4_2/config.py":{"size":9527,"mtimeMs":1789275830314.2026,"hash":"fafdec45f81c2def3798ca7ea86d00cfda7b3ee790c9ba4cc5f00a842ff9a181","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/config.py","name":"config.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/config.py","span":"L1-L201","signature":null,"exported":true,"origin":"ast","body_hash":"fafdec45f81c2def3798ca7ea86d00cfda7b3ee790c9ba4cc5f00a842ff9a181","chars":9527,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py import enum import json from typing import List, Optional, Union from dataclasses import dataclass, field, fields from transformers import PretrainedConfig from vllm.logger import init_logger from vllm.model_executor.layers.quantization import get_quantization_config from vllm.transformers_utils.config import get_hf_text_config from vllm.utils import is_hip # Add for verl from vllm.config import ModelConfig, _get_and_verify_dtype, _get_and_verify_max_len GPTQMarlinConfig = get_quantization_config(\"gptq_marlin\") logger = init_logger(__name__) _GB = 1 << 30 @dataclass"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig","name":"ModelConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/config.py","span":"L37-L144","signature":"class ModelConfig(ModelConfig)","exported":true,"origin":"ast","body_hash":"ee0ec070b3be56e276d2dd111a273d0e711098d9cd58a94c3d9aa1ab2f1225aa","body_text":"class ModelConfig(ModelConfig): \"\"\"Configuration for the model. Args: model: Name or path of the huggingface model to use. tokenizer: Name or path of the huggingface tokenizer to use. tokenizer_mode: Tokenizer mode. \"auto\" will use the fast tokenizer if available, and \"slow\" will always use the slow tokenizer. trust_remote_code: Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer. download_dir: Directory to download and load the weights, default to the default cache directory of huggingface. load_format: The format of the model weights to load: \"auto\" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. \"pt\" will load the weights in the pytorch bin format. \"safetensors\" will load the weights in the safetensors format. \"npcache\" will load the weights in pytorch format and store a numpy cache to speed up the loading. \"dummy\" will initialize the weights with random values, which is mainly for profiling. dtype: Data type for model weights and activations. The \"auto\" option will use FP16 precision for FP32 and FP16 models, and BF16 precision for BF16 models. seed: Random seed for reproducibility. revision: The specific model version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. code_revision: The specific revision to use for the model code on Hugging Face Hub. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. tokenizer_revision: The specific tokenizer version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. max_model_len: Maximum length of a sequence (including prompt and output). If None, will be derived from the model. quantization: Quantization method that was used to quantize the model weights. If None, we assume the model weights are not quantized. quantization_param_path: Path to JSON file containing scaling factors. Used to load KV cache scaling factors into the model when KV cache type is FP8_E4M3 on ROCm (AMD GPU). In the future these will also be used to load activation and weight scaling factors when the model dtype is FP8_E4M3 on ROCm. enforce_eager: Whether to enforce eager execution. If True, we will disable CUDA graph and always execute the model in eager mode. If False, we will use CUDA graph and eager execution in hybrid. max_context_len_to_capture: Maximum context len covered by CUDA graphs. When a sequence has context length larger than this, we fall back to eager mode (DEPRECATED. Use max_seq_len_to_capture instead). max_seq_len_to_capture: Maximum sequence len covered by CUDA graphs. When a sequence has context length larger than this, we fall back to eager mode skip_tokenizer_init: If true, skip initialization of tokenizer and detokenizer. served_model_name: The model name used in metrics tag `model_name`, matches the model name exposed via the APIs. If multiple model names provided, the first name will be used. If not specified, the model name will be the same as `model`. \"\"\" def __init__( self, hf_config: PretrainedConfig, dtype: str, seed: int, revision: Optional[str] = None, code_revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, max_model_len: Optional[int] = None, quantization: Optional[str] = None, quantization_param_path: Optional[str] = None, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: Optional[int] = None, max_logprobs: int = 5, skip_tokenizer_init: bool = False, served_model_name: Optional[Union[str, List[str]]] = None, ) -> None: self.model = hf_config._name_or_path self.tokenizer = hf_config._name_or_path self.seed = seed self.revision = revision self.code_revision = code_revision self.tokenizer_revision = tokenizer_revision self.quantization = quantization self.quantization_param_path = quantization_param_path self.enforce_eager = enforce_eager self.max_context_len_to_capture = max_context_len_to_capture if self.max_context_len_to_capture is not None: raise ValueError(\"`max_context_len_to_capture` is deprecated. \" \"Use `max_seq_len_to_capture` instead.\") self.max_seq_len_to_capture = (max_seq_len_to_capture or max_context_len_to_capture) self.max_logprobs = max_logprobs self.skip_tokenizer_init = skip_tokenizer_init # self.hf_config = get_config(model, trust_remote_code, revision) self.hf_config = hf_config self.hf_text_config = get_hf_text_config(hf_config) # TODO: for multimodal model self.dtype = _get_and_verify_dtype(self.hf_config, dtype) self.max_model_len = _get_and_verify_max_len(self.hf_config, max_model_len) # self.served_model_name = get_served_model_name(model, # served_model_name) # self._verify_load_format() # self._verify_tokenizer_mode() self._verify_quantization() self._verify_cuda_graph()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/config.py","span":"L98-L144","signature":"def __init__( self, hf_config: PretrainedConfig, dtype: str, seed: int, revision: Optional[str] = None, code_revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, max_model_len: Optional[int] = None, quantization: Optional[str] = None, quantization_param_path: Optional[str] = None, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: Optional[int] = None, max_logprobs: int = 5, skip_tokenizer_init: bool = False, served_model_name: Optional[Union[str, List[str]]] = None, ) -> None","exported":false,"origin":"ast","body_hash":"2f037af99496298da9d12bffb83d5f97c78f8c4ef61b21cb177642bdacdec52b","body_text":"def __init__( self, hf_config: PretrainedConfig, dtype: str, seed: int, revision: Optional[str] = None, code_revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, max_model_len: Optional[int] = None, quantization: Optional[str] = None, quantization_param_path: Optional[str] = None, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: Optional[int] = None, max_logprobs: int = 5, skip_tokenizer_init: bool = False, served_model_name: Optional[Union[str, List[str]]] = None, ) -> None: self.model = hf_config._name_or_path self.tokenizer = hf_config._name_or_path self.seed = seed self.revision = revision self.code_revision = code_revision self.tokenizer_revision = tokenizer_revision self.quantization = quantization self.quantization_param_path = quantization_param_path self.enforce_eager = enforce_eager self.max_context_len_to_capture = max_context_len_to_capture if self.max_context_len_to_capture is not None: raise ValueError(\"`max_context_len_to_capture` is deprecated. \" \"Use `max_seq_len_to_capture` instead.\") self.max_seq_len_to_capture = (max_seq_len_to_capture or max_context_len_to_capture) self.max_logprobs = max_logprobs self.skip_tokenizer_init = skip_tokenizer_init # self.hf_config = get_config(model, trust_remote_code, revision) self.hf_config = hf_config self.hf_text_config = get_hf_text_config(hf_config) # TODO: for multimodal model self.dtype = _get_and_verify_dtype(self.hf_config, dtype) self.max_model_len = _get_and_verify_max_len(self.hf_config, max_model_len) # self.served_model_name = get_served_model_name(model, # served_model_name) # self._verify_load_format() # self._verify_tokenizer_mode() self._verify_quantization() self._verify_cuda_graph()","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadFormat","name":"LoadFormat","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/config.py","span":"L147-L154","signature":"class LoadFormat(str, enum.Enum)","exported":true,"origin":"ast","body_hash":"86ea6a79c9ee2abd1dba76e1cbbea5f0cf439e41e0d22f4aa40eae34cca8429b","body_text":"class LoadFormat(str, enum.Enum): AUTO = 'auto' MEGATRON = \"megatron\" HF = \"hf\" DTENSOR = 'dtensor' DUMMY_HF = 'dummy_hf' DUMMY_MEGATRON = 'dummy_megatron' DUMMY_DTENSOR = 'dummy_dtensor'","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig","name":"LoadConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/config.py","span":"L158-L200","signature":"class LoadConfig","exported":true,"origin":"ast","body_hash":"f600abbaffae437a0aa03bd7dc46eddf4c9aa5dd7c50fbd2b44021530e67fcc5","body_text":"class LoadConfig: \"\"\" download_dir: Directory to download and load the weights, default to the default cache directory of huggingface. load_format: The format of the model weights to load: \"auto\" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. \"pt\" will load the weights in the pytorch bin format. \"safetensors\" will load the weights in the safetensors format. \"npcache\" will load the weights in pytorch format and store a numpy cache to speed up the loading. \"dummy\" will initialize the weights with random values, which is mainly for profiling. \"tensorizer\" will use CoreWeave's tensorizer library for fast weight loading. \"\"\" load_format: Union[str, LoadFormat, \"BaseModelLoader\"] = LoadFormat.AUTO download_dir: Optional[str] = None model_loader_extra_config: Optional[Union[str, dict]] = field(default_factory=dict) def __post_init__(self): model_loader_extra_config = self.model_loader_extra_config or {} if isinstance(model_loader_extra_config, str): self.model_loader_extra_config = json.loads(model_loader_extra_config) self._verify_load_format() def _verify_load_format(self) -> None: if not isinstance(self.load_format, str): return load_format = self.load_format.lower() self.load_format = LoadFormat(load_format) rocm_not_supported_load_format: List[str] = [] if is_hip() and load_format in rocm_not_supported_load_format: rocm_supported_load_format = [ f for f in LoadFormat.__members__ if (f not in rocm_not_supported_load_format) ] raise ValueError(f\"load format '{load_format}' is not supported in ROCm. \" f\"Supported load formats are \" f\"{rocm_supported_load_format}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig.__post_init__","name":"__post_init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/config.py","span":"L180-L184","signature":"def __post_init__(self)","exported":false,"origin":"ast","body_hash":"d5238c23b8860df698b7f2ed590be6a658849bcf3ee61618cf5092d9d1ed94a0","body_text":"def __post_init__(self): model_loader_extra_config = self.model_loader_extra_config or {} if isinstance(model_loader_extra_config, str): self.model_loader_extra_config = json.loads(model_loader_extra_config) self._verify_load_format()","summary_state":"pending","summary":null,"crux":null,"owner":"LoadConfig"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig._verify_load_format","name":"_verify_load_format","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/config.py","span":"L186-L200","signature":"def _verify_load_format(self) -> None","exported":false,"origin":"ast","body_hash":"f9f7e30ec352e2c4183434a9e0e078cb7cff2b5e08499624e93b0b4857aff8ab","body_text":"def _verify_load_format(self) -> None: if not isinstance(self.load_format, str): return load_format = self.load_format.lower() self.load_format = LoadFormat(load_format) rocm_not_supported_load_format: List[str] = [] if is_hip() and load_format in rocm_not_supported_load_format: rocm_supported_load_format = [ f for f in LoadFormat.__members__ if (f not in rocm_not_supported_load_format) ] raise ValueError(f\"load format '{load_format}' is not supported in ROCm. \" f\"Supported load formats are \" f\"{rocm_supported_load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"LoadConfig"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"imports","specifier":"enum","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"imports","specifier":"json","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"imports","specifier":"vllm.model_executor.layers.quantization","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"imports","specifier":"vllm.transformers_utils.config","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"calls","name":"get_quantization_config","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig","relation":"extends","name":"ModelConfig","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig.__init__","relation":"calls","name":"get_hf_text_config","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig.__init__","relation":"calls","name":"_get_and_verify_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig.__init__","relation":"calls","name":"_get_and_verify_max_len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig.__init__","relation":"calls","name":"_verify_quantization","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#ModelConfig.__init__","relation":"calls","name":"_verify_cuda_graph","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadFormat","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadFormat","relation":"extends","name":"str","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig","relation":"calls","name":"field","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig.__post_init__","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig.__post_init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig.__post_init__","relation":"calls","name":"loads","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig.__post_init__","relation":"calls","name":"_verify_load_format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py","recvType":"LoadConfig"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig._verify_load_format","file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig._verify_load_format","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig._verify_load_format","relation":"calls","name":"lower","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py","recvType":"LoadFormat"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig._verify_load_format","relation":"calls","name":"LoadFormat","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig._verify_load_format","relation":"calls","name":"is_hip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/config.py#LoadConfig._verify_load_format","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/config.py"}]},"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py":{"size":12378,"mtimeMs":1789275830314.2026,"hash":"8567b6e6a43a709cad9f2ba6d7d3fc34d19c24ba3c9d8c5866d506207aaf64b0","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","name":"dtensor_weight_loaders.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L1-L270","signature":null,"exported":true,"origin":"ast","body_hash":"8567b6e6a43a709cad9f2ba6d7d3fc34d19c24ba3c9d8c5866d506207aaf64b0","chars":12378,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models from typing import Dict, Iterable, Tuple import torch import torch.nn as nn from torch.distributed._tensor import DTensor, Shard, Replicate from vllm.model_executor.layers.linear import * from vllm.model_executor.models import ModelRegistry from vllm.model_executor.model_loader.weight_utils import default_weight_loader __MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__ = { 'GPT2LMHeadModel': gpt2_dtensor_weight_loader, 'LlamaForCausalLM': llama_dtensor_weight_loader, 'LLaMAForCausalLM': llama_dtensor_weight_loader, 'MistralForCausalLM': llama_dtensor_weight_loader, # mistral is the same as llama in vLLM 'InternLMForCausalLM': llama_dtensor_weight_loader, 'AquilaModel': llama_dtensor_weight_loader, 'AquilaForCausalLM': llama_dtensor_weight_loader, 'Phi3ForCausalLM': llama_dtensor_weight_loader, 'GemmaForCausalLM': gemma_dtensor_weight_loader, 'GPTBigCodeForCausalLM': gptbigcode_dtensor_load_weights, 'Starcoder2ForCausalLM': starcoder2_dtensor_load_weights, 'Qwen2ForCausalLM': qwen2_dtensor_weight_loader } # the actor model is .state_dict() # Load dtensor weights # NOTE(sgm): we use per-parameter weight loader in each vllm sub"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","name":"gemma_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L26-L71","signature":"def gemma_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"ec471104a3c35b286365da72de23b6b31e60264e65c561706817ae240ebe4e5c","body_text":"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(vllm_model.named_parameters()) 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 stacked_name = name.replace(shard_name, param_name) # Skip loading extra bias for GPTQ models. if stacked_name.endswith(\".bias\") and stacked_name not in params_dict: continue local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) param = params_dict[stacked_name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) 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 # GemmaRMSNorm is different from Llama's in that it multiplies # (1 + weight) to the output, instead of just weight. if \"norm.weight\" in name: local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) norm_weight = local_loaded_weight + 1.0 param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, norm_weight.to(dtype=param.dtype)) else: 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, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","name":"gptbigcode_dtensor_load_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L74-L86","signature":"def gptbigcode_dtensor_load_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"7494a5e13088c35a8ed94c153c3d247196ed1c1f0c08b38b75cfc6be4046b77e","body_text":"def gptbigcode_dtensor_load_weights(actor_weights: Dict, vllm_model: nn.Module): params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"lm_head.weight\" in name: continue if \".attn.bias\" in name: # Skip attention mask. # NOTE: \"c_attn.bias\" should not be skipped. 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, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","name":"starcoder2_dtensor_load_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L89-L117","signature":"def starcoder2_dtensor_load_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"f7efd018225e19458922d6d47502837ab1afc26da47f5f85e802bd7b89f2a936","body_text":"def starcoder2_dtensor_load_weights(actor_weights: Dict, vllm_model: 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\"), ] params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue for (param_name, weight_name, shard_id) in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in name: continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","name":"llama_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L120-L161","signature":"def llama_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"1736f831264cc292ac39f2e3010167ded6b1db2f01e38512a1834683d157524b","body_text":"def llama_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(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue if (\"rotary_emb.cos_cached\" in name or \"rotary_emb.sin_cached\" in name): # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue # With tie_word_embeddings, we can skip lm_head.weight # The weight might appear unnecessarily in the files if the model is # processed with quantization, LoRA, fine-tuning, etc. if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in name: continue for (param_name, weight_name, shard_id) in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_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, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: # 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, local_loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","name":"qwen2_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L164-L198","signature":"def qwen2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"9b4faed6adddcf37e68aeb25fd9a60eb45af8ef9ef6032e5f1b69683438afb14","body_text":"def qwen2_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(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in name: continue for (param_name, weight_name, shard_id) in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_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, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: # Skip loading extra bias for GPTQ models. if name.endswith(\".bias\") and name not in params_dict: continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gpt2_dtensor_weight_loader","name":"gpt2_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L201-L202","signature":"def gpt2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"e26a0b664180d6dbfeb3ba74e435ba77e0a32726d4fc799cc1282d68edce33d9","body_text":"def gpt2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: pass","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#redistribute_dtensor","name":"redistribute_dtensor","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L205-L215","signature":"def redistribute_dtensor(param_name: str, loaded_weights: DTensor, parallelize_plan: Dict = None)","exported":true,"origin":"ast","body_hash":"8e33f089fd37013b9404ed851cdce475af110e84dc4bc3c6edfbbe7ef9fe5d30","body_text":"def redistribute_dtensor(param_name: str, loaded_weights: DTensor, parallelize_plan: Dict = None): param_name = _process_parameter_names(name=param_name) if parallelize_plan is not None: assert param_name in parallelize_plan.keys(), \\ f\"param name: {param_name} not in parallelize_plan :{parallelize_plan.keys()}\" placement = parallelize_plan[param_name] local_loaded_weights = loaded_weights.redistribute(device_mesh=loaded_weights.device_mesh, placements=placement).to_local() else: local_loaded_weights = loaded_weights.full_tensor() return local_loaded_weights","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#_process_parameter_names","name":"_process_parameter_names","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L218-L231","signature":"def _process_parameter_names(name): # Remove '.weight' if it exists at the end of the string","exported":false,"origin":"ast","body_hash":"074c597f3b319460bbcf2325d017836ddcbf0ce56c58d5d3356ffe01fe6d6547","body_text":"def _process_parameter_names(name): # Remove '.weight' if it exists at the end of the string if name.endswith(\".weight\"): name = name[:-7] # Remove 'model.layers.x.' or 'model.' prefix if \"model.layers\" in name: parts = name.split('.') # Reconstruct the string without 'model.layers.x.' name = '.'.join(parts[3:]) # parts[0] is 'model', parts[1] is 'layers', parts[2] is 'x' elif name.startswith(\"model.\"): name = name[6:] # Remove 'model.' return name","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#load_dtensor_weights","name":"load_dtensor_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L252-L257","signature":"def load_dtensor_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"d614ba13175fcb0a1db90bcd80acdd13cbcddb57eb013b3ce25386d8328656ce","body_text":"def load_dtensor_weights(actor_weights: Dict, vllm_model: nn.Module): weight_loader = _get_model_weight_loader(vllm_model.__class__.__name__) weight_loader(actor_weights, vllm_model) # NOTE(sgm) to reduce peak memory usage, we offload vllm model to cpu # after init, and we need this after sync model weights for in first iter. vllm_model = vllm_model.cuda()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#_get_model_weight_loader","name":"_get_model_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L260-L264","signature":"def _get_model_weight_loader(arch: str)","exported":false,"origin":"ast","body_hash":"2aff4020b0abe87a16fdc5b4ba0f9ed8f437312a1edfce1320e5726f5165c8ae","body_text":"def _get_model_weight_loader(arch: str): if arch in __MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__: return __MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__[arch] raise ValueError(f\"Model architectures {arch} are not supported for now. \" f\"Supported architectures: {__MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__.keys()}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#update_dtensor_weight_loader","name":"update_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","span":"L268-L269","signature":"def update_dtensor_weight_loader()","exported":true,"origin":"ast","body_hash":"3251c12c590cfa8575d836647130786146b1ba578183f484f98d2dac2cce5914","body_text":"def update_dtensor_weight_loader(): pass","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"imports","specifier":"torch.distributed._tensor","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.linear","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.models","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.model_loader.weight_utils","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#gpt2_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#redistribute_dtensor","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"_process_parameter_names","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"to_local","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"redistribute","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"DTensor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"full_tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","recvType":"DTensor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#_process_parameter_names","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"split","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"startswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#load_dtensor_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#load_dtensor_weights","relation":"calls","name":"_get_model_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#load_dtensor_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#load_dtensor_weights","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#_get_model_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py#update_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/dtensor_weight_loaders.py"}]},"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py":{"size":4032,"mtimeMs":1789275830314.2026,"hash":"05fb8e715913bd4e0ac7384bf4e116610639a25700e46c6aafb2fc5df4c81bdc","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","name":"hf_weight_loader.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","span":"L1-L92","signature":null,"exported":true,"origin":"ast","body_hash":"05fb8e715913bd4e0ac7384bf4e116610639a25700e46c6aafb2fc5df4c81bdc","chars":4032,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models from typing import Dict, Union, Optional, Iterable, Tuple import torch import torch.nn as nn from vllm.model_executor.model_loader.utils import set_default_torch_dtype from vllm.model_executor.model_loader.weight_utils import default_weight_loader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#update_hf_weight_loader","name":"update_hf_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","span":"L25-L27","signature":"def update_hf_weight_loader()","exported":true,"origin":"ast","body_hash":"dfd8a23ae8fb6a045ff338e48f7461b621bad2a9c0e6e3c272d0a202f045acfa","body_text":"def update_hf_weight_loader(): from vllm.model_executor.models.gemma import GemmaForCausalLM GemmaForCausalLM.load_weights = gemma_load_weights","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","name":"gemma_load_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","span":"L30-L76","signature":"def gemma_load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]])","exported":true,"origin":"ast","body_hash":"8b5a227288facc2c63197232e3ac07bcd5a2996a116828665d723a162b20f816","body_text":"def gemma_load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): 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()) loaded_params = set() for name, loaded_weight in weights: 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 param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, loaded_weight, 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 # GemmaRMSNorm is different from Llama's in that it multiplies # (1 + weight) to the output, instead of just weight. if \"norm.weight\" in name: norm_weight = loaded_weight + 1.0 # prevent inplace modify actor weights param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, norm_weight) else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight) 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}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","name":"load_hf_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","span":"L79-L91","signature":"def load_hf_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"7d7210564bd93fd001c90777c582c1e6fd88fcf975563b21cf6c14bda243750f","body_text":"def load_hf_weights(actor_weights: Dict, vllm_model: nn.Module): assert isinstance(actor_weights, Dict) with set_default_torch_dtype(next(vllm_model.parameters()).dtype): # TODO vllm_model.load_weights(actor_weights.items()) for _, module in vllm_model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() vllm_model = vllm_model.cuda()","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.utils","file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.weight_utils","file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#update_hf_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","relation":"imports","specifier":"vllm.model_executor.models.gemma","file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"set","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"add","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","recvType":"set"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","recvType":"dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#gemma_load_weights","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"next","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"load_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/hf_weight_loader.py"}]},"verl/third_party/vllm/vllm_v_0_4_2/llm.py":{"size":15147,"mtimeMs":1789275830315.2026,"hash":"8ad6dfe89c1178ebcb764d2cfd9b888d2de1df3ec4a3124935ea6c5bb7a7c4ce","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","name":"llm.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L1-L307","signature":null,"exported":true,"origin":"ast","body_hash":"8ad6dfe89c1178ebcb764d2cfd9b888d2de1df3ec4a3124935ea6c5bb7a7c4ce","chars":15147,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/llm.py from typing import Dict, List, Optional, Tuple, Union from tqdm import tqdm from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from transformers import PretrainedConfig import torch.nn as nn from .arg_utils import EngineArgs from .llm_engine_sp import LLMEngine from vllm.lora.request import LoRARequest from vllm.outputs import RequestOutput from vllm.sampling_params import SamplingParams from vllm.sequence import MultiModalData from vllm.usage.usage_lib import UsageContext from vllm.utils import Counter import torch from torch.nn.utils.rnn import pad_sequence from verl.workers.rollout.tokenizer import HybridEngineBaseTokenizer"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","name":"LLM","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L35-L306","signature":"class LLM","exported":true,"origin":"ast","body_hash":"9b5c30cc4af414ad3c197db02121cd3a299b40699cf06e7b868b3793cbfb3aad","body_text":"class LLM: \"\"\"An LLM for generating texts from given prompts and sampling parameters. This class includes a tokenizer, a language model (possibly distributed across multiple GPUs), and GPU memory space allocated for intermediate states (aka KV cache). Given a batch of prompts and sampling parameters, this class generates texts from the model, using an intelligent batching mechanism and efficient memory management. NOTE: This class is intended to be used for offline inference. For online serving, use the `AsyncLLMEngine` class instead. NOTE: For the comprehensive list of arguments, see `EngineArgs`. Args: model: A HuggingFace Transformers model instance. tokenizer: A HuggingFace Transformers tokenizer instance. tokenizer_mode: The tokenizer mode. \"auto\" will use the fast tokenizer if available, and \"slow\" will always use the slow tokenizer. trust_remote_code: Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer. tensor_parallel_size: The number of GPUs to use for distributed execution with tensor parallelism. dtype: The data type for the model weights and activations. Currently, we support `float32`, `float16`, and `bfloat16`. If `auto`, we use the `torch_dtype` attribute specified in the model config file. However, if the `torch_dtype` in the config is `float32`, we will use `float16` instead. quantization: The method used to quantize the model weights. Currently, we support \"awq\". If None, we assume the model weights are not quantized and use `dtype` to determine the data type of the weights. revision: The specific model version to use. It can be a branch name, a tag name, or a commit id. tokenizer_revision: The specific tokenizer version to use. It can be a branch name, a tag name, or a commit id. seed: The seed to initialize the random number generator for sampling. gpu_memory_utilization: The ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV cache. Higher values will increase the KV cache size and thus improve the model's throughput. However, if the value is too high, it may cause out-of- memory (OOM) errors. swap_space: The size (GiB) of CPU memory per GPU to use as swap space. This can be used for temporarily storing the states of the requests when their `best_of` sampling parameters are larger than 1. If all requests will have `best_of=1`, you can safely set this to 0. Otherwise, too small values may cause out-of-memory (OOM) errors. enforce_eager: Whether to enforce eager execution. If True, we will disable CUDA graph and always execute the model in eager mode. If False, we will use CUDA graph and eager execution in hybrid. max_context_len_to_capture: Maximum context len covered by CUDA graphs. When a sequence has context length larger than this, we fall back to eager mode. disable_custom_all_reduce: See ParallelConfig \"\"\" def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, enforce_eager: bool = False, max_context_len_to_capture: int = None, disable_custom_all_reduce: bool = False, load_format = 'auto', **kwargs, ) -> None: if \"disable_log_stats\" not in kwargs: kwargs[\"disable_log_stats\"] = True engine_args = EngineArgs( model_hf_config=model_hf_config, tensor_parallel_size=tensor_parallel_size, dtype=dtype, quantization=quantization, revision=revision, tokenizer_revision=tokenizer_revision, seed=seed, gpu_memory_utilization=gpu_memory_utilization, swap_space=swap_space, enforce_eager=enforce_eager, max_context_len_to_capture=max_context_len_to_capture, disable_custom_all_reduce=disable_custom_all_reduce, load_format=load_format, **kwargs, ) tokenizer_cls = (PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer) if not isinstance(tokenizer, tokenizer_cls): raise ValueError( f\"Unexpected tokenizer type: {type(tokenizer)}. Must be\" \"one of the following: PreTrainedTokenizer, PreTrainedTokenizerFast, verl.workers.rollout.HybridEngineBaseTokenizer\" ) self.llm_engine = LLMEngine.from_engine_args(model, tokenizer, engine_args) self.request_counter = Counter() def init_cache_engine(self): self.llm_engine.init_cache_engine() def free_cache_engine(self): self.llm_engine.free_cache_engine() def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: return self.llm_engine.tokenizer def set_tokenizer( self, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], ) -> None: self.llm_engine.tokenizer = tokenizer def generate( self, prompts: Optional[Union[str, List[str]]] = None, sampling_params: Optional[Union[SamplingParams, List[SamplingPara","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L89-L135","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, enforce_eager: bool = False, max_context_len_to_capture: int = None, disable_custom_all_reduce: bool = False, load_format = 'auto', **kwargs, ) -> None","exported":false,"origin":"ast","body_hash":"cd4d1ad2207d9374313e50602a343ec6751b15619aa8065c2453f709fab59883","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, enforce_eager: bool = False, max_context_len_to_capture: int = None, disable_custom_all_reduce: bool = False, load_format = 'auto', **kwargs, ) -> None: if \"disable_log_stats\" not in kwargs: kwargs[\"disable_log_stats\"] = True engine_args = EngineArgs( model_hf_config=model_hf_config, tensor_parallel_size=tensor_parallel_size, dtype=dtype, quantization=quantization, revision=revision, tokenizer_revision=tokenizer_revision, seed=seed, gpu_memory_utilization=gpu_memory_utilization, swap_space=swap_space, enforce_eager=enforce_eager, max_context_len_to_capture=max_context_len_to_capture, disable_custom_all_reduce=disable_custom_all_reduce, load_format=load_format, **kwargs, ) tokenizer_cls = (PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer) if not isinstance(tokenizer, tokenizer_cls): raise ValueError( f\"Unexpected tokenizer type: {type(tokenizer)}. Must be\" \"one of the following: PreTrainedTokenizer, PreTrainedTokenizerFast, verl.workers.rollout.HybridEngineBaseTokenizer\" ) self.llm_engine = LLMEngine.from_engine_args(model, tokenizer, engine_args) self.request_counter = Counter()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L137-L138","signature":"def init_cache_engine(self)","exported":true,"origin":"ast","body_hash":"04d68c22194d8053c926485d043932fb2526f0d73a880157acee875717c5937f","body_text":"def init_cache_engine(self): self.llm_engine.init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L140-L141","signature":"def free_cache_engine(self)","exported":true,"origin":"ast","body_hash":"179e8502fa237bb23acb9f06952cf4ae62f3f9010408f63734f8adf6a5163378","body_text":"def free_cache_engine(self): self.llm_engine.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.get_tokenizer","name":"get_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L143-L144","signature":"def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]","exported":true,"origin":"ast","body_hash":"cebad3b0992af1915c8c8a1543b493069d6507f491421f924191ee4b98e60a2f","body_text":"def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: return self.llm_engine.tokenizer","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.set_tokenizer","name":"set_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L146-L150","signature":"def set_tokenizer( self, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], ) -> None","exported":true,"origin":"ast","body_hash":"a9d3cbc026686b5d94c85093062fb18f30a9d4182a5bf7eb9679c456cbc101fa","body_text":"def set_tokenizer( self, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], ) -> None: self.llm_engine.tokenizer = tokenizer","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","name":"generate","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L152-L230","signature":"def generate( self, prompts: Optional[Union[str, List[str]]] = None, sampling_params: Optional[Union[SamplingParams, List[SamplingParams]]] = None, prompt_token_ids: Optional[List[List[int]]] = None, use_tqdm: bool = True, lora_request: Optional[LoRARequest] = None, multi_modal_data: Optional[MultiModalData] = None, ) -> List[RequestOutput]","exported":true,"origin":"ast","body_hash":"cc12349222af253d43bc3c57599dcb1649b16d4e27dce7a5a1aae4df519f26c1","body_text":"def generate( self, prompts: Optional[Union[str, List[str]]] = None, sampling_params: Optional[Union[SamplingParams, List[SamplingParams]]] = None, prompt_token_ids: Optional[List[List[int]]] = None, use_tqdm: bool = True, lora_request: Optional[LoRARequest] = None, multi_modal_data: Optional[MultiModalData] = None, ) -> List[RequestOutput]: \"\"\"Generates the completions for the input prompts. NOTE: This class automatically batches the given prompts, considering the memory constraint. For the best performance, put all of your prompts into a single list and pass it to this method. Args: prompts: A list of prompts to generate completions for. sampling_params: The sampling parameters for text generation. If None, we use the default sampling parameters. When it is a single value, it is applied to every prompt. When it is a list, the list must have the same length as the prompts and it is paired one by one with the prompt. prompt_token_ids: A list of token IDs for the prompts. If None, we use the tokenizer to convert the prompts to token IDs. use_tqdm: Whether to use tqdm to display the progress bar. lora_request: LoRA request to use for generation, if any. multi_modal_data: Multi modal data. Returns: A list of `RequestOutput` objects containing the generated completions in the same order as the input prompts. \"\"\" if prompts is None and prompt_token_ids is None: raise ValueError(\"Either prompts or prompt_token_ids must be \" \"provided.\") if self.llm_engine.model_config.skip_tokenizer_init \\ and prompts is not None: raise ValueError(\"prompts must be None if skip_tokenizer_init \" \"is True\") if isinstance(prompts, str): # Convert a single prompt to a list. prompts = [prompts] if (prompts is not None and prompt_token_ids is not None and len(prompts) != len(prompt_token_ids)): raise ValueError(\"The lengths of prompts and prompt_token_ids \" \"must be the same.\") if prompts is not None: num_requests = len(prompts) else: assert prompt_token_ids is not None num_requests = len(prompt_token_ids) if sampling_params is None: # Use default sampling params. sampling_params = SamplingParams() elif isinstance(sampling_params, list) and len(sampling_params) != num_requests: raise ValueError(\"The lengths of prompts and sampling_params \" \"must be the same.\") if multi_modal_data: multi_modal_data.data = multi_modal_data.data.to(torch.float16) # Add requests to the engine. for i in range(num_requests): prompt = prompts[i] if prompts is not None else None token_ids = None if prompt_token_ids is None else prompt_token_ids[i] if not isinstance(token_ids, list): # NOTE(shengguangming): convert the rollout input into List[str] token_ids = self._pre_process_inputs(token_ids) self._add_request( prompt, sampling_params[i] if isinstance(sampling_params, list) else sampling_params, token_ids, lora_request=lora_request, # Get ith image while maintaining the batch dim. multi_modal_data=MultiModalData(type=multi_modal_data.type, data=multi_modal_data.data[i].unsqueeze(0)) if multi_modal_data else None, ) return self._run_engine(use_tqdm)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._add_request","name":"_add_request","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L232-L246","signature":"def _add_request( self, prompt: Optional[str], sampling_params: SamplingParams, prompt_token_ids: Optional[List[int]], lora_request: Optional[LoRARequest] = None, multi_modal_data: Optional[MultiModalData] = None, ) -> None","exported":false,"origin":"ast","body_hash":"d7d8afe5ee6e1c4480b489b4632f5f6b74025a07b924601576a356231f0ae19c","body_text":"def _add_request( self, prompt: Optional[str], sampling_params: SamplingParams, prompt_token_ids: Optional[List[int]], lora_request: Optional[LoRARequest] = None, multi_modal_data: Optional[MultiModalData] = None, ) -> None: request_id = str(next(self.request_counter)) self.llm_engine.add_request(request_id, prompt, sampling_params, prompt_token_ids, lora_request=lora_request, multi_modal_data=multi_modal_data)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","name":"_run_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L248-L269","signature":"def _run_engine(self, use_tqdm: bool) -> List[RequestOutput]: # Initialize tqdm.","exported":false,"origin":"ast","body_hash":"6d60cc37ae788033d50ebc34560b70894c256736d43beea88514de3b90501dc2","body_text":"def _run_engine(self, use_tqdm: bool) -> List[RequestOutput]: # Initialize tqdm. if use_tqdm: num_requests = self.llm_engine.get_num_unfinished_requests() pbar = tqdm(total=num_requests, desc=\"Processed prompts\", dynamic_ncols=True) # Run the engine. outputs: List[RequestOutput] = [] while self.llm_engine.has_unfinished_requests(): step_outputs = self.llm_engine.step() for output in step_outputs: if output.finished: outputs.append(output) if use_tqdm: pbar.update(1) if use_tqdm: pbar.close() # Sort the outputs by request ID. # This is necessary because some requests may be finished earlier than # its previous requests. outputs = sorted(outputs, key=lambda x: int(x.request_id)) # TODO(shengguangming): maybe we can hack the autoregressive logics without only apply post process for better performance return self._post_process_outputs(outputs)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._pre_process_inputs","name":"_pre_process_inputs","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L273-L278","signature":"def _pre_process_inputs(self, prompt_token_ids: torch.Tensor) -> List[int]: # remove the left padding in the prompt token_id","exported":false,"origin":"ast","body_hash":"fdf721950a6eec9276740f600d40ca1c9190d1b345d858eefb8f2b9f83c0c799","body_text":"def _pre_process_inputs(self, prompt_token_ids: torch.Tensor) -> List[int]: # remove the left padding in the prompt token_id pad_token_id = self.llm_engine.tokenizer.pad_token_id if self.llm_engine.tokenizer.pad_token_id is not None else self.llm_engine.tokenizer.eos_token_id non_pad_index = torch.nonzero(prompt_token_ids != pad_token_id, as_tuple=False)[0][0] token_ids = prompt_token_ids[non_pad_index:].tolist() return token_ids","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","name":"_post_process_outputs","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L281-L300","signature":"def _post_process_outputs(self, request_outputs: List[RequestOutput]) -> Tuple[torch.Tensor, torch.Tensor]","exported":false,"origin":"ast","body_hash":"05b06c04ad38ec323adde48ebace2afb11cf6796e5d16e93e4a46b004820296d","body_text":"def _post_process_outputs(self, request_outputs: List[RequestOutput]) -> Tuple[torch.Tensor, torch.Tensor]: output_token_ids = [] logprobs = [] for request_output in request_outputs: # List[RequestOutput] outputs = request_output.outputs for output in outputs: # List[CompletionOutput], usually len == 1 output_token_ids.append(torch.tensor(output.token_ids)) # TODO(shengguangming): can be optimzied by rewrite the Sampler._get_logprobs() logits logprobs_dicts = output.logprobs if logprobs_dicts is not None: logprob = [] for logprobs_dict, id in zip(logprobs_dicts, output.token_ids): logprob.append(logprobs_dict[id].logprob) logprobs.append(torch.tensor(logprob)) pad_token_id = self.llm_engine.tokenizer.pad_token_id if self.llm_engine.tokenizer.pad_token_id is not None else self.llm_engine.tokenizer.eos_token_id output_token_ids = pad_sequence(output_token_ids, batch_first=True, padding_value=pad_token_id) if len(logprobs) > 0: logprobs = pad_sequence(logprobs, batch_first=True, padding_value=pad_token_id) return output_token_ids, logprobs","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L302-L303","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None","exported":true,"origin":"ast","body_hash":"df77662639c0d1c0008e46c21bdeffd4e84ed28fa0388825031e0577e7d24cf8","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None: self.llm_engine.sync_model_weights(actor_weights=actor_weights, load_format=load_format)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","span":"L305-L306","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"27a73de1e582c115374d023e58aeaf54824c75ca2740c19f48f07d7d67773621","body_text":"def offload_model_weights(self) -> None: self.llm_engine.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"tqdm","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":".arg_utils","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":".llm_engine_sp","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"vllm.outputs","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"vllm.sampling_params","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"vllm.usage.usage_lib","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"torch.nn.utils.rnn","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"imports","specifier":"verl.workers.rollout.tokenizer","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.__init__","relation":"calls","name":"EngineArgs","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.__init__","relation":"calls","name":"type","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.__init__","relation":"calls","name":"from_engine_args","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.__init__","relation":"calls","name":"Counter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.init_cache_engine","relation":"calls","name":"init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.get_tokenizer","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.set_tokenizer","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"SamplingParams","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"_pre_process_inputs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","recvType":"LLM"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"_add_request","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","recvType":"LLM"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"MultiModalData","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.generate","relation":"calls","name":"_run_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","recvType":"LLM"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._add_request","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._add_request","relation":"calls","name":"str","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._add_request","relation":"calls","name":"next","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._add_request","relation":"calls","name":"add_request","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","relation":"calls","name":"get_num_unfinished_requests","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","relation":"calls","name":"tqdm","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","relation":"calls","name":"has_unfinished_requests","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","relation":"calls","name":"step","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","recvType":"sorted"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","relation":"calls","name":"update","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","recvType":"tqdm"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","relation":"calls","name":"close","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","recvType":"tqdm"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","relation":"calls","name":"sorted","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._run_engine","relation":"calls","name":"_post_process_outputs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","recvType":"LLM"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._pre_process_inputs","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._pre_process_inputs","relation":"calls","name":"nonzero","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._pre_process_inputs","relation":"calls","name":"tolist","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","recvType":"pad_sequence"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","relation":"calls","name":"zip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py","recvType":"pad_sequence"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","relation":"calls","name":"pad_sequence","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM._post_process_outputs","relation":"calls","name":"pad_sequence","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm.py#LLM.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm.py"}]},"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py":{"size":12403,"mtimeMs":1789275830315.2026,"hash":"daf7e56923a73e6075ce7e45a648b7d9f2957e26e7c4476be9dec3efb7625556","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","name":"llm_engine_sp.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","span":"L1-L284","signature":null,"exported":true,"origin":"ast","body_hash":"daf7e56923a73e6075ce7e45a648b7d9f2957e26e7c4476be9dec3efb7625556","chars":12403,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/engine/llm_engine.py import torch from typing import Dict, Optional, Union, Type import vllm from vllm.config import (CacheConfig, DecodingConfig, DeviceConfig, LoRAConfig, ParallelConfig, SchedulerConfig, SpeculativeConfig, VisionLanguageConfig) from vllm.core.scheduler import Scheduler from vllm.engine.output_processor.interfaces import (SequenceGroupOutputProcessor) from vllm.engine.output_processor.stop_checker import StopChecker from vllm.executor.executor_base import ExecutorBase from vllm.logger import init_logger from vllm.transformers_utils.detokenizer import Detokenizer from vllm.engine.metrics import StatLogger from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled, usage_message) from vllm.utils import Counter from vllm.engine.llm_engine import _load_generation_config_dict from vllm.engine.llm_engine import LLMEngine import torch.nn as nn from .arg_utils import EngineArgs from .tokenizer import TokenizerGroup from .config import ModelConfig, LoadConfig logger = init_logger(__name__) _LOCAL_LOGGING_INTERVAL_SEC = 5"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine","name":"LLMEngine","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","span":"L43-L283","signature":"class LLMEngine(LLMEngine)","exported":true,"origin":"ast","body_hash":"aaddc29742e0cb683d4c2f5357e2ffd16bad7b9a663a85c0ae7e0e103cdc58bc","body_text":"class LLMEngine(LLMEngine): \"\"\"An LLM engine that receives requests and generates texts. This is the main class for the vLLM engine. It receives requests from clients and generates texts from the LLM. It includes a tokenizer, a language model (possibly distributed across multiple GPUs), and GPU memory space allocated for intermediate states (aka KV cache). This class utilizes iteration-level scheduling and efficient memory management to maximize the serving throughput. The `LLM` class wraps this class for offline batched inference and the `AsyncLLMEngine` class wraps this class for online serving. NOTE: The config arguments are derived from the `EngineArgs` class. For the comprehensive list of arguments, see `EngineArgs`. Args: model: the actor model initialize outside vllm (add for verl) tokenizer: the initialized tokenizer (add for verl) model_config: The configuration related to the LLM model. cache_config: The configuration related to the KV cache memory management. parallel_config: The configuration related to distributed execution. scheduler_config: The configuration related to the request scheduler. distributed_init_method: The initialization method for distributed execution. See `torch.distributed.init_process_group` for details. placement_group: Ray placement group for distributed execution. Required for distributed execution. log_stats: Whether to log statistics. \"\"\" def __init__( self, # NOTE(sgm): first two arguments are added for verl model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, # NOTE(sgm): vllm original arguments model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], speculative_config: Optional[SpeculativeConfig], decoding_config: Optional[DecodingConfig], executor_class: Type[ExecutorBase], log_stats: bool, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, ) -> None: logger.info( \"Initializing an LLM engine (v%s) with config: \" \"model=%r, speculative_config=%r, tokenizer=%r, \" \"skip_tokenizer_init=%s, tokenizer_mode=%s, revision=%s, \" \"tokenizer_revision=%s, trust_remote_code=%s, dtype=%s, \" \"max_seq_len=%d, download_dir=%r, load_format=%s, \" \"tensor_parallel_size=%d, disable_custom_all_reduce=%s, \" \"quantization=%s, enforce_eager=%s, kv_cache_dtype=%s, \" \"quantization_param_path=%s, device_config=%s, \" \"decoding_config=%r, seed=%d, served_model_name=%s)\", vllm.__version__, model_config.model, speculative_config, model_config.tokenizer, model_config.skip_tokenizer_init, # model_config.tokenizer_mode, model_config.revision, model_config.tokenizer_revision, # model_config.trust_remote_code, model_config.dtype, model_config.max_model_len, load_config.download_dir, load_config.load_format, parallel_config.tensor_parallel_size, parallel_config.disable_custom_all_reduce, model_config.quantization, model_config.enforce_eager, cache_config.cache_dtype, model_config.quantization_param_path, device_config.device, decoding_config, model_config.seed, # model_config.served_model_name, ) # TODO(woosuk): Print more configs in debug mode. self.model_config = model_config # TODO: currently is hfconfig self.cache_config = cache_config self.lora_config = lora_config self.vision_language_config = vision_language_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.speculative_config = speculative_config self.load_config = load_config self.decoding_config = decoding_config or DecodingConfig() self.log_stats = log_stats # self.model = model # should not store the model, it should be deleted # TODO(shengguangming): maybe we can choose init here or from arguments if not self.model_config.skip_tokenizer_init: # TODO: check tokenizer class self._init_tokenizer(tokenizer) self.detokenizer = Detokenizer(self.tokenizer) else: self.detokenizer = None self.tokenizer = None self.seq_counter = Counter() # TODO: don't know what's the usage self.generation_config_fields = _load_generation_config_dict(model_config) self.model_executor = executor_class( model=model, # add for spmd_gpu_executor model_config=model_config, cache_config=cache_config, parallel_config=parallel_config, scheduler_config=scheduler_config, device_config=device_config, lora_config=lora_config, vision_language_config=vision_language_config, speculative_config=speculative_config, load_config=load_config, ) # Profile the memory usage and initialize the cache. self._initialize_kv_caches() # If usage stat is enabled, collect relevant info. if is_usage_stats_enabled(): from vllm.model_executor.model_loader import (get_architecture_class_name) usage_message.report_usage( get_architecture_class_name(model_config), usage_context, extra_kvs={ # Common configuration \"dtype\": str(model_config.dtype), \"tensor_parallel_size\": parallel","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","span":"L74-L226","signature":"def __init__( self, # NOTE(sgm): first two arguments are added for verl model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, # NOTE(sgm): vllm original arguments model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], speculative_config: Optional[SpeculativeConfig], decoding_config: Optional[DecodingConfig], executor_class: Type[ExecutorBase], log_stats: bool, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, ) -> None","exported":false,"origin":"ast","body_hash":"d518e0847fa5a0214c49bdde370dfab2915073fa8544a9e250e98603127fc172","body_text":"def __init__( self, # NOTE(sgm): first two arguments are added for verl model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, # NOTE(sgm): vllm original arguments model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], speculative_config: Optional[SpeculativeConfig], decoding_config: Optional[DecodingConfig], executor_class: Type[ExecutorBase], log_stats: bool, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, ) -> None: logger.info( \"Initializing an LLM engine (v%s) with config: \" \"model=%r, speculative_config=%r, tokenizer=%r, \" \"skip_tokenizer_init=%s, tokenizer_mode=%s, revision=%s, \" \"tokenizer_revision=%s, trust_remote_code=%s, dtype=%s, \" \"max_seq_len=%d, download_dir=%r, load_format=%s, \" \"tensor_parallel_size=%d, disable_custom_all_reduce=%s, \" \"quantization=%s, enforce_eager=%s, kv_cache_dtype=%s, \" \"quantization_param_path=%s, device_config=%s, \" \"decoding_config=%r, seed=%d, served_model_name=%s)\", vllm.__version__, model_config.model, speculative_config, model_config.tokenizer, model_config.skip_tokenizer_init, # model_config.tokenizer_mode, model_config.revision, model_config.tokenizer_revision, # model_config.trust_remote_code, model_config.dtype, model_config.max_model_len, load_config.download_dir, load_config.load_format, parallel_config.tensor_parallel_size, parallel_config.disable_custom_all_reduce, model_config.quantization, model_config.enforce_eager, cache_config.cache_dtype, model_config.quantization_param_path, device_config.device, decoding_config, model_config.seed, # model_config.served_model_name, ) # TODO(woosuk): Print more configs in debug mode. self.model_config = model_config # TODO: currently is hfconfig self.cache_config = cache_config self.lora_config = lora_config self.vision_language_config = vision_language_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.speculative_config = speculative_config self.load_config = load_config self.decoding_config = decoding_config or DecodingConfig() self.log_stats = log_stats # self.model = model # should not store the model, it should be deleted # TODO(shengguangming): maybe we can choose init here or from arguments if not self.model_config.skip_tokenizer_init: # TODO: check tokenizer class self._init_tokenizer(tokenizer) self.detokenizer = Detokenizer(self.tokenizer) else: self.detokenizer = None self.tokenizer = None self.seq_counter = Counter() # TODO: don't know what's the usage self.generation_config_fields = _load_generation_config_dict(model_config) self.model_executor = executor_class( model=model, # add for spmd_gpu_executor model_config=model_config, cache_config=cache_config, parallel_config=parallel_config, scheduler_config=scheduler_config, device_config=device_config, lora_config=lora_config, vision_language_config=vision_language_config, speculative_config=speculative_config, load_config=load_config, ) # Profile the memory usage and initialize the cache. self._initialize_kv_caches() # If usage stat is enabled, collect relevant info. if is_usage_stats_enabled(): from vllm.model_executor.model_loader import (get_architecture_class_name) usage_message.report_usage( get_architecture_class_name(model_config), usage_context, extra_kvs={ # Common configuration \"dtype\": str(model_config.dtype), \"tensor_parallel_size\": parallel_config.tensor_parallel_size, \"block_size\": cache_config.block_size, \"gpu_memory_utilization\": cache_config.gpu_memory_utilization, # Quantization \"quantization\": model_config.quantization, \"kv_cache_dtype\": cache_config.cache_dtype, # Feature flags \"enable_lora\": bool(lora_config), \"enable_prefix_caching\": cache_config.enable_prefix_caching, \"enforce_eager\": model_config.enforce_eager, \"disable_custom_all_reduce\": parallel_config.disable_custom_all_reduce, }) if self.tokenizer: # Ping the tokenizer to ensure liveness if it runs in a # different process. self.tokenizer.ping() # Create the scheduler. # NOTE: the cache_config here have been updated with the numbers of # GPU and CPU blocks, which are profiled in the distributed executor. # NOTE(shengguangming): each process will have independent scheduler self.scheduler = Scheduler(scheduler_config, cache_config, lora_config) # Metric Logging. if self.log_stats: self.stat_logger = StatLogger(local_interval=_LOCAL_LOGGING_INTERVAL_SEC, labels=dict(model_name=model_config.served_model_name), max_model_len=self.model_config.max_model_len) self.stat_logger.info(\"cache_config\", self.cache_config) # Create sequence output processor, e.g. for beam search or # speculative decoding. self.output_processor = (SequenceGroupOutputProcessor.create_output_processor( self.scheduler_config, self.detokenizer, self.scheduler, self.seq_counter, self.","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine._init_tokenizer","name":"_init_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","span":"L229-L234","signature":"def _init_tokenizer(self, tokenizer, **tokenizer_init_kwargs)","exported":false,"origin":"ast","body_hash":"012ad286206dcdfe6339e376b29596d7da1ef7041db76fb84ccc7865d8c919fa","body_text":"def _init_tokenizer(self, tokenizer, **tokenizer_init_kwargs): init_kwargs = dict(enable_lora=bool(self.lora_config), max_num_seqs=self.scheduler_config.max_num_seqs, max_input_length=None) init_kwargs.update(tokenizer_init_kwargs) self.tokenizer: TokenizerGroup = TokenizerGroup(tokenizer, **init_kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","span":"L236-L239","signature":"def init_cache_engine(self): # TODO: check whether we should rebuild the CUDAGraph every iter when offload/load KVCache # Re-capture CUDAGraph would be time-consuming","exported":true,"origin":"ast","body_hash":"b8670db02a6509c498d21709df59989b89cdc29c54be45a2684c9a9514906677","body_text":"def init_cache_engine(self): # TODO: check whether we should rebuild the CUDAGraph every iter when offload/load KVCache # Re-capture CUDAGraph would be time-consuming self.model_executor.init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","span":"L241-L242","signature":"def free_cache_engine(self)","exported":true,"origin":"ast","body_hash":"61329183cfad12da6f1cfae074a6cf5526b7ffa477e5571a36f81ab4ce9e6e46","body_text":"def free_cache_engine(self): self.model_executor.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.from_engine_args","name":"from_engine_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","span":"L247-L277","signature":"def from_engine_args( cls, model, tokenizer, engine_args: EngineArgs, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, ) -> \"LLMEngine\"","exported":true,"origin":"ast","body_hash":"974e76ead773d6798bcdc6289722f284cf7b5899283ae98d42e2cc780b9a3625","body_text":"def from_engine_args( cls, model, tokenizer, engine_args: EngineArgs, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, ) -> \"LLMEngine\": \"\"\"Creates an LLM engine from the engine arguments.\"\"\" # Create the engine configs. engine_config = engine_args.create_engine_config() # Initialize the cluster and specify the executor class. assert engine_config.device_config.device_type == \"cuda\", \\ \"Currently, the vllm in verl only support running on GPU\" if engine_config.parallel_config.world_size == 1: engine_config.load_config.load_format = \"dummy_hf\" from .spmd_gpu_executor import SPMDGPUExecutor executor_class = SPMDGPUExecutor # Create the LLM engine. engine = cls( model, tokenizer, **engine_config.to_dict(), executor_class=executor_class, log_stats=not engine_args.disable_log_stats, usage_context=usage_context, ) return engine","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","span":"L279-L280","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None","exported":true,"origin":"ast","body_hash":"3e783d98664a0f1d1e1d5620565f52cce66f192cc38c10123eeb27fe5c8d2776","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None: self.model_executor.sync_model_weights(actor_weights=actor_weights, load_format=load_format)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","span":"L282-L283","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"e77bccf4128d8406c1aae651fad5c639bd60d4b251e0a060f0fdb8196cb26eaf","body_text":"def offload_model_weights(self) -> None: self.model_executor.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.core.scheduler","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.output_processor.interfaces","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.output_processor.stop_checker","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.executor.executor_base","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.transformers_utils.detokenizer","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.metrics","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.usage.usage_lib","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.llm_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.llm_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":".arg_utils","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":".tokenizer","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine","relation":"extends","name":"LLMEngine","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"DecodingConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_init_tokenizer","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Detokenizer","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Counter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_load_generation_config_dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"executor_class","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_initialize_kv_caches","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"is_usage_stats_enabled","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":"vllm.model_executor.model_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"report_usage","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"get_architecture_class_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"str","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"bool","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"ping","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"TokenizerGroup"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Scheduler","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"StatLogger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"StatLogger"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"create_output_processor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"StopChecker","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine._init_tokenizer","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"bool","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"update","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"TokenizerGroup","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.init_cache_engine","relation":"calls","name":"init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"executor_class"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"executor_class"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.from_engine_args","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"create_engine_config","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"EngineArgs"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","relation":"imports","specifier":".spmd_gpu_executor","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"cls","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"to_dict","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"executor_class"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py#LLMEngine.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/llm_engine_sp.py","recvType":"executor_class"}]},"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py":{"size":15726,"mtimeMs":1789275830315.2026,"hash":"8031edaf38d3e80ec53b872b3dae963254d54d49bcbd060db61f2b4037461181","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","name":"megatron_weight_loaders.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L1-L349","signature":null,"exported":true,"origin":"ast","body_hash":"8031edaf38d3e80ec53b872b3dae963254d54d49bcbd060db61f2b4037461181","chars":15726,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models from typing import Dict import torch import torch.nn as nn from vllm.model_executor.layers.linear import * from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding, ParallelLMHead from vllm.model_executor.layers.activation import ScaledActivation from vllm.model_executor.models import ModelRegistry # NOTE(shengguangming): replace the origin weight loader function in the class __LAYER_WEIGHT_MEGATRON_LOADER_REGISTRY__ = { ColumnParallelLinear: parallel_weight_loader, MergedColumnParallelLinear: parallel_weight_loader, QKVParallelLinear: parallel_weight_loader, RowParallelLinear: parallel_weight_loader, VocabParallelEmbedding: parallel_weight_loader, ParallelLMHead: parallel_weight_loader # \"ScaledActivation.weight_loader\": ScaledActivation, # TODO(shengguangming): latest commit in vllm fix awq for this function and add load_weights # \"default_weight_loader\": default_weight_loader } # for layer_class, weight_loader in __LAYER_WEIGHT_MEGATRON_LOADER_REGISTRY__.items(): # # setattr(layer_class, 'megatron_weight_loader', weight_loader) # layer_class.weight_loader = weight_loader __MODEL_MEGATRON_WEIGHT_LOADER_REGISTRY__ = { 'GPT2LMHeadModel': gpt2_weight_loader, 'LlamaForCausalLM': llama_megatron_core_te_weight_loader, # use te backend for open-source megatron 'LLaMAForCausalLM': llama_megatron_core_te_weight_loader, 'MistralForCausalLM': mistral_megatron_weight_loader, } # the actor model is .state_dict() # Load megatron weights # FIXME(shengguangming): the vLLM vocab will pad to 64, which may incur out of bounds # so we need to rewrite the init function of vocab DEFAULT_VOCAB_PADDING_SIZE = 64"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#parallel_weight_loader","name":"parallel_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L27-L34","signature":"def parallel_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None","exported":true,"origin":"ast","body_hash":"0f25a6818304178a092b3b28b97800a19b34e8bb12872e7a788083897c25df2a","body_text":"def parallel_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: \"\"\"Parallel Linear weight loader.\"\"\" assert param.size() == loaded_weight.size( ), 'the parameter size is not align with the loaded weight size, param size: {}, loaded_weight size: {}'.format( param.size(), loaded_weight.size()) assert param.data.dtype == loaded_weight.data.dtype, \"if we want to shared weights, the data type should also be the same\" param.data = loaded_weight.data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#default_weight_loader","name":"default_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L37-L42","signature":"def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None","exported":true,"origin":"ast","body_hash":"27bc592da2cd05ea4807cdcdd1b8feb48351a0a66ebfda38aaa38edd7f0f770b","body_text":"def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: \"\"\"Default weight loader.\"\"\" assert param.size() == loaded_weight.size() assert param.data.dtype == loaded_weight.data.dtype, \"if we want to shared weights, the data type should also be the same\" param.data = loaded_weight.data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#gpt2_weight_loader","name":"gpt2_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L45-L70","signature":"def gpt2_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"d7af0ddffa9a06f7c8b64a06c96f5ce1cbbbd5f2c03114b137268dd5d69048eb","body_text":"def gpt2_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"lm_head.weight\" in name: # GPT-2 ties the weights of the embedding layer and the final # linear layer. continue if \".attn.bias\" in name or \".attn.masked_bias\" in name: # Skip attention mask. # NOTE: \"c_attn.bias\" should not be skipped. continue if not name.startswith(\"transformer.\"): name = \"transformer.\" + name param = params_dict[name] # The HF's GPT-2 implementation uses Conv1D instead of Linear. # Because of this, we need to transpose the weights. # Note(zhuohan): the logic below might break quantized models. for conv1d_weight_name in [\"c_attn\", \"c_proj\", \"c_fc\"]: if conv1d_weight_name not in name: continue if not name.endswith(\".weight\"): continue # TODO: check megatron loaded_weight = loaded_weight.t() weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_weight_loader","name":"llama_megatron_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L73-L82","signature":"def llama_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # NOTE(shengguangming): the megatron llama may have this prefix","exported":true,"origin":"ast","body_hash":"e85f937f0d1b176d431ed32c7fdb4f44e023e98630b81ff16d89c4179292cf97","body_text":"def llama_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","name":"llama_megatron_core_te_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L85-L113","signature":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"6f4f63949dd5590450d23bed267b13c9b246905350e2f2b9fb03a6adfb751030","body_text":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"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.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", 'self_attn.o_proj'), ('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'), ('decoder.final_layernorm', 'model.norm'), ('output_layer', 'lm_head'), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith('.bias') and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader","name":"llama_megatron_core_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L116-L143","signature":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"ebabfc744aea05f3781b4d7fc6869a63a03c65929ca5310a1aa1c2303b787681","body_text":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"self_attention.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", 'self_attn.o_proj'), ( 'input_layernorm', 'input_layernorm', ), ('pre_mlp_layernorm', 'post_attention_layernorm'), ('mlp.linear_fc1', 'mlp.gate_up_proj'), ('mlp.linear_fc2', 'mlp.down_proj'), ('decoder.final_layernorm', 'model.norm'), ('output_layer', 'lm_head'), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith('.bias') and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name","name":"_replace_name","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L146-L166","signature":"def _replace_name(megatron_name, name_mapping)","exported":false,"origin":"ast","body_hash":"0c9ae7affadf31827df88ca4b341b2ef683cd56ad7ed1abc8cc9d9fba2f48f84","body_text":"def _replace_name(megatron_name, name_mapping): for m_name, v_name in name_mapping: if m_name not in megatron_name: continue if 'layers' in megatron_name: # deal with decoder layers megatron_name = megatron_name.replace('decoder', 'model') megatron_name_list = megatron_name.split('.') if 'layer_norm_weight' in megatron_name_list or 'layer_norm_bias' in megatron_name_list: param_name_list = megatron_name_list[:3] param_name_list.append(v_name) param_name = '.'.join(param_name_list) else: param_name_list = megatron_name_list[:3] weight_or_bias = megatron_name_list[-1] param_name_list.append(v_name) param_name_list.append(weight_or_bias) param_name = '.'.join(param_name_list) return param_name else: param_name = megatron_name.replace(m_name, v_name) return param_name","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","name":"llama_megatron_core_te_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L169-L197","signature":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"6f4f63949dd5590450d23bed267b13c9b246905350e2f2b9fb03a6adfb751030","body_text":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"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.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", 'self_attn.o_proj'), ('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'), ('decoder.final_layernorm', 'model.norm'), ('output_layer', 'lm_head'), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith('.bias') and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","name":"llama_megatron_core_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L200-L227","signature":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"ebabfc744aea05f3781b4d7fc6869a63a03c65929ca5310a1aa1c2303b787681","body_text":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"self_attention.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", 'self_attn.o_proj'), ( 'input_layernorm', 'input_layernorm', ), ('pre_mlp_layernorm', 'post_attention_layernorm'), ('mlp.linear_fc1', 'mlp.gate_up_proj'), ('mlp.linear_fc2', 'mlp.down_proj'), ('decoder.final_layernorm', 'model.norm'), ('output_layer', 'lm_head'), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith('.bias') and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name~2","name":"_replace_name","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L230-L250","signature":"def _replace_name(megatron_name, name_mapping)","exported":false,"origin":"ast","body_hash":"0c9ae7affadf31827df88ca4b341b2ef683cd56ad7ed1abc8cc9d9fba2f48f84","body_text":"def _replace_name(megatron_name, name_mapping): for m_name, v_name in name_mapping: if m_name not in megatron_name: continue if 'layers' in megatron_name: # deal with decoder layers megatron_name = megatron_name.replace('decoder', 'model') megatron_name_list = megatron_name.split('.') if 'layer_norm_weight' in megatron_name_list or 'layer_norm_bias' in megatron_name_list: param_name_list = megatron_name_list[:3] param_name_list.append(v_name) param_name = '.'.join(param_name_list) else: param_name_list = megatron_name_list[:3] weight_or_bias = megatron_name_list[-1] param_name_list.append(v_name) param_name_list.append(weight_or_bias) param_name = '.'.join(param_name_list) return param_name else: param_name = megatron_name.replace(m_name, v_name) return param_name","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#mistral_megatron_weight_loader","name":"mistral_megatron_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L253-L262","signature":"def mistral_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # TODO: need to implement a general way to deal with prefix","exported":true,"origin":"ast","body_hash":"da2308149252bdda6caeac91a48a8428bb8f71adb703df0b69a2008805ed02cd","body_text":"def mistral_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # TODO: need to implement a general way to deal with prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#load_megatron_weights","name":"load_megatron_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L290-L295","signature":"def load_megatron_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"308ec912ee64a38977bfc7739ea5f87730fbe66a16585ef28a7e106241f7e5b9","body_text":"def load_megatron_weights(actor_weights: Dict, vllm_model: nn.Module): weight_loader = _get_model_weight_loader(vllm_model.__class__.__name__) weight_loader(actor_weights, vllm_model) # NOTE(sgm) to reduce peak memory usage, we offload vllm model to cpu # after init, and we need this after sync model weights for in first iter. vllm_model = vllm_model.cuda()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_get_model_weight_loader","name":"_get_model_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L298-L302","signature":"def _get_model_weight_loader(arch: str)","exported":false,"origin":"ast","body_hash":"f208266263ddc85ba8ce8716c63c6dcb827fbdab60f6cde510bc807de3ed6ed3","body_text":"def _get_model_weight_loader(arch: str): if arch in __MODEL_MEGATRON_WEIGHT_LOADER_REGISTRY__: return __MODEL_MEGATRON_WEIGHT_LOADER_REGISTRY__[arch] raise ValueError(f\"Model architectures {arch} are not supported for now. \" f\"Supported architectures: {ModelRegistry.get_supported_archs()}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#update_megatron_weight_loader","name":"update_megatron_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L305-L308","signature":"def update_megatron_weight_loader()","exported":true,"origin":"ast","body_hash":"5506b3ead3970687baea9e4b7033107c32cb13245a963493c9a93eacfd5a0845","body_text":"def update_megatron_weight_loader(): for layer_class, weight_loader in __LAYER_WEIGHT_MEGATRON_LOADER_REGISTRY__.items(): layer_class.weight_loader = weight_loader VocabParallelEmbedding.__init__ = vocab_init","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","name":"vocab_init","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","span":"L316-L348","signature":"def vocab_init(self, num_embeddings: int, embedding_dim: int, params_dtype: Optional[torch.dtype] = None, org_num_embeddings: Optional[int] = None, padding_size: int = DEFAULT_VOCAB_PADDING_SIZE)","exported":true,"origin":"ast","body_hash":"522bde8cb981dbade967474fb531553e001ba56c68c425a8ae53c32b21a65abe","body_text":"def vocab_init(self, num_embeddings: int, embedding_dim: int, params_dtype: Optional[torch.dtype] = None, org_num_embeddings: Optional[int] = None, padding_size: int = DEFAULT_VOCAB_PADDING_SIZE): super(VocabParallelEmbedding, self).__init__() # Keep the input dimensions. # TODO (pad to be divided by 4) self.num_embeddings = num_embeddings self.org_vocab_size = org_num_embeddings or num_embeddings # self.num_embeddings_padded = pad_vocab_size(num_embeddings, # padding_size) self.embedding_dim = embedding_dim if params_dtype is None: params_dtype = torch.get_default_dtype() self.tp_size = get_tensor_model_parallel_world_size() # Divide the weight matrix along the vocaburaly dimension. # TODO: remove dependencies from megatron from megatron.core.tensor_parallel.utils import VocabUtility self.vocab_start_index, self.vocab_end_index = (VocabUtility.vocab_range_from_global_vocab_size( self.num_embeddings, get_tensor_model_parallel_rank(), self.tp_size)) self.num_embeddings_per_partition = (self.vocab_end_index - self.vocab_start_index) self.weight = Parameter( torch.empty( self.num_embeddings_per_partition, self.embedding_dim, # device=torch.cuda.current_device(), dtype=params_dtype)) set_weight_attrs(self.weight, {\"parallel_dim\": 0, \"weight_loader\": self.weight_loader})","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.linear","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.vocab_parallel_embedding","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.activation","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.models","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#parallel_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#default_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#default_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#default_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#gpt2_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"startswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"t","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"split","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name~2","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"split","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#mistral_megatron_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#load_megatron_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#load_megatron_weights","relation":"calls","name":"_get_model_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#load_megatron_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#load_megatron_weights","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_get_model_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"get_supported_archs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#update_megatron_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#update_megatron_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","relation":"calls","name":"get_default_dtype","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py","relation":"imports","specifier":"megatron.core.tensor_parallel.utils","file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","relation":"calls","name":"vocab_range_from_global_vocab_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","relation":"calls","name":"Parameter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","relation":"calls","name":"empty","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py#vocab_init","relation":"calls","name":"set_weight_attrs","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/megatron_weight_loaders.py"}]},"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py":{"size":13600,"mtimeMs":1789275830315.2026,"hash":"ec29dcd0785c71e44a3e755ca0d649afc27a98929128f9b76eaa078d76051f13","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","name":"model_loader.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L1-L266","signature":null,"exported":true,"origin":"ast","body_hash":"ec29dcd0785c71e44a3e755ca0d649afc27a98929128f9b76eaa078d76051f13","chars":13600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/model_loader \"\"\"Utilities for selecting and loading models.\"\"\" from typing import Dict, Union, Optional, Iterable, Tuple import torch import torch.nn as nn from transformers import PreTrainedModel from vllm.config import (DeviceConfig, LoRAConfig, ParallelConfig, SchedulerConfig, VisionLanguageConfig) from vllm.model_executor.model_loader import BaseModelLoader from vllm.model_executor.model_loader.loader import _initialize_model from vllm.model_executor.model_loader.utils import set_default_torch_dtype from vllm.distributed.communication_op import tensor_model_parallel_all_gather from .config import ModelConfig, LoadFormat, LoadConfig from .megatron_weight_loaders import load_megatron_weights, update_megatron_weight_loader from .dtensor_weight_loaders import load_dtensor_weights, update_dtensor_weight_loader from .hf_weight_loader import update_hf_weight_loader # FIXME(sgm): hack the _get_logits function in vllm v0.4.2 # as they use ray, the _get_logits result will only need to return to the driver node, # therefore gather is enough. However, we use SPMD instead of a central scheduler, # all_gather is required (aligned with v0.2.6) from vllm.model_executor.layers.logits_processor import LogitsProcessor LogitsProcessor._get_logits = _get_logits"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model","name":"get_model","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L34-L52","signature":"def get_model(actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, load_config: LoadConfig, device_config: DeviceConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig]) -> nn.Module","exported":true,"origin":"ast","body_hash":"b342e3c3d3399c94725bba3fefbc9f0d243e02618c6de676667c3689103af153","body_text":"def get_model(actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, load_config: LoadConfig, device_config: DeviceConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig]) -> nn.Module: loader = get_model_loader(load_config) if load_config.load_format.startswith('dummy'): return loader.load_model(model_config=model_config, device_config=device_config, lora_config=lora_config, vision_language_config=vision_language_config, parallel_config=parallel_config, scheduler_config=scheduler_config) else: return loader.load_model(actor_model=actor_model, model_config=model_config, device_config=device_config, lora_config=lora_config, vision_language_config=vision_language_config, parallel_config=parallel_config, scheduler_config=scheduler_config)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","name":"get_model_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L55-L91","signature":"def get_model_loader(load_config: LoadConfig) -> BaseModelLoader","exported":true,"origin":"ast","body_hash":"5f589ff0ab3f6c557941be091d8e3ac943a019d9ce46843e033c626fc089f67c","body_text":"def get_model_loader(load_config: LoadConfig) -> BaseModelLoader: \"\"\"Get a model loader based on the load format.\"\"\" if isinstance(load_config.load_format, type): return load_config.load_format(load_config) if load_config.load_format == LoadFormat.AUTO: update_megatron_weight_loader() return MegatronLoader(load_config) # NOTE(sgm): change the weight_loader function in runtime if load_config.load_format == LoadFormat.MEGATRON: update_megatron_weight_loader() return MegatronLoader(load_config) if load_config.load_format == LoadFormat.HF: update_hf_weight_loader() return HFLoader(load_config) if load_config.load_format == LoadFormat.DTENSOR: update_dtensor_weight_loader() return DTensorLoader(load_config) if load_config.load_format == LoadFormat.DUMMY_HF: update_hf_weight_loader() return DummyModelLoader(load_config) if load_config.load_format == LoadFormat.DUMMY_MEGATRON: update_megatron_weight_loader() return DummyModelLoader(load_config) if load_config.load_format == LoadFormat.DUMMY_DTENSOR: update_dtensor_weight_loader() return DummyModelLoader(load_config) raise ValueError('load format not supported in verl: {}, only support {} and {}'.format( load_config.load_format, LoadFormat.MEGATRON, LoadFormat.HF))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader","name":"DummyModelLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L94-L112","signature":"class DummyModelLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"8c45acc28337cd3524bac14c8f4409236464069411b10342c65877d7ad0c13d1","body_text":"class DummyModelLoader(BaseModelLoader): \"\"\"Model loader that will set model weights to random values.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def load_model(self, *, model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, vision_language_config) # NOTE(woosuk): For accurate performance evaluation, we assign # random values to the weights. # initialize_dummy_weights(model) return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L97-L101","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"DummyModelLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L103-L112","signature":"def load_model(self, *, model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module","exported":true,"origin":"ast","body_hash":"fc04f5965f9417ae80c3813485bbf9e53b7a4513a13045053295f2581506d578","body_text":"def load_model(self, *, model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, vision_language_config) # NOTE(woosuk): For accurate performance evaluation, we assign # random values to the weights. # initialize_dummy_weights(model) return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"DummyModelLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader","name":"MegatronLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L115-L158","signature":"class MegatronLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"a836098066a711d83e95d139855f792efcc25f584f2504370536cd26853448a9","body_text":"class MegatronLoader(BaseModelLoader): \"\"\"Model loader that can load the model weights from partitioned megatron model.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, vision_language_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_megatron_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_megatron_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L118-L122","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader._get_weights_iterator","name":"_get_weights_iterator","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L124-L131","signature":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model","exported":false,"origin":"ast","body_hash":"694724aa8c3287baa407271e8b4f4dab178819219fb63707f139916d9a3e0340","body_text":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L133-L158","signature":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module","exported":true,"origin":"ast","body_hash":"f98a8e114b1f2ac52e4acdbbb22f4d18000918585809e5d1b3927241ecd2f712","body_text":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, vision_language_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_megatron_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_megatron_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader","name":"HFLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L161-L197","signature":"class HFLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"b66a93690238553251835b6ddaff5b24949d0bf44f8a45aae264f25b5d5178fb","body_text":"class HFLoader(BaseModelLoader): \"\"\"Model loader that can load the model weights from model's full params.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def _get_weights_iterator(self, actor_model: Union[PreTrainedModel, Dict]): if isinstance(actor_model, Dict): return actor_model.items() elif isinstance(actor_model, nn.Module): return dict(actor_model.named_parameters()).items() else: raise ValueError(f'actor model should be Dict or nn.Module, but get {type(actor_model)}') def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): # with torch.device(device_config.device): # NOTE(sgm): init the model in cpu model = _initialize_model(model_config, self.load_config, lora_config, vision_language_config) model.load_weights(self._get_weights_iterator(actor_model)) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L164-L168","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"HFLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader._get_weights_iterator","name":"_get_weights_iterator","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L170-L176","signature":"def _get_weights_iterator(self, actor_model: Union[PreTrainedModel, Dict])","exported":false,"origin":"ast","body_hash":"06a16b0f3b8c8c0cc050f443179f9732fd3c336a501b75e7f33660f421c53568","body_text":"def _get_weights_iterator(self, actor_model: Union[PreTrainedModel, Dict]): if isinstance(actor_model, Dict): return actor_model.items() elif isinstance(actor_model, nn.Module): return dict(actor_model.named_parameters()).items() else: raise ValueError(f'actor model should be Dict or nn.Module, but get {type(actor_model)}')","summary_state":"pending","summary":null,"crux":null,"owner":"HFLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L178-L197","signature":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module","exported":true,"origin":"ast","body_hash":"028c4250266f33dbdf73f8d22523be020440e17e5231229ca99c68d34d266a90","body_text":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): # with torch.device(device_config.device): # NOTE(sgm): init the model in cpu model = _initialize_model(model_config, self.load_config, lora_config, vision_language_config) model.load_weights(self._get_weights_iterator(actor_model)) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"HFLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader","name":"DTensorLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L200-L243","signature":"class DTensorLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"04e159cc81f6b92a284987112e6f52f317d9d6268399dc36bfbe43a2d516a1d1","body_text":"class DTensorLoader(BaseModelLoader): \"\"\"Model loader that can load the model weights from partitioned megatron model.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, vision_language_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_dtensor_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_dtensor_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L203-L207","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"DTensorLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader._get_weights_iterator","name":"_get_weights_iterator","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L209-L216","signature":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model","exported":false,"origin":"ast","body_hash":"694724aa8c3287baa407271e8b4f4dab178819219fb63707f139916d9a3e0340","body_text":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model","summary_state":"pending","summary":null,"crux":null,"owner":"DTensorLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L218-L243","signature":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module","exported":true,"origin":"ast","body_hash":"cabeba1340b3ce26ed880aa9f374f150a0ce0ab7d5836e3e0c39e1b46d571250","body_text":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, vision_language_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_dtensor_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_dtensor_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"DTensorLoader"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#_get_logits","name":"_get_logits","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","span":"L250-L260","signature":"def _get_logits(self, hidden_states: torch.Tensor, embedding: torch.Tensor, embedding_bias: Optional[torch.Tensor]) -> torch.Tensor: # Get the logits for the next tokens.","exported":false,"origin":"ast","body_hash":"d294e846526a09f100321505cbacf5666d45ef972200a51c4e4c25557ec585e0","body_text":"def _get_logits(self, hidden_states: torch.Tensor, embedding: torch.Tensor, embedding_bias: Optional[torch.Tensor]) -> torch.Tensor: # Get the logits for the next tokens. logits = torch.matmul(hidden_states, embedding.t()) if embedding_bias is not None: logits += embedding_bias logits = tensor_model_parallel_all_gather(logits) # Remove paddings in vocab (if any). if logits is not None: logits = logits[:, :self.org_vocab_size] return logits","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.loader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.utils","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":"vllm.distributed.communication_op","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":".megatron_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":".dtensor_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":".hf_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model","relation":"calls","name":"get_model_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model","relation":"calls","name":"startswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"get_model_loader"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"get_model_loader"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"load_format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"LoadConfig"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"update_megatron_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"MegatronLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"update_megatron_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"MegatronLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"update_hf_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"HFLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"update_dtensor_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"DTensorLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"update_hf_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"DummyModelLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"update_megatron_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"DummyModelLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"update_dtensor_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"DummyModelLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#get_model_loader","relation":"calls","name":"format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader._get_weights_iterator","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"load_megatron_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"load_megatron_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader._get_weights_iterator","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"type","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"load_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"_get_weights_iterator","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"HFLoader"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#HFLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader._get_weights_iterator","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"load_dtensor_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"load_dtensor_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#_get_logits","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#_get_logits","relation":"calls","name":"matmul","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#_get_logits","relation":"calls","name":"t","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py#_get_logits","relation":"calls","name":"tensor_model_parallel_all_gather","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py","relation":"imports","specifier":"vllm.model_executor.layers.logits_processor","file":"verl/third_party/vllm/vllm_v_0_4_2/model_loader.py"}]},"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py":{"size":12581,"mtimeMs":1789275830315.2026,"hash":"9a7a637e14192a5253c708ba86584995a4f6fecc9a0ec6935c60957f6ca94750","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","name":"model_runner.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","span":"L1-L282","signature":null,"exported":true,"origin":"ast","body_hash":"9a7a637e14192a5253c708ba86584995a4f6fecc9a0ec6935c60957f6ca94750","chars":12581,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/worker/model_runner.py import torch import torch.nn as nn from enum import IntEnum from typing import Dict, List, Optional, Set, Tuple, Union from vllm.attention import (AttentionMetadata, get_attn_backend) from vllm.config import (DeviceConfig, LoRAConfig, ParallelConfig, SchedulerConfig, VisionLanguageConfig) from vllm.logger import init_logger from vllm.lora.layers import LoRAMapping from vllm.lora.request import LoRARequest from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager from vllm.model_executor import SamplingMetadata from vllm.sequence import (MultiModalData, SamplerOutput, SequenceData, SequenceGroupMetadata) from vllm.utils import (CudaMemoryProfiler, is_hip, is_pin_memory_available) from vllm.worker.model_runner import ModelRunner, CUDAGraphRunner from .model_loader import get_model from .config import ModelConfig, LoadConfig logger = init_logger(__name__) # How batches are constructed."},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#BatchType","name":"BatchType","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","span":"L39-L45","signature":"class BatchType(IntEnum): # Every batch is prefill.","exported":true,"origin":"ast","body_hash":"2936fda8ae00fc0b54ec6c67c640b8e048c86af85365d8366a5dd10212f1a15d","body_text":"class BatchType(IntEnum): # Every batch is prefill. PREFILL = 0 # Every batch is decode. DECODE = 1 # Batch is a mixture of prefill and decode. MIXED = 2","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner","name":"ModelRunner","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","span":"L48-L281","signature":"class ModelRunner(ModelRunner)","exported":true,"origin":"ast","body_hash":"e4775dca4eaacc4fe94a2dab94d71405af6f401e273a10036a542dfde156f298","body_text":"class ModelRunner(ModelRunner): def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", vision_language_config: Optional[VisionLanguageConfig] = None, ): self.model_config = model_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.lora_config = lora_config self.load_config = load_config # model_config can be None in tests/samplers/test_sampler.py. # FIXME(woosuk): This is a hack to make the tests work. Refactor this. self.sliding_window = (model_config.get_sliding_window() if model_config is not None else None) self.device_config = (device_config if device_config is not None else DeviceConfig()) self.device = self.device_config.device # NOTE(sgm): add for verl self.model = model # this will be replaced by get_model() # Set after load_model. self.lora_manager: LRUCacheWorkerLoRAManager = None self.graph_runners: Dict[int, CUDAGraphRunner] = {} self.graph_memory_pool: Optional[Tuple[int, int]] = None # Set during graph capture. self.max_seq_len_to_capture = (self.model_config.max_seq_len_to_capture if self.model_config is not None else 0) self.pin_memory = is_pin_memory_available() self.kv_cache_dtype = kv_cache_dtype self.vision_language_config = vision_language_config self.attn_backend = get_attn_backend(self.model_config.dtype if model_config is not None else None) # Lazy initialization self.block_size: int # Set after initial profiling. # When using CUDA graph, the input block tables must be padded to # max_seq_len_to_capture. However, creating the block table in # Python can be expensive. To optimize this, we cache the block table # in numpy and only copy the actual input content at every iteration. # The shape of the cached block table will be # (max batch size to capture, max context len to capture / block size). self.graph_block_tables: torch.Tensor # Set after initial profiling. # Set if the backend is flashinfer. self.flashinfer_workspace_buffer: torch.Tensor # NOTE(sgm): initialize model using the actor model def load_model(self) -> None: with CudaMemoryProfiler() as m: self.model = get_model(actor_model=self.model, model_config=self.model_config, device_config=self.device_config, lora_config=self.lora_config, load_config=self.load_config, parallel_config=self.parallel_config, scheduler_config=self.scheduler_config, vision_language_config=self.vision_language_config) self.model_memory_usage = m.consumed_memory logger.info(\"Loading model weights took %.4f GB\", self.model_memory_usage / float(2**30)) if self.lora_config: assert hasattr(self.model, \"supported_lora_modules\") and self.model.supported_lora_modules, ( \"Model does not support LoRA\") assert hasattr(self.model, \"embedding_modules\"), \"Model does not have embedding_modules\" assert hasattr(self.model, \"embedding_padding_modules\"), \"Model does not have embedding_padding_modules\" self.lora_manager = LRUCacheWorkerLoRAManager(self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens, self.vocab_size, self.lora_config, self.device, self.model.embedding_modules, self.model.embedding_padding_modules) self.model = self.lora_manager.create_lora_manager(self.model) if self.kv_cache_dtype == \"fp8\" and is_hip(): # Currently scaled KV cache is only enabled on ROCm if self.model_config.quantization_param_path is not None: if callable(getattr(self.model, \"load_kv_cache_scales\", None)): self.model.load_kv_cache_scales(self.model_config.quantization_param_path) else: raise RuntimeError( \"Using FP8 KV cache and scaling factors provided but \" \"model %s does not support loading scaling factors.\", self.model.__class__) else: logger.warning(\"Using FP8 KV cache but no scaling factors \" \"provided. Defaulting to scaling factors of 1.0. \" \"This may lead to less accurate results!\") elif self.model_config.quantization_param_path is not None: logger.warning(\"KV cache scaling factors provided, \" \"but the KV cache data type is not FP8. \" \"KV cache scaling factors will not be used.\") def prepare_input_tensors( self, seq_group_metadata_list: List[SequenceGroupMetadata], ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, SamplingMetadata, Set[LoRARequest], LoRAMapping, torch.Tensor]: # NOTE(sgm): all workers prepare the input in the same way prefill_reqs = [] decode_reqs = [] for seq_group_meta in seq_group_metadata_list: if seq_group_meta.is_prompt: prefill_reqs.append(seq_group_meta) else: decode_reqs.append(seq_group_meta) # Prepare input tensors. ( input_tokens, input_positions, prefill_attn_metadata, seq_lens, query_lens, lora_index_mapping, lora_prompt_mapping, lora_requests, multi_modal_input, slot_mapping, ) = self._prepare_prompt(prefill_reqs) ( decode_input_tokens, decode_input_positions, decode_attn_metadata, decode_lora_index_mapping, decode_lo","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","span":"L50-L102","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", vision_language_config: Optional[VisionLanguageConfig] = None, )","exported":false,"origin":"ast","body_hash":"c0f049d28c76018aa981b2ff6233a53d7cd5d7f7849f66414f178de6223f9642","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", vision_language_config: Optional[VisionLanguageConfig] = None, ): self.model_config = model_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.lora_config = lora_config self.load_config = load_config # model_config can be None in tests/samplers/test_sampler.py. # FIXME(woosuk): This is a hack to make the tests work. Refactor this. self.sliding_window = (model_config.get_sliding_window() if model_config is not None else None) self.device_config = (device_config if device_config is not None else DeviceConfig()) self.device = self.device_config.device # NOTE(sgm): add for verl self.model = model # this will be replaced by get_model() # Set after load_model. self.lora_manager: LRUCacheWorkerLoRAManager = None self.graph_runners: Dict[int, CUDAGraphRunner] = {} self.graph_memory_pool: Optional[Tuple[int, int]] = None # Set during graph capture. self.max_seq_len_to_capture = (self.model_config.max_seq_len_to_capture if self.model_config is not None else 0) self.pin_memory = is_pin_memory_available() self.kv_cache_dtype = kv_cache_dtype self.vision_language_config = vision_language_config self.attn_backend = get_attn_backend(self.model_config.dtype if model_config is not None else None) # Lazy initialization self.block_size: int # Set after initial profiling. # When using CUDA graph, the input block tables must be padded to # max_seq_len_to_capture. However, creating the block table in # Python can be expensive. To optimize this, we cache the block table # in numpy and only copy the actual input content at every iteration. # The shape of the cached block table will be # (max batch size to capture, max context len to capture / block size). self.graph_block_tables: torch.Tensor # Set after initial profiling. # Set if the backend is flashinfer. self.flashinfer_workspace_buffer: torch.Tensor","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","span":"L105-L145","signature":"def load_model(self) -> None","exported":true,"origin":"ast","body_hash":"a2cdeee42e2e15871137993ec5785b0124a6ed0dba62e0139c60103d38fb8576","body_text":"def load_model(self) -> None: with CudaMemoryProfiler() as m: self.model = get_model(actor_model=self.model, model_config=self.model_config, device_config=self.device_config, lora_config=self.lora_config, load_config=self.load_config, parallel_config=self.parallel_config, scheduler_config=self.scheduler_config, vision_language_config=self.vision_language_config) self.model_memory_usage = m.consumed_memory logger.info(\"Loading model weights took %.4f GB\", self.model_memory_usage / float(2**30)) if self.lora_config: assert hasattr(self.model, \"supported_lora_modules\") and self.model.supported_lora_modules, ( \"Model does not support LoRA\") assert hasattr(self.model, \"embedding_modules\"), \"Model does not have embedding_modules\" assert hasattr(self.model, \"embedding_padding_modules\"), \"Model does not have embedding_padding_modules\" self.lora_manager = LRUCacheWorkerLoRAManager(self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens, self.vocab_size, self.lora_config, self.device, self.model.embedding_modules, self.model.embedding_padding_modules) self.model = self.lora_manager.create_lora_manager(self.model) if self.kv_cache_dtype == \"fp8\" and is_hip(): # Currently scaled KV cache is only enabled on ROCm if self.model_config.quantization_param_path is not None: if callable(getattr(self.model, \"load_kv_cache_scales\", None)): self.model.load_kv_cache_scales(self.model_config.quantization_param_path) else: raise RuntimeError( \"Using FP8 KV cache and scaling factors provided but \" \"model %s does not support loading scaling factors.\", self.model.__class__) else: logger.warning(\"Using FP8 KV cache but no scaling factors \" \"provided. Defaulting to scaling factors of 1.0. \" \"This may lead to less accurate results!\") elif self.model_config.quantization_param_path is not None: logger.warning(\"KV cache scaling factors provided, \" \"but the KV cache data type is not FP8. \" \"KV cache scaling factors will not be used.\")","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","name":"prepare_input_tensors","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","span":"L147-L235","signature":"def prepare_input_tensors( self, seq_group_metadata_list: List[SequenceGroupMetadata], ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, SamplingMetadata, Set[LoRARequest], LoRAMapping, torch.Tensor]: # NOTE(sgm): all workers prepare the input in the same way","exported":true,"origin":"ast","body_hash":"35cb4aa86492f67727dba19a8ab065b4fd9b82c0e7fdc14d263864e273c42385","body_text":"def prepare_input_tensors( self, seq_group_metadata_list: List[SequenceGroupMetadata], ) -> Tuple[torch.Tensor, torch.Tensor, AttentionMetadata, SamplingMetadata, Set[LoRARequest], LoRAMapping, torch.Tensor]: # NOTE(sgm): all workers prepare the input in the same way prefill_reqs = [] decode_reqs = [] for seq_group_meta in seq_group_metadata_list: if seq_group_meta.is_prompt: prefill_reqs.append(seq_group_meta) else: decode_reqs.append(seq_group_meta) # Prepare input tensors. ( input_tokens, input_positions, prefill_attn_metadata, seq_lens, query_lens, lora_index_mapping, lora_prompt_mapping, lora_requests, multi_modal_input, slot_mapping, ) = self._prepare_prompt(prefill_reqs) ( decode_input_tokens, decode_input_positions, decode_attn_metadata, decode_lora_index_mapping, decode_lora_prompt_mapping, decode_lora_requests, decode_slot_mapping, ) = self._prepare_decode(decode_reqs) sampling_metadata = SamplingMetadata.prepare(seq_group_metadata_list, seq_lens, query_lens, self.device, self.pin_memory) if not self.scheduler_config.chunked_prefill_enabled: assert (len(prefill_reqs) and len(decode_reqs)) == 0 num_prefills = len(seq_lens) num_prefill_tokens = len(input_tokens) num_decode_tokens = len(decode_input_tokens) # Coalesce tensors. Note that attn_metadata is currently not # coalesced for simplicity. input_tokens.extend(decode_input_tokens) input_positions.extend(decode_input_positions) slot_mapping.extend(decode_slot_mapping) lora_index_mapping.extend(decode_lora_index_mapping) lora_prompt_mapping.extend(decode_lora_prompt_mapping) lora_requests.update(decode_lora_requests) input_tokens = torch.tensor(input_tokens, dtype=torch.long, device=self.device) input_positions = torch.tensor(input_positions, dtype=torch.long, device=self.device) slot_mapping = torch.tensor(slot_mapping, dtype=torch.long, device=self.device) if self.lora_config: lora_mapping = LoRAMapping( lora_index_mapping, lora_prompt_mapping, ) else: lora_mapping = None # Broadcast the metadata. # If batch contains both prefill and decode, it sends 2 broadcasts. # If it only contains 1 type, it triggers a single broadcast. if (prefill_attn_metadata is not None and decode_attn_metadata is not None): batch_type = BatchType.MIXED elif prefill_attn_metadata is not None: batch_type = BatchType.PREFILL else: batch_type = BatchType.DECODE attn_metadata = AttentionMetadata( num_prefills=num_prefills, slot_mapping=slot_mapping, num_prefill_tokens=num_prefill_tokens, num_decode_tokens=num_decode_tokens, prefill_metadata=prefill_attn_metadata, decode_metadata=decode_attn_metadata, kv_cache_dtype=self.kv_cache_dtype, ) return (input_tokens, input_positions, attn_metadata, sampling_metadata, lora_requests, lora_mapping, multi_modal_input)","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.execute_model","name":"execute_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","span":"L238-L281","signature":"def execute_model( self, seq_group_metadata_list: List[SequenceGroupMetadata], kv_caches: List[torch.Tensor], ) -> Optional[SamplerOutput]","exported":true,"origin":"ast","body_hash":"9624e78131091280e6008d2fc1382df030c552f784152ec37ff9ab53828653bb","body_text":"def execute_model( self, seq_group_metadata_list: List[SequenceGroupMetadata], kv_caches: List[torch.Tensor], ) -> Optional[SamplerOutput]: (input_tokens, input_positions, attn_metadata, sampling_metadata, lora_requests, lora_mapping, multi_modal_input) = self.prepare_input_tensors(seq_group_metadata_list) if self.lora_config: self.set_active_loras(lora_requests, lora_mapping) # Currently cuda graph is only supported by the decode phase. prefill_meta = attn_metadata.prefill_metadata decode_meta = attn_metadata.decode_metadata if prefill_meta is None and decode_meta.use_cuda_graph: graph_batch_size = input_tokens.shape[0] model_executable = self.graph_runners[graph_batch_size] else: model_executable = self.model execute_model_kwargs = { \"input_ids\": input_tokens, \"positions\": input_positions, \"kv_caches\": kv_caches, \"attn_metadata\": attn_metadata, } if self.vision_language_config: execute_model_kwargs.update({\"image_input\": multi_modal_input}) hidden_states = model_executable(**execute_model_kwargs) # Compute the logits. logits = self.model.compute_logits(hidden_states, sampling_metadata) # Only perform sampling in the driver worker. # if not self.is_driver_worker: # return None # TODO(sgm): perform sampling on rank 0 # Sample the next token. output = self.model.sample( logits=logits, sampling_metadata=sampling_metadata, ) return output","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"enum","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"vllm.attention","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"vllm.lora.layers","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"vllm.lora.worker_manager","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"vllm.model_executor","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":"vllm.worker.model_runner","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":".model_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#BatchType","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#BatchType","relation":"extends","name":"IntEnum","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner","relation":"extends","name":"ModelRunner","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.__init__","relation":"calls","name":"get_sliding_window","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.__init__","relation":"calls","name":"DeviceConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.__init__","relation":"calls","name":"is_pin_memory_available","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.__init__","relation":"calls","name":"get_attn_backend","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"CudaMemoryProfiler","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"get_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"float","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"LRUCacheWorkerLoRAManager","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"create_lora_manager","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"LRUCacheWorkerLoRAManager"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"is_hip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"callable","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"load_kv_cache_scales","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"get_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.load_model","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"_prepare_prompt","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"_prepare_decode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"prepare","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"extend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"update","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"LoRAMapping","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.prepare_input_tensors","relation":"calls","name":"AttentionMetadata","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner","relation":"calls","name":"inference_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.execute_model","file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.execute_model","relation":"calls","name":"prepare_input_tensors","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.execute_model","relation":"calls","name":"set_active_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.execute_model","relation":"calls","name":"update","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.execute_model","relation":"calls","name":"model_executable","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.execute_model","relation":"calls","name":"compute_logits","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"get_model"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py#ModelRunner.execute_model","relation":"calls","name":"sample","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/model_runner.py","recvType":"get_model"}]},"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py":{"size":12798,"mtimeMs":1789275830315.2026,"hash":"29aa956ffc9b828b1d60fd47975b575c455ff82053fec1c6b7bed9d273040bc7","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","name":"parallel_state.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L1-L295","signature":null,"exported":true,"origin":"ast","body_hash":"29aa956ffc9b828b1d60fd47975b575c455ff82053fec1c6b7bed9d273040bc7","chars":12798,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM team. # Adapted from # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. \"\"\"Model and data parallel groups.\"\"\" import os import torch import torch.distributed from typing import Optional import vllm.distributed.parallel_state as ps import vllm.envs as envs from vllm.logger import init_logger from torch.distributed.device_mesh import init_device_mesh logger = init_logger(__name__) \"\"\" This version is strongly tied with Megatron to implement HybridEngine and weight sharing between vllm and Megatron. - We assume the Megatron tp+dp+pp world is already established before calling this function. \"\"\" # Device mesh for using DTensor _DEVICE_MESH = None # Tensor model parallel group that the current rank belongs to. _TP_DEVICE_GROUP = None _TP_CPU_GROUP = None # This method is for initializing the ParallelGroup when using HybridEngine \"\"\" Device mesh utilities \"\"\" \"\"\" Tensor model parallel utilities \"\"\""},{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","name":"initialize_parallel_state","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L35-L63","signature":"def initialize_parallel_state( distributed_init_method: str = \"env://\", backend: str = \"nccl\", tensor_model_parallel_size: int = 1, num_tp_per_train_tp: int = 1, pipeline_model_parallel_size: int = 1, ): # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573","exported":true,"origin":"ast","body_hash":"a68773165dfe0a095f1b61d726ddffe974fb8435572db06055b0e4d17116e923","body_text":"def initialize_parallel_state( distributed_init_method: str = \"env://\", backend: str = \"nccl\", tensor_model_parallel_size: int = 1, num_tp_per_train_tp: int = 1, pipeline_model_parallel_size: int = 1, ): # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # NOTE(sgm): Modify for verl, Env vars will be set by TORCHRUN. rank = int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) # Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" ps.init_distributed_environment(world_size, rank, distributed_init_method, local_rank, backend) if torch.distributed.get_world_size() > 1: # NOTE: build a sepearate inference group with infer tp & micro dp initialize_model_parallel_for_vllm(tensor_model_parallel_size=tensor_model_parallel_size, num_tensor_model_parallel_groups_per_train_tp=num_tp_per_train_tp) else: initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size, backend)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#ensure_model_parallel_initialized","name":"ensure_model_parallel_initialized","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L66-L89","signature":"def ensure_model_parallel_initialized( tensor_model_parallel_size: int, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None","exported":true,"origin":"ast","body_hash":"43261a5affa6ed31f4dfe49d2f0d326528f30ebeed291b5306fb641ac4577259","body_text":"def ensure_model_parallel_initialized( tensor_model_parallel_size: int, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None: \"\"\"Helper to initialize model parallel groups if they are not initialized, or ensure tensor-parallel and pipeline-parallel sizes are equal to expected values if the model parallel groups are initialized. \"\"\" # get the backend of _DEVICE_WORLD_GROUP backend = backend or torch.distributed.get_backend() if not model_parallel_is_initialized(): initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size, backend) return assert (get_tensor_model_parallel_world_size() == tensor_model_parallel_size), ( \"tensor parallel group already initialized, but of unexpected size: \" f\"{get_tensor_model_parallel_world_size()=} vs. \" f\"{tensor_model_parallel_size=}\") # assert (get_pipeline_model_parallel_world_size( # ) == pipeline_model_parallel_size), ( # \"pipeline parallel group already initialized, but of unexpected size: \" # f\"{get_pipeline_model_parallel_world_size()=} vs. \" # f\"{pipeline_model_parallel_size=}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#model_parallel_is_initialized","name":"model_parallel_is_initialized","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L92-L95","signature":"def model_parallel_is_initialized()","exported":true,"origin":"ast","body_hash":"51f8fa1bdf7cc5d89fbe81e04a15d77c74e462055731e35cb6466802589a3719","body_text":"def model_parallel_is_initialized(): \"\"\"Check if tensor and pipeline parallel groups are initialized.\"\"\" return (ps._TP_DEVICE_GROUP is not None) # and _PIPELINE_MODEL_PARALLEL_GROUP is not None)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","name":"initialize_model_parallel_for_vllm","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L98-L169","signature":"def initialize_model_parallel_for_vllm(tensor_model_parallel_size: int, num_tensor_model_parallel_groups_per_train_tp: int = 1) -> None","exported":true,"origin":"ast","body_hash":"039c4084636359bcce3aa87a42f40c7ec786df5e371f180f3ac7cf03fe3d0363","body_text":"def initialize_model_parallel_for_vllm(tensor_model_parallel_size: int, num_tensor_model_parallel_groups_per_train_tp: int = 1) -> None: from torch.distributed import new_group # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() assert isinstance(tensor_model_parallel_size, int) # assert num_tensor_model_parallel_groups_per_train_tp == 1 and not different_tp_group # assert num_tensor_model_parallel_groups_per_train_tp > 1 and different_tp_group # Build the tensor model-parallel groups. assert ps._TP_DEVICE_GROUP is None, (\"tensor model parallel group is already initialized\") global _TP_DEVICE_GROUP global _TP_CPU_GROUP global _DEVICE_MESH world_size: int = torch.distributed.get_world_size() rank = torch.distributed.get_rank() backend = torch.distributed.get_backend() num_tensor_model_parallel_groups = world_size // tensor_model_parallel_size if num_tensor_model_parallel_groups_per_train_tp == 1: # if tensor_model_parallel_size == train_tensor_parallel_size: # using the same tp group as Megatron/vllm for i in range(num_tensor_model_parallel_groups): ranks = range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) group = torch.distributed.new_group(ranks, backend=backend) cpu_group = torch.distributed.new_group(ranks, backend=\"gloo\") if rank in ranks: _TP_DEVICE_GROUP = group _TP_CPU_GROUP = cpu_group ps._TP_DEVICE_GROUP = group ps._TP_CPU_GROUP = cpu_group # no _MICRO_DATA_PARALLEL_GROUP else: # initialize a micro_dp group and a tp group # assume training tp=4, infer tp=2, then, weight is partitioned as # [1], [2], [3], [4] for training and [1,2], [1,2], [3,4], [3,4] for inference # Build the inference tp groups # train_tp = train_tensor_parallel_size train_tp = num_tensor_model_parallel_groups_per_train_tp * tensor_model_parallel_size # num_tensor_model_parallel_groups_per_train_tp = train_tp // tensor_model_parallel_size assert _TP_DEVICE_GROUP is None, (\"tensor model parallel group is already initialized\") for i in range(num_tensor_model_parallel_groups // num_tensor_model_parallel_groups_per_train_tp): start = train_tp * i end = train_tp * (i + 1) for j in range(num_tensor_model_parallel_groups_per_train_tp): ranks = list(range(start, end, num_tensor_model_parallel_groups_per_train_tp)) for i in range(len(ranks)): ranks[i] += j group = torch.distributed.new_group(ranks) cpu_group = torch.distributed.new_group(ranks, backend='gloo') if rank in ranks: _TP_DEVICE_GROUP = group _TP_CPU_GROUP = cpu_group ps._TP_DEVICE_GROUP = _TP_DEVICE_GROUP ps._TP_CPU_GROUP = cpu_group # Build the pipeline model-parallel groups. # global _PIPELINE_MODEL_PARALLEL_GROUP # global _PIPELINE_GLOBAL_RANKS # assert ps._PIPELINE_MODEL_PARALLEL_GROUP is None, (\"pipeline model parallel group is already initialized\") # ps._PIPELINE_MODEL_PARALLEL_GROUP = mpu.get_pipeline_model_parallel_group() # ps._PIPELINE_GLOBAL_RANKS = mpu.get_pipeline_model_parallel_ranks()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","name":"initialize_model_parallel","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L172-L255","signature":"def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None","exported":true,"origin":"ast","body_hash":"cab35072c63cb9926bde087c486f453ed166c39c1c45598ad2f7ca554912d6b5","body_text":"def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None: \"\"\" NOTE: This method is a hack from the open-sourced version without asertion of world_size = tp * pp Initialize model parallel groups. Arguments: tensor_model_parallel_size: number of GPUs used for tensor model parallelism. pipeline_model_parallel_size: number of GPUs used for pipeline model parallelism. Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize the model pipeline. The present function will create 4 tensor model-parallel groups and 2 pipeline model-parallel groups: 4 tensor model-parallel groups: [g0, g1], [g2, g3], [g4, g5], [g6, g7] 2 pipeline model-parallel groups: [g0, g2, g4, g6], [g1, g3, g5, g7] Note that for efficiency, the caller should make sure adjacent ranks are on the same DGX box. For example if we are using 2 DGX-1 boxes with a total of 16 GPUs, rank 0 to 7 belong to the first box and ranks 8 to 15 belong to the second box. \"\"\" # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() world_size: int = torch.distributed.get_world_size() # get the backend of _DEVICE_WORLD_GROUP backend = backend or torch.distributed.get_backend() # NOTE(sgm) we don't assert world_size == tp * pp # DP is not managed by vllm but by the veRL WorkerGroup num_tensor_model_parallel_groups: int = (world_size // tensor_model_parallel_size) num_pipeline_model_parallel_groups: int = (world_size // pipeline_model_parallel_size) rank = torch.distributed.get_rank() # Build device mesh for TP if num_tensor_model_parallel_groups > 1: device_mesh = init_device_mesh(\"cuda\", (num_tensor_model_parallel_groups, tensor_model_parallel_size), mesh_dim_names=(\"replicate\", \"tp_shard\")) else: device_mesh = init_device_mesh(\"cuda\", (tensor_model_parallel_size,), mesh_dim_names=[\"tp_shard\"]) shard_group = device_mesh.get_group(mesh_dim=\"tp_shard\") # Build the tensor model-parallel groups. global _TP_DEVICE_GROUP, _TP_CPU_GROUP global _DEVICE_MESH assert _TP_DEVICE_GROUP is None, (\"tensor model parallel group is already initialized\") assert _DEVICE_MESH is None, (\"device mesh in vllm is already initialized\") _DEVICE_MESH = device_mesh # for i in range(num_tensor_model_parallel_groups): # ranks = range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) # group = torch.distributed.new_group(ranks, backend=backend) # cpu_group = torch.distributed.new_group(ranks, backend=\"gloo\") # assert torch.distributed.get_process_group_ranks(shard_group) == torch.distributed.get_process_group_ranks(cpu_group) # ranks = torch.distributed.get_process_group_ranks(shard_group) # cpu_group = torch.distributed.new_group(ranks, backend=\"gloo\") # TODO: this will hang # cpu_group = torch.distributed.new_group(, backend=\"gloo\") # if rank == 0: # print(f'rank: {rank}') # print(f'ranks: {ranks}') # print(f'torch.distributed.get_process_group_ranks(shard_group): {torch.distributed.get_process_group_ranks(shard_group)}') # if rank in ranks: _TP_DEVICE_GROUP = shard_group ps._TP_DEVICE_GROUP = _TP_DEVICE_GROUP # ps._TP_CPU_GROUP = cpu_group # TODO: will hang when used with device mesh # TODO: init using device mesh # Build the pipeline model-parallel groups. assert ps._PIPELINE_MODEL_PARALLEL_GROUP is None, (\"pipeline model parallel group is already initialized\") for i in range(num_pipeline_model_parallel_groups): ranks = range(i, world_size, num_pipeline_model_parallel_groups) group = torch.distributed.new_group(ranks, backend=backend) if rank in ranks: ps._PIPELINE_MODEL_PARALLEL_GROUP = group ps._PIPELINE_GLOBAL_RANKS = ranks","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_device_mesh","name":"get_device_mesh","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L263-L265","signature":"def get_device_mesh()","exported":true,"origin":"ast","body_hash":"2ac9c1236bccce8a6f34fa68c3df23bd33dcc07a6041cb0f4ff8d7cff2dcd147","body_text":"def get_device_mesh(): assert _DEVICE_MESH is not None, (\"device mesh is not initialized\") return _DEVICE_MESH","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_group","name":"get_tensor_model_parallel_group","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L273-L276","signature":"def get_tensor_model_parallel_group()","exported":true,"origin":"ast","body_hash":"096ae53877a5693bfa12299a7abf8ba9736f421c899722297c9cd46dfd9e2d88","body_text":"def get_tensor_model_parallel_group(): \"\"\"Get the tensor model parallel group the caller rank belongs to.\"\"\" assert _TP_DEVICE_GROUP is not None, (\"tensor model parallel group is not initialized\") return _TP_DEVICE_GROUP","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_world_size","name":"get_tensor_model_parallel_world_size","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L279-L281","signature":"def get_tensor_model_parallel_world_size()","exported":true,"origin":"ast","body_hash":"c1af0d2aaf0c6f73d3dc5df62d1ec18f42133f1075e5262abc2daf509ff9dc98","body_text":"def get_tensor_model_parallel_world_size(): \"\"\"Return world size for the tensor model parallel group.\"\"\" return torch.distributed.get_world_size(group=get_tensor_model_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_rank","name":"get_tensor_model_parallel_rank","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L284-L286","signature":"def get_tensor_model_parallel_rank()","exported":true,"origin":"ast","body_hash":"7d7e3527eb0fbc1e8f6ca37c5449d0a44e76ded85a0b824cf4f92bda58558364","body_text":"def get_tensor_model_parallel_rank(): \"\"\"Return my rank for the tensor model parallel group.\"\"\" return torch.distributed.get_rank(group=get_tensor_model_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_src_rank","name":"get_tensor_model_parallel_src_rank","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","span":"L289-L294","signature":"def get_tensor_model_parallel_src_rank()","exported":true,"origin":"ast","body_hash":"cfa50af80e4b8151b634ff31e0fe23018462a6f55018e05a8cc9c4918ef4e48e","body_text":"def get_tensor_model_parallel_src_rank(): \"\"\"Calculate the global rank corresponding to the first local rank in the tensor model parallel group.\"\"\" global_rank = torch.distributed.get_rank() local_world_size = get_tensor_model_parallel_world_size() return (global_rank // local_world_size) * local_world_size","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"imports","specifier":"torch.distributed","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"imports","specifier":"torch.distributed.device_mesh","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","relation":"calls","name":"init_distributed_environment","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","relation":"calls","name":"initialize_model_parallel_for_vllm","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_parallel_state","relation":"calls","name":"initialize_model_parallel","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#ensure_model_parallel_initialized","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_backend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"model_parallel_is_initialized","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"initialize_model_parallel","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#model_parallel_is_initialized","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"imports","specifier":"torch.distributed","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_backend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"new_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"new_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"new_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"new_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_backend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","relation":"calls","name":"init_device_mesh","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","relation":"calls","name":"init_device_mesh","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","recvType":"init_device_mesh"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#initialize_model_parallel","relation":"calls","name":"new_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_device_mesh","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_group","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_world_size","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_world_size","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_world_size","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_rank","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_rank","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_src_rank","file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_src_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py#get_tensor_model_parallel_src_rank","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/parallel_state.py"}]},"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py":{"size":8647,"mtimeMs":1789275830316.2026,"hash":"29254e03f754ed706f1265f7ab839dd9c465cd48cbd8549ad6cdf8a4db702a63","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","name":"spmd_gpu_executor.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L1-L219","signature":null,"exported":true,"origin":"ast","body_hash":"29254e03f754ed706f1265f7ab839dd9c465cd48cbd8549ad6cdf8a4db702a63","chars":8647,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/executor/gpu_executor.py import os import socket from typing import Any, Dict, List, Optional, Set, Tuple import torch import vllm.envs as envs from vllm.executor.executor_base import ExecutorBase, ExecutorAsyncBase from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.sequence import SamplerOutput, ExecuteModelRequest from vllm.config import (CacheConfig, DeviceConfig, LoRAConfig, ParallelConfig, SchedulerConfig, SpeculativeConfig, VisionLanguageConfig) from .config import ModelConfig, LoadConfig logger = init_logger(__name__) # TODO(sgm): not implemented async executor yet"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","name":"SPMDGPUExecutor","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L33-L175","signature":"class SPMDGPUExecutor(ExecutorBase)","exported":true,"origin":"ast","body_hash":"31a698f406dc5af7191ad6234cedc1967249770f5c6550307772ebe03ddbc68e","body_text":"class SPMDGPUExecutor(ExecutorBase): \"\"\"SPMD-based multi-GPU executor implementations.\"\"\" def __init__( self, model, # pytorch model itself or its parameter dict model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], speculative_config: Optional[SpeculativeConfig], ) -> None: self.model_config = model_config self.cache_config = cache_config self.lora_config = lora_config self.load_config = load_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.vision_language_config = vision_language_config self.speculative_config = speculative_config distributed_init_method = initialize_cluster(parallel_config) self._init_executor(model, distributed_init_method) # TODO(sgm): verl not support speculative decode now def _init_executor(self, model, distributed_init_method) -> None: assert (not self.speculative_config), \"Speculative decoding not yet supported for multi-GPU backend.\" # Create the parallel worker for each GPU. self._init_workers_sp(model, distributed_init_method) def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker from .worker import Worker # pylint: disable=import-outside-toplevel rank = int(os.getenv(\"RANK\")) local_rank = int(os.getenv(\"LOCAL_RANK\")) print(f'local rank {local_rank}') self.worker = Worker( model, self.model_config, self.parallel_config, self.scheduler_config, self.device_config, self.cache_config, self.load_config, local_rank, rank, distributed_init_method, lora_config=self.lora_config, vision_language_config=self.vision_language_config, ) # NOTE(shengguangming): torch.distributed.init_process_group will be called inside the init_model() self.worker.init_device() self.worker.load_model() def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Determine the number of available KV blocks. This invokes `determine_num_available_blocks` on each worker and takes the min of the results, guaranteeing that the selected cache sizes are compatible with all workers. Returns: - tuple[num_gpu_blocks, num_cpu_blocks] \"\"\" # Get the maximum number of blocks that can be allocated on GPU and CPU. num_blocks = self.worker.determine_num_available_blocks() # NOTE(shengguangming): Now we don't use a shared centralized controler but each process will # have its own scheduler num_gpu_blocks = num_blocks[0] num_cpu_blocks = num_blocks[1] return num_gpu_blocks, num_cpu_blocks def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: \"\"\"Initialize the KV cache in all workers. \"\"\" # NOTE: We log here to avoid multiple logs when number of workers is # greater than one. We could log in the engine, but not all executors # have GPUs. logger.info(\"# GPU blocks: %d, # CPU blocks: %d\", num_gpu_blocks, num_cpu_blocks) self.cache_config.num_gpu_blocks = num_gpu_blocks self.cache_config.num_cpu_blocks = num_cpu_blocks if torch.distributed.get_rank() == 0: print( f'before init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB' ) self.worker.initialize_cache(num_gpu_blocks=num_gpu_blocks, num_cpu_blocks=num_cpu_blocks) if torch.distributed.get_rank() == 0: print( f'after init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB' ) # NOTE(sgm): This will not profile & capture the model(CUDAGraph) when rebuilding KVCache def init_cache_engine(self) -> None: self.worker._init_cache_engine() def free_cache_engine(self) -> None: self.worker.free_cache_engine() def execute_model(self, execute_model_req) -> List[SamplerOutput]: all_outputs = self.worker.execute_model(execute_model_req=execute_model_req) # NOTE(sgm): # Each GPU in vllm under verl has its own spmd_gpu_executor, therefore all GPUs should return the outputs # In vllm with ray, only the driver worker returns the sampling results. return all_outputs def add_lora(self, lora_request: LoRARequest) -> bool: assert lora_request.lora_int_id > 0, \"lora_id must be greater than 0.\" return self.worker.add_lora(lora_request=lora_request) def remove_lora(self, lora_id: int) -> bool: assert lora_id > 0, \"lora_id must be greater than 0.\" return self.worker.remove_lora(lora_id=lora_id) def list_loras(self) -> Set[int]: return self.worker.list_loras() def check_health(self) -> None: # SPMDExecutor will always be healthy as long as # it's running. return # NOTE(sgm): add for verl def offload_model_weights(self) -> None: self.worker.offload_model_weights() def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None: self.worker.sync_model_weights(actor_weights=actor_weights, load_format=load_format)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L36-L60","signature":"def __init__( self, model, # pytorch model itself or its parameter dict model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], speculative_config: Optional[SpeculativeConfig], ) -> None","exported":false,"origin":"ast","body_hash":"b29bdea4a61cc30804fdd87f5e8fad471bf147ccafeb2cb60ee5cdbd63e0dfd1","body_text":"def __init__( self, model, # pytorch model itself or its parameter dict model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], vision_language_config: Optional[VisionLanguageConfig], speculative_config: Optional[SpeculativeConfig], ) -> None: self.model_config = model_config self.cache_config = cache_config self.lora_config = lora_config self.load_config = load_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.vision_language_config = vision_language_config self.speculative_config = speculative_config distributed_init_method = initialize_cluster(parallel_config) self._init_executor(model, distributed_init_method)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_executor","name":"_init_executor","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L63-L67","signature":"def _init_executor(self, model, distributed_init_method) -> None","exported":false,"origin":"ast","body_hash":"250ab123bb88f594c2ad692273d68d9f61a0601c0730510375dc528bc3d4e8d4","body_text":"def _init_executor(self, model, distributed_init_method) -> None: assert (not self.speculative_config), \"Speculative decoding not yet supported for multi-GPU backend.\" # Create the parallel worker for each GPU. self._init_workers_sp(model, distributed_init_method)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","name":"_init_workers_sp","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L69-L95","signature":"def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker","exported":false,"origin":"ast","body_hash":"01f57ef0b4bb1a450228e6370b9f17d70ae66bc3f4a47c3d20cb1fe47f59e918","body_text":"def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker from .worker import Worker # pylint: disable=import-outside-toplevel rank = int(os.getenv(\"RANK\")) local_rank = int(os.getenv(\"LOCAL_RANK\")) print(f'local rank {local_rank}') self.worker = Worker( model, self.model_config, self.parallel_config, self.scheduler_config, self.device_config, self.cache_config, self.load_config, local_rank, rank, distributed_init_method, lora_config=self.lora_config, vision_language_config=self.vision_language_config, ) # NOTE(shengguangming): torch.distributed.init_process_group will be called inside the init_model() self.worker.init_device() self.worker.load_model()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.determine_num_available_blocks","name":"determine_num_available_blocks","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L97-L115","signature":"def determine_num_available_blocks(self) -> Tuple[int, int]","exported":true,"origin":"ast","body_hash":"dbbbc1bd7aa827704263dacd159512c096e5693fe6fa928309c9ed5451aa5514","body_text":"def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Determine the number of available KV blocks. This invokes `determine_num_available_blocks` on each worker and takes the min of the results, guaranteeing that the selected cache sizes are compatible with all workers. Returns: - tuple[num_gpu_blocks, num_cpu_blocks] \"\"\" # Get the maximum number of blocks that can be allocated on GPU and CPU. num_blocks = self.worker.determine_num_available_blocks() # NOTE(shengguangming): Now we don't use a shared centralized controler but each process will # have its own scheduler num_gpu_blocks = num_blocks[0] num_cpu_blocks = num_blocks[1] return num_gpu_blocks, num_cpu_blocks","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","name":"initialize_cache","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L117-L137","signature":"def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None","exported":true,"origin":"ast","body_hash":"1a1bc217eeed8020b0c93c5ea547950d65edd6448ee5de0cfa6514fc704ef066","body_text":"def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: \"\"\"Initialize the KV cache in all workers. \"\"\" # NOTE: We log here to avoid multiple logs when number of workers is # greater than one. We could log in the engine, but not all executors # have GPUs. logger.info(\"# GPU blocks: %d, # CPU blocks: %d\", num_gpu_blocks, num_cpu_blocks) self.cache_config.num_gpu_blocks = num_gpu_blocks self.cache_config.num_cpu_blocks = num_cpu_blocks if torch.distributed.get_rank() == 0: print( f'before init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB' ) self.worker.initialize_cache(num_gpu_blocks=num_gpu_blocks, num_cpu_blocks=num_cpu_blocks) if torch.distributed.get_rank() == 0: print( f'after init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB' )","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L140-L141","signature":"def init_cache_engine(self) -> None","exported":true,"origin":"ast","body_hash":"135998160b42b08111bb719b769019409c044dd0da5bd18e3cf1819bf579782e","body_text":"def init_cache_engine(self) -> None: self.worker._init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L143-L144","signature":"def free_cache_engine(self) -> None","exported":true,"origin":"ast","body_hash":"869ecde616def986458e35749106985c797ce877ae77ec9bab2ddf6490106436","body_text":"def free_cache_engine(self) -> None: self.worker.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.execute_model","name":"execute_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L146-L152","signature":"def execute_model(self, execute_model_req) -> List[SamplerOutput]","exported":true,"origin":"ast","body_hash":"aae188cc605ed5322f94d274292f73007610b372cfc40fda0aa9a0800811b7a5","body_text":"def execute_model(self, execute_model_req) -> List[SamplerOutput]: all_outputs = self.worker.execute_model(execute_model_req=execute_model_req) # NOTE(sgm): # Each GPU in vllm under verl has its own spmd_gpu_executor, therefore all GPUs should return the outputs # In vllm with ray, only the driver worker returns the sampling results. return all_outputs","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.add_lora","name":"add_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L154-L156","signature":"def add_lora(self, lora_request: LoRARequest) -> bool","exported":true,"origin":"ast","body_hash":"22a01766fd3ba5823ac7f31e7164043c28514e0346dceaf1705ede572c20a6c0","body_text":"def add_lora(self, lora_request: LoRARequest) -> bool: assert lora_request.lora_int_id > 0, \"lora_id must be greater than 0.\" return self.worker.add_lora(lora_request=lora_request)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.remove_lora","name":"remove_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L158-L160","signature":"def remove_lora(self, lora_id: int) -> bool","exported":true,"origin":"ast","body_hash":"13816aebed74124d7222df89c015c122a2dace11bbe74355d4b3e2b0d545e089","body_text":"def remove_lora(self, lora_id: int) -> bool: assert lora_id > 0, \"lora_id must be greater than 0.\" return self.worker.remove_lora(lora_id=lora_id)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.list_loras","name":"list_loras","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L162-L163","signature":"def list_loras(self) -> Set[int]","exported":true,"origin":"ast","body_hash":"84207e642ba5b1e08cab7f7c4d2b2014f03cbfebd74f378d55838bda1d0c31b7","body_text":"def list_loras(self) -> Set[int]: return self.worker.list_loras()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.check_health","name":"check_health","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L165-L168","signature":"def check_health(self) -> None: # SPMDExecutor will always be healthy as long as # it's running.","exported":true,"origin":"ast","body_hash":"48ca06f8f34b874ea336f29bb675c8616e07b51643c0ee9e4708d53fc9b3fa9a","body_text":"def check_health(self) -> None: # SPMDExecutor will always be healthy as long as # it's running. return","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L171-L172","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"c346b8cbee5a5fe63905573385e113fb7b59c46e58b26ecc132e331ba728ea49","body_text":"def offload_model_weights(self) -> None: self.worker.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L174-L175","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None","exported":true,"origin":"ast","body_hash":"e7853eafaa9c2dd18668c019f3a34f0b873eb14ebc971798a00af8501b8de047","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None: self.worker.sync_model_weights(actor_weights=actor_weights, load_format=load_format)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#initialize_cluster","name":"initialize_cluster","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L178-L199","signature":"def initialize_cluster( parallel_config: ParallelConfig, engine_use_ray: bool = False, ray_address: Optional[str] = None, ) -> Tuple[str, Optional[None]]","exported":true,"origin":"ast","body_hash":"d737f1830326ce53e8589f097411c5f17acdd325ad8b62f3817f6527312dfe79","body_text":"def initialize_cluster( parallel_config: ParallelConfig, engine_use_ray: bool = False, ray_address: Optional[str] = None, ) -> Tuple[str, Optional[None]]: \"\"\"Initialize the distributed cluster probably with Ray. Args: parallel_config: The configurations for parallel execution. Returns: The `distributed_init_method` is the address for initializing the distributed backend. \"\"\" # Initialize cluster locally. port = get_open_port() # We need to setup the distributed init method to make sure # the distributed megatron code (e.g., get world size) works correctly. # distributed_init_method = f\"tcp://localhost:{port}\" distributed_init_method = 'env://' return distributed_init_method","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#get_open_port","name":"get_open_port","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L202-L205","signature":"def get_open_port()","exported":true,"origin":"ast","body_hash":"d88039216a61e6ec2a0e1328331a7b83f0af92fe7306e2625e8d4f34b557c26e","body_text":"def get_open_port(): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind((\"\", 0)) return s.getsockname()[1]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync","name":"SPMDGPUExecutorAsync","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L209-L218","signature":"class SPMDGPUExecutorAsync(SPMDGPUExecutor, ExecutorAsyncBase)","exported":true,"origin":"ast","body_hash":"a4473dfeed05058503ca8f50986836a978e7445454d17a15c9a1a71dc1c1ae0c","body_text":"class SPMDGPUExecutorAsync(SPMDGPUExecutor, ExecutorAsyncBase): async def execute_model_async(self, execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]: \"\"\"Executes one model step on the given sequences.\"\"\" raise NotImplementedError async def check_health_async(self) -> None: \"\"\"Checks if the executor is healthy. If not, it should raise an exception.\"\"\" self.check_health()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync.execute_model_async","name":"execute_model_async","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L211-L213","signature":"async def execute_model_async(self, execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]","exported":true,"origin":"ast","body_hash":"832d82699df49a6c8509c3c7f505f16afd4d4b985ff4fd27e7e1f5bdf2d96c0c","body_text":"async def execute_model_async(self, execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]: \"\"\"Executes one model step on the given sequences.\"\"\" raise NotImplementedError","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutorAsync"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync.check_health_async","name":"check_health_async","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","span":"L215-L218","signature":"async def check_health_async(self) -> None","exported":true,"origin":"ast","body_hash":"fe8c9a207c0e656691830f7feb9005f0015b9529086d8913552ff6134a65430b","body_text":"async def check_health_async(self) -> None: \"\"\"Checks if the executor is healthy. If not, it should raise an exception.\"\"\" self.check_health()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutorAsync"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":"socket","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.executor.executor_base","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"extends","name":"ExecutorBase","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","relation":"calls","name":"initialize_cluster","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","relation":"calls","name":"_init_executor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"SPMDGPUExecutor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_executor","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_executor","relation":"calls","name":"_init_workers_sp","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"SPMDGPUExecutor"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"imports","specifier":".worker","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"Worker","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"init_device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.determine_num_available_blocks","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.determine_num_available_blocks","relation":"calls","name":"determine_num_available_blocks","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_allocated","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_reserved","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"initialize_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_allocated","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_reserved","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.init_cache_engine","relation":"calls","name":"_init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.execute_model","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.execute_model","relation":"calls","name":"execute_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.add_lora","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.add_lora","relation":"calls","name":"add_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.remove_lora","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.remove_lora","relation":"calls","name":"remove_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.list_loras","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.list_loras","relation":"calls","name":"list_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.check_health","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutor.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#initialize_cluster","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#initialize_cluster","relation":"calls","name":"get_open_port","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#get_open_port","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#get_open_port","relation":"calls","name":"socket","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#get_open_port","relation":"calls","name":"bind","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#get_open_port","relation":"calls","name":"getsockname","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"extends","name":"SPMDGPUExecutor","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"extends","name":"ExecutorAsyncBase","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync.execute_model_async","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync.check_health_async","file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py#SPMDGPUExecutorAsync.check_health_async","relation":"calls","name":"check_health","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/spmd_gpu_executor.py","recvType":"SPMDGPUExecutorAsync"}]},"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py":{"size":3310,"mtimeMs":1789275830316.2026,"hash":"497a401f78f5c164308c4d48df0144e8b0c7a47bd4dfb3212681c8d511fe98bb","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","name":"tokenizer.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","span":"L1-L78","signature":null,"exported":true,"origin":"ast","body_hash":"497a401f78f5c164308c4d48df0144e8b0c7a47bd4dfb3212681c8d511fe98bb","chars":3310,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/tokenizer_group/tokenizer_group.py from typing import List, Optional, Tuple, Union from transformers import (AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast) from vllm.lora.request import LoRARequest from vllm.utils import make_async, LRUCache from vllm.transformers_utils.tokenizers import *"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup","name":"TokenizerGroup","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","span":"L25-L77","signature":"class TokenizerGroup","exported":true,"origin":"ast","body_hash":"c107b245077e329b7be7c53602c2fad707f0231d9dcb1df954ad8280adebefb6","body_text":"class TokenizerGroup: \"\"\"A group of tokenizers that can be used for LoRA adapters.\"\"\" def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int]): self.enable_lora = enable_lora self.max_input_length = max_input_length self.tokenizer = tokenizer self.lora_tokenizers = LRUCache[PreTrainedTokenizer](capacity=max_num_seqs) if enable_lora else None def ping(self) -> bool: \"\"\"Check if the tokenizer group is alive.\"\"\" return True def get_max_input_len(self, lora_request: Optional[LoRARequest] = None) -> Optional[int]: \"\"\"Get the maximum input length for the LoRA request.\"\"\" return self.max_input_length def encode(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = self.get_lora_tokenizer(lora_request) return tokenizer.encode(prompt) async def encode_async(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = await self.get_lora_tokenizer_async(lora_request) return tokenizer.encode(prompt) def get_lora_tokenizer(self, lora_request: Optional[LoRARequest]) -> \"PreTrainedTokenizer\": if not lora_request or not self.enable_lora: return self.tokenizer if lora_request.lora_int_id not in self.lora_tokenizers: # TODO(sgm): the lora tokenizer is also passed, but may be different tokenizer = self.tokenizer # tokenizer = (get_lora_tokenizer( # lora_request, **self.tokenizer_config) or self.tokenizer) self.lora_tokenizers.put(lora_request.lora_int_id, tokenizer) return tokenizer else: return self.lora_tokenizers.get(lora_request.lora_int_id) # FIXME(sgm): for simplicity, we assign the special token here @property def pad_token_id(self): return self.tokenizer.pad_token_id @property def eos_token_id(self): return self.tokenizer.eos_token_id","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","span":"L28-L33","signature":"def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int])","exported":false,"origin":"ast","body_hash":"ac7258b8cb647fbec870bbffdc1633018f85352db29d6a00fd04abe0a3a94186","body_text":"def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int]): self.enable_lora = enable_lora self.max_input_length = max_input_length self.tokenizer = tokenizer self.lora_tokenizers = LRUCache[PreTrainedTokenizer](capacity=max_num_seqs) if enable_lora else None","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.ping","name":"ping","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","span":"L35-L37","signature":"def ping(self) -> bool","exported":true,"origin":"ast","body_hash":"97d31a1c22925222b76fe44142930b66980680e4e167b8a2be973980797507c3","body_text":"def ping(self) -> bool: \"\"\"Check if the tokenizer group is alive.\"\"\" return True","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.get_max_input_len","name":"get_max_input_len","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","span":"L39-L41","signature":"def get_max_input_len(self, lora_request: Optional[LoRARequest] = None) -> Optional[int]","exported":true,"origin":"ast","body_hash":"232098bcc1725540b28685d973367034b497ccefd004f560ceab1bf3cae8ae28","body_text":"def get_max_input_len(self, lora_request: Optional[LoRARequest] = None) -> Optional[int]: \"\"\"Get the maximum input length for the LoRA request.\"\"\" return self.max_input_length","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.encode","name":"encode","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","span":"L43-L48","signature":"def encode(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]","exported":true,"origin":"ast","body_hash":"c7cd151da6e5d1fb514dcc275e13f4c289111810f56df69c36b64c0d3f2cbc9c","body_text":"def encode(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = self.get_lora_tokenizer(lora_request) return tokenizer.encode(prompt)","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.encode_async","name":"encode_async","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","span":"L50-L55","signature":"async def encode_async(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]","exported":true,"origin":"ast","body_hash":"e0ccf679cf9c6ab3aeebd6ff19c82fdfdb9328b3357f16319f00fadca9561589","body_text":"async def encode_async(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = await self.get_lora_tokenizer_async(lora_request) return tokenizer.encode(prompt)","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.get_lora_tokenizer","name":"get_lora_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","span":"L57-L68","signature":"def get_lora_tokenizer(self, lora_request: Optional[LoRARequest]) -> \"PreTrainedTokenizer\"","exported":true,"origin":"ast","body_hash":"66dbb721b6874d42fc9c39c04fd34dd635d82d307c22dcbff76c9464a47c01ca","body_text":"def get_lora_tokenizer(self, lora_request: Optional[LoRARequest]) -> \"PreTrainedTokenizer\": if not lora_request or not self.enable_lora: return self.tokenizer if lora_request.lora_int_id not in self.lora_tokenizers: # TODO(sgm): the lora tokenizer is also passed, but may be different tokenizer = self.tokenizer # tokenizer = (get_lora_tokenizer( # lora_request, **self.tokenizer_config) or self.tokenizer) self.lora_tokenizers.put(lora_request.lora_int_id, tokenizer) return tokenizer else: return self.lora_tokenizers.get(lora_request.lora_int_id)","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.pad_token_id","name":"pad_token_id","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","span":"L72-L73","signature":"def pad_token_id(self)","exported":true,"origin":"ast","body_hash":"09c75a3902b0bb0ca3d7cec97618ef037c84e09c3213f62cd5ee9aa99c2b6669","body_text":"def pad_token_id(self): return self.tokenizer.pad_token_id","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.eos_token_id","name":"eos_token_id","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","span":"L76-L77","signature":"def eos_token_id(self)","exported":true,"origin":"ast","body_hash":"57455769f69da68b62a323b65045e6baf5ecb3dde6a703266578aeeb1bf096a2","body_text":"def eos_token_id(self): return self.tokenizer.eos_token_id","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","relation":"imports","specifier":"vllm.transformers_utils.tokenizers","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.ping","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.get_max_input_len","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.encode","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.encode","relation":"calls","name":"get_lora_tokenizer","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","recvType":"TokenizerGroup"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.encode","relation":"calls","name":"encode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.encode_async","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.encode_async","relation":"calls","name":"get_lora_tokenizer_async","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py","recvType":"TokenizerGroup"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.encode_async","relation":"calls","name":"encode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.get_lora_tokenizer","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.get_lora_tokenizer","relation":"calls","name":"put","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.get_lora_tokenizer","relation":"calls","name":"get","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.pad_token_id","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py#TokenizerGroup.eos_token_id","file":"verl/third_party/vllm/vllm_v_0_4_2/tokenizer.py"}]},"verl/third_party/vllm/vllm_v_0_4_2/worker.py":{"size":13310,"mtimeMs":1789275830316.2026,"hash":"5de27b2ac756955564514b6e0e097b2e7d1c98c93f6a495ffd5ae29a240865a1","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","name":"worker.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L1-L293","signature":null,"exported":true,"origin":"ast","body_hash":"5de27b2ac756955564514b6e0e097b2e7d1c98c93f6a495ffd5ae29a240865a1","chars":13310,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/worker/worker.py \"\"\"A GPU worker class.\"\"\" import os import gc from typing import Dict, List, Tuple, Optional, Union import torch import torch.distributed import torch.nn as nn from vllm.config import (CacheConfig, DeviceConfig, LoRAConfig, ParallelConfig, SchedulerConfig, VisionLanguageConfig) from vllm.model_executor import set_random_seed from vllm.sequence import SamplerOutput, ExecuteModelRequest from vllm.worker.cache_engine import CacheEngine from vllm.distributed.device_communicators import pynccl_utils from vllm.distributed.device_communicators.custom_all_reduce import (init_custom_ar) # TODO(sgm): check why vllm has similar file in vllm.model_executor.parallel_utils.parallel_state from vllm.distributed import get_tensor_model_parallel_cpu_group, init_distributed_environment, get_tensor_model_parallel_group from vllm.worker.worker import Worker, _check_if_gpu_supports_dtype from .model_runner import ModelRunner from .megatron_weight_loaders import load_megatron_weights from .hf_weight_loader import load_hf_weights from .dtensor_weight_loaders import load_dtensor_weights from .parallel_state import (ensure_model_parallel_initialized) from .config import ModelConfig, LoadConfig, LoadFormat"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","name":"Worker","kind":"class","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L42-L254","signature":"class Worker(Worker)","exported":true,"origin":"ast","body_hash":"ee31d3952912b6a30b33de06c1b38c029a8598a9047fe11e3a823258b75a2c98","body_text":"class Worker(Worker): \"\"\"A worker class that executes (a partition of) the model on a GPU. Each worker is associated with a single GPU. The worker is responsible for maintaining the KV cache and executing the model on the GPU. In case of distributed inference, each worker is assigned a partition of the model. \"\"\" def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, local_rank: int, rank: int, distributed_init_method: str, lora_config: Optional[LoRAConfig] = None, vision_language_config: Optional[VisionLanguageConfig] = None, is_driver_worker: bool = False, ) -> None: # self.model = model # will be replaced in the init_model self.model_config = model_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.cache_config = cache_config self.local_rank = local_rank self.rank = rank self.distributed_init_method = distributed_init_method self.lora_config = lora_config self.load_config = load_config self.is_driver_worker = is_driver_worker if self.is_driver_worker: assert self.rank == 0, \"The driver worker must have rank 0.\" self.vision_language_config = vision_language_config if self.vision_language_config: assert not self.lora_config, (\"To be tested: vision language model with LoRA settings.\") self.model_runner = ModelRunner( model, model_config, parallel_config, scheduler_config, device_config, load_config=load_config, lora_config=self.lora_config, kv_cache_dtype=self.cache_config.cache_dtype, vision_language_config=vision_language_config, ) # Uninitialized cache engine. Will be initialized by # init_cache_engine. self.cache_engine: CacheEngine = None self.gpu_cache: List[torch.Tensor] = None # NOTE(sgm): For offloading inference engine params self.cpu_model = None def init_device(self) -> None: if self.device_config.device.type == \"cuda\": # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # NOTE(sgm): Modify for verl, Env vars will be set by TORCHRUN. self.rank = self.rank if self.rank is not None else int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) self.device = torch.device(f\"cuda:{local_rank}\") if self.rank < 0: raise ValueError(\"Invalid or unspecified rank.\") torch.cuda.set_device(self.device) # Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" self.parallel_config.world_size = world_size _check_if_gpu_supports_dtype(self.model_config.dtype) torch.cuda.empty_cache() self.init_gpu_memory = torch.cuda.mem_get_info()[0] else: raise RuntimeError(f\"Not support device type: {self.device_config.device}\") # Initialize the distributed environment. init_worker_distributed_environment(self.parallel_config, self.rank, self.distributed_init_method, self.local_rank) # Set random seed. set_random_seed(self.model_config.seed) # self.model = get_model(actor_model=self.model, model_config=self.model_config) @torch.inference_mode() def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Profiles the peak memory usage of the model to determine how many KV blocks may be allocated without OOMs. The engine will first conduct a profiling of the existing memory usage. Then, it calculate the maximum possible number of GPU and CPU blocks that can be allocated with the remaining free memory. .. tip:: You may limit the usage of GPU memory by adjusting the `gpu_memory_utilization` parameter. \"\"\" # Profile the memory usage of the model and get the maximum number of # cache blocks that can be allocated with the remaining free memory. torch.cuda.empty_cache() # torch.cuda.reset_peak_memory_stats() # Execute a forward pass with dummy inputs to profile the memory usage # of the model. self.model_runner.profile_run() # Calculate the number of blocks that can be allocated with the # profiled peak memory. torch.cuda.synchronize() free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info() peak_memory = total_gpu_memory - free_gpu_memory assert peak_memory > 0, (\"Error in memory profiling. This happens when the GPU memory was \" \"not properly cleaned up before initializing the vLLM instance.\") cache_block_size = self.get_cache_block_size_bytes() # NOTE(sgm) use the remaining memory num_gpu_blocks = int((free_gpu_memory * self.cache_config.gpu_memory_utilization) // cache_block_size) # num_gpu_blocks = int((total_gpu_memory * self.cache_config.gpu_memory_utilization - peak_memory) // ca","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L50-L103","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, local_rank: int, rank: int, distributed_init_method: str, lora_config: Optional[LoRAConfig] = None, vision_language_config: Optional[VisionLanguageConfig] = None, is_driver_worker: bool = False, ) -> None: # self.model = model # will be replaced in the init_model","exported":false,"origin":"ast","body_hash":"e08729c22a1706ad3f87c3a8652579c5329f51bbc19def93e3d6c384eaae6c37","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, local_rank: int, rank: int, distributed_init_method: str, lora_config: Optional[LoRAConfig] = None, vision_language_config: Optional[VisionLanguageConfig] = None, is_driver_worker: bool = False, ) -> None: # self.model = model # will be replaced in the init_model self.model_config = model_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.cache_config = cache_config self.local_rank = local_rank self.rank = rank self.distributed_init_method = distributed_init_method self.lora_config = lora_config self.load_config = load_config self.is_driver_worker = is_driver_worker if self.is_driver_worker: assert self.rank == 0, \"The driver worker must have rank 0.\" self.vision_language_config = vision_language_config if self.vision_language_config: assert not self.lora_config, (\"To be tested: vision language model with LoRA settings.\") self.model_runner = ModelRunner( model, model_config, parallel_config, scheduler_config, device_config, load_config=load_config, lora_config=self.lora_config, kv_cache_dtype=self.cache_config.cache_dtype, vision_language_config=vision_language_config, ) # Uninitialized cache engine. Will be initialized by # init_cache_engine. self.cache_engine: CacheEngine = None self.gpu_cache: List[torch.Tensor] = None # NOTE(sgm): For offloading inference engine params self.cpu_model = None","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","name":"init_device","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L105-L139","signature":"def init_device(self) -> None","exported":true,"origin":"ast","body_hash":"762272295f8ebaa0d8a603527f1c7088967ebd4ff4562d0c715a5656dcddc70b","body_text":"def init_device(self) -> None: if self.device_config.device.type == \"cuda\": # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # NOTE(sgm): Modify for verl, Env vars will be set by TORCHRUN. self.rank = self.rank if self.rank is not None else int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) self.device = torch.device(f\"cuda:{local_rank}\") if self.rank < 0: raise ValueError(\"Invalid or unspecified rank.\") torch.cuda.set_device(self.device) # Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" self.parallel_config.world_size = world_size _check_if_gpu_supports_dtype(self.model_config.dtype) torch.cuda.empty_cache() self.init_gpu_memory = torch.cuda.mem_get_info()[0] else: raise RuntimeError(f\"Not support device type: {self.device_config.device}\") # Initialize the distributed environment. init_worker_distributed_environment(self.parallel_config, self.rank, self.distributed_init_method, self.local_rank) # Set random seed. set_random_seed(self.model_config.seed) # self.model = get_model(actor_model=self.model, model_config=self.model_config)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","name":"determine_num_available_blocks","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L142-L197","signature":"def determine_num_available_blocks(self) -> Tuple[int, int]","exported":true,"origin":"ast","body_hash":"1233b161a20029928b361ad4ad59b985ddd8b5df2b31130501eebe33c16f9b0d","body_text":"def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Profiles the peak memory usage of the model to determine how many KV blocks may be allocated without OOMs. The engine will first conduct a profiling of the existing memory usage. Then, it calculate the maximum possible number of GPU and CPU blocks that can be allocated with the remaining free memory. .. tip:: You may limit the usage of GPU memory by adjusting the `gpu_memory_utilization` parameter. \"\"\" # Profile the memory usage of the model and get the maximum number of # cache blocks that can be allocated with the remaining free memory. torch.cuda.empty_cache() # torch.cuda.reset_peak_memory_stats() # Execute a forward pass with dummy inputs to profile the memory usage # of the model. self.model_runner.profile_run() # Calculate the number of blocks that can be allocated with the # profiled peak memory. torch.cuda.synchronize() free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info() peak_memory = total_gpu_memory - free_gpu_memory assert peak_memory > 0, (\"Error in memory profiling. This happens when the GPU memory was \" \"not properly cleaned up before initializing the vLLM instance.\") cache_block_size = self.get_cache_block_size_bytes() # NOTE(sgm) use the remaining memory num_gpu_blocks = int((free_gpu_memory * self.cache_config.gpu_memory_utilization) // cache_block_size) # num_gpu_blocks = int((total_gpu_memory * self.cache_config.gpu_memory_utilization - peak_memory) // cache_block_size) num_cpu_blocks = int(self.cache_config.swap_space_bytes // cache_block_size) num_gpu_blocks = max(num_gpu_blocks, 0) num_cpu_blocks = max(num_cpu_blocks, 0) if self.model_runner.lora_manager: self.model_runner.remove_all_loras() # NOTE(sgm): Add for verl, synchronize number of blocks with all the rank num_gpu_blocks = torch.tensor([num_gpu_blocks], device='cuda') num_cpu_blocks = torch.tensor([num_cpu_blocks], device='cuda') torch.distributed.all_reduce(num_gpu_blocks, op=torch.distributed.ReduceOp.MIN, group=get_tensor_model_parallel_group()) torch.distributed.all_reduce(num_cpu_blocks, op=torch.distributed.ReduceOp.MIN, group=get_tensor_model_parallel_group()) num_gpu_blocks = num_gpu_blocks.item() num_cpu_blocks = num_cpu_blocks.item() gc.collect() torch.cuda.empty_cache() return num_gpu_blocks, num_cpu_blocks","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker._init_cache_engine","name":"_init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L199-L201","signature":"def _init_cache_engine(self)","exported":false,"origin":"ast","body_hash":"ddbf730d543fa64f637055bf2b4047182f388f84f7375434e0765a7cd602b012","body_text":"def _init_cache_engine(self): if self.cache_engine is None and self.gpu_cache is None: super()._init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L203-L206","signature":"def free_cache_engine(self): # ensure `enforce_eager=True`","exported":true,"origin":"ast","body_hash":"d384948b80b2c04ee44319bcd88c9b941e3967b68e2050a10073ade04af52223","body_text":"def free_cache_engine(self): # ensure `enforce_eager=True` self.cache_engine = None self.gpu_cache = None","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.execute_model","name":"execute_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L209-L234","signature":"def execute_model(self, execute_model_req: Optional[ExecuteModelRequest] = None) -> List[SamplerOutput]","exported":true,"origin":"ast","body_hash":"c7343a5088911cba6c0576b11b347e447cbf260e7e2c6e7468877c123f7ef914","body_text":"def execute_model(self, execute_model_req: Optional[ExecuteModelRequest] = None) -> List[SamplerOutput]: if execute_model_req is None: seq_group_metadata_list = None else: seq_group_metadata_list = execute_model_req.seq_group_metadata_list # NOTE(sgm): each SPMD rank will have identical input assert seq_group_metadata_list is not None assert execute_model_req is not None num_seq_groups = len(seq_group_metadata_list) blocks_to_swap_in = execute_model_req.blocks_to_swap_in blocks_to_swap_out = execute_model_req.blocks_to_swap_out blocks_to_copy = execute_model_req.blocks_to_copy self.cache_swap(blocks_to_swap_in, blocks_to_swap_out, blocks_to_copy) # If there is no input, we don't need to execute the model. if num_seq_groups == 0: return [] output = self.model_runner.execute_model(seq_group_metadata_list, self.gpu_cache) # Worker only supports single-step execution. Wrap the output in a list # to conform to interface. return [output]","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L237-L244","signature":"def sync_model_weights(self, actor_weights: Dict, load_format: str)","exported":true,"origin":"ast","body_hash":"61f049debd75c5e97bb95de1fbaf9d769e14978f2c8fd9b2ab6143756010d1d1","body_text":"def sync_model_weights(self, actor_weights: Dict, load_format: str): if load_format in [LoadFormat.MEGATRON, LoadFormat.AUTO]: load_megatron_weights(actor_weights, self.model_runner.model) elif load_format == LoadFormat.HF: # full model state dict without no sharding load_hf_weights(actor_weights, self.model_runner.model) elif load_format == LoadFormat.DTENSOR: load_dtensor_weights(actor_weights, self.model_runner.model)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L246-L254","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"58d8d5e64a0765436d968d3a9641e3432be0a43a9ca286f99b1f827937e19e31","body_text":"def offload_model_weights(self) -> None: if self.cpu_model == None: self.cpu_model = {} for name, params in self.model_runner.model.named_parameters(): self.cpu_model[name] = torch.empty_like(params, device='cpu') params.data = self.cpu_model[name] else: for name, params in self.model_runner.model.named_parameters(): params.data = self.cpu_model[name]","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#init_worker_distributed_environment","name":"init_worker_distributed_environment","kind":"function","path":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","span":"L257-L292","signature":"def init_worker_distributed_environment( parallel_config: ParallelConfig, rank: int, distributed_init_method: Optional[str] = \"env://\", local_rank: int = -1, ) -> None","exported":true,"origin":"ast","body_hash":"e22c29f3cd53bbef8f4b049f5bcb2f631f04a33dbda0e1347cf3a1f070f5f87b","body_text":"def init_worker_distributed_environment( parallel_config: ParallelConfig, rank: int, distributed_init_method: Optional[str] = \"env://\", local_rank: int = -1, ) -> None: \"\"\"Initialize the distributed environment.\"\"\" # NOTE(sgm) use tcp://localhost:xxxx will hang in HF setting without megatron init_distributed_environment(parallel_config.world_size, rank, distributed_init_method, local_rank) ensure_model_parallel_initialized(tensor_model_parallel_size=parallel_config.tensor_parallel_size, pipeline_model_parallel_size=parallel_config.pipeline_parallel_size) # TODO(sgm): check whether need this # if pynccl_utils.is_initialized(): # pynccl_world_size = pynccl_utils.get_world_size() # if pynccl_world_size != parallel_config.world_size: # raise RuntimeError( # \"pynccl is already initialized but the pynccl world \" # \"size does not match parallel_config.world_size \" # f\"({pynccl_world_size} vs. {parallel_config.world_size}).\") # elif parallel_config.world_size > 1: # # NOTE(woosuk): We don't initialize pynccl process group when world size # # is 1. # # NOTE(kaichao): By default, pynccl is initialized for tp group. # pynccl_utils.init_process_group( # group=get_tensor_model_parallel_cpu_group()) # # Initialize a custom fast all-reduce implementation. # if not parallel_config.disable_custom_all_reduce: # init_custom_ar() # A small all_reduce for warmup. torch.distributed.all_reduce(torch.zeros(1).cuda()) # if pynccl_utils.is_initialized(): # pynccl_utils.all_reduce(torch.zeros(1).cuda())","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"gc","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"torch.distributed","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"vllm.model_executor","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"vllm.worker.cache_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"vllm.distributed.device_communicators","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"vllm.distributed.device_communicators.custom_all_reduce","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"vllm.distributed","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":"vllm.worker.worker","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":".model_runner","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":".megatron_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":".hf_weight_loader","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":".dtensor_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":".parallel_state","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"extends","name":"Worker","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.__init__","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.__init__","relation":"calls","name":"ModelRunner","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"set_device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"_check_if_gpu_supports_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"mem_get_info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"init_worker_distributed_environment","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.init_device","relation":"calls","name":"set_random_seed","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"calls","name":"inference_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"profile_run","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"synchronize","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"mem_get_info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"get_cache_block_size_bytes","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"remove_all_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"item","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","recvType":"max"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"item","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","recvType":"max"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"collect","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker._init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker._init_cache_engine","relation":"calls","name":"_init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker._init_cache_engine","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"calls","name":"inference_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.execute_model","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.execute_model","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.execute_model","relation":"calls","name":"cache_swap","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.execute_model","relation":"calls","name":"execute_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","recvType":"ModelRunner"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.sync_model_weights","relation":"calls","name":"load_megatron_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.sync_model_weights","relation":"calls","name":"load_hf_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.sync_model_weights","relation":"calls","name":"load_dtensor_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.offload_model_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.offload_model_weights","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#Worker.offload_model_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#init_worker_distributed_environment","file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#init_worker_distributed_environment","relation":"calls","name":"init_distributed_environment","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#init_worker_distributed_environment","relation":"calls","name":"ensure_model_parallel_initialized","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#init_worker_distributed_environment","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#init_worker_distributed_environment","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_4_2/worker.py#init_worker_distributed_environment","relation":"calls","name":"zeros","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_4_2/worker.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/__init__.py":{"size":600,"mtimeMs":1789275830316.2026,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/__init__.py","name":"__init__.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py":{"size":23746,"mtimeMs":1789275830316.2026,"hash":"f4fe86d839cefc374e4606f13ce53a7f224e13c2221bca80d775eabb6c03990c","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","name":"arg_utils.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","span":"L1-L454","signature":null,"exported":true,"origin":"ast","body_hash":"f4fe86d839cefc374e4606f13ce53a7f224e13c2221bca80d775eabb6c03990c","chars":23746,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/engine/arg_utils.py import os import argparse import dataclasses import json from dataclasses import dataclass from typing import TYPE_CHECKING, List, Optional, Tuple, Type, Union import torch.nn as nn from transformers import PretrainedConfig from .config import ModelConfig, LoadConfig from vllm.config import (CacheConfig, DecodingConfig, DeviceConfig, EngineConfig, LoRAConfig, MultiModalConfig, ObservabilityConfig, ParallelConfig, PromptAdapterConfig, SchedulerConfig, SpeculativeConfig, TokenizerPoolConfig) from vllm.executor.executor_base import ExecutorBase from vllm.logger import init_logger from vllm.utils import FlexibleArgumentParser from vllm.model_executor.layers.quantization import QUANTIZATION_METHODS from vllm.utils import str_to_int_tuple if TYPE_CHECKING: from vllm.transformers_utils.tokenizer_group.base_tokenizer_group import (BaseTokenizerGroup) logger = init_logger(__name__) @dataclass"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#nullable_str","name":"nullable_str","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","span":"L43-L46","signature":"def nullable_str(val: str)","exported":true,"origin":"ast","body_hash":"900c08104995d0c284480fe8066923076c935ca69209e3a86885fbcc330eac6b","body_text":"def nullable_str(val: str): if not val or val == \"None\": return None return val","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs","name":"EngineArgs","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","span":"L50-L453","signature":"class EngineArgs","exported":true,"origin":"ast","body_hash":"081dac2fd7648cd5406ef24ac35f445b9e5b157be035844b0cf6a7be81848942","body_text":"class EngineArgs: \"\"\"Arguments for vLLM engine.\"\"\" model_hf_config: PretrainedConfig = None # for verl served_model_name = None # TODO(sgm): check this # tokenizer: Optional[str] = None # TODO(sgm): check this skip_tokenizer_init: bool = False tokenizer_mode: str = 'auto' trust_remote_code: bool = False download_dir: Optional[str] = None load_format: str = 'auto' dtype: str = 'auto' kv_cache_dtype: str = 'auto' quantization_param_path: Optional[str] = None seed: int = 0 max_model_len: Optional[int] = None worker_use_ray: bool = False # Note: Specifying a custom executor backend by passing a class # is intended for expert use only. The API may change without # notice. distributed_executor_backend: Optional[Union[str, Type[ExecutorBase]]] = None pipeline_parallel_size: int = 1 tensor_parallel_size: int = 1 max_parallel_loading_workers: Optional[int] = None block_size: int = 16 enable_prefix_caching: bool = False disable_sliding_window: bool = False use_v2_block_manager: bool = False swap_space: int = 4 # GiB cpu_offload_gb: int = 0 # GiB gpu_memory_utilization: float = 0.90 max_num_batched_tokens: Optional[int] = None max_num_seqs: int = 256 max_logprobs: int = 20 # Default value for OpenAI Chat Completions API disable_log_stats: bool = False revision: Optional[str] = None code_revision: Optional[str] = None rope_scaling: Optional[dict] = None rope_theta: Optional[float] = None tokenizer_revision: Optional[str] = None quantization: Optional[str] = None enforce_eager: bool = False max_context_len_to_capture: Optional[int] = None max_seq_len_to_capture: int = 8192 disable_custom_all_reduce: bool = False tokenizer_pool_size: int = 0 # Note: Specifying a tokenizer pool by passing a class # is intended for expert use only. The API may change without # notice. tokenizer_pool_type: Union[str, Type[\"BaseTokenizerGroup\"]] = \"ray\" tokenizer_pool_extra_config: Optional[dict] = None enable_lora: bool = False max_loras: int = 1 max_lora_rank: int = 16 enable_prompt_adapter: bool = False max_prompt_adapters: int = 1 max_prompt_adapter_token: int = 0 fully_sharded_loras: bool = False lora_extra_vocab_size: int = 256 long_lora_scaling_factors: Optional[Tuple[float]] = None lora_dtype: str = 'auto' max_cpu_loras: Optional[int] = None device: str = 'auto' ray_workers_use_nsight: bool = False num_gpu_blocks_override: Optional[int] = None num_lookahead_slots: int = 0 model_loader_extra_config: Optional[dict] = None ignore_patterns: Optional[Union[str, List[str]]] = None preemption_mode: Optional[str] = None scheduler_delay_factor: float = 0.0 enable_chunked_prefill: Optional[bool] = None guided_decoding_backend: str = 'outlines' # Speculative decoding configuration. speculative_model: Optional[str] = None speculative_draft_tensor_parallel_size: Optional[int] = None num_speculative_tokens: Optional[int] = None speculative_max_model_len: Optional[int] = None speculative_disable_by_batch_size: Optional[int] = None ngram_prompt_lookup_max: Optional[int] = None ngram_prompt_lookup_min: Optional[int] = None spec_decoding_acceptance_method: str = 'rejection_sampler' typical_acceptance_sampler_posterior_threshold: Optional[float] = None typical_acceptance_sampler_posterior_alpha: Optional[float] = None qlora_adapter_name_or_path: Optional[str] = None disable_logprobs_during_spec_decoding: Optional[bool] = None otlp_traces_endpoint: Optional[str] = None @staticmethod def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: \"\"\"Shared CLI arguments for vLLM engine.\"\"\" # Model arguments # TODO(shengguangming): delete the unused args parser.add_argument('--model', type=str, default='facebook/opt-125m', help='name or path of the huggingface model to use') parser.add_argument('--tokenizer', type=str, default=EngineArgs.tokenizer, help='name or path of the huggingface tokenizer to use') parser.add_argument('--revision', type=str, default=None, help='the specific model version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-revision', type=str, default=None, help='the specific tokenizer version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-mode', type=str, default=EngineArgs.tokenizer_mode, choices=['auto', 'slow'], help='tokenizer mode. \"auto\" will use the fast ' 'tokenizer if available, and \"slow\" will ' 'always use the slow tokenizer.') parser.add_argument('--trust-remote-code', action='store_true', help='trust remote code from huggingface') parser.add_argument('--download-dir', type=str, default=EngineArgs.download_dir, help='directory to download and load the weights, ' 'default to the default cache dir of ' 'huggingface') parser.add_argument('--load-format', type=str, default=EngineArgs.load_format, choices=['auto', 'pt', 'safetensors', 'npcache', 'dummy'], help='The format of the model weights to load. ' '\"auto\" will tr","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","name":"add_cli_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","span":"L140-L254","signature":"def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser","exported":true,"origin":"ast","body_hash":"cf1908f71aee396b15d02e5aad05a2803adf4e9fd0905ff5fcb9c8727ba16fb5","body_text":"def add_cli_args(parser: argparse.ArgumentParser) -> argparse.ArgumentParser: \"\"\"Shared CLI arguments for vLLM engine.\"\"\" # Model arguments # TODO(shengguangming): delete the unused args parser.add_argument('--model', type=str, default='facebook/opt-125m', help='name or path of the huggingface model to use') parser.add_argument('--tokenizer', type=str, default=EngineArgs.tokenizer, help='name or path of the huggingface tokenizer to use') parser.add_argument('--revision', type=str, default=None, help='the specific model version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-revision', type=str, default=None, help='the specific tokenizer version to use. It can be a branch ' 'name, a tag name, or a commit id. If unspecified, will use ' 'the default version.') parser.add_argument('--tokenizer-mode', type=str, default=EngineArgs.tokenizer_mode, choices=['auto', 'slow'], help='tokenizer mode. \"auto\" will use the fast ' 'tokenizer if available, and \"slow\" will ' 'always use the slow tokenizer.') parser.add_argument('--trust-remote-code', action='store_true', help='trust remote code from huggingface') parser.add_argument('--download-dir', type=str, default=EngineArgs.download_dir, help='directory to download and load the weights, ' 'default to the default cache dir of ' 'huggingface') parser.add_argument('--load-format', type=str, default=EngineArgs.load_format, choices=['auto', 'pt', 'safetensors', 'npcache', 'dummy'], help='The format of the model weights to load. ' '\"auto\" will try to load the weights in the safetensors format ' 'and fall back to the pytorch bin format if safetensors format ' 'is not available. ' '\"pt\" will load the weights in the pytorch bin format. ' '\"safetensors\" will load the weights in the safetensors format. ' '\"npcache\" will load the weights in pytorch format and store ' 'a numpy cache to speed up the loading. ' '\"dummy\" will initialize the weights with random values, ' 'which is mainly for profiling.') parser.add_argument('--dtype', type=str, default=EngineArgs.dtype, choices=['auto', 'half', 'float16', 'bfloat16', 'float', 'float32'], help='data type for model weights and activations. ' 'The \"auto\" option will use FP16 precision ' 'for FP32 and FP16 models, and BF16 precision ' 'for BF16 models.') parser.add_argument('--max-model-len', type=int, default=None, help='model context length. If unspecified, ' 'will be automatically derived from the model.') # Parallel arguments parser.add_argument('--worker-use-ray', action='store_true', help='use Ray for distributed serving, will be ' 'automatically set when using more than 1 GPU') parser.add_argument('--pipeline-parallel-size', '-pp', type=int, default=EngineArgs.pipeline_parallel_size, help='number of pipeline stages') parser.add_argument('--tensor-parallel-size', '-tp', type=int, default=EngineArgs.tensor_parallel_size, help='number of tensor parallel replicas') # KV cache arguments parser.add_argument('--block-size', type=int, default=EngineArgs.block_size, choices=[8, 16, 32], help='token block size') # TODO(woosuk): Support fine-grained seeds (e.g., seed per request). parser.add_argument('--seed', type=int, default=EngineArgs.seed, help='random seed') parser.add_argument('--swap-space', type=int, default=EngineArgs.swap_space, help='CPU swap space size (GiB) per GPU') parser.add_argument('--gpu-memory-utilization', type=float, default=EngineArgs.gpu_memory_utilization, help='the percentage of GPU memory to be used for' 'the model executor') parser.add_argument('--max-num-batched-tokens', type=int, default=EngineArgs.max_num_batched_tokens, help='maximum number of batched tokens per ' 'iteration') parser.add_argument('--max-num-seqs', type=int, default=EngineArgs.max_num_seqs, help='maximum number of sequences per iteration') parser.add_argument('--disable-log-stats', action='store_true', help='disable logging statistics') # Quantization settings. parser.add_argument('--quantization', '-q', type=str, choices=['awq', None], default=None, help='Method used to quantize the weights') return parser","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.from_cli_args","name":"from_cli_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","span":"L257-L262","signature":"def from_cli_args(cls, args: argparse.Namespace) -> 'EngineArgs': # Get the list of attributes of this dataclass.","exported":true,"origin":"ast","body_hash":"21089f327dd4125f76a223a109fdee39de77f84d10cd1fad1f24de8837c1b6f1","body_text":"def from_cli_args(cls, args: argparse.Namespace) -> 'EngineArgs': # Get the list of attributes of this dataclass. attrs = [attr.name for attr in dataclasses.fields(cls)] # Set the attributes from the parsed arguments. engine_args = cls(**{attr: getattr(args, attr) for attr in attrs}) return engine_args","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","name":"create_engine_config","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","span":"L264-L453","signature":"def create_engine_config( self, ) -> EngineConfig: # bitsandbytes quantization needs a specific model loader # so we make sure the quant method and the load format are consistent","exported":true,"origin":"ast","body_hash":"33dbc9b097ecb51a5cbba60c913724b4f204f33854dde9fa998dd2e1b4343328","body_text":"def create_engine_config( self, ) -> EngineConfig: # bitsandbytes quantization needs a specific model loader # so we make sure the quant method and the load format are consistent if (self.quantization == \"bitsandbytes\" or self.qlora_adapter_name_or_path is not None) and \\ self.load_format != \"bitsandbytes\": raise ValueError(\"BitsAndBytes quantization and QLoRA adapter only support \" f\"'bitsandbytes' load format, but got {self.load_format}\") if (self.load_format == \"bitsandbytes\" or self.qlora_adapter_name_or_path is not None) and \\ self.quantization != \"bitsandbytes\": raise ValueError(\"BitsAndBytes load format and QLoRA adapter only support \" f\"'bitsandbytes' quantization, but got {self.quantization}\") assert self.cpu_offload_gb >= 0, (\"CPU offload space must be non-negative\" f\", but got {self.cpu_offload_gb}\") multimodal_config = MultiModalConfig() device_config = DeviceConfig(self.device) # NOTE(sgm): we only modify ModelConfig, other configs are import from vllm model_config = ModelConfig(hf_config=self.model_hf_config, tokenizer_mode=self.tokenizer_mode, trust_remote_code=self.trust_remote_code, dtype=self.dtype, seed=self.seed, revision=self.revision, code_revision=self.code_revision, rope_scaling=self.rope_scaling, rope_theta=self.rope_theta, tokenizer_revision=self.tokenizer_revision, max_model_len=self.max_model_len, quantization=self.quantization, quantization_param_path=self.quantization_param_path, enforce_eager=self.enforce_eager, max_context_len_to_capture=self.max_context_len_to_capture, max_seq_len_to_capture=self.max_seq_len_to_capture, max_logprobs=self.max_logprobs, disable_sliding_window=self.disable_sliding_window, skip_tokenizer_init=self.skip_tokenizer_init, served_model_name=self.served_model_name, multimodal_config=multimodal_config) cache_config = CacheConfig( block_size=self.block_size, gpu_memory_utilization=self.gpu_memory_utilization, swap_space=self.swap_space, cache_dtype=self.kv_cache_dtype, num_gpu_blocks_override=self.num_gpu_blocks_override, sliding_window=model_config.get_sliding_window(), enable_prefix_caching=self.enable_prefix_caching, cpu_offload_gb=self.cpu_offload_gb, ) parallel_config = ParallelConfig(pipeline_parallel_size=self.pipeline_parallel_size, tensor_parallel_size=self.tensor_parallel_size, worker_use_ray=self.worker_use_ray, max_parallel_loading_workers=self.max_parallel_loading_workers, disable_custom_all_reduce=self.disable_custom_all_reduce, tokenizer_pool_config=TokenizerPoolConfig.create_config( self.tokenizer_pool_size, self.tokenizer_pool_type, self.tokenizer_pool_extra_config, ), ray_workers_use_nsight=self.ray_workers_use_nsight, distributed_executor_backend=self.distributed_executor_backend) # NOTE[VERL]: Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" parallel_config.world_size = world_size max_model_len = model_config.max_model_len use_long_context = max_model_len > 32768 if self.enable_chunked_prefill is None: # If not explicitly set, enable chunked prefill by default for # long context (> 32K) models. This is to avoid OOM errors in the # initial memory profiling phase. if use_long_context: is_gpu = device_config.device_type == \"cuda\" use_sliding_window = (model_config.get_sliding_window() is not None) use_spec_decode = self.speculative_model is not None has_seqlen_agnostic_layers = (model_config.contains_seqlen_agnostic_layers(parallel_config)) if (is_gpu and not use_sliding_window and not use_spec_decode and not self.enable_lora and not self.enable_prompt_adapter and not self.enable_prefix_caching and not has_seqlen_agnostic_layers): self.enable_chunked_prefill = True logger.warning(\"Chunked prefill is enabled by default for models with \" \"max_model_len > 32K. Currently, chunked prefill might \" \"not work with some features or models. If you \" \"encounter any issues, please disable chunked prefill \" \"by setting --enable-chunked-prefill=False.\") if self.enable_chunked_prefill is None: self.enable_chunked_prefill = False if not self.enable_chunked_prefill and use_long_context: logger.warning( \"The model has a long context length (%s). This may cause OOM \" \"errors during the initial memory profiling phase, or result \" \"in low performance due to small KV cache space. Consider \" \"setting --max-model-len to a smaller value.\", max_model_len) # TODO: spec config speculative_config = SpeculativeConfig.maybe_create_spec_config( target_model_config=model_config, target_parallel_config=parallel_config, target_dtype=self.dtype, speculative_model=self.speculative_model, speculative_draft_tensor_parallel_size = \\ self.speculative_draft_tensor_parallel_size, num_speculative_tokens=self.num_speculative_tokens, speculative_disable_by_batch_size=self. speculative_disable_by_batch_size, speculative_max_model_len=self.speculative_max_model_len, enable_chunked_prefill=self.enable_chunked_prefill, use_v2_block_manager=self.use_v2_block_manager, dis","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"argparse","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"json","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"vllm.executor.executor_base","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"vllm.model_executor.layers.quantization","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"imports","specifier":"vllm.transformers_utils.tokenizer_group.base_tokenizer_group","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#nullable_str","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.add_cli_args","relation":"calls","name":"add_argument","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.from_cli_args","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.from_cli_args","relation":"calls","name":"fields","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.from_cli_args","relation":"calls","name":"cls","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.from_cli_args","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"MultiModalConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"DeviceConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"ModelConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"CacheConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"get_sliding_window","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"ParallelConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"create_config","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"get_sliding_window","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"contains_seqlen_agnostic_layers","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"maybe_create_spec_config","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"SchedulerConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"LoRAConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"LoadConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"PromptAdapterConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"DecodingConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"ObservabilityConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"get_sliding_window","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"EngineConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/arg_utils.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/config.py":{"size":11948,"mtimeMs":1789275830316.2026,"hash":"5dac246763fa9ba0331f36ee72bbb43fae1a702bb354adc156e08cbaa6f0fa5e","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/config.py","name":"config.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/config.py","span":"L1-L247","signature":null,"exported":true,"origin":"ast","body_hash":"5dac246763fa9ba0331f36ee72bbb43fae1a702bb354adc156e08cbaa6f0fa5e","chars":11948,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py import enum import json from typing import List, Optional, Union from dataclasses import dataclass, field, fields import torch from transformers import PretrainedConfig from vllm.logger import init_logger from vllm.model_executor.layers.quantization import get_quantization_config from vllm.transformers_utils.config import get_hf_text_config from vllm.utils import is_hip, print_warning_once # Add for verl from vllm.config import ModelConfig, _get_and_verify_dtype, _get_and_verify_max_len, get_served_model_name GPTQMarlinConfig = get_quantization_config(\"gptq_marlin\") logger = init_logger(__name__) _GB = 1 << 30 # TODO: check whether this is necessary @dataclass"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig","name":"ModelConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/config.py","span":"L38-L178","signature":"class ModelConfig(ModelConfig)","exported":true,"origin":"ast","body_hash":"0598c4dcb658a8521595364cf97d26579bd7962fc9b7cc3863085b33ee12f834","body_text":"class ModelConfig(ModelConfig): \"\"\"Configuration for the model. Args: model: Name or path of the huggingface model to use. tokenizer: Name or path of the huggingface tokenizer to use. tokenizer_mode: Tokenizer mode. \"auto\" will use the fast tokenizer if available, and \"slow\" will always use the slow tokenizer. trust_remote_code: Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer. download_dir: Directory to download and load the weights, default to the default cache directory of huggingface. load_format: The format of the model weights to load: \"auto\" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. \"pt\" will load the weights in the pytorch bin format. \"safetensors\" will load the weights in the safetensors format. \"npcache\" will load the weights in pytorch format and store a numpy cache to speed up the loading. \"dummy\" will initialize the weights with random values, which is mainly for profiling. dtype: Data type for model weights and activations. The \"auto\" option will use FP16 precision for FP32 and FP16 models, and BF16 precision for BF16 models. seed: Random seed for reproducibility. revision: The specific model version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. code_revision: The specific revision to use for the model code on Hugging Face Hub. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. tokenizer_revision: The specific tokenizer version to use. It can be a branch name, a tag name, or a commit id. If unspecified, will use the default version. max_model_len: Maximum length of a sequence (including prompt and output). If None, will be derived from the model. quantization: Quantization method that was used to quantize the model weights. If None, we assume the model weights are not quantized. quantization_param_path: Path to JSON file containing scaling factors. Used to load KV cache scaling factors into the model when KV cache type is FP8_E4M3 on ROCm (AMD GPU). In the future these will also be used to load activation and weight scaling factors when the model dtype is FP8_E4M3 on ROCm. enforce_eager: Whether to enforce eager execution. If True, we will disable CUDA graph and always execute the model in eager mode. If False, we will use CUDA graph and eager execution in hybrid. max_context_len_to_capture: Maximum context len covered by CUDA graphs. When a sequence has context length larger than this, we fall back to eager mode (DEPRECATED. Use max_seq_len_to_capture instead). max_seq_len_to_capture: Maximum sequence len covered by CUDA graphs. When a sequence has context length larger than this, we fall back to eager mode skip_tokenizer_init: If true, skip initialization of tokenizer and detokenizer. served_model_name: The model name used in metrics tag `model_name`, matches the model name exposed via the APIs. If multiple model names provided, the first name will be used. If not specified, the model name will be the same as `model`. \"\"\" def __init__( self, hf_config: PretrainedConfig, tokenizer_mode: str, trust_remote_code: bool, dtype: Union[str, torch.dtype], seed: int, revision: Optional[str] = None, code_revision: Optional[str] = None, rope_scaling: Optional[dict] = None, rope_theta: Optional[float] = None, tokenizer_revision: Optional[str] = None, max_model_len: Optional[int] = None, quantization: Optional[str] = None, quantization_param_path: Optional[str] = None, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: Optional[int] = None, max_logprobs: int = 20, disable_sliding_window: bool = False, skip_tokenizer_init: bool = False, served_model_name: Optional[Union[str, List[str]]] = None, multimodal_config: Optional[\"MultiModalConfig\"] = None, ) -> None: self.model = hf_config._name_or_path self.tokenizer = hf_config._name_or_path # NOTE(sgm): same as open-sourced self.tokenizer_mode = tokenizer_mode self.trust_remote_code = trust_remote_code self.seed = seed self.revision = revision self.code_revision = code_revision self.rope_scaling = rope_scaling self.rope_theta = rope_theta # The tokenizer version is consistent with the model version by default. if tokenizer_revision is None: self.tokenizer_revision = revision else: self.tokenizer_revision = tokenizer_revision self.quantization = quantization self.quantization_param_path = quantization_param_path self.enforce_eager = enforce_eager if max_context_len_to_capture is not None: raise ValueError(\"`max_context_len_to_capture` is deprecated. \" \"Use `max_seq_len_to_capture` instead.\") self.max_seq_len_to_capture = max_seq_len_to_capture self.max_logprobs = max_logprobs self.disable_sliding_window = disable_sliding_window self.skip_tokenizer_init = skip_tokenizer_init # self.hf_config = get_config(model, trust_remote_code, revision) self.hf_config = hf_config self.hf_text","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/config.py","span":"L99-L178","signature":"def __init__( self, hf_config: PretrainedConfig, tokenizer_mode: str, trust_remote_code: bool, dtype: Union[str, torch.dtype], seed: int, revision: Optional[str] = None, code_revision: Optional[str] = None, rope_scaling: Optional[dict] = None, rope_theta: Optional[float] = None, tokenizer_revision: Optional[str] = None, max_model_len: Optional[int] = None, quantization: Optional[str] = None, quantization_param_path: Optional[str] = None, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: Optional[int] = None, max_logprobs: int = 20, disable_sliding_window: bool = False, skip_tokenizer_init: bool = False, served_model_name: Optional[Union[str, List[str]]] = None, multimodal_config: Optional[\"MultiModalConfig\"] = None, ) -> None","exported":false,"origin":"ast","body_hash":"c60660421929282a7b22893fb6361867e08fd09c123202a351b3ee6f2d0ed967","body_text":"def __init__( self, hf_config: PretrainedConfig, tokenizer_mode: str, trust_remote_code: bool, dtype: Union[str, torch.dtype], seed: int, revision: Optional[str] = None, code_revision: Optional[str] = None, rope_scaling: Optional[dict] = None, rope_theta: Optional[float] = None, tokenizer_revision: Optional[str] = None, max_model_len: Optional[int] = None, quantization: Optional[str] = None, quantization_param_path: Optional[str] = None, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: Optional[int] = None, max_logprobs: int = 20, disable_sliding_window: bool = False, skip_tokenizer_init: bool = False, served_model_name: Optional[Union[str, List[str]]] = None, multimodal_config: Optional[\"MultiModalConfig\"] = None, ) -> None: self.model = hf_config._name_or_path self.tokenizer = hf_config._name_or_path # NOTE(sgm): same as open-sourced self.tokenizer_mode = tokenizer_mode self.trust_remote_code = trust_remote_code self.seed = seed self.revision = revision self.code_revision = code_revision self.rope_scaling = rope_scaling self.rope_theta = rope_theta # The tokenizer version is consistent with the model version by default. if tokenizer_revision is None: self.tokenizer_revision = revision else: self.tokenizer_revision = tokenizer_revision self.quantization = quantization self.quantization_param_path = quantization_param_path self.enforce_eager = enforce_eager if max_context_len_to_capture is not None: raise ValueError(\"`max_context_len_to_capture` is deprecated. \" \"Use `max_seq_len_to_capture` instead.\") self.max_seq_len_to_capture = max_seq_len_to_capture self.max_logprobs = max_logprobs self.disable_sliding_window = disable_sliding_window self.skip_tokenizer_init = skip_tokenizer_init # self.hf_config = get_config(model, trust_remote_code, revision) self.hf_config = hf_config self.hf_text_config = get_hf_text_config(hf_config) self.dtype = _get_and_verify_dtype(self.hf_text_config, dtype) # self.served_model_name = get_served_model_name(model, # served_model_name) # self._verify_load_format() # self._verify_tokenizer_mode() if (not self.disable_sliding_window and self.hf_text_config.model_type == \"gemma2\" and self.hf_text_config.sliding_window is not None): print_warning_once(\"Gemma 2 uses sliding window attention for every odd layer, \" \"which is currently not supported by vLLM. Disabling sliding \" \"window and capping the max length to the sliding window size \" f\"({self.hf_text_config.sliding_window}).\") self.disable_sliding_window = True self.max_model_len = _get_and_verify_max_len(hf_config=self.hf_text_config, max_model_len=max_model_len, disable_sliding_window=self.disable_sliding_window, sliding_window_len=self.get_hf_config_sliding_window()) self.served_model_name = get_served_model_name( self.model, # str served_model_name) self.multimodal_config = multimodal_config if not self.skip_tokenizer_init: self._verify_tokenizer_mode() self._verify_embedding_mode() self._verify_quantization() self._verify_cuda_graph()","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadFormat","name":"LoadFormat","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/config.py","span":"L181-L188","signature":"class LoadFormat(str, enum.Enum)","exported":true,"origin":"ast","body_hash":"86ea6a79c9ee2abd1dba76e1cbbea5f0cf439e41e0d22f4aa40eae34cca8429b","body_text":"class LoadFormat(str, enum.Enum): AUTO = 'auto' MEGATRON = \"megatron\" HF = \"hf\" DTENSOR = 'dtensor' DUMMY_HF = 'dummy_hf' DUMMY_MEGATRON = 'dummy_megatron' DUMMY_DTENSOR = 'dummy_dtensor'","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig","name":"LoadConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/config.py","span":"L193-L246","signature":"class LoadConfig","exported":true,"origin":"ast","body_hash":"17fef7183b587bc8293d96f97d8214d995039d878a918f7ad3cf562e74abdc12","body_text":"class LoadConfig: \"\"\" download_dir: Directory to download and load the weights, default to the default cache directory of huggingface. load_format: The format of the model weights to load: \"auto\" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. \"pt\" will load the weights in the pytorch bin format. \"safetensors\" will load the weights in the safetensors format. \"npcache\" will load the weights in pytorch format and store a numpy cache to speed up the loading. \"dummy\" will initialize the weights with random values, which is mainly for profiling. \"tensorizer\" will use CoreWeave's tensorizer library for fast weight loading. \"bitsandbytes\" will load nf4 type weights. ignore_patterns: The list of patterns to ignore when loading the model. Default to \"original/**/*\" to avoid repeated loading of llama's checkpoints. \"\"\" load_format: Union[str, LoadFormat, \"BaseModelLoader\"] = LoadFormat.AUTO download_dir: Optional[str] = None model_loader_extra_config: Optional[Union[str, dict]] = field(default_factory=dict) ignore_patterns: Optional[Union[List[str], str]] = None def __post_init__(self): model_loader_extra_config = self.model_loader_extra_config or {} if isinstance(model_loader_extra_config, str): self.model_loader_extra_config = json.loads(model_loader_extra_config) self._verify_load_format() if self.ignore_patterns is not None and len(self.ignore_patterns) > 0: logger.info(\"Ignoring the following patterns when downloading weights: %s\", self.ignore_patterns) else: self.ignore_patterns = [\"original/**/*\"] def _verify_load_format(self) -> None: if not isinstance(self.load_format, str): return load_format = self.load_format.lower() self.load_format = LoadFormat(load_format) rocm_not_supported_load_format: List[str] = [] if is_hip() and load_format in rocm_not_supported_load_format: rocm_supported_load_format = [ f for f in LoadFormat.__members__ if (f not in rocm_not_supported_load_format) ] raise ValueError(f\"load format '{load_format}' is not supported in ROCm. \" f\"Supported load formats are \" f\"{rocm_supported_load_format}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig.__post_init__","name":"__post_init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/config.py","span":"L221-L230","signature":"def __post_init__(self)","exported":false,"origin":"ast","body_hash":"5094b5416a54d00d4a1013fa86e20a0199223b86997c6d96b2410245cb2b7f25","body_text":"def __post_init__(self): model_loader_extra_config = self.model_loader_extra_config or {} if isinstance(model_loader_extra_config, str): self.model_loader_extra_config = json.loads(model_loader_extra_config) self._verify_load_format() if self.ignore_patterns is not None and len(self.ignore_patterns) > 0: logger.info(\"Ignoring the following patterns when downloading weights: %s\", self.ignore_patterns) else: self.ignore_patterns = [\"original/**/*\"]","summary_state":"pending","summary":null,"crux":null,"owner":"LoadConfig"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig._verify_load_format","name":"_verify_load_format","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/config.py","span":"L232-L246","signature":"def _verify_load_format(self) -> None","exported":false,"origin":"ast","body_hash":"f9f7e30ec352e2c4183434a9e0e078cb7cff2b5e08499624e93b0b4857aff8ab","body_text":"def _verify_load_format(self) -> None: if not isinstance(self.load_format, str): return load_format = self.load_format.lower() self.load_format = LoadFormat(load_format) rocm_not_supported_load_format: List[str] = [] if is_hip() and load_format in rocm_not_supported_load_format: rocm_supported_load_format = [ f for f in LoadFormat.__members__ if (f not in rocm_not_supported_load_format) ] raise ValueError(f\"load format '{load_format}' is not supported in ROCm. \" f\"Supported load formats are \" f\"{rocm_supported_load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"LoadConfig"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"enum","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"json","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"vllm.model_executor.layers.quantization","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"vllm.transformers_utils.config","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"calls","name":"get_quantization_config","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig","relation":"extends","name":"ModelConfig","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"get_hf_text_config","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"_get_and_verify_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"print_warning_once","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"_get_and_verify_max_len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"get_hf_config_sliding_window","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"get_served_model_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"_verify_tokenizer_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"_verify_embedding_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"_verify_quantization","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#ModelConfig.__init__","relation":"calls","name":"_verify_cuda_graph","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py","recvType":"ModelConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadFormat","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadFormat","relation":"extends","name":"str","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig","relation":"calls","name":"field","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig.__post_init__","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig.__post_init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig.__post_init__","relation":"calls","name":"loads","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig.__post_init__","relation":"calls","name":"_verify_load_format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py","recvType":"LoadConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig.__post_init__","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig.__post_init__","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig._verify_load_format","file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig._verify_load_format","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig._verify_load_format","relation":"calls","name":"lower","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py","recvType":"LoadFormat"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig._verify_load_format","relation":"calls","name":"LoadFormat","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig._verify_load_format","relation":"calls","name":"is_hip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/config.py#LoadConfig._verify_load_format","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/config.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py":{"size":15827,"mtimeMs":1789275830317.203,"hash":"3f71614c586d78588d8ee69d66e1533896c51e985e0cb8719cb8bd124e35c3f3","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","name":"dtensor_weight_loaders.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L1-L341","signature":null,"exported":true,"origin":"ast","body_hash":"3f71614c586d78588d8ee69d66e1533896c51e985e0cb8719cb8bd124e35c3f3","chars":15827,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models from typing import Dict, Iterable, Tuple import torch import torch.nn as nn from torch.distributed._tensor import DTensor, Shard, Replicate from vllm.model_executor.layers.linear import * from vllm.model_executor.models import ModelRegistry from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.utils import is_pp_missing_parameter from vllm.model_executor.layers.fused_moe import FusedMoE __MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__ = { 'GPT2LMHeadModel': gpt2_dtensor_weight_loader, 'LlamaForCausalLM': llama_dtensor_weight_loader, 'LLaMAForCausalLM': llama_dtensor_weight_loader, 'MistralForCausalLM': llama_dtensor_weight_loader, # mistral is the same as llama in vLLM 'InternLMForCausalLM': llama_dtensor_weight_loader, 'AquilaModel': llama_dtensor_weight_loader, 'AquilaForCausalLM': llama_dtensor_weight_loader, 'Phi3ForCausalLM': llama_dtensor_weight_loader, 'GemmaForCausalLM': gemma_dtensor_weight_loader, 'Gemma2ForCausalLM': gemma_dtensor_weight_loader, 'GPTBigCodeForCausalLM': gptbigcode_dtensor_load_weights, 'Starcoder2ForCausalLM': starcoder2_dtensor_load_weights, 'Qwen2ForCausalLM': qwen2_dtensor_weight_loader, 'DeepseekV2ForCausalLM': deepseekv2_dtensor_weight_loader } # the actor model is .state_dict() # Load dtensor weights # NOTE(sgm): we use per-parameter weight loader in each vllm sub"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","name":"gemma_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L27-L61","signature":"def gemma_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"04c523a72722155ecd45cb012e7917ae37de4cdc4e82bedaa5ff8ad292cf1db1","body_text":"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(vllm_model.named_parameters()) 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 stacked_name = name.replace(shard_name, param_name) # Skip loading extra bias for GPTQ models. if stacked_name.endswith(\".bias\") and stacked_name not in params_dict: continue local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) param = params_dict[stacked_name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) 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, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","name":"gptbigcode_dtensor_load_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L64-L76","signature":"def gptbigcode_dtensor_load_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"7494a5e13088c35a8ed94c153c3d247196ed1c1f0c08b38b75cfc6be4046b77e","body_text":"def gptbigcode_dtensor_load_weights(actor_weights: Dict, vllm_model: nn.Module): params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"lm_head.weight\" in name: continue if \".attn.bias\" in name: # Skip attention mask. # NOTE: \"c_attn.bias\" should not be skipped. 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, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","name":"starcoder2_dtensor_load_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L79-L107","signature":"def starcoder2_dtensor_load_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"f7efd018225e19458922d6d47502837ab1afc26da47f5f85e802bd7b89f2a936","body_text":"def starcoder2_dtensor_load_weights(actor_weights: Dict, vllm_model: 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\"), ] params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue for (param_name, weight_name, shard_id) in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in name: continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","name":"llama_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L110-L151","signature":"def llama_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"1736f831264cc292ac39f2e3010167ded6b1db2f01e38512a1834683d157524b","body_text":"def llama_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(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue if (\"rotary_emb.cos_cached\" in name or \"rotary_emb.sin_cached\" in name): # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue # With tie_word_embeddings, we can skip lm_head.weight # The weight might appear unnecessarily in the files if the model is # processed with quantization, LoRA, fine-tuning, etc. if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in name: continue for (param_name, weight_name, shard_id) in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_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, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: # 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, local_loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","name":"qwen2_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L154-L188","signature":"def qwen2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"9b4faed6adddcf37e68aeb25fd9a60eb45af8ef9ef6032e5f1b69683438afb14","body_text":"def qwen2_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(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in name: continue for (param_name, weight_name, shard_id) in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_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, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: # Skip loading extra bias for GPTQ models. if name.endswith(\".bias\") and name not in params_dict: continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","name":"deepseekv2_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L194-L267","signature":"def deepseekv2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"a466c5c9d6fe0a68232cf5aebdc766014622f31452e0ecf9df26a1cf2803439d","body_text":"def deepseekv2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: stacked_params_mapping = [ # (param_name, shard_name, shard_id) (\"gate_up_proj\", \"gate_proj\", 0), (\"gate_up_proj\", \"up_proj\", 1), ] # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) expert_params_mapping = FusedMoE.make_expert_params_mapping(ckpt_gate_proj_name=\"gate_proj\", ckpt_down_proj_name=\"down_proj\", ckpt_up_proj_name=\"up_proj\", num_experts=vllm_model.config.n_routed_experts) params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue for (param_name, weight_name, shard_id) in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: continue # We have mlp.experts[0].gate_proj in the checkpoint. # Since we handle the experts below in expert_params_mapping, # we need to skip here BEFORE we update the name, otherwise # name will be updated to mlp.experts[0].gate_up_proj, which # will then be updated below in expert_params_mapping # for mlp.experts[0].gate_gate_up_proj, which breaks load. if ((\"mlp.experts.\" in name) and name not in params_dict): continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(\".bias\") and name not in params_dict: continue if is_pp_missing_parameter(name, vllm_model): continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue name = name.replace(weight_name, param_name) if is_pp_missing_parameter(name, vllm_model): continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype), weight_name, shard_id=shard_id, expert_id=expert_id) break else: # Skip loading extra bias for GPTQ models. if name.endswith(\".bias\") and name not in params_dict: continue if is_pp_missing_parameter(name, vllm_model): continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gpt2_dtensor_weight_loader","name":"gpt2_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L270-L271","signature":"def gpt2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"e26a0b664180d6dbfeb3ba74e435ba77e0a32726d4fc799cc1282d68edce33d9","body_text":"def gpt2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: pass","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#redistribute_dtensor","name":"redistribute_dtensor","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L274-L284","signature":"def redistribute_dtensor(param_name: str, loaded_weights: DTensor, parallelize_plan: Dict = None)","exported":true,"origin":"ast","body_hash":"8e33f089fd37013b9404ed851cdce475af110e84dc4bc3c6edfbbe7ef9fe5d30","body_text":"def redistribute_dtensor(param_name: str, loaded_weights: DTensor, parallelize_plan: Dict = None): param_name = _process_parameter_names(name=param_name) if parallelize_plan is not None: assert param_name in parallelize_plan.keys(), \\ f\"param name: {param_name} not in parallelize_plan :{parallelize_plan.keys()}\" placement = parallelize_plan[param_name] local_loaded_weights = loaded_weights.redistribute(device_mesh=loaded_weights.device_mesh, placements=placement).to_local() else: local_loaded_weights = loaded_weights.full_tensor() return local_loaded_weights","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#_process_parameter_names","name":"_process_parameter_names","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L287-L300","signature":"def _process_parameter_names(name): # Remove '.weight' if it exists at the end of the string","exported":false,"origin":"ast","body_hash":"074c597f3b319460bbcf2325d017836ddcbf0ce56c58d5d3356ffe01fe6d6547","body_text":"def _process_parameter_names(name): # Remove '.weight' if it exists at the end of the string if name.endswith(\".weight\"): name = name[:-7] # Remove 'model.layers.x.' or 'model.' prefix if \"model.layers\" in name: parts = name.split('.') # Reconstruct the string without 'model.layers.x.' name = '.'.join(parts[3:]) # parts[0] is 'model', parts[1] is 'layers', parts[2] is 'x' elif name.startswith(\"model.\"): name = name[6:] # Remove 'model.' return name","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#load_dtensor_weights","name":"load_dtensor_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L323-L328","signature":"def load_dtensor_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"d614ba13175fcb0a1db90bcd80acdd13cbcddb57eb013b3ce25386d8328656ce","body_text":"def load_dtensor_weights(actor_weights: Dict, vllm_model: nn.Module): weight_loader = _get_model_weight_loader(vllm_model.__class__.__name__) weight_loader(actor_weights, vllm_model) # NOTE(sgm) to reduce peak memory usage, we offload vllm model to cpu # after init, and we need this after sync model weights for in first iter. vllm_model = vllm_model.cuda()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#_get_model_weight_loader","name":"_get_model_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L331-L335","signature":"def _get_model_weight_loader(arch: str)","exported":false,"origin":"ast","body_hash":"2aff4020b0abe87a16fdc5b4ba0f9ed8f437312a1edfce1320e5726f5165c8ae","body_text":"def _get_model_weight_loader(arch: str): if arch in __MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__: return __MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__[arch] raise ValueError(f\"Model architectures {arch} are not supported for now. \" f\"Supported architectures: {__MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__.keys()}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#update_dtensor_weight_loader","name":"update_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","span":"L339-L340","signature":"def update_dtensor_weight_loader()","exported":true,"origin":"ast","body_hash":"3251c12c590cfa8575d836647130786146b1ba578183f484f98d2dac2cce5914","body_text":"def update_dtensor_weight_loader(): pass","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"imports","specifier":"torch.distributed._tensor","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.linear","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.models","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.model_loader.weight_utils","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.models.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.fused_moe","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"make_expert_params_mapping","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"is_pp_missing_parameter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"is_pp_missing_parameter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"is_pp_missing_parameter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#gpt2_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#redistribute_dtensor","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"_process_parameter_names","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"to_local","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"redistribute","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"DTensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"full_tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","recvType":"DTensor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#_process_parameter_names","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"split","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"startswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#load_dtensor_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#load_dtensor_weights","relation":"calls","name":"_get_model_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#load_dtensor_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#load_dtensor_weights","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#_get_model_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py#update_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/dtensor_weight_loaders.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py":{"size":1895,"mtimeMs":1789275830317.203,"hash":"8c819582f3ad7b23df842087cdf4a1992392e931cd419bcdc19879742ee828fe","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","name":"hf_weight_loader.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","span":"L1-L45","signature":null,"exported":true,"origin":"ast","body_hash":"8c819582f3ad7b23df842087cdf4a1992392e931cd419bcdc19879742ee828fe","chars":1895,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models from typing import Dict, Union, Optional, Iterable, Tuple import torch import torch.nn as nn from vllm.model_executor.model_loader.utils import set_default_torch_dtype from vllm.model_executor.model_loader.weight_utils import default_weight_loader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#update_hf_weight_loader","name":"update_hf_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","span":"L25-L27","signature":"def update_hf_weight_loader()","exported":true,"origin":"ast","body_hash":"61f93704751e75420cb304d950a62061b6f5959c50e3d81c6635f7f5c32f461b","body_text":"def update_hf_weight_loader(): print('no hf weight loader need to be updated') return","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","name":"load_hf_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","span":"L30-L44","signature":"def load_hf_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"ba5f0f2cc08e7ff77c5fbacded07d517f7cee3cee904129e40a6bb641a06de22","body_text":"def load_hf_weights(actor_weights: Dict, vllm_model: nn.Module): assert isinstance(actor_weights, Dict) with set_default_torch_dtype(next(vllm_model.parameters()).dtype): # TODO if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in actor_weights.keys(): del actor_weights[\"lm_head.weight\"] vllm_model.load_weights(actor_weights.items()) for _, module in vllm_model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() vllm_model = vllm_model.cuda()","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.weight_utils","file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#update_hf_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#update_hf_weight_loader","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"next","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"load_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/hf_weight_loader.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/llm.py":{"size":12271,"mtimeMs":1789275830317.203,"hash":"47be604307df31d16239191df3051ded9a7a0d7d7a69f3ee88539686556c2f30","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","name":"llm.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L1-L240","signature":null,"exported":true,"origin":"ast","body_hash":"47be604307df31d16239191df3051ded9a7a0d7d7a69f3ee88539686556c2f30","chars":12271,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/llm.py from contextlib import contextmanager from typing import ClassVar, List, Optional, Sequence, Union, cast, overload, Dict, Tuple from tqdm import tqdm from transformers import PreTrainedTokenizer, PreTrainedTokenizerFast from transformers import PretrainedConfig import torch.nn as nn from .arg_utils import EngineArgs from .llm_engine_sp import LLMEngine from vllm import LLM from vllm.inputs import (PromptInputs, TextPrompt, TokensPrompt, parse_and_batch_prompt) from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.model_executor.guided_decoding import (GuidedDecodingRequest, get_local_guided_decoding_logits_processor) from vllm.model_executor.guided_decoding.guided_fields import LLMGuidedOptions from vllm.outputs import EmbeddingRequestOutput, RequestOutput from vllm.pooling_params import PoolingParams from vllm.prompt_adapter.request import PromptAdapterRequest from vllm.sampling_params import SamplingParams from vllm.transformers_utils.tokenizer import get_cached_tokenizer from vllm.usage.usage_lib import UsageContext from vllm.utils import Counter, deprecate_kwargs import torch from torch.nn.utils.rnn import pad_sequence from verl.workers.rollout.tokenizer import HybridEngineBaseTokenizer"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","name":"LLM","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L43-L239","signature":"class LLM(LLM)","exported":true,"origin":"ast","body_hash":"f1bc8ce4d3368fe834cf2f2898dba8bc953936aea8ccb93f06cd9bbec22bd120","body_text":"class LLM(LLM): \"\"\"An LLM for generating texts from given prompts and sampling parameters. This class includes a tokenizer, a language model (possibly distributed across multiple GPUs), and GPU memory space allocated for intermediate states (aka KV cache). Given a batch of prompts and sampling parameters, this class generates texts from the model, using an intelligent batching mechanism and efficient memory management. NOTE: This class is intended to be used for offline inference. For online serving, use the `AsyncLLMEngine` class instead. NOTE: For the comprehensive list of arguments, see `EngineArgs`. Args: model: A HuggingFace Transformers model instance. tokenizer: A HuggingFace Transformers tokenizer instance. tokenizer_mode: The tokenizer mode. \"auto\" will use the fast tokenizer if available, and \"slow\" will always use the slow tokenizer. trust_remote_code: Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer. tensor_parallel_size: The number of GPUs to use for distributed execution with tensor parallelism. dtype: The data type for the model weights and activations. Currently, we support `float32`, `float16`, and `bfloat16`. If `auto`, we use the `torch_dtype` attribute specified in the model config file. However, if the `torch_dtype` in the config is `float32`, we will use `float16` instead. quantization: The method used to quantize the model weights. Currently, we support \"awq\". If None, we assume the model weights are not quantized and use `dtype` to determine the data type of the weights. revision: The specific model version to use. It can be a branch name, a tag name, or a commit id. tokenizer_revision: The specific tokenizer version to use. It can be a branch name, a tag name, or a commit id. seed: The seed to initialize the random number generator for sampling. gpu_memory_utilization: The ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV cache. Higher values will increase the KV cache size and thus improve the model's throughput. However, if the value is too high, it may cause out-of- memory (OOM) errors. swap_space: The size (GiB) of CPU memory per GPU to use as swap space. This can be used for temporarily storing the states of the requests when their `best_of` sampling parameters are larger than 1. If all requests will have `best_of=1`, you can safely set this to 0. Otherwise, too small values may cause out-of-memory (OOM) errors. enforce_eager: Whether to enforce eager execution. If True, we will disable CUDA graph and always execute the model in eager mode. If False, we will use CUDA graph and eager execution in hybrid. max_context_len_to_capture: Maximum context len covered by CUDA graphs. When a sequence has context length larger than this, we fall back to eager mode. disable_custom_all_reduce: See ParallelConfig \"\"\" def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, skip_tokenizer_init: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, cpu_offload_gb: float = 0, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: int = 8192, disable_custom_all_reduce: bool = False, load_format = 'auto', **kwargs, ) -> None: if \"disable_log_stats\" not in kwargs: kwargs[\"disable_log_stats\"] = True engine_args = EngineArgs( model_hf_config=model_hf_config, tensor_parallel_size=tensor_parallel_size, dtype=dtype, quantization=quantization, revision=revision, tokenizer_revision=tokenizer_revision, seed=seed, gpu_memory_utilization=gpu_memory_utilization, swap_space=swap_space, cpu_offload_gb=cpu_offload_gb, enforce_eager=enforce_eager, max_context_len_to_capture=max_context_len_to_capture, max_seq_len_to_capture=max_seq_len_to_capture, disable_custom_all_reduce=disable_custom_all_reduce, load_format=load_format, skip_tokenizer_init=skip_tokenizer_init, **kwargs, ) tokenizer_cls = (PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer) if not isinstance(tokenizer, tokenizer_cls): raise ValueError( f\"Unexpected tokenizer type: {type(tokenizer)}. Must be\" \"one of the following: PreTrainedTokenizer, PreTrainedTokenizerFast, verl.workers.rollout.HybridEngineBaseTokenizer\" ) self.llm_engine = LLMEngine.from_engine_args(model, tokenizer, engine_args) # TODO: check usagecontext self.request_counter = Counter() def init_cache_engine(self): self.llm_engine.init_cache_engine() def free_cache_engine(self): self.llm_engine.free_cache_engine() def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: return self.llm_engine.tokenizer def set_token","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L97-L149","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, skip_tokenizer_init: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, cpu_offload_gb: float = 0, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: int = 8192, disable_custom_all_reduce: bool = False, load_format = 'auto', **kwargs, ) -> None","exported":false,"origin":"ast","body_hash":"5c142c4ea1a705e5417c0d42614aa0c806f149e0e62c4ee5cc139695e06b8e38","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, skip_tokenizer_init: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, cpu_offload_gb: float = 0, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: int = 8192, disable_custom_all_reduce: bool = False, load_format = 'auto', **kwargs, ) -> None: if \"disable_log_stats\" not in kwargs: kwargs[\"disable_log_stats\"] = True engine_args = EngineArgs( model_hf_config=model_hf_config, tensor_parallel_size=tensor_parallel_size, dtype=dtype, quantization=quantization, revision=revision, tokenizer_revision=tokenizer_revision, seed=seed, gpu_memory_utilization=gpu_memory_utilization, swap_space=swap_space, cpu_offload_gb=cpu_offload_gb, enforce_eager=enforce_eager, max_context_len_to_capture=max_context_len_to_capture, max_seq_len_to_capture=max_seq_len_to_capture, disable_custom_all_reduce=disable_custom_all_reduce, load_format=load_format, skip_tokenizer_init=skip_tokenizer_init, **kwargs, ) tokenizer_cls = (PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer) if not isinstance(tokenizer, tokenizer_cls): raise ValueError( f\"Unexpected tokenizer type: {type(tokenizer)}. Must be\" \"one of the following: PreTrainedTokenizer, PreTrainedTokenizerFast, verl.workers.rollout.HybridEngineBaseTokenizer\" ) self.llm_engine = LLMEngine.from_engine_args(model, tokenizer, engine_args) # TODO: check usagecontext self.request_counter = Counter()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L151-L152","signature":"def init_cache_engine(self)","exported":true,"origin":"ast","body_hash":"04d68c22194d8053c926485d043932fb2526f0d73a880157acee875717c5937f","body_text":"def init_cache_engine(self): self.llm_engine.init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L154-L155","signature":"def free_cache_engine(self)","exported":true,"origin":"ast","body_hash":"179e8502fa237bb23acb9f06952cf4ae62f3f9010408f63734f8adf6a5163378","body_text":"def free_cache_engine(self): self.llm_engine.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.get_tokenizer","name":"get_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L157-L158","signature":"def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]","exported":true,"origin":"ast","body_hash":"cebad3b0992af1915c8c8a1543b493069d6507f491421f924191ee4b98e60a2f","body_text":"def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: return self.llm_engine.tokenizer","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.set_tokenizer","name":"set_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L160-L164","signature":"def set_tokenizer( self, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], ) -> None","exported":true,"origin":"ast","body_hash":"a9d3cbc026686b5d94c85093062fb18f30a9d4182a5bf7eb9679c456cbc101fa","body_text":"def set_tokenizer( self, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], ) -> None: self.llm_engine.tokenizer = tokenizer","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","name":"_run_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L166-L202","signature":"def _run_engine(self, *, use_tqdm: bool) -> List[Union[RequestOutput, EmbeddingRequestOutput]]: # Initialize tqdm.","exported":false,"origin":"ast","body_hash":"82e45cb5c2525cdc4c73f8d391a366c1c8b1791146ecf4e8840e7dd9f2a75a35","body_text":"def _run_engine(self, *, use_tqdm: bool) -> List[Union[RequestOutput, EmbeddingRequestOutput]]: # Initialize tqdm. if use_tqdm: num_requests = self.llm_engine.get_num_unfinished_requests() pbar = tqdm( total=num_requests, desc=\"Processed prompts\", dynamic_ncols=True, postfix=(f\"est. speed input: {0:.2f} toks/s, \" f\"output: {0:.2f} toks/s\"), ) # Run the engine. outputs: List[Union[RequestOutput, EmbeddingRequestOutput]] = [] total_in_toks = 0 total_out_toks = 0 while self.llm_engine.has_unfinished_requests(): step_outputs = self.llm_engine.step() for output in step_outputs: if output.finished: outputs.append(output) if use_tqdm: if isinstance(output, RequestOutput): # Calculate tokens only for RequestOutput total_in_toks += len(output.prompt_token_ids) in_spd = total_in_toks / pbar.format_dict[\"elapsed\"] total_out_toks += sum(len(stp.token_ids) for stp in output.outputs) out_spd = total_out_toks / pbar.format_dict[\"elapsed\"] pbar.postfix = (f\"est. speed input: {in_spd:.2f} toks/s, \" f\"output: {out_spd:.2f} toks/s\") pbar.update(1) if use_tqdm: pbar.close() # Sort the outputs by request ID. # This is necessary because some requests may be finished earlier than # its previous requests. outputs = sorted(outputs, key=lambda x: int(x.request_id)) return self._post_process_outputs(outputs)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","name":"_post_process_outputs","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L214-L233","signature":"def _post_process_outputs(self, request_outputs: List[RequestOutput]) -> Tuple[torch.Tensor, torch.Tensor]","exported":false,"origin":"ast","body_hash":"05b06c04ad38ec323adde48ebace2afb11cf6796e5d16e93e4a46b004820296d","body_text":"def _post_process_outputs(self, request_outputs: List[RequestOutput]) -> Tuple[torch.Tensor, torch.Tensor]: output_token_ids = [] logprobs = [] for request_output in request_outputs: # List[RequestOutput] outputs = request_output.outputs for output in outputs: # List[CompletionOutput], usually len == 1 output_token_ids.append(torch.tensor(output.token_ids)) # TODO(shengguangming): can be optimzied by rewrite the Sampler._get_logprobs() logits logprobs_dicts = output.logprobs if logprobs_dicts is not None: logprob = [] for logprobs_dict, id in zip(logprobs_dicts, output.token_ids): logprob.append(logprobs_dict[id].logprob) logprobs.append(torch.tensor(logprob)) pad_token_id = self.llm_engine.tokenizer.pad_token_id if self.llm_engine.tokenizer.pad_token_id is not None else self.llm_engine.tokenizer.eos_token_id output_token_ids = pad_sequence(output_token_ids, batch_first=True, padding_value=pad_token_id) if len(logprobs) > 0: logprobs = pad_sequence(logprobs, batch_first=True, padding_value=pad_token_id) return output_token_ids, logprobs","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L235-L236","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None","exported":true,"origin":"ast","body_hash":"df77662639c0d1c0008e46c21bdeffd4e84ed28fa0388825031e0577e7d24cf8","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None: self.llm_engine.sync_model_weights(actor_weights=actor_weights, load_format=load_format)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","span":"L238-L239","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"27a73de1e582c115374d023e58aeaf54824c75ca2740c19f48f07d7d67773621","body_text":"def offload_model_weights(self) -> None: self.llm_engine.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"contextlib","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"tqdm","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":".arg_utils","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":".llm_engine_sp","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.inputs","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.model_executor.guided_decoding","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.model_executor.guided_decoding.guided_fields","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.outputs","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.pooling_params","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.prompt_adapter.request","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.sampling_params","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.transformers_utils.tokenizer","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.usage.usage_lib","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"torch.nn.utils.rnn","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"imports","specifier":"verl.workers.rollout.tokenizer","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","relation":"extends","name":"LLM","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.__init__","relation":"calls","name":"EngineArgs","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.__init__","relation":"calls","name":"type","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.__init__","relation":"calls","name":"from_engine_args","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.__init__","relation":"calls","name":"Counter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.init_cache_engine","relation":"calls","name":"init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.get_tokenizer","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.set_tokenizer","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"get_num_unfinished_requests","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"tqdm","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"has_unfinished_requests","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"step","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","recvType":"sorted"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"sum","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"update","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","recvType":"tqdm"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"close","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","recvType":"tqdm"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"sorted","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._run_engine","relation":"calls","name":"_post_process_outputs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","recvType":"LLM"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","recvType":"pad_sequence"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","relation":"calls","name":"zip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py","recvType":"pad_sequence"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","relation":"calls","name":"pad_sequence","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM._post_process_outputs","relation":"calls","name":"pad_sequence","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm.py#LLM.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py":{"size":14689,"mtimeMs":1789275830317.203,"hash":"85fe01207028cd8b23f86be73179e2795eb7f45420a7f48db0d27835c9dcd1d1","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","name":"llm_engine_sp.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","span":"L1-L329","signature":null,"exported":true,"origin":"ast","body_hash":"85fe01207028cd8b23f86be73179e2795eb7f45420a7f48db0d27835c9dcd1d1","chars":14689,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/engine/llm_engine.py import torch from typing import Dict, Optional, Union, Type import vllm.envs as envs from vllm.config import (CacheConfig, DecodingConfig, DeviceConfig, EngineConfig, LoRAConfig, MultiModalConfig, ObservabilityConfig, ParallelConfig, PromptAdapterConfig, SchedulerConfig, SpeculativeConfig) from vllm.core.scheduler import Scheduler from vllm.engine.output_processor.interfaces import (SequenceGroupOutputProcessor) from vllm.engine.output_processor.stop_checker import StopChecker from vllm.executor.executor_base import ExecutorBase from vllm.inputs import INPUT_REGISTRY, LLMInputs, PromptInputs from vllm.logger import init_logger from vllm.transformers_utils.detokenizer import Detokenizer from vllm.engine.metrics import (LoggingStatLogger, PrometheusStatLogger, StatLoggerBase, Stats) from vllm.tracing import (SpanAttributes, SpanKind, extract_trace_context, init_tracer) from vllm.usage.usage_lib import (UsageContext, is_usage_stats_enabled, usage_message) from vllm.utils import Counter from vllm.engine.llm_engine import _load_generation_config_dict from vllm.engine.llm_engine import LLMEngine from vllm.version import __version__ as VLLM_VERSION import torch.nn as nn from .arg_utils import EngineArgs from .tokenizer import TokenizerGroup from .config import ModelConfig, LoadConfig logger = init_logger(__name__) _LOCAL_LOGGING_INTERVAL_SEC = 5"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","name":"LLMEngine","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","span":"L46-L328","signature":"class LLMEngine(LLMEngine)","exported":true,"origin":"ast","body_hash":"3b712bf4170c87f321d05015231d14a5e54e8207687c5f679b3b2f2df3c1bd03","body_text":"class LLMEngine(LLMEngine): \"\"\"An LLM engine that receives requests and generates texts. This is the main class for the vLLM engine. It receives requests from clients and generates texts from the LLM. It includes a tokenizer, a language model (possibly distributed across multiple GPUs), and GPU memory space allocated for intermediate states (aka KV cache). This class utilizes iteration-level scheduling and efficient memory management to maximize the serving throughput. The `LLM` class wraps this class for offline batched inference and the `AsyncLLMEngine` class wraps this class for online serving. NOTE: The config arguments are derived from the `EngineArgs` class. For the comprehensive list of arguments, see `EngineArgs`. Args: model: the actor model initialize outside vllm (add for verl) tokenizer: the initialized tokenizer (add for verl) model_config: The configuration related to the LLM model. cache_config: The configuration related to the KV cache memory management. parallel_config: The configuration related to distributed execution. scheduler_config: The configuration related to the request scheduler. distributed_init_method: The initialization method for distributed execution. See `torch.distributed.init_process_group` for details. placement_group: Ray placement group for distributed execution. Required for distributed execution. log_stats: Whether to log statistics. \"\"\" def __init__( self, # NOTE(sgm): first two arguments are added for verl model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, # NOTE(sgm): vllm original arguments model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], speculative_config: Optional[SpeculativeConfig], decoding_config: Optional[DecodingConfig], observability_config: Optional[ObservabilityConfig], prompt_adapter_config: Optional[PromptAdapterConfig], executor_class: Type[ExecutorBase], log_stats: bool, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, stat_loggers: Optional[Dict[str, StatLoggerBase]] = None, ) -> None: logger.info( \"Initializing an LLM engine (v%s) with config: \" \"model=%r, speculative_config=%r, tokenizer=%r, \" \"skip_tokenizer_init=%s, revision=%s, \" \"rope_scaling=%r, rope_theta=%r, tokenizer_revision=%s, \" \"trust_remote_code=%s, dtype=%s, max_seq_len=%d, \" \"download_dir=%r, load_format=%s, tensor_parallel_size=%d, \" \"pipeline_parallel_size=%d, \" \"disable_custom_all_reduce=%s, quantization=%s, \" \"enforce_eager=%s, kv_cache_dtype=%s, \" \"quantization_param_path=%s, device_config=%s, \" \"decoding_config=%r, observability_config=%r, \" \"seed=%d, served_model_name=%s, use_v2_block_manager=%s, \" \"enable_prefix_caching=%s)\", VLLM_VERSION, model_config.model, speculative_config, model_config.tokenizer, model_config.skip_tokenizer_init, model_config.revision, model_config.rope_scaling, model_config.rope_theta, model_config.tokenizer_revision, model_config.trust_remote_code, model_config.dtype, model_config.max_model_len, load_config.download_dir, load_config.load_format, parallel_config.tensor_parallel_size, parallel_config.pipeline_parallel_size, parallel_config.disable_custom_all_reduce, model_config.quantization, model_config.enforce_eager, cache_config.cache_dtype, model_config.quantization_param_path, device_config.device, decoding_config, observability_config, model_config.seed, model_config.served_model_name, scheduler_config.use_v2_block_manager, cache_config.enable_prefix_caching, ) # TODO(woosuk): Print more configs in debug mode. self.model_config = model_config self.cache_config = cache_config self.lora_config = lora_config self.multimodal_config = multimodal_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.speculative_config = speculative_config self.load_config = load_config self.decoding_config = decoding_config or DecodingConfig() self.prompt_adapter_config = prompt_adapter_config self.observability_config = observability_config or ObservabilityConfig() self.log_stats = log_stats # self.model = model # should not store the model, it should be deleted # TODO(shengguangming): maybe we can choose init here or from arguments if not self.model_config.skip_tokenizer_init: self.tokenizer = self._init_tokenizer(tokenizer) self.detokenizer = Detokenizer(self.tokenizer) else: self.tokenizer = None self.detokenizer = None self.seq_counter = Counter() self.generation_config_fields = _load_generation_config_dict(model_config) self.input_processor = INPUT_REGISTRY.create_input_processor(self.model_config) self.model_executor = executor_class( model=model, # add for spmd_gpu_executor model_config=model_config, cache_config=cache_config, parallel_config=parallel_config, scheduler_config=scheduler_config, device_config=device_config, lora_config=","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","span":"L77-L260","signature":"def __init__( self, # NOTE(sgm): first two arguments are added for verl model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, # NOTE(sgm): vllm original arguments model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], speculative_config: Optional[SpeculativeConfig], decoding_config: Optional[DecodingConfig], observability_config: Optional[ObservabilityConfig], prompt_adapter_config: Optional[PromptAdapterConfig], executor_class: Type[ExecutorBase], log_stats: bool, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, stat_loggers: Optional[Dict[str, StatLoggerBase]] = None, ) -> None","exported":false,"origin":"ast","body_hash":"b1a21a71b19deb379fa4c007dcacb46aef595d3b944309376914f3453a0181b7","body_text":"def __init__( self, # NOTE(sgm): first two arguments are added for verl model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, # NOTE(sgm): vllm original arguments model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], speculative_config: Optional[SpeculativeConfig], decoding_config: Optional[DecodingConfig], observability_config: Optional[ObservabilityConfig], prompt_adapter_config: Optional[PromptAdapterConfig], executor_class: Type[ExecutorBase], log_stats: bool, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, stat_loggers: Optional[Dict[str, StatLoggerBase]] = None, ) -> None: logger.info( \"Initializing an LLM engine (v%s) with config: \" \"model=%r, speculative_config=%r, tokenizer=%r, \" \"skip_tokenizer_init=%s, revision=%s, \" \"rope_scaling=%r, rope_theta=%r, tokenizer_revision=%s, \" \"trust_remote_code=%s, dtype=%s, max_seq_len=%d, \" \"download_dir=%r, load_format=%s, tensor_parallel_size=%d, \" \"pipeline_parallel_size=%d, \" \"disable_custom_all_reduce=%s, quantization=%s, \" \"enforce_eager=%s, kv_cache_dtype=%s, \" \"quantization_param_path=%s, device_config=%s, \" \"decoding_config=%r, observability_config=%r, \" \"seed=%d, served_model_name=%s, use_v2_block_manager=%s, \" \"enable_prefix_caching=%s)\", VLLM_VERSION, model_config.model, speculative_config, model_config.tokenizer, model_config.skip_tokenizer_init, model_config.revision, model_config.rope_scaling, model_config.rope_theta, model_config.tokenizer_revision, model_config.trust_remote_code, model_config.dtype, model_config.max_model_len, load_config.download_dir, load_config.load_format, parallel_config.tensor_parallel_size, parallel_config.pipeline_parallel_size, parallel_config.disable_custom_all_reduce, model_config.quantization, model_config.enforce_eager, cache_config.cache_dtype, model_config.quantization_param_path, device_config.device, decoding_config, observability_config, model_config.seed, model_config.served_model_name, scheduler_config.use_v2_block_manager, cache_config.enable_prefix_caching, ) # TODO(woosuk): Print more configs in debug mode. self.model_config = model_config self.cache_config = cache_config self.lora_config = lora_config self.multimodal_config = multimodal_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.speculative_config = speculative_config self.load_config = load_config self.decoding_config = decoding_config or DecodingConfig() self.prompt_adapter_config = prompt_adapter_config self.observability_config = observability_config or ObservabilityConfig() self.log_stats = log_stats # self.model = model # should not store the model, it should be deleted # TODO(shengguangming): maybe we can choose init here or from arguments if not self.model_config.skip_tokenizer_init: self.tokenizer = self._init_tokenizer(tokenizer) self.detokenizer = Detokenizer(self.tokenizer) else: self.tokenizer = None self.detokenizer = None self.seq_counter = Counter() self.generation_config_fields = _load_generation_config_dict(model_config) self.input_processor = INPUT_REGISTRY.create_input_processor(self.model_config) self.model_executor = executor_class( model=model, # add for spmd_gpu_executor model_config=model_config, cache_config=cache_config, parallel_config=parallel_config, scheduler_config=scheduler_config, device_config=device_config, lora_config=lora_config, multimodal_config=multimodal_config, speculative_config=speculative_config, load_config=load_config, prompt_adapter_config=prompt_adapter_config, ) # Profile the memory usage and initialize the cache. if not self.model_config.embedding_mode: self._initialize_kv_caches() # If usage stat is enabled, collect relevant info. if is_usage_stats_enabled(): from vllm.model_executor.model_loader import (get_architecture_class_name) usage_message.report_usage( get_architecture_class_name(model_config), usage_context, extra_kvs={ # Common configuration \"dtype\": str(model_config.dtype), \"tensor_parallel_size\": parallel_config.tensor_parallel_size, \"block_size\": cache_config.block_size, \"gpu_memory_utilization\": cache_config.gpu_memory_utilization, # Quantization \"quantization\": model_config.quantization, \"kv_cache_dtype\": str(cache_config.cache_dtype), # Feature flags \"enable_lora\": bool(lora_config), \"enable_prompt_adapter\": bool(prompt_adapter_config), \"enable_prefix_caching\": cache_config.enable_prefix_caching, \"enforce_eager\": model_config.enforce_eager, \"disable_custom_all_reduce\": parallel_config.disable_custom_all_reduce, }) if self.tokenizer: # Ping the tokenizer to ensure liveness if it runs in a # different process. self.tokenizer.ping() # Create the scheduler. # NOTE: the cache_config here have been updated with the numbers of # GPU and CPU blocks, which are profiled","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine._init_tokenizer","name":"_init_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","span":"L263-L268","signature":"def _init_tokenizer(self, tokenizer, **tokenizer_init_kwargs)","exported":false,"origin":"ast","body_hash":"14caea76c616136989b92122bb1c73ca013e47bed311a2b9dbc88616d313fe4a","body_text":"def _init_tokenizer(self, tokenizer, **tokenizer_init_kwargs): init_kwargs = dict(enable_lora=bool(self.lora_config), max_num_seqs=self.scheduler_config.max_num_seqs, max_input_length=None) init_kwargs.update(tokenizer_init_kwargs) return TokenizerGroup(tokenizer, **init_kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","span":"L270-L273","signature":"def init_cache_engine(self): # TODO: check whether we should rebuild the CUDAGraph every iter when offload/load KVCache # Re-capture CUDAGraph would be time-consuming","exported":true,"origin":"ast","body_hash":"b8670db02a6509c498d21709df59989b89cdc29c54be45a2684c9a9514906677","body_text":"def init_cache_engine(self): # TODO: check whether we should rebuild the CUDAGraph every iter when offload/load KVCache # Re-capture CUDAGraph would be time-consuming self.model_executor.init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","span":"L275-L276","signature":"def free_cache_engine(self)","exported":true,"origin":"ast","body_hash":"61329183cfad12da6f1cfae074a6cf5526b7ffa477e5571a36f81ab4ce9e6e46","body_text":"def free_cache_engine(self): self.model_executor.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine._get_executor_cls","name":"_get_executor_cls","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","span":"L281-L290","signature":"def _get_executor_cls(cls, engine_config: EngineConfig) -> Type[ExecutorBase]","exported":false,"origin":"ast","body_hash":"f92fc7878f1c72670451b401e5fd15b852dcc92a46c4d285535ee8e924d8a777","body_text":"def _get_executor_cls(cls, engine_config: EngineConfig) -> Type[ExecutorBase]: assert engine_config.device_config.device_type == \"cuda\", \\ \"Currently, the vllm in verl only support running on GPU\" if engine_config.parallel_config.world_size == 1: engine_config.load_config.load_format = \"dummy_hf\" from .spmd_gpu_executor import SPMDGPUExecutor executor_class = SPMDGPUExecutor return executor_class","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.from_engine_args","name":"from_engine_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","span":"L293-L322","signature":"def from_engine_args( cls, model, tokenizer, engine_args: EngineArgs, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, stat_loggers: Optional[Dict[str, StatLoggerBase]] = None, ) -> \"LLMEngine\"","exported":true,"origin":"ast","body_hash":"996b7b8515a3aef99cd0d0cfd7364940fd96ef82cff9edb8cb9d7e66c903e58c","body_text":"def from_engine_args( cls, model, tokenizer, engine_args: EngineArgs, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, stat_loggers: Optional[Dict[str, StatLoggerBase]] = None, ) -> \"LLMEngine\": \"\"\"Creates an LLM engine from the engine arguments.\"\"\" # Create the engine configs. engine_config = engine_args.create_engine_config() executor_class = cls._get_executor_cls(engine_config) # Initialize the cluster and specify the executor class. assert engine_config.device_config.device_type == \"cuda\", \\ \"Currently, the vllm in verl only support running on GPU\" from .spmd_gpu_executor import SPMDGPUExecutor executor_class = SPMDGPUExecutor # Create the LLM engine. engine = cls( model, tokenizer, **engine_config.to_dict(), executor_class=executor_class, log_stats=not engine_args.disable_log_stats, usage_context=usage_context, stat_loggers=stat_loggers, ) return engine","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","span":"L324-L325","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None","exported":true,"origin":"ast","body_hash":"3e783d98664a0f1d1e1d5620565f52cce66f192cc38c10123eeb27fe5c8d2776","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None: self.model_executor.sync_model_weights(actor_weights=actor_weights, load_format=load_format)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","span":"L327-L328","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"e77bccf4128d8406c1aae651fad5c639bd60d4b251e0a060f0fdb8196cb26eaf","body_text":"def offload_model_weights(self) -> None: self.model_executor.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.core.scheduler","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.output_processor.interfaces","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.output_processor.stop_checker","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.executor.executor_base","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.inputs","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.transformers_utils.detokenizer","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.metrics","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.tracing","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.usage.usage_lib","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.llm_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.llm_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.version","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":".arg_utils","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":".tokenizer","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","relation":"extends","name":"LLMEngine","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"DecodingConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"ObservabilityConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_init_tokenizer","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Detokenizer","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Counter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_load_generation_config_dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"create_input_processor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"executor_class","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_initialize_kv_caches","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"is_usage_stats_enabled","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":"vllm.model_executor.model_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"report_usage","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"get_architecture_class_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"str","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"str","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"bool","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"bool","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"ping","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Scheduler","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"LoggingStatLogger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"PrometheusStatLogger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"init_tracer","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"create_output_processor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"StopChecker","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine._init_tokenizer","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"bool","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"update","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","recvType":"dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"TokenizerGroup","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.init_cache_engine","relation":"calls","name":"init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","recvType":"executor_class"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","recvType":"executor_class"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine._get_executor_cls","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":".spmd_gpu_executor","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.from_engine_args","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"create_engine_config","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","recvType":"EngineArgs"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"_get_executor_cls","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","relation":"imports","specifier":".spmd_gpu_executor","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"cls","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"to_dict","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","recvType":"executor_class"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py#LLMEngine.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/llm_engine_sp.py","recvType":"executor_class"}]},"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py":{"size":13939,"mtimeMs":1789275830317.203,"hash":"08dccb70b440541f1f3c68e5f519c1c419bce66360885af1084c4fee2b972079","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","name":"megatron_weight_loaders.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L1-L308","signature":null,"exported":true,"origin":"ast","body_hash":"08dccb70b440541f1f3c68e5f519c1c419bce66360885af1084c4fee2b972079","chars":13939,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models from typing import Dict import torch import torch.nn as nn from vllm.model_executor.layers.linear import * from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding, ParallelLMHead from vllm.model_executor.layers.activation import ScaledActivation from vllm.model_executor.models import ModelRegistry # NOTE(shengguangming): replace the origin weight loader function in the class __LAYER_WEIGHT_MEGATRON_LOADER_REGISTRY__ = { ColumnParallelLinear: parallel_weight_loader, MergedColumnParallelLinear: parallel_weight_loader, QKVParallelLinear: parallel_weight_loader, RowParallelLinear: parallel_weight_loader, VocabParallelEmbedding: parallel_weight_loader, ParallelLMHead: parallel_weight_loader # \"ScaledActivation.weight_loader\": ScaledActivation, # TODO(shengguangming): latest commit in vllm fix awq for this function and add load_weights # \"default_weight_loader\": default_weight_loader } # for layer_class, weight_loader in __LAYER_WEIGHT_MEGATRON_LOADER_REGISTRY__.items(): # # setattr(layer_class, 'megatron_weight_loader', weight_loader) # layer_class.weight_loader = weight_loader __MODEL_MEGATRON_WEIGHT_LOADER_REGISTRY__ = { 'GPT2LMHeadModel': gpt2_weight_loader, 'LlamaForCausalLM': llama_megatron_weight_loader, # use te backend for open-source megatron 'LLaMAForCausalLM': llama_megatron_weight_loader, 'MistralForCausalLM': mistral_megatron_weight_loader, } # the actor model is .state_dict() # Load megatron weights"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#parallel_weight_loader","name":"parallel_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L27-L34","signature":"def parallel_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None","exported":true,"origin":"ast","body_hash":"0f25a6818304178a092b3b28b97800a19b34e8bb12872e7a788083897c25df2a","body_text":"def parallel_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: \"\"\"Parallel Linear weight loader.\"\"\" assert param.size() == loaded_weight.size( ), 'the parameter size is not align with the loaded weight size, param size: {}, loaded_weight size: {}'.format( param.size(), loaded_weight.size()) assert param.data.dtype == loaded_weight.data.dtype, \"if we want to shared weights, the data type should also be the same\" param.data = loaded_weight.data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#default_weight_loader","name":"default_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L37-L42","signature":"def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None","exported":true,"origin":"ast","body_hash":"27bc592da2cd05ea4807cdcdd1b8feb48351a0a66ebfda38aaa38edd7f0f770b","body_text":"def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: \"\"\"Default weight loader.\"\"\" assert param.size() == loaded_weight.size() assert param.data.dtype == loaded_weight.data.dtype, \"if we want to shared weights, the data type should also be the same\" param.data = loaded_weight.data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#gpt2_weight_loader","name":"gpt2_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L45-L70","signature":"def gpt2_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"d7af0ddffa9a06f7c8b64a06c96f5ce1cbbbd5f2c03114b137268dd5d69048eb","body_text":"def gpt2_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"lm_head.weight\" in name: # GPT-2 ties the weights of the embedding layer and the final # linear layer. continue if \".attn.bias\" in name or \".attn.masked_bias\" in name: # Skip attention mask. # NOTE: \"c_attn.bias\" should not be skipped. continue if not name.startswith(\"transformer.\"): name = \"transformer.\" + name param = params_dict[name] # The HF's GPT-2 implementation uses Conv1D instead of Linear. # Because of this, we need to transpose the weights. # Note(zhuohan): the logic below might break quantized models. for conv1d_weight_name in [\"c_attn\", \"c_proj\", \"c_fc\"]: if conv1d_weight_name not in name: continue if not name.endswith(\".weight\"): continue # TODO: check megatron loaded_weight = loaded_weight.t() weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_weight_loader","name":"llama_megatron_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L73-L82","signature":"def llama_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # NOTE(shengguangming): the megatron llama may have this prefix","exported":true,"origin":"ast","body_hash":"e85f937f0d1b176d431ed32c7fdb4f44e023e98630b81ff16d89c4179292cf97","body_text":"def llama_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","name":"llama_megatron_core_te_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L85-L113","signature":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"6f4f63949dd5590450d23bed267b13c9b246905350e2f2b9fb03a6adfb751030","body_text":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"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.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", 'self_attn.o_proj'), ('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'), ('decoder.final_layernorm', 'model.norm'), ('output_layer', 'lm_head'), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith('.bias') and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader","name":"llama_megatron_core_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L116-L143","signature":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"ebabfc744aea05f3781b4d7fc6869a63a03c65929ca5310a1aa1c2303b787681","body_text":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"self_attention.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", 'self_attn.o_proj'), ( 'input_layernorm', 'input_layernorm', ), ('pre_mlp_layernorm', 'post_attention_layernorm'), ('mlp.linear_fc1', 'mlp.gate_up_proj'), ('mlp.linear_fc2', 'mlp.down_proj'), ('decoder.final_layernorm', 'model.norm'), ('output_layer', 'lm_head'), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith('.bias') and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name","name":"_replace_name","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L146-L166","signature":"def _replace_name(megatron_name, name_mapping)","exported":false,"origin":"ast","body_hash":"0c9ae7affadf31827df88ca4b341b2ef683cd56ad7ed1abc8cc9d9fba2f48f84","body_text":"def _replace_name(megatron_name, name_mapping): for m_name, v_name in name_mapping: if m_name not in megatron_name: continue if 'layers' in megatron_name: # deal with decoder layers megatron_name = megatron_name.replace('decoder', 'model') megatron_name_list = megatron_name.split('.') if 'layer_norm_weight' in megatron_name_list or 'layer_norm_bias' in megatron_name_list: param_name_list = megatron_name_list[:3] param_name_list.append(v_name) param_name = '.'.join(param_name_list) else: param_name_list = megatron_name_list[:3] weight_or_bias = megatron_name_list[-1] param_name_list.append(v_name) param_name_list.append(weight_or_bias) param_name = '.'.join(param_name_list) return param_name else: param_name = megatron_name.replace(m_name, v_name) return param_name","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","name":"llama_megatron_core_te_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L169-L197","signature":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"6f4f63949dd5590450d23bed267b13c9b246905350e2f2b9fb03a6adfb751030","body_text":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"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.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", 'self_attn.o_proj'), ('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'), ('decoder.final_layernorm', 'model.norm'), ('output_layer', 'lm_head'), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith('.bias') and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","name":"llama_megatron_core_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L200-L227","signature":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"ebabfc744aea05f3781b4d7fc6869a63a03c65929ca5310a1aa1c2303b787681","body_text":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"self_attention.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", 'self_attn.o_proj'), ( 'input_layernorm', 'input_layernorm', ), ('pre_mlp_layernorm', 'post_attention_layernorm'), ('mlp.linear_fc1', 'mlp.gate_up_proj'), ('mlp.linear_fc2', 'mlp.down_proj'), ('decoder.final_layernorm', 'model.norm'), ('output_layer', 'lm_head'), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith('.bias') and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name~2","name":"_replace_name","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L230-L250","signature":"def _replace_name(megatron_name, name_mapping)","exported":false,"origin":"ast","body_hash":"0c9ae7affadf31827df88ca4b341b2ef683cd56ad7ed1abc8cc9d9fba2f48f84","body_text":"def _replace_name(megatron_name, name_mapping): for m_name, v_name in name_mapping: if m_name not in megatron_name: continue if 'layers' in megatron_name: # deal with decoder layers megatron_name = megatron_name.replace('decoder', 'model') megatron_name_list = megatron_name.split('.') if 'layer_norm_weight' in megatron_name_list or 'layer_norm_bias' in megatron_name_list: param_name_list = megatron_name_list[:3] param_name_list.append(v_name) param_name = '.'.join(param_name_list) else: param_name_list = megatron_name_list[:3] weight_or_bias = megatron_name_list[-1] param_name_list.append(v_name) param_name_list.append(weight_or_bias) param_name = '.'.join(param_name_list) return param_name else: param_name = megatron_name.replace(m_name, v_name) return param_name","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#mistral_megatron_weight_loader","name":"mistral_megatron_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L253-L262","signature":"def mistral_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # TODO: need to implement a general way to deal with prefix","exported":true,"origin":"ast","body_hash":"da2308149252bdda6caeac91a48a8428bb8f71adb703df0b69a2008805ed02cd","body_text":"def mistral_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # TODO: need to implement a general way to deal with prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#load_megatron_weights","name":"load_megatron_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L290-L295","signature":"def load_megatron_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"308ec912ee64a38977bfc7739ea5f87730fbe66a16585ef28a7e106241f7e5b9","body_text":"def load_megatron_weights(actor_weights: Dict, vllm_model: nn.Module): weight_loader = _get_model_weight_loader(vllm_model.__class__.__name__) weight_loader(actor_weights, vllm_model) # NOTE(sgm) to reduce peak memory usage, we offload vllm model to cpu # after init, and we need this after sync model weights for in first iter. vllm_model = vllm_model.cuda()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_get_model_weight_loader","name":"_get_model_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L298-L302","signature":"def _get_model_weight_loader(arch: str)","exported":false,"origin":"ast","body_hash":"f208266263ddc85ba8ce8716c63c6dcb827fbdab60f6cde510bc807de3ed6ed3","body_text":"def _get_model_weight_loader(arch: str): if arch in __MODEL_MEGATRON_WEIGHT_LOADER_REGISTRY__: return __MODEL_MEGATRON_WEIGHT_LOADER_REGISTRY__[arch] raise ValueError(f\"Model architectures {arch} are not supported for now. \" f\"Supported architectures: {ModelRegistry.get_supported_archs()}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#update_megatron_weight_loader","name":"update_megatron_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","span":"L305-L307","signature":"def update_megatron_weight_loader()","exported":true,"origin":"ast","body_hash":"795025216be7d96096150c3271a691f3f03257fc38134d8fe0313242e0025e67","body_text":"def update_megatron_weight_loader(): for layer_class, weight_loader in __LAYER_WEIGHT_MEGATRON_LOADER_REGISTRY__.items(): layer_class.weight_loader = weight_loader","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.linear","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.vocab_parallel_embedding","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.activation","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.models","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#parallel_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#default_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#default_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#default_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#gpt2_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"startswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"t","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"split","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name~2","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"split","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#mistral_megatron_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#load_megatron_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#load_megatron_weights","relation":"calls","name":"_get_model_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#load_megatron_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#load_megatron_weights","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_get_model_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"get_supported_archs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#update_megatron_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py#update_megatron_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/megatron_weight_loaders.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py":{"size":14974,"mtimeMs":1789275830317.203,"hash":"0e43f762258581069c1d20717591a6cbeb42288c033fa2aabc92c77f536768f4","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","name":"model_loader.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L1-L303","signature":null,"exported":true,"origin":"ast","body_hash":"0e43f762258581069c1d20717591a6cbeb42288c033fa2aabc92c77f536768f4","chars":14974,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/model_loader from typing import Dict, Union, Optional, Iterable, Tuple import torch import torch.nn as nn from transformers import PreTrainedModel from vllm.config import (CacheConfig, DeviceConfig, LoadConfig, LoRAConfig, ModelConfig, MultiModalConfig, ParallelConfig, SchedulerConfig) from vllm.model_executor.model_loader import BaseModelLoader from vllm.model_executor.model_loader.loader import _initialize_model from vllm.model_executor.model_loader.utils import set_default_torch_dtype from vllm.distributed.communication_op import tensor_model_parallel_all_gather from .config import ModelConfig, LoadFormat, LoadConfig from .megatron_weight_loaders import load_megatron_weights, update_megatron_weight_loader from .dtensor_weight_loaders import load_dtensor_weights, update_dtensor_weight_loader from .hf_weight_loader import update_hf_weight_loader # FIXME(sgm): hack the _get_logits function in vllm v0.4.2 # as they use ray, the _get_logits result will only need to return to the driver node, # therefore gather is enough. However, we use SPMD instead of a central scheduler, # all_gather is required (aligned with v0.2.6) from vllm.model_executor.layers.logits_processor import LogitsProcessor LogitsProcessor.__init__ = logitsprocessor_init # use all_gather"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model","name":"get_model","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L35-L61","signature":"def get_model(actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, load_config: LoadConfig, device_config: DeviceConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], cache_config: CacheConfig = None) -> nn.Module","exported":true,"origin":"ast","body_hash":"174e42e2ae6c88c4ef58eaa628ebbd25fe66a0320d48116663adef662b4b5267","body_text":"def get_model(actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, load_config: LoadConfig, device_config: DeviceConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], cache_config: CacheConfig = None) -> nn.Module: loader = get_model_loader(load_config) if load_config.load_format.startswith('dummy'): return loader.load_model(model_config=model_config, device_config=device_config, lora_config=lora_config, multimodal_config=multimodal_config, parallel_config=parallel_config, scheduler_config=scheduler_config, cache_config=cache_config) else: return loader.load_model(actor_model=actor_model, model_config=model_config, device_config=device_config, lora_config=lora_config, multimodal_config=multimodal_config, parallel_config=parallel_config, scheduler_config=scheduler_config, cache_config=cache_config)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","name":"get_model_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L64-L100","signature":"def get_model_loader(load_config: LoadConfig) -> BaseModelLoader","exported":true,"origin":"ast","body_hash":"5f589ff0ab3f6c557941be091d8e3ac943a019d9ce46843e033c626fc089f67c","body_text":"def get_model_loader(load_config: LoadConfig) -> BaseModelLoader: \"\"\"Get a model loader based on the load format.\"\"\" if isinstance(load_config.load_format, type): return load_config.load_format(load_config) if load_config.load_format == LoadFormat.AUTO: update_megatron_weight_loader() return MegatronLoader(load_config) # NOTE(sgm): change the weight_loader function in runtime if load_config.load_format == LoadFormat.MEGATRON: update_megatron_weight_loader() return MegatronLoader(load_config) if load_config.load_format == LoadFormat.HF: update_hf_weight_loader() return HFLoader(load_config) if load_config.load_format == LoadFormat.DTENSOR: update_dtensor_weight_loader() return DTensorLoader(load_config) if load_config.load_format == LoadFormat.DUMMY_HF: update_hf_weight_loader() return DummyModelLoader(load_config) if load_config.load_format == LoadFormat.DUMMY_MEGATRON: update_megatron_weight_loader() return DummyModelLoader(load_config) if load_config.load_format == LoadFormat.DUMMY_DTENSOR: update_dtensor_weight_loader() return DummyModelLoader(load_config) raise ValueError('load format not supported in verl: {}, only support {} and {}'.format( load_config.load_format, LoadFormat.MEGATRON, LoadFormat.HF))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader","name":"DummyModelLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L103-L122","signature":"class DummyModelLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"71401300b687313a6081fe930ebf26d8e546e2d5cbd54a8017aaf0a601988feb","body_text":"class DummyModelLoader(BaseModelLoader): \"\"\"Model loader that will set model weights to random values.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def load_model(self, *, model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, multimodal_config, cache_config, scheduler_config) # NOTE(woosuk): For accurate performance evaluation, we assign # random values to the weights. # initialize_dummy_weights(model) return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L106-L110","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"DummyModelLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L112-L122","signature":"def load_model(self, *, model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module","exported":true,"origin":"ast","body_hash":"11cb314e365f2376b03be4c8b4c90eda1f90f4b38fba2f35b42e81d3b4497495","body_text":"def load_model(self, *, model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, multimodal_config, cache_config, scheduler_config) # NOTE(woosuk): For accurate performance evaluation, we assign # random values to the weights. # initialize_dummy_weights(model) return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"DummyModelLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader","name":"MegatronLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L125-L169","signature":"class MegatronLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"f40c4aa02899afa502171f28c534847b9fb34eef4fae39f631581d20e5b4cf19","body_text":"class MegatronLoader(BaseModelLoader): \"\"\"Model loader that can load the model weights from partitioned megatron model.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, multimodal_config, cache_config, scheduler_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_megatron_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_megatron_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L128-L132","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader._get_weights_iterator","name":"_get_weights_iterator","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L134-L141","signature":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model","exported":false,"origin":"ast","body_hash":"694724aa8c3287baa407271e8b4f4dab178819219fb63707f139916d9a3e0340","body_text":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L143-L169","signature":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module","exported":true,"origin":"ast","body_hash":"64ac7df89985b55daa0d49cca34ce2cd21e367f6ea1b03458bf65c116625948e","body_text":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, multimodal_config, cache_config, scheduler_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_megatron_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_megatron_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader","name":"HFLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L172-L209","signature":"class HFLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"ccc707265e95b5f46a05ce46d92f8f68dbacb686ee5d2cc591b376426b4374f1","body_text":"class HFLoader(BaseModelLoader): \"\"\"Model loader that can load the model weights from model's full params.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def _get_weights_iterator(self, actor_model: Union[PreTrainedModel, Dict]): if isinstance(actor_model, Dict): return actor_model.items() elif isinstance(actor_model, nn.Module): return dict(actor_model.named_parameters()).items() else: raise ValueError(f'actor model should be Dict or nn.Module, but get {type(actor_model)}') def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): # with torch.device(device_config.device): # NOTE(sgm): init the model in cpu model = _initialize_model(model_config, self.load_config, lora_config, multimodal_config, cache_config, scheduler_config) model.load_weights(self._get_weights_iterator(actor_model)) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L175-L179","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"HFLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader._get_weights_iterator","name":"_get_weights_iterator","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L181-L187","signature":"def _get_weights_iterator(self, actor_model: Union[PreTrainedModel, Dict])","exported":false,"origin":"ast","body_hash":"06a16b0f3b8c8c0cc050f443179f9732fd3c336a501b75e7f33660f421c53568","body_text":"def _get_weights_iterator(self, actor_model: Union[PreTrainedModel, Dict]): if isinstance(actor_model, Dict): return actor_model.items() elif isinstance(actor_model, nn.Module): return dict(actor_model.named_parameters()).items() else: raise ValueError(f'actor model should be Dict or nn.Module, but get {type(actor_model)}')","summary_state":"pending","summary":null,"crux":null,"owner":"HFLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L189-L209","signature":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module","exported":true,"origin":"ast","body_hash":"d7d03beea1d0923172e88b8d117f320d43ee36ae9950744857743cf8afbe8b0b","body_text":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): # with torch.device(device_config.device): # NOTE(sgm): init the model in cpu model = _initialize_model(model_config, self.load_config, lora_config, multimodal_config, cache_config, scheduler_config) model.load_weights(self._get_weights_iterator(actor_model)) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"HFLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader","name":"DTensorLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L212-L256","signature":"class DTensorLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"1d7f295e3c6e641e8711e68a20138396008a1cf04d2c23d15ebd942236af24fc","body_text":"class DTensorLoader(BaseModelLoader): \"\"\"Model loader that can load the model weights from partitioned megatron model.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, multimodal_config, cache_config, scheduler_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_dtensor_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_dtensor_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L215-L219","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"DTensorLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader._get_weights_iterator","name":"_get_weights_iterator","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L221-L228","signature":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model","exported":false,"origin":"ast","body_hash":"694724aa8c3287baa407271e8b4f4dab178819219fb63707f139916d9a3e0340","body_text":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model","summary_state":"pending","summary":null,"crux":null,"owner":"DTensorLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L230-L256","signature":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module","exported":true,"origin":"ast","body_hash":"5833f75ea773bbd7b8db0c10c6d55b4a3d6a588c5bb13254768fa4a08952b1dd","body_text":"def load_model(self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, multimodal_config, cache_config, scheduler_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_dtensor_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_dtensor_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"DTensorLoader"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#_get_logits","name":"_get_logits","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L263-L273","signature":"def _get_logits(self, hidden_states: torch.Tensor, embedding: torch.Tensor, embedding_bias: Optional[torch.Tensor]) -> torch.Tensor: # Get the logits for the next tokens.","exported":false,"origin":"ast","body_hash":"d294e846526a09f100321505cbacf5666d45ef972200a51c4e4c25557ec585e0","body_text":"def _get_logits(self, hidden_states: torch.Tensor, embedding: torch.Tensor, embedding_bias: Optional[torch.Tensor]) -> torch.Tensor: # Get the logits for the next tokens. logits = torch.matmul(hidden_states, embedding.t()) if embedding_bias is not None: logits += embedding_bias logits = tensor_model_parallel_all_gather(logits) # Remove paddings in vocab (if any). if logits is not None: logits = logits[:, :self.org_vocab_size] return logits","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#logitsprocessor_init","name":"logitsprocessor_init","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","span":"L279-L299","signature":"def logitsprocessor_init(self, vocab_size: int, org_vocab_size: Optional[int] = None, scale: float = 1.0, logits_as_input: bool = False, soft_cap: Optional[float] = None) -> None","exported":true,"origin":"ast","body_hash":"526dc28c5304194f58b7f1678afd08bf8e66553eb3f1dc3ac00f7ac9adfa0ce2","body_text":"def logitsprocessor_init(self, vocab_size: int, org_vocab_size: Optional[int] = None, scale: float = 1.0, logits_as_input: bool = False, soft_cap: Optional[float] = None) -> None: \"\"\" Args: scale: A scaling factor to apply to the logits. \"\"\" super(LogitsProcessor, self).__init__() self.scale = scale self.vocab_size = vocab_size # Whether the input is logits (default is hidden states). self.logits_as_input = logits_as_input # original vocabulary size (without LoRA). self.org_vocab_size = org_vocab_size or vocab_size # Soft cap the logits. Used in Gemma 2. self.soft_cap = soft_cap # Whether to use gather or all-gather to gather the logits. self.use_gather = False","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.loader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":"vllm.distributed.communication_op","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":".megatron_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":".dtensor_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":".hf_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model","relation":"calls","name":"get_model_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model","relation":"calls","name":"startswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"get_model_loader"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"get_model_loader"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"load_format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"LoadConfig"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"update_megatron_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"MegatronLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"update_megatron_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"MegatronLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"update_hf_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"HFLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"update_dtensor_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"DTensorLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"update_hf_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"DummyModelLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"update_megatron_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"DummyModelLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"update_dtensor_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"DummyModelLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#get_model_loader","relation":"calls","name":"format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader._get_weights_iterator","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"load_megatron_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"load_megatron_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader._get_weights_iterator","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"type","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"load_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"_get_weights_iterator","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"HFLoader"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#HFLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader._get_weights_iterator","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"load_dtensor_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"load_dtensor_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#_get_logits","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#_get_logits","relation":"calls","name":"matmul","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#_get_logits","relation":"calls","name":"t","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#_get_logits","relation":"calls","name":"tensor_model_parallel_all_gather","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"imports","specifier":"vllm.model_executor.layers.logits_processor","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#logitsprocessor_init","file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#logitsprocessor_init","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py#logitsprocessor_init","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_loader.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py":{"size":6857,"mtimeMs":1789275830317.203,"hash":"5f3ed36a09a9ae7c3ce2036c77f955a6c1b5ed1999dcf7e4df65a2b3ce606c1d","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","name":"model_runner.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","span":"L1-L151","signature":null,"exported":true,"origin":"ast","body_hash":"5f3ed36a09a9ae7c3ce2036c77f955a6c1b5ed1999dcf7e4df65a2b3ce606c1d","chars":6857,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/worker/model_runner.py import torch import torch.nn as nn from enum import IntEnum from typing import Dict, List, Optional, Set, Tuple, Union import warnings import vllm.envs as envs from vllm.attention import (AttentionMetadata, get_attn_backend) from vllm.config import (CacheConfig, DeviceConfig, LoRAConfig, MultiModalConfig, ParallelConfig, PromptAdapterConfig, SchedulerConfig) from vllm.logger import init_logger from vllm.lora.layers import LoRAMapping from vllm.lora.request import LoRARequest from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager from vllm.model_executor import SamplingMetadata from vllm.model_executor.models.interfaces import (supports_lora, supports_vision) from vllm.utils import (CudaMemoryProfiler, is_hip, is_pin_memory_available) from vllm.worker.model_runner import ModelRunner, CUDAGraphRunner from vllm.prompt_adapter.worker_manager import (LRUCacheWorkerPromptAdapterManager) from .model_loader import get_model from .config import ModelConfig, LoadConfig logger = init_logger(__name__) # How batches are constructed."},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#BatchType","name":"BatchType","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","span":"L43-L49","signature":"class BatchType(IntEnum): # Every batch is prefill.","exported":true,"origin":"ast","body_hash":"2936fda8ae00fc0b54ec6c67c640b8e048c86af85365d8366a5dd10212f1a15d","body_text":"class BatchType(IntEnum): # Every batch is prefill. PREFILL = 0 # Every batch is decode. DECODE = 1 # Batch is a mixture of prefill and decode. MIXED = 2","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner","name":"ModelRunner","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","span":"L52-L150","signature":"class ModelRunner(ModelRunner)","exported":true,"origin":"ast","body_hash":"17d5f99a9058e42254b4517a3e757e82a172400aaeaa8b8460cbb33313b5770f","body_text":"class ModelRunner(ModelRunner): def __init__( self, model: Union[nn.Module, Dict], # [verl] model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", is_driver_worker: bool = False, prompt_adapter_config: Optional[PromptAdapterConfig] = None, multimodal_config: Optional[MultiModalConfig] = None, return_hidden_states: bool = False, ): super().__init__( model_config, parallel_config, scheduler_config, device_config, cache_config, load_config, lora_config, kv_cache_dtype, is_driver_worker=True, # a hack prompt_adapter_config=prompt_adapter_config, multimodal_config=multimodal_config, return_hidden_states=return_hidden_states) # NOTE(sgm): add for verl self.model = model # this will be replaced by get_model() # NOTE(sgm): initialize model using the actor model def load_model(self) -> None: logger.info(\"Starting to load model %s...\", self.model_config.model) with CudaMemoryProfiler() as m: self.model = get_model(actor_model=self.model, model_config=self.model_config, device_config=self.device_config, lora_config=self.lora_config, load_config=self.load_config, parallel_config=self.parallel_config, scheduler_config=self.scheduler_config, multimodal_config=self.multimodal_config, cache_config=self.cache_config) self.model_memory_usage = m.consumed_memory logger.info(\"Loading model weights took %.4f GB\", self.model_memory_usage / float(2**30)) if self.lora_config: assert supports_lora(self.model), \"Model does not support LoRA\" assert not supports_vision(self.model), \"To be tested: vision language model with LoRA settings.\" self.lora_manager = LRUCacheWorkerLoRAManager( self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens, self.vocab_size, self.lora_config, self.device, self.model.embedding_modules, self.model.embedding_padding_modules, max_position_embeddings=self.model.config.max_position_embeddings, ) self.model = self.lora_manager.create_lora_manager(self.model) if self.prompt_adapter_config: self.prompt_adapter_manager = LRUCacheWorkerPromptAdapterManager( self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens, self.device, self.prompt_adapter_config) self.model = (self.prompt_adapter_manager.create_prompt_adapter_manager(self.model)) if self.kv_cache_dtype == \"fp8\" and is_hip(): # Currently only ROCm accepts kv-cache scaling factors # via quantization_param_path and this will be deprecated # in the future. if self.model_config.quantization_param_path is not None: if callable(getattr(self.model, \"load_kv_cache_scales\", None)): warnings.warn( \"Loading kv cache scaling factor from JSON is \" \"deprecated and will be removed. Please include \" \"kv cache scaling factors in the model checkpoint.\", FutureWarning, stacklevel=2) self.model.load_kv_cache_scales(self.model_config.quantization_param_path) logger.info(\"Loaded KV cache scaling factors from %s\", self.model_config.quantization_param_path) else: raise RuntimeError( \"Using FP8 KV cache and scaling factors provided but \" \"model %s does not support loading scaling factors.\", self.model.__class__) else: logger.warning(\"Using FP8 KV cache but no scaling factors \" \"provided. Defaulting to scaling factors of 1.0. \" \"This may lead to less accurate results!\") if envs.VLLM_TEST_DYNAMO_GRAPH_CAPTURE: self.model = torch.compile(self.model, fullgraph=True, backend=\"eager\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","span":"L54-L86","signature":"def __init__( self, model: Union[nn.Module, Dict], # [verl] model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", is_driver_worker: bool = False, prompt_adapter_config: Optional[PromptAdapterConfig] = None, multimodal_config: Optional[MultiModalConfig] = None, return_hidden_states: bool = False, )","exported":false,"origin":"ast","body_hash":"d7550ba499f4f84657dda3cfaae57c56ff42fcb17223ad9e221ad009f2317683","body_text":"def __init__( self, model: Union[nn.Module, Dict], # [verl] model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", is_driver_worker: bool = False, prompt_adapter_config: Optional[PromptAdapterConfig] = None, multimodal_config: Optional[MultiModalConfig] = None, return_hidden_states: bool = False, ): super().__init__( model_config, parallel_config, scheduler_config, device_config, cache_config, load_config, lora_config, kv_cache_dtype, is_driver_worker=True, # a hack prompt_adapter_config=prompt_adapter_config, multimodal_config=multimodal_config, return_hidden_states=return_hidden_states) # NOTE(sgm): add for verl self.model = model # this will be replaced by get_model()","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","span":"L89-L150","signature":"def load_model(self) -> None","exported":true,"origin":"ast","body_hash":"010c8088ed5802d4ce3450238bc7b4646e127163b45093678c939f3013e26af3","body_text":"def load_model(self) -> None: logger.info(\"Starting to load model %s...\", self.model_config.model) with CudaMemoryProfiler() as m: self.model = get_model(actor_model=self.model, model_config=self.model_config, device_config=self.device_config, lora_config=self.lora_config, load_config=self.load_config, parallel_config=self.parallel_config, scheduler_config=self.scheduler_config, multimodal_config=self.multimodal_config, cache_config=self.cache_config) self.model_memory_usage = m.consumed_memory logger.info(\"Loading model weights took %.4f GB\", self.model_memory_usage / float(2**30)) if self.lora_config: assert supports_lora(self.model), \"Model does not support LoRA\" assert not supports_vision(self.model), \"To be tested: vision language model with LoRA settings.\" self.lora_manager = LRUCacheWorkerLoRAManager( self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens, self.vocab_size, self.lora_config, self.device, self.model.embedding_modules, self.model.embedding_padding_modules, max_position_embeddings=self.model.config.max_position_embeddings, ) self.model = self.lora_manager.create_lora_manager(self.model) if self.prompt_adapter_config: self.prompt_adapter_manager = LRUCacheWorkerPromptAdapterManager( self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens, self.device, self.prompt_adapter_config) self.model = (self.prompt_adapter_manager.create_prompt_adapter_manager(self.model)) if self.kv_cache_dtype == \"fp8\" and is_hip(): # Currently only ROCm accepts kv-cache scaling factors # via quantization_param_path and this will be deprecated # in the future. if self.model_config.quantization_param_path is not None: if callable(getattr(self.model, \"load_kv_cache_scales\", None)): warnings.warn( \"Loading kv cache scaling factor from JSON is \" \"deprecated and will be removed. Please include \" \"kv cache scaling factors in the model checkpoint.\", FutureWarning, stacklevel=2) self.model.load_kv_cache_scales(self.model_config.quantization_param_path) logger.info(\"Loaded KV cache scaling factors from %s\", self.model_config.quantization_param_path) else: raise RuntimeError( \"Using FP8 KV cache and scaling factors provided but \" \"model %s does not support loading scaling factors.\", self.model.__class__) else: logger.warning(\"Using FP8 KV cache but no scaling factors \" \"provided. Defaulting to scaling factors of 1.0. \" \"This may lead to less accurate results!\") if envs.VLLM_TEST_DYNAMO_GRAPH_CAPTURE: self.model = torch.compile(self.model, fullgraph=True, backend=\"eager\")","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"enum","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"warnings","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.attention","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.lora.layers","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.lora.worker_manager","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.model_executor","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.model_executor.models.interfaces","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.worker.model_runner","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":"vllm.prompt_adapter.worker_manager","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":".model_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#BatchType","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#BatchType","relation":"extends","name":"IntEnum","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner","relation":"extends","name":"ModelRunner","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"CudaMemoryProfiler","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"get_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"float","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"supports_lora","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"supports_vision","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"LRUCacheWorkerLoRAManager","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"create_lora_manager","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","recvType":"LRUCacheWorkerLoRAManager"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"LRUCacheWorkerPromptAdapterManager","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"create_prompt_adapter_manager","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","recvType":"LRUCacheWorkerPromptAdapterManager"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"is_hip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"callable","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"warn","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"load_kv_cache_scales","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","recvType":"get_model"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py#ModelRunner.load_model","relation":"calls","name":"compile","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/model_runner.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py":{"size":13091,"mtimeMs":1789275830317.203,"hash":"f9b0ec025d5523972e62f343f5a80c35dc6f68199a735ee1dc6f45e68e731fbb","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","name":"parallel_state.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L1-L304","signature":null,"exported":true,"origin":"ast","body_hash":"f9b0ec025d5523972e62f343f5a80c35dc6f68199a735ee1dc6f45e68e731fbb","chars":13091,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM team. # Adapted from # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. \"\"\"Model and data parallel groups.\"\"\" import os import torch import torch.distributed from typing import Optional import vllm.distributed.parallel_state as ps from vllm.distributed.parallel_state import get_pp_group, get_world_group, init_distributed_environment, init_model_parallel_group import vllm.envs as envs from vllm.logger import init_logger from torch.distributed.device_mesh import init_device_mesh logger = init_logger(__name__) \"\"\" This version is strongly tied with Megatron to implement HybridEngine and weight sharing between vllm and Megatron. - We assume the Megatron tp+dp+pp world is already established before calling this function. \"\"\" # Device mesh for using DTensor _DEVICE_MESH = None # Tensor model parallel group that the current rank belongs to. _TP = None # Pipeline model parallel group that the current rank belongs to. _PP = None # This method is for initializing the ParallelGroup when using HybridEngine # TODO(sgm): deviate from the v0.5.4, not pp now \"\"\" Device mesh utilities \"\"\" \"\"\" Tensor model parallel utilities \"\"\""},{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","name":"initialize_parallel_state","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L37-L65","signature":"def initialize_parallel_state( distributed_init_method: str = \"env://\", backend: str = \"nccl\", tensor_model_parallel_size: int = 1, num_tp_per_train_tp: int = 1, pipeline_model_parallel_size: int = 1, ): # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573","exported":true,"origin":"ast","body_hash":"217b72043f1a4eb925270bf6070f0f1469d3376d7a2f0ff825165e00247c2cbb","body_text":"def initialize_parallel_state( distributed_init_method: str = \"env://\", backend: str = \"nccl\", tensor_model_parallel_size: int = 1, num_tp_per_train_tp: int = 1, pipeline_model_parallel_size: int = 1, ): # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # NOTE(sgm): Modify for verl, Env vars will be set by TORCHRUN. rank = int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) # Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" init_distributed_environment(world_size, rank, distributed_init_method, local_rank, backend) if torch.distributed.get_world_size() > 1: # NOTE: build a sepearate inference group with infer tp & micro dp initialize_model_parallel_for_vllm(tensor_model_parallel_size=tensor_model_parallel_size, num_tensor_model_parallel_groups_per_train_tp=num_tp_per_train_tp) else: initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size, backend)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#ensure_model_parallel_initialized","name":"ensure_model_parallel_initialized","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L68-L91","signature":"def ensure_model_parallel_initialized( tensor_model_parallel_size: int, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None","exported":true,"origin":"ast","body_hash":"ee2d7fbd303c49788967d4627d0dc1178465269367c334e800181dde28d329e1","body_text":"def ensure_model_parallel_initialized( tensor_model_parallel_size: int, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None: \"\"\"Helper to initialize model parallel groups if they are not initialized, or ensure tensor-parallel and pipeline-parallel sizes are equal to expected values if the model parallel groups are initialized. \"\"\" # get the backend of _DEVICE_WORLD_GROUP backend = backend or torch.distributed.get_backend(get_world_group().device_group) if not model_parallel_is_initialized(): initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size, backend) return assert (get_tensor_model_parallel_world_size() == tensor_model_parallel_size), ( \"tensor parallel group already initialized, but of unexpected size: \" f\"{get_tensor_model_parallel_world_size()=} vs. \" f\"{tensor_model_parallel_size=}\") pp_world_size = get_pp_group().world_size assert (pp_world_size == pipeline_model_parallel_size), ( \"pipeline parallel group already initialized, but of unexpected size: \" f\"{pp_world_size=} vs. \" f\"{pipeline_model_parallel_size=}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#model_parallel_is_initialized","name":"model_parallel_is_initialized","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L95-L98","signature":"def model_parallel_is_initialized()","exported":true,"origin":"ast","body_hash":"ae02ee0f5a445b448f56207018602740fdc482a8e5383166665c368c8421bccc","body_text":"def model_parallel_is_initialized(): \"\"\"Check if tensor and pipeline parallel groups are initialized.\"\"\" return (ps._TP is not None) # and _PIPELINE_MODEL_PARALLEL_GROUP is not None)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","name":"initialize_model_parallel_for_vllm","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L101-L188","signature":"def initialize_model_parallel_for_vllm(tensor_model_parallel_size: int, num_tensor_model_parallel_groups_per_train_tp: int = 1, pipeline_model_parallel_size: int = 1) -> None","exported":true,"origin":"ast","body_hash":"b70decce5a554cc9a2ad5469a0758d0e197ff442eddd5996386554b08f22c068","body_text":"def initialize_model_parallel_for_vllm(tensor_model_parallel_size: int, num_tensor_model_parallel_groups_per_train_tp: int = 1, pipeline_model_parallel_size: int = 1) -> None: from torch.distributed import new_group # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() assert isinstance(tensor_model_parallel_size, int) # assert num_tensor_model_parallel_groups_per_train_tp == 1 and not different_tp_group # assert num_tensor_model_parallel_groups_per_train_tp > 1 and different_tp_group # Build the tensor model-parallel groups. assert ps._TP is None, (\"tensor model parallel group is already initialized\") global _TP world_size: int = torch.distributed.get_world_size() rank = torch.distributed.get_rank() backend = torch.distributed.get_backend() num_tensor_model_parallel_groups = world_size // tensor_model_parallel_size if num_tensor_model_parallel_groups_per_train_tp == 1: # if tensor_model_parallel_size == train_tensor_parallel_size: # using the same tp group as Megatron/vllm assert _TP is None, (\"tensor model parallel group is already initialized\") group_ranks = [] for i in range(num_tensor_model_parallel_groups): ranks = range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) group_ranks.append(ranks) _TP = init_model_parallel_group( group_ranks=group_ranks, local_rank=get_world_group().local_rank, backend=backend, use_custom_allreduce=False, # TODO: check why True is not work in Ray trainer use_message_queue_broadcaster=True) ps._TP = _TP # _MICRO_DATA_PARALLEL_GROUP is move to hybrid engine else: # initialize a micro_dp group and a tp group # assume training tp=4, infer tp=2, then, weight is partitioned as # [1], [2], [3], [4] for training and [1,2], [1,2], [3,4], [3,4] for inference # Build the inference tp groups # train_tp = train_tensor_parallel_size train_tp = num_tensor_model_parallel_groups_per_train_tp * tensor_model_parallel_size # num_tensor_model_parallel_groups_per_train_tp = train_tp // tensor_model_parallel_size assert _TP is None, (\"tensor model parallel group is already initialized\") group_ranks = [] for i in range(num_tensor_model_parallel_groups // num_tensor_model_parallel_groups_per_train_tp): start = train_tp * i end = train_tp * (i + 1) for j in range(num_tensor_model_parallel_groups_per_train_tp): ranks = list(range(start, end, num_tensor_model_parallel_groups_per_train_tp)) for i in range(len(ranks)): ranks[i] += j group_ranks.append(ranks) _TP = init_model_parallel_group( group_ranks=group_ranks, local_rank=get_world_group().local_rank, backend=backend, use_custom_allreduce=False, # TODO: check why True is not work in Ray trainer use_message_queue_broadcaster=True) ps._TP = _TP # Build the pipeline model-parallel groups. # global _PIPELINE_MODEL_PARALLEL_GROUP # global _PIPELINE_GLOBAL_RANKS # assert ps._PIPELINE_MODEL_PARALLEL_GROUP is None, (\"pipeline model parallel group is already initialized\") # ps._PIPELINE_MODEL_PARALLEL_GROUP = mpu.get_pipeline_model_parallel_group() # ps._PIPELINE_GLOBAL_RANKS = mpu.get_pipeline_model_parallel_ranks() # TODO: init using device mesh (not support hybrid engine now) # Build the pipeline model-parallel groups. num_pipeline_model_parallel_groups: int = (world_size // pipeline_model_parallel_size) global _PP assert _PP is None, (\"pipeline model parallel group is already initialized\") group_ranks = [] for i in range(num_pipeline_model_parallel_groups): ranks = list(range(i, world_size, num_pipeline_model_parallel_groups)) group_ranks.append(ranks) # pipeline parallel does not need custom allreduce _PP = init_model_parallel_group(group_ranks, get_world_group().local_rank, backend, use_custom_allreduce=False) ps._PP = _PP # for verl","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","name":"initialize_model_parallel","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L191-L264","signature":"def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None","exported":true,"origin":"ast","body_hash":"6df5843fd3529c4eacf993225112ff7760bed9868f386ce7cf2163e01af1a4e0","body_text":"def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None: \"\"\" NOTE: This method is a hack from the open-sourced version without asertion of world_size = tp * pp Initialize model parallel groups. Arguments: tensor_model_parallel_size: number of GPUs used for tensor model parallelism. pipeline_model_parallel_size: number of GPUs used for pipeline model parallelism. Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize the model pipeline. The present function will create 4 tensor model-parallel groups and 2 pipeline model-parallel groups: 4 tensor model-parallel groups: [g0, g1], [g2, g3], [g4, g5], [g6, g7] 2 pipeline model-parallel groups: [g0, g2, g4, g6], [g1, g3, g5, g7] Note that for efficiency, the caller should make sure adjacent ranks are on the same DGX box. For example if we are using 2 DGX-1 boxes with a total of 16 GPUs, rank 0 to 7 belong to the first box and ranks 8 to 15 belong to the second box. \"\"\" # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() world_size: int = torch.distributed.get_world_size() backend = backend or torch.distributed.get_backend(ps.get_world_group().device_group) # NOTE(sgm) we don't assert world_size == tp * pp # DP is not managed by vllm but by the veRL WorkerGroup # 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})\") num_tensor_model_parallel_groups: int = (world_size // tensor_model_parallel_size) rank = torch.distributed.get_rank() global _TP assert _TP is None, (\"tensor model parallel group is already initialized\") group_ranks = [] for i in range(num_tensor_model_parallel_groups): ranks = list(range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size)) group_ranks.append(ranks) # message queue broadcaster is only used in tensor model parallel group _TP = init_model_parallel_group( group_ranks, get_world_group().local_rank, backend, use_custom_allreduce=False, # TODO: check why True is not work in Ray trainer use_message_queue_broadcaster=True) ps._TP = _TP # TODO: init using device mesh (not support hybrid engine now) # Build the pipeline model-parallel groups. num_pipeline_model_parallel_groups: int = (world_size // pipeline_model_parallel_size) global _PP assert _PP is None, (\"pipeline model parallel group is already initialized\") group_ranks = [] for i in range(num_pipeline_model_parallel_groups): ranks = list(range(i, world_size, num_pipeline_model_parallel_groups)) group_ranks.append(ranks) # pipeline parallel does not need custom allreduce _PP = init_model_parallel_group(group_ranks, get_world_group().local_rank, backend, use_custom_allreduce=False) ps._PP = _PP # for verl","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_device_mesh","name":"get_device_mesh","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L272-L274","signature":"def get_device_mesh()","exported":true,"origin":"ast","body_hash":"2ac9c1236bccce8a6f34fa68c3df23bd33dcc07a6041cb0f4ff8d7cff2dcd147","body_text":"def get_device_mesh(): assert _DEVICE_MESH is not None, (\"device mesh is not initialized\") return _DEVICE_MESH","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_group","name":"get_tensor_model_parallel_group","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L282-L285","signature":"def get_tensor_model_parallel_group()","exported":true,"origin":"ast","body_hash":"cce2f525ea01566a94c641ca09b83e3d049800e965231549a1522b424bc74fbd","body_text":"def get_tensor_model_parallel_group(): \"\"\"Get the tensor model parallel group the caller rank belongs to.\"\"\" assert _TP is not None, (\"tensor model parallel group is not initialized\") return _TP.device_group","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_world_size","name":"get_tensor_model_parallel_world_size","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L288-L290","signature":"def get_tensor_model_parallel_world_size()","exported":true,"origin":"ast","body_hash":"c1af0d2aaf0c6f73d3dc5df62d1ec18f42133f1075e5262abc2daf509ff9dc98","body_text":"def get_tensor_model_parallel_world_size(): \"\"\"Return world size for the tensor model parallel group.\"\"\" return torch.distributed.get_world_size(group=get_tensor_model_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_rank","name":"get_tensor_model_parallel_rank","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L293-L295","signature":"def get_tensor_model_parallel_rank()","exported":true,"origin":"ast","body_hash":"7d7e3527eb0fbc1e8f6ca37c5449d0a44e76ded85a0b824cf4f92bda58558364","body_text":"def get_tensor_model_parallel_rank(): \"\"\"Return my rank for the tensor model parallel group.\"\"\" return torch.distributed.get_rank(group=get_tensor_model_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_src_rank","name":"get_tensor_model_parallel_src_rank","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","span":"L298-L303","signature":"def get_tensor_model_parallel_src_rank()","exported":true,"origin":"ast","body_hash":"cfa50af80e4b8151b634ff31e0fe23018462a6f55018e05a8cc9c4918ef4e48e","body_text":"def get_tensor_model_parallel_src_rank(): \"\"\"Calculate the global rank corresponding to the first local rank in the tensor model parallel group.\"\"\" global_rank = torch.distributed.get_rank() local_world_size = get_tensor_model_parallel_world_size() return (global_rank // local_world_size) * local_world_size","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"imports","specifier":"torch.distributed","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"imports","specifier":"vllm.distributed.parallel_state","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"imports","specifier":"torch.distributed.device_mesh","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","relation":"calls","name":"init_distributed_environment","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","relation":"calls","name":"initialize_model_parallel_for_vllm","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_parallel_state","relation":"calls","name":"initialize_model_parallel","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#ensure_model_parallel_initialized","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_backend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"model_parallel_is_initialized","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"initialize_model_parallel","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_pp_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#model_parallel_is_initialized","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"imports","specifier":"torch.distributed","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_backend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"init_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"init_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"init_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_backend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_world_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"init_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"init_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_device_mesh","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_group","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_world_size","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_world_size","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_world_size","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_rank","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_rank","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_src_rank","file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_src_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py#get_tensor_model_parallel_src_rank","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/parallel_state.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py":{"size":10184,"mtimeMs":1789275830317.203,"hash":"3ce0cddfb329f0898287fcef40209e88cd8d8f6e4dc1e7ad78e31ac2bacfce7f","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","name":"spmd_gpu_executor.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L1-L254","signature":null,"exported":true,"origin":"ast","body_hash":"3ce0cddfb329f0898287fcef40209e88cd8d8f6e4dc1e7ad78e31ac2bacfce7f","chars":10184,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/executor/gpu_executor.py import os import socket from typing import Any, Dict, List, Optional, Set, Tuple import torch import vllm.envs as envs from vllm.executor.executor_base import ExecutorBase, ExecutorAsyncBase from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.sequence import SamplerOutput, ExecuteModelRequest from vllm.config import (CacheConfig, DeviceConfig, LoRAConfig, MultiModalConfig, ParallelConfig, PromptAdapterConfig, SchedulerConfig, SpeculativeConfig) from .config import ModelConfig, LoadConfig logger = init_logger(__name__) # TODO(sgm): not implemented async executor yet"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","name":"SPMDGPUExecutor","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L34-L210","signature":"class SPMDGPUExecutor(ExecutorBase)","exported":true,"origin":"ast","body_hash":"d85266b0ca1574187cff585716c3b40fc2037574d120ad3646c9d5ecacb91c9b","body_text":"class SPMDGPUExecutor(ExecutorBase): \"\"\"SPMD-based multi-GPU executor implementations.\"\"\" def __init__( self, model, # pytorch model itself or its parameter dict model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], speculative_config: Optional[SpeculativeConfig], prompt_adapter_config: Optional[PromptAdapterConfig], ) -> None: self.model_config = model_config self.cache_config = cache_config self.lora_config = lora_config self.load_config = load_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.multimodal_config = multimodal_config self.speculative_config = speculative_config self.prompt_adapter_config = prompt_adapter_config distributed_init_method = initialize_cluster(parallel_config) self._init_executor(model, distributed_init_method) # TODO(sgm): verl not support speculative decode now def _init_executor(self, model, distributed_init_method) -> None: assert (not self.speculative_config), \"Speculative decoding not yet supported for multi-GPU backend.\" # Create the parallel worker for each GPU. self._init_workers_sp(model, distributed_init_method) def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker from .worker import Worker # pylint: disable=import-outside-toplevel rank = int(os.getenv(\"RANK\")) local_rank = int(os.getenv(\"LOCAL_RANK\")) print(f'local rank {local_rank}') # see https://github.com/NVIDIA/nccl/issues/1234 os.environ['NCCL_CUMEM_ENABLE'] = '0' self.worker = Worker( model, self.model_config, self.parallel_config, self.scheduler_config, self.device_config, self.cache_config, self.load_config, local_rank, rank, distributed_init_method, lora_config=self.lora_config, multimodal_config=self.multimodal_config, speculative_config=None, prompt_adapter_config=self.speculative_config, is_driver_worker=True, model_runner_cls=None, # use the default one ) # NOTE(shengguangming): torch.distributed.init_process_group will be called inside the init_model() self.worker.init_device() self.worker.load_model() def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Determine the number of available KV blocks. This invokes `determine_num_available_blocks` on each worker and takes the min of the results, guaranteeing that the selected cache sizes are compatible with all workers. Returns: - tuple[num_gpu_blocks, num_cpu_blocks] \"\"\" # Get the maximum number of blocks that can be allocated on GPU and CPU. num_blocks = self.worker.determine_num_available_blocks() # NOTE(shengguangming): Now we don't use a shared centralized controler but each process will # have its own scheduler num_gpu_blocks = num_blocks[0] num_cpu_blocks = num_blocks[1] return num_gpu_blocks, num_cpu_blocks def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: \"\"\"Initialize the KV cache in all workers. \"\"\" # NOTE: We log here to avoid multiple logs when number of workers is # greater than one. We could log in the engine, but not all executors # have GPUs. logger.info(\"# GPU blocks: %d, # CPU blocks: %d\", num_gpu_blocks, num_cpu_blocks) self.cache_config.num_gpu_blocks = num_gpu_blocks self.cache_config.num_cpu_blocks = num_cpu_blocks if torch.distributed.get_rank() == 0: print( f'before init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB' ) self.worker.initialize_cache(num_gpu_blocks=num_gpu_blocks, num_cpu_blocks=num_cpu_blocks) if torch.distributed.get_rank() == 0: print( f'after init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB' ) # NOTE(sgm): This will not profile & capture the model(CUDAGraph) when rebuilding KVCache def init_cache_engine(self) -> None: self.worker._init_cache_engine() def free_cache_engine(self) -> None: self.worker.free_cache_engine() def execute_model(self, execute_model_req) -> List[SamplerOutput]: all_outputs = self.worker.execute_model(execute_model_req=execute_model_req) # NOTE(sgm): # Each GPU in vllm under verl has its own spmd_gpu_executor, therefore all GPUs should return the outputs # In vllm with ray, only the driver worker returns the sampling results. return all_outputs def add_lora(self, lora_request: LoRARequest) -> bool: assert lora_request.lora_int_id > 0, \"lora_id must be greater than 0.\" return self.worker.add_lora(lora_request=lora_request) def remove_lora(self, lora_id: int) -> bool: assert lora_id > 0, \"lora_id must be greater than 0.\" return self.worker.remove_lora(lora_id=lora_id) def list_loras(self) -> Set[int]: return self.worker.list_loras() def check_health(self) -> None: # SPMDExecutor will always be healthy as long as # it's ru","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L37-L63","signature":"def __init__( self, model, # pytorch model itself or its parameter dict model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], speculative_config: Optional[SpeculativeConfig], prompt_adapter_config: Optional[PromptAdapterConfig], ) -> None","exported":false,"origin":"ast","body_hash":"efb6aa3aeb4240545c5fc41e1941e39fc71a9712ba0bcc180ce3aa9be0c799c2","body_text":"def __init__( self, model, # pytorch model itself or its parameter dict model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], multimodal_config: Optional[MultiModalConfig], speculative_config: Optional[SpeculativeConfig], prompt_adapter_config: Optional[PromptAdapterConfig], ) -> None: self.model_config = model_config self.cache_config = cache_config self.lora_config = lora_config self.load_config = load_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.multimodal_config = multimodal_config self.speculative_config = speculative_config self.prompt_adapter_config = prompt_adapter_config distributed_init_method = initialize_cluster(parallel_config) self._init_executor(model, distributed_init_method)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_executor","name":"_init_executor","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L66-L70","signature":"def _init_executor(self, model, distributed_init_method) -> None","exported":false,"origin":"ast","body_hash":"250ab123bb88f594c2ad692273d68d9f61a0601c0730510375dc528bc3d4e8d4","body_text":"def _init_executor(self, model, distributed_init_method) -> None: assert (not self.speculative_config), \"Speculative decoding not yet supported for multi-GPU backend.\" # Create the parallel worker for each GPU. self._init_workers_sp(model, distributed_init_method)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","name":"_init_workers_sp","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L72-L105","signature":"def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker","exported":false,"origin":"ast","body_hash":"60040a790ec4c9c49c5e73de8f677654382901bb4c8161e7b24032dcc75c42f5","body_text":"def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker from .worker import Worker # pylint: disable=import-outside-toplevel rank = int(os.getenv(\"RANK\")) local_rank = int(os.getenv(\"LOCAL_RANK\")) print(f'local rank {local_rank}') # see https://github.com/NVIDIA/nccl/issues/1234 os.environ['NCCL_CUMEM_ENABLE'] = '0' self.worker = Worker( model, self.model_config, self.parallel_config, self.scheduler_config, self.device_config, self.cache_config, self.load_config, local_rank, rank, distributed_init_method, lora_config=self.lora_config, multimodal_config=self.multimodal_config, speculative_config=None, prompt_adapter_config=self.speculative_config, is_driver_worker=True, model_runner_cls=None, # use the default one ) # NOTE(shengguangming): torch.distributed.init_process_group will be called inside the init_model() self.worker.init_device() self.worker.load_model()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.determine_num_available_blocks","name":"determine_num_available_blocks","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L107-L125","signature":"def determine_num_available_blocks(self) -> Tuple[int, int]","exported":true,"origin":"ast","body_hash":"dbbbc1bd7aa827704263dacd159512c096e5693fe6fa928309c9ed5451aa5514","body_text":"def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Determine the number of available KV blocks. This invokes `determine_num_available_blocks` on each worker and takes the min of the results, guaranteeing that the selected cache sizes are compatible with all workers. Returns: - tuple[num_gpu_blocks, num_cpu_blocks] \"\"\" # Get the maximum number of blocks that can be allocated on GPU and CPU. num_blocks = self.worker.determine_num_available_blocks() # NOTE(shengguangming): Now we don't use a shared centralized controler but each process will # have its own scheduler num_gpu_blocks = num_blocks[0] num_cpu_blocks = num_blocks[1] return num_gpu_blocks, num_cpu_blocks","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","name":"initialize_cache","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L127-L147","signature":"def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None","exported":true,"origin":"ast","body_hash":"1a1bc217eeed8020b0c93c5ea547950d65edd6448ee5de0cfa6514fc704ef066","body_text":"def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: \"\"\"Initialize the KV cache in all workers. \"\"\" # NOTE: We log here to avoid multiple logs when number of workers is # greater than one. We could log in the engine, but not all executors # have GPUs. logger.info(\"# GPU blocks: %d, # CPU blocks: %d\", num_gpu_blocks, num_cpu_blocks) self.cache_config.num_gpu_blocks = num_gpu_blocks self.cache_config.num_cpu_blocks = num_cpu_blocks if torch.distributed.get_rank() == 0: print( f'before init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB' ) self.worker.initialize_cache(num_gpu_blocks=num_gpu_blocks, num_cpu_blocks=num_cpu_blocks) if torch.distributed.get_rank() == 0: print( f'after init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB' )","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L150-L151","signature":"def init_cache_engine(self) -> None","exported":true,"origin":"ast","body_hash":"135998160b42b08111bb719b769019409c044dd0da5bd18e3cf1819bf579782e","body_text":"def init_cache_engine(self) -> None: self.worker._init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L153-L154","signature":"def free_cache_engine(self) -> None","exported":true,"origin":"ast","body_hash":"869ecde616def986458e35749106985c797ce877ae77ec9bab2ddf6490106436","body_text":"def free_cache_engine(self) -> None: self.worker.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.execute_model","name":"execute_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L156-L162","signature":"def execute_model(self, execute_model_req) -> List[SamplerOutput]","exported":true,"origin":"ast","body_hash":"aae188cc605ed5322f94d274292f73007610b372cfc40fda0aa9a0800811b7a5","body_text":"def execute_model(self, execute_model_req) -> List[SamplerOutput]: all_outputs = self.worker.execute_model(execute_model_req=execute_model_req) # NOTE(sgm): # Each GPU in vllm under verl has its own spmd_gpu_executor, therefore all GPUs should return the outputs # In vllm with ray, only the driver worker returns the sampling results. return all_outputs","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.add_lora","name":"add_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L164-L166","signature":"def add_lora(self, lora_request: LoRARequest) -> bool","exported":true,"origin":"ast","body_hash":"22a01766fd3ba5823ac7f31e7164043c28514e0346dceaf1705ede572c20a6c0","body_text":"def add_lora(self, lora_request: LoRARequest) -> bool: assert lora_request.lora_int_id > 0, \"lora_id must be greater than 0.\" return self.worker.add_lora(lora_request=lora_request)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.remove_lora","name":"remove_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L168-L170","signature":"def remove_lora(self, lora_id: int) -> bool","exported":true,"origin":"ast","body_hash":"13816aebed74124d7222df89c015c122a2dace11bbe74355d4b3e2b0d545e089","body_text":"def remove_lora(self, lora_id: int) -> bool: assert lora_id > 0, \"lora_id must be greater than 0.\" return self.worker.remove_lora(lora_id=lora_id)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.list_loras","name":"list_loras","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L172-L173","signature":"def list_loras(self) -> Set[int]","exported":true,"origin":"ast","body_hash":"84207e642ba5b1e08cab7f7c4d2b2014f03cbfebd74f378d55838bda1d0c31b7","body_text":"def list_loras(self) -> Set[int]: return self.worker.list_loras()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.check_health","name":"check_health","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L175-L178","signature":"def check_health(self) -> None: # SPMDExecutor will always be healthy as long as # it's running.","exported":true,"origin":"ast","body_hash":"48ca06f8f34b874ea336f29bb675c8616e07b51643c0ee9e4708d53fc9b3fa9a","body_text":"def check_health(self) -> None: # SPMDExecutor will always be healthy as long as # it's running. return","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.add_prompt_adapter","name":"add_prompt_adapter","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L183-L186","signature":"def add_prompt_adapter(self, prompt_adapter_request: PromptAdapterRequest) -> bool","exported":true,"origin":"ast","body_hash":"ca414add4567bf4da32f457533bc207bedb920d26ee303818d284649fb8c9cc4","body_text":"def add_prompt_adapter(self, prompt_adapter_request: PromptAdapterRequest) -> bool: assert prompt_adapter_request.prompt_adapter_id > 0, \\ \"prompt_adapter_id must be greater than 0.\" return self.worker.add_prompt_adapter(prompt_adapter_request)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.list_prompt_adapters","name":"list_prompt_adapters","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L188-L189","signature":"def list_prompt_adapters(self) -> Set[int]","exported":true,"origin":"ast","body_hash":"2f22e71786a397ac8276f938821afca6e4319a157e7b64045de915ac7d3c7c9f","body_text":"def list_prompt_adapters(self) -> Set[int]: return self.worker.list_prompt_adapters()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.pin_lora","name":"pin_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L191-L193","signature":"def pin_lora(self, lora_id: int) -> bool","exported":true,"origin":"ast","body_hash":"c86bee657fbd9fd8a3dc326489d09c0027c46e8d655205b5ddc8af2122ea076d","body_text":"def pin_lora(self, lora_id: int) -> bool: assert lora_id > 0, \"lora_id must be greater than 0.\" return self.worker.pin_lora(lora_id)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.pin_prompt_adapter","name":"pin_prompt_adapter","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L195-L198","signature":"def pin_prompt_adapter(self, prompt_adapter_id: int) -> bool","exported":true,"origin":"ast","body_hash":"e28685e28d6b9d35343e7d9249656fb52d328daac664ce0b1266b32f888c0c3f","body_text":"def pin_prompt_adapter(self, prompt_adapter_id: int) -> bool: assert prompt_adapter_id > 0, \\ \"prompt_adapter_id must be greater than 0.\" return self.worker.pin_prompt_adapter(prompt_adapter_id)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.remove_prompt_adapter","name":"remove_prompt_adapter","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L200-L203","signature":"def remove_prompt_adapter(self, prompt_adapter_id: int) -> bool","exported":true,"origin":"ast","body_hash":"206a88ff5a5553261eec7bd5906de4ab2cb79e75a4d2d166d5f14eddc5bd7e49","body_text":"def remove_prompt_adapter(self, prompt_adapter_id: int) -> bool: assert prompt_adapter_id > 0, \\ \"prompt_adapter_id must be greater than 0.\" return self.worker.remove_prompt_adapter(prompt_adapter_id)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L206-L207","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"c346b8cbee5a5fe63905573385e113fb7b59c46e58b26ecc132e331ba728ea49","body_text":"def offload_model_weights(self) -> None: self.worker.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L209-L210","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None","exported":true,"origin":"ast","body_hash":"e7853eafaa9c2dd18668c019f3a34f0b873eb14ebc971798a00af8501b8de047","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None: self.worker.sync_model_weights(actor_weights=actor_weights, load_format=load_format)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#initialize_cluster","name":"initialize_cluster","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L213-L234","signature":"def initialize_cluster( parallel_config: ParallelConfig, engine_use_ray: bool = False, ray_address: Optional[str] = None, ) -> Tuple[str, Optional[None]]","exported":true,"origin":"ast","body_hash":"d737f1830326ce53e8589f097411c5f17acdd325ad8b62f3817f6527312dfe79","body_text":"def initialize_cluster( parallel_config: ParallelConfig, engine_use_ray: bool = False, ray_address: Optional[str] = None, ) -> Tuple[str, Optional[None]]: \"\"\"Initialize the distributed cluster probably with Ray. Args: parallel_config: The configurations for parallel execution. Returns: The `distributed_init_method` is the address for initializing the distributed backend. \"\"\" # Initialize cluster locally. port = get_open_port() # We need to setup the distributed init method to make sure # the distributed megatron code (e.g., get world size) works correctly. # distributed_init_method = f\"tcp://localhost:{port}\" distributed_init_method = 'env://' return distributed_init_method","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#get_open_port","name":"get_open_port","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L237-L240","signature":"def get_open_port()","exported":true,"origin":"ast","body_hash":"d88039216a61e6ec2a0e1328331a7b83f0af92fe7306e2625e8d4f34b557c26e","body_text":"def get_open_port(): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind((\"\", 0)) return s.getsockname()[1]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync","name":"SPMDGPUExecutorAsync","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L244-L253","signature":"class SPMDGPUExecutorAsync(SPMDGPUExecutor, ExecutorAsyncBase)","exported":true,"origin":"ast","body_hash":"a4473dfeed05058503ca8f50986836a978e7445454d17a15c9a1a71dc1c1ae0c","body_text":"class SPMDGPUExecutorAsync(SPMDGPUExecutor, ExecutorAsyncBase): async def execute_model_async(self, execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]: \"\"\"Executes one model step on the given sequences.\"\"\" raise NotImplementedError async def check_health_async(self) -> None: \"\"\"Checks if the executor is healthy. If not, it should raise an exception.\"\"\" self.check_health()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync.execute_model_async","name":"execute_model_async","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L246-L248","signature":"async def execute_model_async(self, execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]","exported":true,"origin":"ast","body_hash":"832d82699df49a6c8509c3c7f505f16afd4d4b985ff4fd27e7e1f5bdf2d96c0c","body_text":"async def execute_model_async(self, execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]: \"\"\"Executes one model step on the given sequences.\"\"\" raise NotImplementedError","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutorAsync"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync.check_health_async","name":"check_health_async","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","span":"L250-L253","signature":"async def check_health_async(self) -> None","exported":true,"origin":"ast","body_hash":"fe8c9a207c0e656691830f7feb9005f0015b9529086d8913552ff6134a65430b","body_text":"async def check_health_async(self) -> None: \"\"\"Checks if the executor is healthy. If not, it should raise an exception.\"\"\" self.check_health()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutorAsync"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":"socket","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.executor.executor_base","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"extends","name":"ExecutorBase","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","relation":"calls","name":"initialize_cluster","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","relation":"calls","name":"_init_executor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"SPMDGPUExecutor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_executor","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_executor","relation":"calls","name":"_init_workers_sp","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"SPMDGPUExecutor"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":".worker","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"Worker","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"init_device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.determine_num_available_blocks","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.determine_num_available_blocks","relation":"calls","name":"determine_num_available_blocks","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_allocated","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_reserved","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"initialize_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_allocated","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_reserved","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.init_cache_engine","relation":"calls","name":"_init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.execute_model","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.execute_model","relation":"calls","name":"execute_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.add_lora","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.add_lora","relation":"calls","name":"add_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.remove_lora","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.remove_lora","relation":"calls","name":"remove_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.list_loras","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.list_loras","relation":"calls","name":"list_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.check_health","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.prompt_adapter.request","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.add_prompt_adapter","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.add_prompt_adapter","relation":"calls","name":"add_prompt_adapter","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.list_prompt_adapters","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.list_prompt_adapters","relation":"calls","name":"list_prompt_adapters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.pin_lora","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.pin_lora","relation":"calls","name":"pin_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.pin_prompt_adapter","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.pin_prompt_adapter","relation":"calls","name":"pin_prompt_adapter","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.remove_prompt_adapter","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.remove_prompt_adapter","relation":"calls","name":"remove_prompt_adapter","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutor.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#initialize_cluster","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#initialize_cluster","relation":"calls","name":"get_open_port","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#get_open_port","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#get_open_port","relation":"calls","name":"socket","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#get_open_port","relation":"calls","name":"bind","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#get_open_port","relation":"calls","name":"getsockname","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"extends","name":"SPMDGPUExecutor","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"extends","name":"ExecutorAsyncBase","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync.execute_model_async","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync.check_health_async","file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py#SPMDGPUExecutorAsync.check_health_async","relation":"calls","name":"check_health","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/spmd_gpu_executor.py","recvType":"SPMDGPUExecutorAsync"}]},"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py":{"size":3310,"mtimeMs":1789275830318.203,"hash":"497a401f78f5c164308c4d48df0144e8b0c7a47bd4dfb3212681c8d511fe98bb","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","name":"tokenizer.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","span":"L1-L78","signature":null,"exported":true,"origin":"ast","body_hash":"497a401f78f5c164308c4d48df0144e8b0c7a47bd4dfb3212681c8d511fe98bb","chars":3310,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/tokenizer_group/tokenizer_group.py from typing import List, Optional, Tuple, Union from transformers import (AutoTokenizer, PreTrainedTokenizer, PreTrainedTokenizerFast) from vllm.lora.request import LoRARequest from vllm.utils import make_async, LRUCache from vllm.transformers_utils.tokenizers import *"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup","name":"TokenizerGroup","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","span":"L25-L77","signature":"class TokenizerGroup","exported":true,"origin":"ast","body_hash":"c107b245077e329b7be7c53602c2fad707f0231d9dcb1df954ad8280adebefb6","body_text":"class TokenizerGroup: \"\"\"A group of tokenizers that can be used for LoRA adapters.\"\"\" def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int]): self.enable_lora = enable_lora self.max_input_length = max_input_length self.tokenizer = tokenizer self.lora_tokenizers = LRUCache[PreTrainedTokenizer](capacity=max_num_seqs) if enable_lora else None def ping(self) -> bool: \"\"\"Check if the tokenizer group is alive.\"\"\" return True def get_max_input_len(self, lora_request: Optional[LoRARequest] = None) -> Optional[int]: \"\"\"Get the maximum input length for the LoRA request.\"\"\" return self.max_input_length def encode(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = self.get_lora_tokenizer(lora_request) return tokenizer.encode(prompt) async def encode_async(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = await self.get_lora_tokenizer_async(lora_request) return tokenizer.encode(prompt) def get_lora_tokenizer(self, lora_request: Optional[LoRARequest]) -> \"PreTrainedTokenizer\": if not lora_request or not self.enable_lora: return self.tokenizer if lora_request.lora_int_id not in self.lora_tokenizers: # TODO(sgm): the lora tokenizer is also passed, but may be different tokenizer = self.tokenizer # tokenizer = (get_lora_tokenizer( # lora_request, **self.tokenizer_config) or self.tokenizer) self.lora_tokenizers.put(lora_request.lora_int_id, tokenizer) return tokenizer else: return self.lora_tokenizers.get(lora_request.lora_int_id) # FIXME(sgm): for simplicity, we assign the special token here @property def pad_token_id(self): return self.tokenizer.pad_token_id @property def eos_token_id(self): return self.tokenizer.eos_token_id","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","span":"L28-L33","signature":"def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int])","exported":false,"origin":"ast","body_hash":"ac7258b8cb647fbec870bbffdc1633018f85352db29d6a00fd04abe0a3a94186","body_text":"def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int]): self.enable_lora = enable_lora self.max_input_length = max_input_length self.tokenizer = tokenizer self.lora_tokenizers = LRUCache[PreTrainedTokenizer](capacity=max_num_seqs) if enable_lora else None","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.ping","name":"ping","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","span":"L35-L37","signature":"def ping(self) -> bool","exported":true,"origin":"ast","body_hash":"97d31a1c22925222b76fe44142930b66980680e4e167b8a2be973980797507c3","body_text":"def ping(self) -> bool: \"\"\"Check if the tokenizer group is alive.\"\"\" return True","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.get_max_input_len","name":"get_max_input_len","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","span":"L39-L41","signature":"def get_max_input_len(self, lora_request: Optional[LoRARequest] = None) -> Optional[int]","exported":true,"origin":"ast","body_hash":"232098bcc1725540b28685d973367034b497ccefd004f560ceab1bf3cae8ae28","body_text":"def get_max_input_len(self, lora_request: Optional[LoRARequest] = None) -> Optional[int]: \"\"\"Get the maximum input length for the LoRA request.\"\"\" return self.max_input_length","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.encode","name":"encode","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","span":"L43-L48","signature":"def encode(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]","exported":true,"origin":"ast","body_hash":"c7cd151da6e5d1fb514dcc275e13f4c289111810f56df69c36b64c0d3f2cbc9c","body_text":"def encode(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = self.get_lora_tokenizer(lora_request) return tokenizer.encode(prompt)","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.encode_async","name":"encode_async","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","span":"L50-L55","signature":"async def encode_async(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]","exported":true,"origin":"ast","body_hash":"e0ccf679cf9c6ab3aeebd6ff19c82fdfdb9328b3357f16319f00fadca9561589","body_text":"async def encode_async(self, prompt: str, request_id: Optional[str] = None, lora_request: Optional[LoRARequest] = None) -> List[int]: tokenizer = await self.get_lora_tokenizer_async(lora_request) return tokenizer.encode(prompt)","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.get_lora_tokenizer","name":"get_lora_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","span":"L57-L68","signature":"def get_lora_tokenizer(self, lora_request: Optional[LoRARequest]) -> \"PreTrainedTokenizer\"","exported":true,"origin":"ast","body_hash":"66dbb721b6874d42fc9c39c04fd34dd635d82d307c22dcbff76c9464a47c01ca","body_text":"def get_lora_tokenizer(self, lora_request: Optional[LoRARequest]) -> \"PreTrainedTokenizer\": if not lora_request or not self.enable_lora: return self.tokenizer if lora_request.lora_int_id not in self.lora_tokenizers: # TODO(sgm): the lora tokenizer is also passed, but may be different tokenizer = self.tokenizer # tokenizer = (get_lora_tokenizer( # lora_request, **self.tokenizer_config) or self.tokenizer) self.lora_tokenizers.put(lora_request.lora_int_id, tokenizer) return tokenizer else: return self.lora_tokenizers.get(lora_request.lora_int_id)","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.pad_token_id","name":"pad_token_id","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","span":"L72-L73","signature":"def pad_token_id(self)","exported":true,"origin":"ast","body_hash":"09c75a3902b0bb0ca3d7cec97618ef037c84e09c3213f62cd5ee9aa99c2b6669","body_text":"def pad_token_id(self): return self.tokenizer.pad_token_id","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.eos_token_id","name":"eos_token_id","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","span":"L76-L77","signature":"def eos_token_id(self)","exported":true,"origin":"ast","body_hash":"57455769f69da68b62a323b65045e6baf5ecb3dde6a703266578aeeb1bf096a2","body_text":"def eos_token_id(self): return self.tokenizer.eos_token_id","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","relation":"imports","specifier":"vllm.transformers_utils.tokenizers","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.ping","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.get_max_input_len","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.encode","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.encode","relation":"calls","name":"get_lora_tokenizer","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","recvType":"TokenizerGroup"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.encode","relation":"calls","name":"encode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.encode_async","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.encode_async","relation":"calls","name":"get_lora_tokenizer_async","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py","recvType":"TokenizerGroup"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.encode_async","relation":"calls","name":"encode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.get_lora_tokenizer","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.get_lora_tokenizer","relation":"calls","name":"put","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.get_lora_tokenizer","relation":"calls","name":"get","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.pad_token_id","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py#TokenizerGroup.eos_token_id","file":"verl/third_party/vllm/vllm_v_0_5_4/tokenizer.py"}]},"verl/third_party/vllm/vllm_v_0_5_4/worker.py":{"size":15305,"mtimeMs":1789275830318.203,"hash":"4235c9a8d985c1c4c4b6db086b3515c38681517a11f77004a488c6790383897c","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","name":"worker.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L1-L324","signature":null,"exported":true,"origin":"ast","body_hash":"4235c9a8d985c1c4c4b6db086b3515c38681517a11f77004a488c6790383897c","chars":15305,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/worker/worker.py \"\"\"A GPU worker class.\"\"\" import os import gc from typing import Dict, List, Tuple, Optional, Union, Type import torch import torch.distributed import torch.nn as nn from vllm.config import (CacheConfig, DeviceConfig, LoRAConfig, MultiModalConfig, ParallelConfig, PromptAdapterConfig, SchedulerConfig, SpeculativeConfig) from vllm.model_executor import set_random_seed from vllm.sequence import (ExecuteModelRequest, IntermediateTensors, SamplerOutput) from vllm.worker.cache_engine import CacheEngine # TODO(sgm): check why vllm has similar file in vllm.model_executor.parallel_utils.parallel_state from vllm.distributed import (init_distributed_environment, set_custom_all_reduce, get_tensor_model_parallel_group) from vllm.worker.worker_base import WorkerInput from vllm.worker.worker import Worker, _check_if_gpu_supports_dtype from vllm.worker.model_runner_base import ModelRunnerBase, ModelRunnerInputBase from vllm.worker.embedding_model_runner import EmbeddingModelRunner from vllm.worker.model_runner import GPUModelRunnerBase from .model_runner import ModelRunner from .megatron_weight_loaders import load_megatron_weights from .hf_weight_loader import load_hf_weights from .dtensor_weight_loaders import load_dtensor_weights from .parallel_state import (ensure_model_parallel_initialized) from .config import ModelConfig, LoadConfig, LoadFormat"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","name":"Worker","kind":"class","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L44-L283","signature":"class Worker(Worker)","exported":true,"origin":"ast","body_hash":"7a9d27b5df666f5c1c2c4ac2eda6a86754725cf04a0d74392813f446294c5366","body_text":"class Worker(Worker): \"\"\"A worker class that executes (a partition of) the model on a GPU. Each worker is associated with a single GPU. The worker is responsible for maintaining the KV cache and executing the model on the GPU. In case of distributed inference, each worker is assigned a partition of the model. \"\"\" def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, local_rank: int, rank: int, distributed_init_method: str, lora_config: Optional[LoRAConfig] = None, multimodal_config: Optional[MultiModalConfig] = None, speculative_config: Optional[SpeculativeConfig] = None, prompt_adapter_config: Optional[PromptAdapterConfig] = None, is_driver_worker: bool = False, model_runner_cls: Optional[Type[GPUModelRunnerBase]] = None, ) -> None: # self.model = model # will be replaced in the init_model self.model_config = model_config self.parallel_config = parallel_config self.parallel_config.rank = rank self.scheduler_config = scheduler_config self.device_config = device_config self.cache_config = cache_config self.local_rank = local_rank self.rank = rank self.distributed_init_method = distributed_init_method self.lora_config = lora_config self.load_config = load_config self.prompt_adapter_config = prompt_adapter_config self.is_driver_worker = is_driver_worker # TODO: we don't need driver # if parallel_config and is_driver_worker: # assert rank % parallel_config.tensor_parallel_size == 0, \\ # \"Driver worker should be rank 0 of tensor parallel group.\" if self.model_config.trust_remote_code: # note: lazy import to avoid importing torch before initializing from vllm.utils import init_cached_hf_modules init_cached_hf_modules() self.multimodal_config = multimodal_config # Return hidden states from target model if the draft model is an # mlp_speculator speculative_args = {} if speculative_config is None \\ or (speculative_config.draft_model_config.model == model_config.model) \\ or (speculative_config.draft_model_config.hf_config.model_type not in [\"medusa\", \"mlp_speculator\"]) \\ else {\"return_hidden_states\": True} # TODO(sgm): set correct model runner class ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner if model_runner_cls is not None: ModelRunnerClass = model_runner_cls elif self.model_config.embedding_mode: ModelRunnerClass = EmbeddingModelRunner self.model_runner: GPUModelRunnerBase = ModelRunnerClass( model, # [VERL]: add for verl model_config, parallel_config, scheduler_config, device_config, cache_config, load_config=load_config, lora_config=self.lora_config, kv_cache_dtype=self.cache_config.cache_dtype, is_driver_worker=is_driver_worker, prompt_adapter_config=prompt_adapter_config, multimodal_config=multimodal_config, **speculative_args, ) # Uninitialized cache engine. Will be initialized by # initialize_cache. self.cache_engine: List[CacheEngine] = None # Initialize gpu_cache as embedding models don't initialize kv_caches self.gpu_cache: Optional[List[List[torch.Tensor]]] = None # NOTE(sgm): [VERL] For offloading inference engine params self.cpu_model = None def init_device(self) -> None: if self.device_config.device.type == \"cuda\": # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # NOTE(sgm): Modify for verl, Env vars will be set by TORCHRUN. self.rank = self.rank if self.rank is not None else int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) self.device = torch.device(f\"cuda:{local_rank}\") if self.rank < 0: raise ValueError(\"Invalid or unspecified rank.\") torch.cuda.set_device(self.device) # Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" self.parallel_config.world_size = world_size _check_if_gpu_supports_dtype(self.model_config.dtype) torch.cuda.empty_cache() self.init_gpu_memory = torch.cuda.mem_get_info()[0] else: raise RuntimeError(f\"Not support device type: {self.device_config.device}\") # Initialize the distributed environment. init_worker_distributed_environment(self.parallel_config, self.rank, self.distributed_init_method, self.local_rank) # Set random seed. set_random_seed(self.model_config.seed) # self.model = get_model(actor_model=self.model, model_config=self.model_config) @torch.inference_mode() def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Profiles the peak memory usage of the model to determine how many KV blocks may be allocated without OOMs. The engine will first conduct a profiling of the","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L52-L132","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, local_rank: int, rank: int, distributed_init_method: str, lora_config: Optional[LoRAConfig] = None, multimodal_config: Optional[MultiModalConfig] = None, speculative_config: Optional[SpeculativeConfig] = None, prompt_adapter_config: Optional[PromptAdapterConfig] = None, is_driver_worker: bool = False, model_runner_cls: Optional[Type[GPUModelRunnerBase]] = None, ) -> None: # self.model = model # will be replaced in the init_model","exported":false,"origin":"ast","body_hash":"d9d5d136a2d17429b68b1dc3148742f9200c28cffc863a231fcba48ad935df1a","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, local_rank: int, rank: int, distributed_init_method: str, lora_config: Optional[LoRAConfig] = None, multimodal_config: Optional[MultiModalConfig] = None, speculative_config: Optional[SpeculativeConfig] = None, prompt_adapter_config: Optional[PromptAdapterConfig] = None, is_driver_worker: bool = False, model_runner_cls: Optional[Type[GPUModelRunnerBase]] = None, ) -> None: # self.model = model # will be replaced in the init_model self.model_config = model_config self.parallel_config = parallel_config self.parallel_config.rank = rank self.scheduler_config = scheduler_config self.device_config = device_config self.cache_config = cache_config self.local_rank = local_rank self.rank = rank self.distributed_init_method = distributed_init_method self.lora_config = lora_config self.load_config = load_config self.prompt_adapter_config = prompt_adapter_config self.is_driver_worker = is_driver_worker # TODO: we don't need driver # if parallel_config and is_driver_worker: # assert rank % parallel_config.tensor_parallel_size == 0, \\ # \"Driver worker should be rank 0 of tensor parallel group.\" if self.model_config.trust_remote_code: # note: lazy import to avoid importing torch before initializing from vllm.utils import init_cached_hf_modules init_cached_hf_modules() self.multimodal_config = multimodal_config # Return hidden states from target model if the draft model is an # mlp_speculator speculative_args = {} if speculative_config is None \\ or (speculative_config.draft_model_config.model == model_config.model) \\ or (speculative_config.draft_model_config.hf_config.model_type not in [\"medusa\", \"mlp_speculator\"]) \\ else {\"return_hidden_states\": True} # TODO(sgm): set correct model runner class ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner if model_runner_cls is not None: ModelRunnerClass = model_runner_cls elif self.model_config.embedding_mode: ModelRunnerClass = EmbeddingModelRunner self.model_runner: GPUModelRunnerBase = ModelRunnerClass( model, # [VERL]: add for verl model_config, parallel_config, scheduler_config, device_config, cache_config, load_config=load_config, lora_config=self.lora_config, kv_cache_dtype=self.cache_config.cache_dtype, is_driver_worker=is_driver_worker, prompt_adapter_config=prompt_adapter_config, multimodal_config=multimodal_config, **speculative_args, ) # Uninitialized cache engine. Will be initialized by # initialize_cache. self.cache_engine: List[CacheEngine] = None # Initialize gpu_cache as embedding models don't initialize kv_caches self.gpu_cache: Optional[List[List[torch.Tensor]]] = None # NOTE(sgm): [VERL] For offloading inference engine params self.cpu_model = None","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","name":"init_device","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L134-L168","signature":"def init_device(self) -> None","exported":true,"origin":"ast","body_hash":"762272295f8ebaa0d8a603527f1c7088967ebd4ff4562d0c715a5656dcddc70b","body_text":"def init_device(self) -> None: if self.device_config.device.type == \"cuda\": # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # NOTE(sgm): Modify for verl, Env vars will be set by TORCHRUN. self.rank = self.rank if self.rank is not None else int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) self.device = torch.device(f\"cuda:{local_rank}\") if self.rank < 0: raise ValueError(\"Invalid or unspecified rank.\") torch.cuda.set_device(self.device) # Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" self.parallel_config.world_size = world_size _check_if_gpu_supports_dtype(self.model_config.dtype) torch.cuda.empty_cache() self.init_gpu_memory = torch.cuda.mem_get_info()[0] else: raise RuntimeError(f\"Not support device type: {self.device_config.device}\") # Initialize the distributed environment. init_worker_distributed_environment(self.parallel_config, self.rank, self.distributed_init_method, self.local_rank) # Set random seed. set_random_seed(self.model_config.seed) # self.model = get_model(actor_model=self.model, model_config=self.model_config)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","name":"determine_num_available_blocks","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L171-L227","signature":"def determine_num_available_blocks(self) -> Tuple[int, int]","exported":true,"origin":"ast","body_hash":"ce0397c0e7f9c76ce5efbc55529fb217297905a3f39186e5a5e620cf7b7a39df","body_text":"def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Profiles the peak memory usage of the model to determine how many KV blocks may be allocated without OOMs. The engine will first conduct a profiling of the existing memory usage. Then, it calculate the maximum possible number of GPU and CPU blocks that can be allocated with the remaining free memory. .. tip:: You may limit the usage of GPU memory by adjusting the `gpu_memory_utilization` parameter. \"\"\" # Profile the memory usage of the model and get the maximum number of # cache blocks that can be allocated with the remaining free memory. torch.cuda.empty_cache() # torch.cuda.reset_peak_memory_stats() # Execute a forward pass with dummy inputs to profile the memory usage # of the model. self.model_runner.profile_run() # Calculate the number of blocks that can be allocated with the # profiled peak memory. torch.cuda.synchronize() free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info() peak_memory = total_gpu_memory - free_gpu_memory assert peak_memory > 0, (\"Error in memory profiling. This happens when the GPU memory was \" \"not properly cleaned up before initializing the vLLM instance.\") cache_block_size = self.get_cache_block_size_bytes() # NOTE(sgm) [VERL] use the remaining memory num_gpu_blocks = int((free_gpu_memory * self.cache_config.gpu_memory_utilization) // cache_block_size) # num_gpu_blocks = int((total_gpu_memory * self.cache_config.gpu_memory_utilization - peak_memory) // cache_block_size) num_cpu_blocks = int(self.cache_config.swap_space_bytes // cache_block_size) num_gpu_blocks = max(num_gpu_blocks, 0) num_cpu_blocks = max(num_cpu_blocks, 0) if self.model_runner.lora_manager: self.model_runner.remove_all_loras() # NOTE(sgm): Add for [VERL], synchronize number of blocks with all the rank num_gpu_blocks = torch.tensor([num_gpu_blocks], device='cuda') num_cpu_blocks = torch.tensor([num_cpu_blocks], device='cuda') torch.distributed.all_reduce(num_gpu_blocks, op=torch.distributed.ReduceOp.MIN, group=get_tensor_model_parallel_group().device_group) torch.distributed.all_reduce(num_cpu_blocks, op=torch.distributed.ReduceOp.MIN, group=get_tensor_model_parallel_group().device_group) num_gpu_blocks = num_gpu_blocks.item() num_cpu_blocks = num_cpu_blocks.item() gc.collect() torch.cuda.empty_cache() return num_gpu_blocks, num_cpu_blocks","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker._init_cache_engine","name":"_init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L229-L231","signature":"def _init_cache_engine(self)","exported":false,"origin":"ast","body_hash":"ddbf730d543fa64f637055bf2b4047182f388f84f7375434e0765a7cd602b012","body_text":"def _init_cache_engine(self): if self.cache_engine is None and self.gpu_cache is None: super()._init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L233-L236","signature":"def free_cache_engine(self): # ensure `enforce_eager=True`","exported":true,"origin":"ast","body_hash":"d384948b80b2c04ee44319bcd88c9b941e3967b68e2050a10073ade04af52223","body_text":"def free_cache_engine(self): # ensure `enforce_eager=True` self.cache_engine = None self.gpu_cache = None","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.execute_model","name":"execute_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L239-L263","signature":"def execute_model(self, execute_model_req: ExecuteModelRequest, intermediate_tensors: Optional[IntermediateTensors] = None) -> Optional[List[SamplerOutput]]","exported":true,"origin":"ast","body_hash":"a38f711b02ba25f57599a500d004e227b5d8737742fc8c7d1945c61a8232f91a","body_text":"def execute_model(self, execute_model_req: ExecuteModelRequest, intermediate_tensors: Optional[IntermediateTensors] = None) -> Optional[List[SamplerOutput]]: \"\"\" Execute model in Single Program Multiple Data (SPMD) fashion. All workers take the same request, prepare the input and execute the model. \"\"\" assert execute_model_req is not None, (\"_execute_model_spmd() requires each worker to take in an \" \"ExecuteModelRequest\") worker_input: WorkerInput = self.prepare_worker_input(execute_model_req=execute_model_req) model_input: ModelRunnerInputBase = (self.model_runner.prepare_model_input( execute_model_req.seq_group_metadata_list)) # verl.worker.workerbase.WorkerBase # swap cache super().execute_worker(worker_input) # If there is no input, we don't need to execute the model. if worker_input.num_seq_groups == 0: return [] return self.model_runner.execute_model( model_input, self.kv_cache[worker_input.virtual_engine] if self.kv_cache is not None else None, intermediate_tensors)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L266-L273","signature":"def sync_model_weights(self, actor_weights: Dict, load_format: str)","exported":true,"origin":"ast","body_hash":"61f049debd75c5e97bb95de1fbaf9d769e14978f2c8fd9b2ab6143756010d1d1","body_text":"def sync_model_weights(self, actor_weights: Dict, load_format: str): if load_format in [LoadFormat.MEGATRON, LoadFormat.AUTO]: load_megatron_weights(actor_weights, self.model_runner.model) elif load_format == LoadFormat.HF: # full model state dict without no sharding load_hf_weights(actor_weights, self.model_runner.model) elif load_format == LoadFormat.DTENSOR: load_dtensor_weights(actor_weights, self.model_runner.model)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L275-L283","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"58d8d5e64a0765436d968d3a9641e3432be0a43a9ca286f99b1f827937e19e31","body_text":"def offload_model_weights(self) -> None: if self.cpu_model == None: self.cpu_model = {} for name, params in self.model_runner.model.named_parameters(): self.cpu_model[name] = torch.empty_like(params, device='cpu') params.data = self.cpu_model[name] else: for name, params in self.model_runner.model.named_parameters(): params.data = self.cpu_model[name]","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#init_worker_distributed_environment","name":"init_worker_distributed_environment","kind":"function","path":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","span":"L286-L323","signature":"def init_worker_distributed_environment( parallel_config: ParallelConfig, rank: int, distributed_init_method: Optional[str] = \"env://\", local_rank: int = -1, ) -> None","exported":true,"origin":"ast","body_hash":"47aae7bc783a62330d323b54287bc9471810f2cc94d258383dc2c3a8cc25c467","body_text":"def init_worker_distributed_environment( parallel_config: ParallelConfig, rank: int, distributed_init_method: Optional[str] = \"env://\", local_rank: int = -1, ) -> None: \"\"\"Initialize the distributed environment.\"\"\" set_custom_all_reduce(not parallel_config.disable_custom_all_reduce) # NOTE(sgm) use tcp://localhost:xxxx will hang in HF setting without megatron init_distributed_environment(parallel_config.world_size, rank, distributed_init_method, local_rank) ensure_model_parallel_initialized(tensor_model_parallel_size=parallel_config.tensor_parallel_size, pipeline_model_parallel_size=parallel_config.pipeline_parallel_size) # TODO(sgm): check whether need this # if pynccl_utils.is_initialized(): # pynccl_world_size = pynccl_utils.get_world_size() # if pynccl_world_size != parallel_config.world_size: # raise RuntimeError( # \"pynccl is already initialized but the pynccl world \" # \"size does not match parallel_config.world_size \" # f\"({pynccl_world_size} vs. {parallel_config.world_size}).\") # elif parallel_config.world_size > 1: # # NOTE(woosuk): We don't initialize pynccl process group when world size # # is 1. # # NOTE(kaichao): By default, pynccl is initialized for tp group. # pynccl_utils.init_process_group( # group=get_tensor_model_parallel_cpu_group()) # # Initialize a custom fast all-reduce implementation. # if not parallel_config.disable_custom_all_reduce: # init_custom_ar() # A small all_reduce for warmup. torch.distributed.all_reduce(torch.zeros(1).cuda()) # if pynccl_utils.is_initialized(): # pynccl_utils.all_reduce(torch.zeros(1).cuda())","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"gc","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"torch.distributed","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.model_executor","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.worker.cache_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.distributed","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.worker.worker_base","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.worker.worker","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.worker.model_runner_base","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.worker.embedding_model_runner","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.worker.model_runner","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":".model_runner","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":".megatron_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":".hf_weight_loader","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":".dtensor_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":".parallel_state","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","relation":"extends","name":"Worker","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.__init__","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.__init__","relation":"calls","name":"init_cached_hf_modules","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.__init__","relation":"calls","name":"ModelRunnerClass","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"set_device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"_check_if_gpu_supports_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"mem_get_info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"init_worker_distributed_environment","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.init_device","relation":"calls","name":"set_random_seed","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","relation":"calls","name":"inference_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"profile_run","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","recvType":"ModelRunnerClass"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"synchronize","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"mem_get_info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"get_cache_block_size_bytes","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"remove_all_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","recvType":"ModelRunnerClass"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"item","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","recvType":"max"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"item","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","recvType":"max"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"collect","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker._init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker._init_cache_engine","relation":"calls","name":"_init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker._init_cache_engine","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.execute_model","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.execute_model","relation":"calls","name":"prepare_worker_input","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.execute_model","relation":"calls","name":"prepare_model_input","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","recvType":"ModelRunnerClass"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.execute_model","relation":"calls","name":"execute_worker","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.execute_model","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.execute_model","relation":"calls","name":"execute_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","recvType":"ModelRunnerClass"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.sync_model_weights","relation":"calls","name":"load_megatron_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.sync_model_weights","relation":"calls","name":"load_hf_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.sync_model_weights","relation":"calls","name":"load_dtensor_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.offload_model_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.offload_model_weights","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#Worker.offload_model_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#init_worker_distributed_environment","file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#init_worker_distributed_environment","relation":"calls","name":"set_custom_all_reduce","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#init_worker_distributed_environment","relation":"calls","name":"init_distributed_environment","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#init_worker_distributed_environment","relation":"calls","name":"ensure_model_parallel_initialized","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#init_worker_distributed_environment","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#init_worker_distributed_environment","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_5_4/worker.py#init_worker_distributed_environment","relation":"calls","name":"zeros","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_5_4/worker.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/__init__.py":{"size":600,"mtimeMs":1789275830318.203,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/__init__.py","name":"__init__.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py":{"size":3142,"mtimeMs":1789275830318.203,"hash":"fa211803d2705904fa27b4af7c4db629e60623046cc675bc04f5786a20675efc","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","name":"arg_utils.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","span":"L1-L79","signature":null,"exported":true,"origin":"ast","body_hash":"fa211803d2705904fa27b4af7c4db629e60623046cc675bc04f5786a20675efc","chars":3142,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/engine/arg_utils.py import os from dataclasses import dataclass from transformers import PretrainedConfig from vllm.config import EngineConfig from vllm.engine.arg_utils import EngineArgs from .config import LoadConfig, ModelConfig @dataclass"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs","name":"EngineArgs","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","span":"L27-L78","signature":"class EngineArgs(EngineArgs)","exported":true,"origin":"ast","body_hash":"0536f116cad1dbbd6f44493fe7c712af4a85b622a2fb2f63de00ee04edc64384","body_text":"class EngineArgs(EngineArgs): model_hf_config: PretrainedConfig = None # for verl def __post_init__(self): pass def create_model_config(self) -> ModelConfig: return ModelConfig( hf_config=self.model_hf_config, tokenizer_mode=self.tokenizer_mode, trust_remote_code=self.trust_remote_code, dtype=self.dtype, seed=self.seed, revision=self.revision, code_revision=self.code_revision, rope_scaling=self.rope_scaling, rope_theta=self.rope_theta, tokenizer_revision=self.tokenizer_revision, max_model_len=self.max_model_len, quantization=self.quantization, quantization_param_path=self.quantization_param_path, enforce_eager=self.enforce_eager, max_context_len_to_capture=self.max_context_len_to_capture, max_seq_len_to_capture=self.max_seq_len_to_capture, max_logprobs=self.max_logprobs, disable_sliding_window=self.disable_sliding_window, skip_tokenizer_init=self.skip_tokenizer_init, served_model_name=self.served_model_name, limit_mm_per_prompt=self.limit_mm_per_prompt, use_async_output_proc=not self.disable_async_output_proc, override_neuron_config=self.override_neuron_config, config_format=self.config_format, mm_processor_kwargs=self.mm_processor_kwargs, ) def create_load_config(self) -> LoadConfig: return LoadConfig( load_format=self.load_format, download_dir=self.download_dir, model_loader_extra_config=self.model_loader_extra_config, ignore_patterns=self.ignore_patterns, ) def create_engine_config(self) -> EngineConfig: engine_config = super().create_engine_config() # NOTE[VERL]: Use the world_size set by torchrun world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" engine_config.parallel_config.world_size = world_size return engine_config","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.__post_init__","name":"__post_init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","span":"L30-L31","signature":"def __post_init__(self)","exported":false,"origin":"ast","body_hash":"46a2c1b4cda8162afa0e53406e74c3fa2773a3d3cf49cbfad2f0e012ad34962d","body_text":"def __post_init__(self): pass","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_model_config","name":"create_model_config","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","span":"L33-L60","signature":"def create_model_config(self) -> ModelConfig","exported":true,"origin":"ast","body_hash":"3cdaa2eb583755cc7ee52f6c06e4d3a6d6cbbadea6bec6481ab4afc4988e3016","body_text":"def create_model_config(self) -> ModelConfig: return ModelConfig( hf_config=self.model_hf_config, tokenizer_mode=self.tokenizer_mode, trust_remote_code=self.trust_remote_code, dtype=self.dtype, seed=self.seed, revision=self.revision, code_revision=self.code_revision, rope_scaling=self.rope_scaling, rope_theta=self.rope_theta, tokenizer_revision=self.tokenizer_revision, max_model_len=self.max_model_len, quantization=self.quantization, quantization_param_path=self.quantization_param_path, enforce_eager=self.enforce_eager, max_context_len_to_capture=self.max_context_len_to_capture, max_seq_len_to_capture=self.max_seq_len_to_capture, max_logprobs=self.max_logprobs, disable_sliding_window=self.disable_sliding_window, skip_tokenizer_init=self.skip_tokenizer_init, served_model_name=self.served_model_name, limit_mm_per_prompt=self.limit_mm_per_prompt, use_async_output_proc=not self.disable_async_output_proc, override_neuron_config=self.override_neuron_config, config_format=self.config_format, mm_processor_kwargs=self.mm_processor_kwargs, )","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_load_config","name":"create_load_config","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","span":"L62-L68","signature":"def create_load_config(self) -> LoadConfig","exported":true,"origin":"ast","body_hash":"1b0b41bbcb50a943e065875cb081e44cb92a60fb8100089a863d88fb919ac17c","body_text":"def create_load_config(self) -> LoadConfig: return LoadConfig( load_format=self.load_format, download_dir=self.download_dir, model_loader_extra_config=self.model_loader_extra_config, ignore_patterns=self.ignore_patterns, )","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_engine_config","name":"create_engine_config","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","span":"L70-L78","signature":"def create_engine_config(self) -> EngineConfig","exported":true,"origin":"ast","body_hash":"205ff65a867a66c0ac8f07ed62910d43068e986f8a190227b9e20ff535354901","body_text":"def create_engine_config(self) -> EngineConfig: engine_config = super().create_engine_config() # NOTE[VERL]: Use the world_size set by torchrun world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" engine_config.parallel_config.world_size = world_size return engine_config","summary_state":"pending","summary":null,"crux":null,"owner":"EngineArgs"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","relation":"imports","specifier":"vllm.engine.arg_utils","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs","relation":"extends","name":"EngineArgs","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.__post_init__","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_model_config","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_model_config","relation":"calls","name":"ModelConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_load_config","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_load_config","relation":"calls","name":"LoadConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_engine_config","file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"create_engine_config","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py#EngineArgs.create_engine_config","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/arg_utils.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/config.py":{"size":4167,"mtimeMs":1789275830318.203,"hash":"9380fdc1024ef0ec327c75c35195a009edfdcf39a4e92f53a45149fe246cec6d","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/config.py","name":"config.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/config.py","span":"L1-L106","signature":null,"exported":true,"origin":"ast","body_hash":"9380fdc1024ef0ec327c75c35195a009edfdcf39a4e92f53a45149fe246cec6d","chars":4167,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py import enum import json from dataclasses import dataclass, field from typing import TYPE_CHECKING, List, Optional, Union from transformers import PretrainedConfig # Add for verl from vllm.config import ModelConfig from vllm.logger import init_logger from vllm.utils import is_hip if TYPE_CHECKING: from vllm.model_executor.model_loader.loader import BaseModelLoader logger = init_logger(__name__) @dataclass"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadFormat","name":"LoadFormat","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/config.py","span":"L34-L41","signature":"class LoadFormat(str, enum.Enum)","exported":true,"origin":"ast","body_hash":"28b7fada87a4f663b8593985bb82bd6681e6a4ab7ef047ee5236f7d78e149941","body_text":"class LoadFormat(str, enum.Enum): AUTO = \"auto\" MEGATRON = \"megatron\" HF = \"hf\" DTENSOR = \"dtensor\" DUMMY_HF = \"dummy_hf\" DUMMY_MEGATRON = \"dummy_megatron\" DUMMY_DTENSOR = \"dummy_dtensor\"","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/config.py#ModelConfig","name":"ModelConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/config.py","span":"L44-L48","signature":"class ModelConfig(ModelConfig)","exported":true,"origin":"ast","body_hash":"f2f9a35098d45fd4869cab61ef089d692a84897e0142cfa16fb6669771862fd0","body_text":"class ModelConfig(ModelConfig): def __init__(self, hf_config: PretrainedConfig, *args, **kwargs) -> None: super().__init__(model=hf_config._name_or_path, tokenizer=hf_config._name_or_path, *args, **kwargs) self.hf_config = hf_config","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/config.py#ModelConfig.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/config.py","span":"L46-L48","signature":"def __init__(self, hf_config: PretrainedConfig, *args, **kwargs) -> None","exported":false,"origin":"ast","body_hash":"e06bead7676f178cc098a39e55e14bdd7ae6617c87158b97104c36ebcc38f7d0","body_text":"def __init__(self, hf_config: PretrainedConfig, *args, **kwargs) -> None: super().__init__(model=hf_config._name_or_path, tokenizer=hf_config._name_or_path, *args, **kwargs) self.hf_config = hf_config","summary_state":"pending","summary":null,"crux":null,"owner":"ModelConfig"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig","name":"LoadConfig","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/config.py","span":"L52-L105","signature":"class LoadConfig","exported":true,"origin":"ast","body_hash":"67e0f3183826a100694398b0aaccda5cf1e6a6238607523ec3ba473dc971aca6","body_text":"class LoadConfig: \"\"\" download_dir: Directory to download and load the weights, default to the default cache directory of huggingface. load_format: The format of the model weights to load: \"auto\" will try to load the weights in the safetensors format and fall back to the pytorch bin format if safetensors format is not available. \"pt\" will load the weights in the pytorch bin format. \"safetensors\" will load the weights in the safetensors format. \"npcache\" will load the weights in pytorch format and store a numpy cache to speed up the loading. \"dummy\" will initialize the weights with random values, which is mainly for profiling. \"tensorizer\" will use CoreWeave's tensorizer library for fast weight loading. \"bitsandbytes\" will load nf4 type weights. ignore_patterns: The list of patterns to ignore when loading the model. Default to \"original/**/*\" to avoid repeated loading of llama's checkpoints. \"\"\" load_format: Union[str, LoadFormat, \"BaseModelLoader\"] = LoadFormat.AUTO download_dir: Optional[str] = None model_loader_extra_config: Optional[Union[str, dict]] = field(default_factory=dict) ignore_patterns: Optional[Union[List[str], str]] = None def __post_init__(self): model_loader_extra_config = self.model_loader_extra_config or {} if isinstance(model_loader_extra_config, str): self.model_loader_extra_config = json.loads(model_loader_extra_config) self._verify_load_format() if self.ignore_patterns is not None and len(self.ignore_patterns) > 0: logger.info(\"Ignoring the following patterns when downloading weights: %s\", self.ignore_patterns) else: self.ignore_patterns = [\"original/**/*\"] def _verify_load_format(self) -> None: if not isinstance(self.load_format, str): return load_format = self.load_format.lower() self.load_format = LoadFormat(load_format) rocm_not_supported_load_format: List[str] = [] if is_hip() and load_format in rocm_not_supported_load_format: rocm_supported_load_format = [ f for f in LoadFormat.__members__ if (f not in rocm_not_supported_load_format) ] raise ValueError(f\"load format '{load_format}' is not supported in ROCm. \" f\"Supported load formats are \" f\"{rocm_supported_load_format}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig.__post_init__","name":"__post_init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/config.py","span":"L80-L89","signature":"def __post_init__(self)","exported":false,"origin":"ast","body_hash":"5094b5416a54d00d4a1013fa86e20a0199223b86997c6d96b2410245cb2b7f25","body_text":"def __post_init__(self): model_loader_extra_config = self.model_loader_extra_config or {} if isinstance(model_loader_extra_config, str): self.model_loader_extra_config = json.loads(model_loader_extra_config) self._verify_load_format() if self.ignore_patterns is not None and len(self.ignore_patterns) > 0: logger.info(\"Ignoring the following patterns when downloading weights: %s\", self.ignore_patterns) else: self.ignore_patterns = [\"original/**/*\"]","summary_state":"pending","summary":null,"crux":null,"owner":"LoadConfig"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig._verify_load_format","name":"_verify_load_format","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/config.py","span":"L91-L105","signature":"def _verify_load_format(self) -> None","exported":false,"origin":"ast","body_hash":"f9f7e30ec352e2c4183434a9e0e078cb7cff2b5e08499624e93b0b4857aff8ab","body_text":"def _verify_load_format(self) -> None: if not isinstance(self.load_format, str): return load_format = self.load_format.lower() self.load_format = LoadFormat(load_format) rocm_not_supported_load_format: List[str] = [] if is_hip() and load_format in rocm_not_supported_load_format: rocm_supported_load_format = [ f for f in LoadFormat.__members__ if (f not in rocm_not_supported_load_format) ] raise ValueError(f\"load format '{load_format}' is not supported in ROCm. \" f\"Supported load formats are \" f\"{rocm_supported_load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"LoadConfig"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"imports","specifier":"enum","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"imports","specifier":"json","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"imports","specifier":"dataclasses","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"imports","specifier":"vllm.model_executor.model_loader.loader","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadFormat","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadFormat","relation":"extends","name":"str","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/config.py#ModelConfig","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#ModelConfig","relation":"extends","name":"ModelConfig","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#ModelConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/config.py#ModelConfig.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#ModelConfig.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#ModelConfig.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig","relation":"calls","name":"field","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig.__post_init__","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig.__post_init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig.__post_init__","relation":"calls","name":"loads","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig.__post_init__","relation":"calls","name":"_verify_load_format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py","recvType":"LoadConfig"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig.__post_init__","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig.__post_init__","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig._verify_load_format","file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig._verify_load_format","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig._verify_load_format","relation":"calls","name":"lower","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py","recvType":"LoadFormat"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig._verify_load_format","relation":"calls","name":"LoadFormat","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig._verify_load_format","relation":"calls","name":"is_hip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/config.py#LoadConfig._verify_load_format","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/config.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py":{"size":17356,"mtimeMs":1789275830318.203,"hash":"280ba6eacd660cb208f6e0fe1b427cd634c67927121328d6b3bc1226d59d0280","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","name":"dtensor_weight_loaders.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L1-L381","signature":null,"exported":true,"origin":"ast","body_hash":"280ba6eacd660cb208f6e0fe1b427cd634c67927121328d6b3bc1226d59d0280","chars":17356,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/model_loader from typing import Dict import torch.nn as nn from torch.distributed._tensor import DTensor from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.utils import is_pp_missing_parameter from vllm.model_executor.layers.fused_moe import FusedMoE __MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__ = { \"GPT2LMHeadModel\": gpt2_dtensor_weight_loader, \"LlamaForCausalLM\": llama_dtensor_weight_loader, \"LLaMAForCausalLM\": llama_dtensor_weight_loader, \"MistralForCausalLM\": llama_dtensor_weight_loader, # mistral is the same as llama in vLLM \"InternLMForCausalLM\": llama_dtensor_weight_loader, \"AquilaModel\": llama_dtensor_weight_loader, \"AquilaForCausalLM\": llama_dtensor_weight_loader, \"Phi3ForCausalLM\": llama_dtensor_weight_loader, \"GemmaForCausalLM\": gemma_dtensor_weight_loader, \"Gemma2ForCausalLM\": gemma_dtensor_weight_loader, \"GPTBigCodeForCausalLM\": gptbigcode_dtensor_load_weights, \"Starcoder2ForCausalLM\": starcoder2_dtensor_load_weights, \"Qwen2ForCausalLM\": qwen2_dtensor_weight_loader, \"DeepseekV2ForCausalLM\": deepseekv2_dtensor_weight_loader, \"Qwen2VLForConditionalGeneration\": qwen2vl_dtensor_weight_loader, } # the actor model is .state_dict() # Load dtensor weights # NOTE(sgm): we use per-parameter weight loader in each vllm sub"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","name":"gemma_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L24-L58","signature":"def gemma_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"35350e6cb89b3798ab1d6f2f9440feb37eb4da95d6cd8b05e563f390cdfc4e5e","body_text":"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(vllm_model.named_parameters()) 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 stacked_name = name.replace(shard_name, param_name) # Skip loading extra bias for GPTQ models. if stacked_name.endswith(\".bias\") and stacked_name not in params_dict: continue local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) param = params_dict[stacked_name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) 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, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","name":"gptbigcode_dtensor_load_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L61-L73","signature":"def gptbigcode_dtensor_load_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"7494a5e13088c35a8ed94c153c3d247196ed1c1f0c08b38b75cfc6be4046b77e","body_text":"def gptbigcode_dtensor_load_weights(actor_weights: Dict, vllm_model: nn.Module): params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"lm_head.weight\" in name: continue if \".attn.bias\" in name: # Skip attention mask. # NOTE: \"c_attn.bias\" should not be skipped. 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, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","name":"starcoder2_dtensor_load_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L76-L104","signature":"def starcoder2_dtensor_load_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"e13c1e9a0944634e39269ff65f9fecf9fad5ef957c93d6ac35727465476f3896","body_text":"def starcoder2_dtensor_load_weights(actor_weights: Dict, vllm_model: 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\"), ] params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_name, param_name) local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) param = params_dict[name] weight_loader = param.weight_loader weight_loader(param, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in name: continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","name":"llama_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L107-L148","signature":"def llama_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"4e7619de64eee88249884eb893161ea6c86e46f48c0b4ff71dd5f1efb67fa07f","body_text":"def llama_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(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue if \"rotary_emb.cos_cached\" in name or \"rotary_emb.sin_cached\" in name: # Models trained using ColossalAI may include these tensors in # the checkpoint. Skip them. continue # With tie_word_embeddings, we can skip lm_head.weight # The weight might appear unnecessarily in the files if the model is # processed with quantization, LoRA, fine-tuning, etc. if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in name: continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_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, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: # 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, local_loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","name":"qwen2_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L151-L185","signature":"def qwen2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"d634800a458d316c8c20305017f2eb526adabb69ada10e7f834c8db536aa50dc","body_text":"def qwen2_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(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in name: continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_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, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: # Skip loading extra bias for GPTQ models. if name.endswith(\".bias\") and name not in params_dict: continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","name":"qwen2vl_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L188-L222","signature":"def qwen2vl_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"f4caef72eead6bfb4b996cab616abec20ec679a1990f029d0aacadc36f37e301","body_text":"def qwen2vl_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(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in name: continue for param_name, weight_name, shard_id in stacked_params_mapping: if weight_name not in name: continue name = name.replace(weight_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, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: # Skip loading extra bias for GPTQ models. if name.endswith(\".bias\") and name not in params_dict: continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","name":"deepseekv2_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L228-L305","signature":"def deepseekv2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"6ec264e61d6a93fe5c3bf388847ba367d47ce4d0865b52155a42fadce0b51137","body_text":"def deepseekv2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: stacked_params_mapping = [ # (param_name, shard_name, shard_id) (\"gate_up_proj\", \"gate_proj\", 0), (\"gate_up_proj\", \"up_proj\", 1), ] # Params for weights, fp8 weight scales, fp8 activation scales # (param_name, weight_name, expert_id, shard_id) expert_params_mapping = FusedMoE.make_expert_params_mapping( ckpt_gate_proj_name=\"gate_proj\", ckpt_down_proj_name=\"down_proj\", ckpt_up_proj_name=\"up_proj\", num_experts=vllm_model.config.n_routed_experts, ) params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue for param_name, weight_name, shard_id in stacked_params_mapping: # Skip non-stacked layers and experts (experts handled below). if weight_name not in name: continue # We have mlp.experts[0].gate_proj in the checkpoint. # Since we handle the experts below in expert_params_mapping, # we need to skip here BEFORE we update the name, otherwise # name will be updated to mlp.experts[0].gate_up_proj, which # will then be updated below in expert_params_mapping # for mlp.experts[0].gate_gate_up_proj, which breaks load. if (\"mlp.experts.\" in name) and name not in params_dict: continue name = name.replace(weight_name, param_name) # Skip loading extra bias for GPTQ models. if name.endswith(\".bias\") and name not in params_dict: continue if is_pp_missing_parameter(name, vllm_model): continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype), shard_id) break else: for mapping in expert_params_mapping: param_name, weight_name, expert_id, shard_id = mapping if weight_name not in name: continue name = name.replace(weight_name, param_name) if is_pp_missing_parameter(name, vllm_model): continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader( param, local_loaded_weight.to(dtype=param.dtype), weight_name, shard_id=shard_id, expert_id=expert_id, ) break else: # Skip loading extra bias for GPTQ models. if name.endswith(\".bias\") and name not in params_dict: continue if is_pp_missing_parameter(name, vllm_model): continue param = params_dict[name] local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight) weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, local_loaded_weight.to(dtype=param.dtype))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gpt2_dtensor_weight_loader","name":"gpt2_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L308-L309","signature":"def gpt2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"e26a0b664180d6dbfeb3ba74e435ba77e0a32726d4fc799cc1282d68edce33d9","body_text":"def gpt2_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: pass","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#redistribute_dtensor","name":"redistribute_dtensor","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L312-L323","signature":"def redistribute_dtensor(param_name: str, loaded_weights: DTensor, parallelize_plan: Dict = None)","exported":true,"origin":"ast","body_hash":"32cdf551169eca64177afa6234b3b8b4738adef11e09cf64d2298ed7166aecc6","body_text":"def redistribute_dtensor(param_name: str, loaded_weights: DTensor, parallelize_plan: Dict = None): param_name = _process_parameter_names(name=param_name) if parallelize_plan is not None: assert ( param_name in parallelize_plan.keys()), f\"param name: {param_name} not in parallelize_plan :{parallelize_plan.keys()}\" placement = parallelize_plan[param_name] local_loaded_weights = loaded_weights.redistribute(device_mesh=loaded_weights.device_mesh, placements=placement).to_local() else: local_loaded_weights = loaded_weights.full_tensor() return local_loaded_weights","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#_process_parameter_names","name":"_process_parameter_names","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L326-L339","signature":"def _process_parameter_names(name): # Remove '.weight' if it exists at the end of the string","exported":false,"origin":"ast","body_hash":"32ced0f6904a96973b48620ad345f314a91ddd2e87c876246a75d308ca0226c5","body_text":"def _process_parameter_names(name): # Remove '.weight' if it exists at the end of the string if name.endswith(\".weight\"): name = name[:-7] # Remove 'model.layers.x.' or 'model.' prefix if \"model.layers\" in name: parts = name.split(\".\") # Reconstruct the string without 'model.layers.x.' name = \".\".join(parts[3:]) # parts[0] is 'model', parts[1] is 'layers', parts[2] is 'x' elif name.startswith(\"model.\"): name = name[6:] # Remove 'model.' return name","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#load_dtensor_weights","name":"load_dtensor_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L363-L368","signature":"def load_dtensor_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"d614ba13175fcb0a1db90bcd80acdd13cbcddb57eb013b3ce25386d8328656ce","body_text":"def load_dtensor_weights(actor_weights: Dict, vllm_model: nn.Module): weight_loader = _get_model_weight_loader(vllm_model.__class__.__name__) weight_loader(actor_weights, vllm_model) # NOTE(sgm) to reduce peak memory usage, we offload vllm model to cpu # after init, and we need this after sync model weights for in first iter. vllm_model = vllm_model.cuda()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#_get_model_weight_loader","name":"_get_model_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L371-L375","signature":"def _get_model_weight_loader(arch: str)","exported":false,"origin":"ast","body_hash":"2aff4020b0abe87a16fdc5b4ba0f9ed8f437312a1edfce1320e5726f5165c8ae","body_text":"def _get_model_weight_loader(arch: str): if arch in __MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__: return __MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__[arch] raise ValueError(f\"Model architectures {arch} are not supported for now. \" f\"Supported architectures: {__MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__.keys()}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#update_dtensor_weight_loader","name":"update_dtensor_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","span":"L379-L380","signature":"def update_dtensor_weight_loader()","exported":true,"origin":"ast","body_hash":"3251c12c590cfa8575d836647130786146b1ba578183f484f98d2dac2cce5914","body_text":"def update_dtensor_weight_loader(): pass","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"imports","specifier":"torch.distributed._tensor","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.model_loader.weight_utils","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.models.utils","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gemma_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gptbigcode_dtensor_load_weights","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#starcoder2_dtensor_load_weights","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#llama_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#qwen2vl_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.fused_moe","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"make_expert_params_mapping","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"is_pp_missing_parameter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"is_pp_missing_parameter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"is_pp_missing_parameter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"redistribute_dtensor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#deepseekv2_dtensor_weight_loader","relation":"calls","name":"to","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"redistribute_dtensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#gpt2_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#redistribute_dtensor","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"_process_parameter_names","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"to_local","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"redistribute","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"DTensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#redistribute_dtensor","relation":"calls","name":"full_tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","recvType":"DTensor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#_process_parameter_names","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"split","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#_process_parameter_names","relation":"calls","name":"startswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#load_dtensor_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#load_dtensor_weights","relation":"calls","name":"_get_model_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#load_dtensor_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#load_dtensor_weights","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#_get_model_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py#update_dtensor_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/dtensor_weight_loaders.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py":{"size":1773,"mtimeMs":1789275830318.203,"hash":"01e371e44c45c0a4ecc2a757f3244a7220f6baa31c52a57d3b6c58f07a367752","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","name":"hf_weight_loader.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","span":"L1-L42","signature":null,"exported":true,"origin":"ast","body_hash":"01e371e44c45c0a4ecc2a757f3244a7220f6baa31c52a57d3b6c58f07a367752","chars":1773,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/model_loader from typing import Dict import torch.nn as nn from vllm.model_executor.model_loader.utils import set_default_torch_dtype"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#update_hf_weight_loader","name":"update_hf_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","span":"L22-L24","signature":"def update_hf_weight_loader()","exported":true,"origin":"ast","body_hash":"0cfa434a2a336b6f473cf6dc29e094817916dd89411d0ff741380d2aacf82d44","body_text":"def update_hf_weight_loader(): print(\"no hf weight loader need to be updated\") return","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","name":"load_hf_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","span":"L27-L41","signature":"def load_hf_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"ba5f0f2cc08e7ff77c5fbacded07d517f7cee3cee904129e40a6bb641a06de22","body_text":"def load_hf_weights(actor_weights: Dict, vllm_model: nn.Module): assert isinstance(actor_weights, Dict) with set_default_torch_dtype(next(vllm_model.parameters()).dtype): # TODO if vllm_model.config.tie_word_embeddings and \"lm_head.weight\" in actor_weights.keys(): del actor_weights[\"lm_head.weight\"] vllm_model.load_weights(actor_weights.items()) for _, module in vllm_model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() vllm_model = vllm_model.cuda()","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.utils","file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#update_hf_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#update_hf_weight_loader","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"next","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"keys","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"load_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py#load_hf_weights","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/hf_weight_loader.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/llm.py":{"size":10236,"mtimeMs":1789275830318.203,"hash":"5bc262b93ac5e938c82cc51f0e48b22a33cf13d0f5966f02afa15b60445203ff","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","name":"llm.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L1-L201","signature":null,"exported":true,"origin":"ast","body_hash":"5bc262b93ac5e938c82cc51f0e48b22a33cf13d0f5966f02afa15b60445203ff","chars":10236,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/entrypoints/llm.py from typing import Dict, List, Optional, Tuple, Union import torch import torch.nn as nn from torch.nn.utils.rnn import pad_sequence from transformers import PretrainedConfig, PreTrainedTokenizer, PreTrainedTokenizerFast from verl.workers.rollout.tokenizer import HybridEngineBaseTokenizer from vllm import LLM from vllm.outputs import EmbeddingRequestOutput, RequestOutput from vllm.utils import Counter from .arg_utils import EngineArgs from .llm_engine_sp import LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","name":"LLM","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L31-L200","signature":"class LLM(LLM)","exported":true,"origin":"ast","body_hash":"0692ad9846e8627239ff929ab2b635ee128bd199d394becb264faed570666ebf","body_text":"class LLM(LLM): \"\"\"An LLM for generating texts from given prompts and sampling parameters. This class includes a tokenizer, a language model (possibly distributed across multiple GPUs), and GPU memory space allocated for intermediate states (aka KV cache). Given a batch of prompts and sampling parameters, this class generates texts from the model, using an intelligent batching mechanism and efficient memory management. NOTE: This class is intended to be used for offline inference. For online serving, use the `AsyncLLMEngine` class instead. NOTE: For the comprehensive list of arguments, see `EngineArgs`. Args: model: A HuggingFace Transformers model instance. tokenizer: A HuggingFace Transformers tokenizer instance. tokenizer_mode: The tokenizer mode. \"auto\" will use the fast tokenizer if available, and \"slow\" will always use the slow tokenizer. trust_remote_code: Trust remote code (e.g., from HuggingFace) when downloading the model and tokenizer. tensor_parallel_size: The number of GPUs to use for distributed execution with tensor parallelism. dtype: The data type for the model weights and activations. Currently, we support `float32`, `float16`, and `bfloat16`. If `auto`, we use the `torch_dtype` attribute specified in the model config file. However, if the `torch_dtype` in the config is `float32`, we will use `float16` instead. quantization: The method used to quantize the model weights. Currently, we support \"awq\". If None, we assume the model weights are not quantized and use `dtype` to determine the data type of the weights. revision: The specific model version to use. It can be a branch name, a tag name, or a commit id. tokenizer_revision: The specific tokenizer version to use. It can be a branch name, a tag name, or a commit id. seed: The seed to initialize the random number generator for sampling. gpu_memory_utilization: The ratio (between 0 and 1) of GPU memory to reserve for the model weights, activations, and KV cache. Higher values will increase the KV cache size and thus improve the model's throughput. However, if the value is too high, it may cause out-of- memory (OOM) errors. swap_space: The size (GiB) of CPU memory per GPU to use as swap space. This can be used for temporarily storing the states of the requests when their `best_of` sampling parameters are larger than 1. If all requests will have `best_of=1`, you can safely set this to 0. Otherwise, too small values may cause out-of-memory (OOM) errors. enforce_eager: Whether to enforce eager execution. If True, we will disable CUDA graph and always execute the model in eager mode. If False, we will use CUDA graph and eager execution in hybrid. max_context_len_to_capture: Maximum context len covered by CUDA graphs. When a sequence has context length larger than this, we fall back to eager mode. disable_custom_all_reduce: See ParallelConfig \"\"\" def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, skip_tokenizer_init: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, cpu_offload_gb: float = 0, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: int = 8192, disable_custom_all_reduce: bool = False, load_format=\"auto\", **kwargs, ) -> None: if \"disable_log_stats\" not in kwargs: kwargs[\"disable_log_stats\"] = True removed_vision_keys = (\"image_token_id\", \"image_feature_size\", \"image_input_shape\", \"image_input_type\") if any(k in kwargs for k in removed_vision_keys): raise TypeError(\"There is no need to pass vision-related arguments anymore.\") engine_args = EngineArgs( model_hf_config=model_hf_config, # tokenizer=tokenizer, tokenizer_mode=tokenizer_mode, skip_tokenizer_init=skip_tokenizer_init, trust_remote_code=trust_remote_code, tensor_parallel_size=tensor_parallel_size, dtype=dtype, quantization=quantization, revision=revision, tokenizer_revision=tokenizer_revision, seed=seed, gpu_memory_utilization=gpu_memory_utilization, swap_space=swap_space, cpu_offload_gb=cpu_offload_gb, enforce_eager=enforce_eager, max_context_len_to_capture=max_context_len_to_capture, max_seq_len_to_capture=max_seq_len_to_capture, disable_custom_all_reduce=disable_custom_all_reduce, load_format=load_format, **kwargs, ) tokenizer_cls = (PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer) if not isinstance(tokenizer, tokenizer_cls): raise ValueError( f\"Unexpected tokenizer type: {type(tokenizer)}. Must be\" \"one of the following: PreTrainedTokenizer, PreTrainedTokenizerFast, verl.workers.rollout.HybridEngineBaseTokenizer\" ) self.llm_engine = LLMEngine.from_engine_args(model, tokenizer, engine_a","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L85-L143","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, skip_tokenizer_init: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, cpu_offload_gb: float = 0, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: int = 8192, disable_custom_all_reduce: bool = False, load_format=\"auto\", **kwargs, ) -> None","exported":false,"origin":"ast","body_hash":"705dad1dce8b2d1b964e1380b10624fdb2034b731634d8ce98ca3aaad33f47db","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer], model_hf_config: PretrainedConfig, tokenizer_mode: str = \"auto\", trust_remote_code: bool = False, skip_tokenizer_init: bool = False, tensor_parallel_size: int = 1, dtype: str = \"auto\", quantization: Optional[str] = None, revision: Optional[str] = None, tokenizer_revision: Optional[str] = None, seed: int = 0, gpu_memory_utilization: float = 0.9, swap_space: int = 4, cpu_offload_gb: float = 0, enforce_eager: bool = False, max_context_len_to_capture: Optional[int] = None, max_seq_len_to_capture: int = 8192, disable_custom_all_reduce: bool = False, load_format=\"auto\", **kwargs, ) -> None: if \"disable_log_stats\" not in kwargs: kwargs[\"disable_log_stats\"] = True removed_vision_keys = (\"image_token_id\", \"image_feature_size\", \"image_input_shape\", \"image_input_type\") if any(k in kwargs for k in removed_vision_keys): raise TypeError(\"There is no need to pass vision-related arguments anymore.\") engine_args = EngineArgs( model_hf_config=model_hf_config, # tokenizer=tokenizer, tokenizer_mode=tokenizer_mode, skip_tokenizer_init=skip_tokenizer_init, trust_remote_code=trust_remote_code, tensor_parallel_size=tensor_parallel_size, dtype=dtype, quantization=quantization, revision=revision, tokenizer_revision=tokenizer_revision, seed=seed, gpu_memory_utilization=gpu_memory_utilization, swap_space=swap_space, cpu_offload_gb=cpu_offload_gb, enforce_eager=enforce_eager, max_context_len_to_capture=max_context_len_to_capture, max_seq_len_to_capture=max_seq_len_to_capture, disable_custom_all_reduce=disable_custom_all_reduce, load_format=load_format, **kwargs, ) tokenizer_cls = (PreTrainedTokenizer, PreTrainedTokenizerFast, HybridEngineBaseTokenizer) if not isinstance(tokenizer, tokenizer_cls): raise ValueError( f\"Unexpected tokenizer type: {type(tokenizer)}. Must be\" \"one of the following: PreTrainedTokenizer, PreTrainedTokenizerFast, verl.workers.rollout.HybridEngineBaseTokenizer\" ) self.llm_engine = LLMEngine.from_engine_args(model, tokenizer, engine_args) # TODO: check usagecontext self.request_counter = Counter()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L145-L146","signature":"def init_cache_engine(self)","exported":true,"origin":"ast","body_hash":"04d68c22194d8053c926485d043932fb2526f0d73a880157acee875717c5937f","body_text":"def init_cache_engine(self): self.llm_engine.init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L148-L149","signature":"def free_cache_engine(self)","exported":true,"origin":"ast","body_hash":"179e8502fa237bb23acb9f06952cf4ae62f3f9010408f63734f8adf6a5163378","body_text":"def free_cache_engine(self): self.llm_engine.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.get_tokenizer","name":"get_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L151-L152","signature":"def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]","exported":true,"origin":"ast","body_hash":"cebad3b0992af1915c8c8a1543b493069d6507f491421f924191ee4b98e60a2f","body_text":"def get_tokenizer(self) -> Union[PreTrainedTokenizer, PreTrainedTokenizerFast]: return self.llm_engine.tokenizer","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.set_tokenizer","name":"set_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L154-L158","signature":"def set_tokenizer( self, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], ) -> None","exported":true,"origin":"ast","body_hash":"a9d3cbc026686b5d94c85093062fb18f30a9d4182a5bf7eb9679c456cbc101fa","body_text":"def set_tokenizer( self, tokenizer: Union[PreTrainedTokenizer, PreTrainedTokenizerFast], ) -> None: self.llm_engine.tokenizer = tokenizer","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._run_engine","name":"_run_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L160-L162","signature":"def _run_engine(self, *, use_tqdm: bool) -> List[Union[RequestOutput, EmbeddingRequestOutput]]","exported":false,"origin":"ast","body_hash":"6fb42d07bbd58fca78525619af30637007ab0e9f639953f061759317bac8cc22","body_text":"def _run_engine(self, *, use_tqdm: bool) -> List[Union[RequestOutput, EmbeddingRequestOutput]]: outputs = super()._run_engine(use_tqdm=use_tqdm) return self._post_process_outputs(outputs)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","name":"_post_process_outputs","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L174-L194","signature":"def _post_process_outputs(self, request_outputs: List[RequestOutput]) -> Tuple[torch.Tensor, torch.Tensor]","exported":false,"origin":"ast","body_hash":"2f059d479cd20191c8fd2da82dfd4cf2ed401588749f45661f35ac8827ad465a","body_text":"def _post_process_outputs(self, request_outputs: List[RequestOutput]) -> Tuple[torch.Tensor, torch.Tensor]: output_token_ids = [] logprobs = [] for request_output in request_outputs: # List[RequestOutput] outputs = request_output.outputs for output in outputs: # List[CompletionOutput], usually len == 1 output_token_ids.append(torch.tensor(output.token_ids)) # TODO(shengguangming): can be optimzied by rewrite the Sampler._get_logprobs() logits logprobs_dicts = output.logprobs if logprobs_dicts is not None: logprob = [] for logprobs_dict, id in zip(logprobs_dicts, output.token_ids): logprob.append(logprobs_dict[id].logprob) logprobs.append(torch.tensor(logprob)) pad_token_id = (self.llm_engine.tokenizer.pad_token_id if self.llm_engine.tokenizer.pad_token_id is not None else self.llm_engine.tokenizer.eos_token_id) output_token_ids = pad_sequence(output_token_ids, batch_first=True, padding_value=pad_token_id) if len(logprobs) > 0: logprobs = pad_sequence(logprobs, batch_first=True, padding_value=pad_token_id) return output_token_ids, logprobs","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L196-L197","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None","exported":true,"origin":"ast","body_hash":"df77662639c0d1c0008e46c21bdeffd4e84ed28fa0388825031e0577e7d24cf8","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None: self.llm_engine.sync_model_weights(actor_weights=actor_weights, load_format=load_format)","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","span":"L199-L200","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"27a73de1e582c115374d023e58aeaf54824c75ca2740c19f48f07d7d67773621","body_text":"def offload_model_weights(self) -> None: self.llm_engine.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"LLM"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"imports","specifier":"torch.nn.utils.rnn","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"imports","specifier":"verl.workers.rollout.tokenizer","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"imports","specifier":"vllm","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"imports","specifier":"vllm.outputs","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"imports","specifier":".arg_utils","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"imports","specifier":".llm_engine_sp","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","relation":"extends","name":"LLM","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.__init__","relation":"calls","name":"any","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.__init__","relation":"calls","name":"TypeError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.__init__","relation":"calls","name":"EngineArgs","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.__init__","relation":"calls","name":"type","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.__init__","relation":"calls","name":"from_engine_args","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.__init__","relation":"calls","name":"Counter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.init_cache_engine","relation":"calls","name":"init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.get_tokenizer","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.set_tokenizer","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._run_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._run_engine","relation":"calls","name":"_run_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._run_engine","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._run_engine","relation":"calls","name":"_post_process_outputs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","recvType":"LLM"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","recvType":"pad_sequence"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","relation":"calls","name":"zip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py","recvType":"pad_sequence"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","relation":"calls","name":"pad_sequence","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM._post_process_outputs","relation":"calls","name":"pad_sequence","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm.py#LLM.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py":{"size":17823,"mtimeMs":1789275830318.203,"hash":"618c464b9b80f595863c9c831c65a53a854f695012cc506f63c691073bf86e53","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","name":"llm_engine_sp.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L1-L409","signature":null,"exported":true,"origin":"ast","body_hash":"618c464b9b80f595863c9c831c65a53a854f695012cc506f63c691073bf86e53","chars":17823,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/engine/llm_engine.py from functools import partial from typing import Callable, Dict, Optional, Type, Union import torch import torch.nn as nn from vllm.config import ( CacheConfig, DecodingConfig, DeviceConfig, EngineConfig, LoadConfig, LoRAConfig, ModelConfig, ObservabilityConfig, ParallelConfig, PromptAdapterConfig, SchedulerConfig, SpeculativeConfig, ) from vllm.core.scheduler import Scheduler from vllm.engine.arg_utils import EngineArgs from vllm.engine.llm_engine import LLMEngine, SchedulerContext, SchedulerOutputState, _load_generation_config_dict from vllm.engine.metrics_types import StatLoggerBase from vllm.engine.output_processor.interfaces import SequenceGroupOutputProcessor from vllm.engine.output_processor.stop_checker import StopChecker from vllm.executor.executor_base import ExecutorBase from vllm.inputs import INPUT_REGISTRY, InputRegistry from vllm.inputs.preprocess import InputPreprocessor from vllm.logger import init_logger from vllm.sequence import Sequence from vllm.tracing import init_tracer from vllm.transformers_utils.detokenizer import Detokenizer from vllm.transformers_utils.tokenizer import AnyTokenizer from vllm.usage.usage_lib import UsageContext, is_usage_stats_enabled, usage_message from vllm.utils import Counter, weak_bind from vllm.version import __version__ as VLLM_VERSION from .arg_utils import EngineArgs from .config import LoadConfig, ModelConfig from .tokenizer import TokenizerGroup logger = init_logger(__name__) _LOCAL_LOGGING_INTERVAL_SEC = 5"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","name":"LLMEngine","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L61-L408","signature":"class LLMEngine(LLMEngine)","exported":true,"origin":"ast","body_hash":"46fb09d7c2276880c6e9da3797b31c380c1d1261c99ac5324a77546afc773c51","body_text":"class LLMEngine(LLMEngine): \"\"\"An LLM engine that receives requests and generates texts. This is the main class for the vLLM engine. It receives requests from clients and generates texts from the LLM. It includes a tokenizer, a language model (possibly distributed across multiple GPUs), and GPU memory space allocated for intermediate states (aka KV cache). This class utilizes iteration-level scheduling and efficient memory management to maximize the serving throughput. The :class:`~vllm.LLM` class wraps this class for offline batched inference and the :class:`AsyncLLMEngine` class wraps this class for online serving. The config arguments are derived from :class:`~vllm.EngineArgs`. (See :ref:`engine_args`) Args: model_config: The configuration related to the LLM model. cache_config: The configuration related to the KV cache memory management. parallel_config: The configuration related to distributed execution. scheduler_config: The configuration related to the request scheduler. device_config: The configuration related to the device. lora_config (Optional): The configuration related to serving multi-LoRA. speculative_config (Optional): The configuration related to speculative decoding. executor_class: The model executor class for managing distributed execution. prompt_adapter_config (Optional): The configuration related to serving prompt adapters. log_stats: Whether to log statistics. usage_context: Specified entry point, used for usage info collection. \"\"\" def __init__( self, # NOTE(sgm): first two arguments are added for verl model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, # NOTE(sgm): vllm original arguments model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], speculative_config: Optional[SpeculativeConfig], decoding_config: Optional[DecodingConfig], observability_config: Optional[ObservabilityConfig], prompt_adapter_config: Optional[PromptAdapterConfig], executor_class: Type[ExecutorBase], log_stats: bool, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, stat_loggers: Optional[Dict[str, StatLoggerBase]] = None, input_registry: InputRegistry = INPUT_REGISTRY, use_cached_outputs: bool = False, ) -> None: logger.info( \"Initializing an LLM engine (v%s) with config: \" \"model=%r, speculative_config=%r, tokenizer=%r, \" \"skip_tokenizer_init=%s, tokenizer_mode=%s, revision=%s, \" \"override_neuron_config=%s, \" \"rope_scaling=%r, rope_theta=%r, tokenizer_revision=%s, \" \"trust_remote_code=%s, dtype=%s, max_seq_len=%d, \" \"download_dir=%r, load_format=%s, tensor_parallel_size=%d, \" \"pipeline_parallel_size=%d, \" \"disable_custom_all_reduce=%s, quantization=%s, \" \"enforce_eager=%s, kv_cache_dtype=%s, \" \"quantization_param_path=%s, device_config=%s, \" \"decoding_config=%r, observability_config=%r, \" \"seed=%d, served_model_name=%s, use_v2_block_manager=%s, \" \"num_scheduler_steps=%d, chunked_prefill_enabled=%s \" \"multi_step_stream_outputs=%s, enable_prefix_caching=%s, \" \"use_async_output_proc=%s, use_cached_outputs=%s, \" \"mm_processor_kwargs=%s)\", VLLM_VERSION, model_config.model, speculative_config, model_config.tokenizer, model_config.skip_tokenizer_init, model_config.tokenizer_mode, model_config.revision, model_config.override_neuron_config, model_config.rope_scaling, model_config.rope_theta, model_config.tokenizer_revision, model_config.trust_remote_code, model_config.dtype, model_config.max_model_len, load_config.download_dir, load_config.load_format, parallel_config.tensor_parallel_size, parallel_config.pipeline_parallel_size, parallel_config.disable_custom_all_reduce, model_config.quantization, model_config.enforce_eager, cache_config.cache_dtype, model_config.quantization_param_path, device_config.device, decoding_config, observability_config, model_config.seed, model_config.served_model_name, scheduler_config.use_v2_block_manager, scheduler_config.num_scheduler_steps, scheduler_config.chunked_prefill_enabled, scheduler_config.multi_step_stream_outputs, cache_config.enable_prefix_caching, model_config.use_async_output_proc, use_cached_outputs, model_config.mm_processor_kwargs, ) # TODO(woosuk): Print more configs in debug mode. self.model_config = model_config self.cache_config = cache_config self.lora_config = lora_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.speculative_config = speculative_config self.load_config = load_config self.decoding_config = decoding_config or DecodingConfig() self.prompt_adapter_config = prompt_adapter_config self.observability_config = observability_config or ObservabilityConfig() self.log_stats = log_stats self.use_cached_outputs = use_cached_outputs if not self.model_config.skip_tokenizer_init: self.tokenizer = self._init_tokenizer(tokenizer) self.detokenizer = Detokenizer(self.tokenizer) tokenizer_gr","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L95-L334","signature":"def __init__( self, # NOTE(sgm): first two arguments are added for verl model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, # NOTE(sgm): vllm original arguments model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], speculative_config: Optional[SpeculativeConfig], decoding_config: Optional[DecodingConfig], observability_config: Optional[ObservabilityConfig], prompt_adapter_config: Optional[PromptAdapterConfig], executor_class: Type[ExecutorBase], log_stats: bool, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, stat_loggers: Optional[Dict[str, StatLoggerBase]] = None, input_registry: InputRegistry = INPUT_REGISTRY, use_cached_outputs: bool = False, ) -> None","exported":false,"origin":"ast","body_hash":"4e2f68da0e010f58512e1c1af8389c22423738af3689cdcdf309ad49b9296a8b","body_text":"def __init__( self, # NOTE(sgm): first two arguments are added for verl model: Union[nn.Module, Dict], # model itself or its parameter dict tokenizer: nn.Module, # NOTE(sgm): vllm original arguments model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], speculative_config: Optional[SpeculativeConfig], decoding_config: Optional[DecodingConfig], observability_config: Optional[ObservabilityConfig], prompt_adapter_config: Optional[PromptAdapterConfig], executor_class: Type[ExecutorBase], log_stats: bool, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, stat_loggers: Optional[Dict[str, StatLoggerBase]] = None, input_registry: InputRegistry = INPUT_REGISTRY, use_cached_outputs: bool = False, ) -> None: logger.info( \"Initializing an LLM engine (v%s) with config: \" \"model=%r, speculative_config=%r, tokenizer=%r, \" \"skip_tokenizer_init=%s, tokenizer_mode=%s, revision=%s, \" \"override_neuron_config=%s, \" \"rope_scaling=%r, rope_theta=%r, tokenizer_revision=%s, \" \"trust_remote_code=%s, dtype=%s, max_seq_len=%d, \" \"download_dir=%r, load_format=%s, tensor_parallel_size=%d, \" \"pipeline_parallel_size=%d, \" \"disable_custom_all_reduce=%s, quantization=%s, \" \"enforce_eager=%s, kv_cache_dtype=%s, \" \"quantization_param_path=%s, device_config=%s, \" \"decoding_config=%r, observability_config=%r, \" \"seed=%d, served_model_name=%s, use_v2_block_manager=%s, \" \"num_scheduler_steps=%d, chunked_prefill_enabled=%s \" \"multi_step_stream_outputs=%s, enable_prefix_caching=%s, \" \"use_async_output_proc=%s, use_cached_outputs=%s, \" \"mm_processor_kwargs=%s)\", VLLM_VERSION, model_config.model, speculative_config, model_config.tokenizer, model_config.skip_tokenizer_init, model_config.tokenizer_mode, model_config.revision, model_config.override_neuron_config, model_config.rope_scaling, model_config.rope_theta, model_config.tokenizer_revision, model_config.trust_remote_code, model_config.dtype, model_config.max_model_len, load_config.download_dir, load_config.load_format, parallel_config.tensor_parallel_size, parallel_config.pipeline_parallel_size, parallel_config.disable_custom_all_reduce, model_config.quantization, model_config.enforce_eager, cache_config.cache_dtype, model_config.quantization_param_path, device_config.device, decoding_config, observability_config, model_config.seed, model_config.served_model_name, scheduler_config.use_v2_block_manager, scheduler_config.num_scheduler_steps, scheduler_config.chunked_prefill_enabled, scheduler_config.multi_step_stream_outputs, cache_config.enable_prefix_caching, model_config.use_async_output_proc, use_cached_outputs, model_config.mm_processor_kwargs, ) # TODO(woosuk): Print more configs in debug mode. self.model_config = model_config self.cache_config = cache_config self.lora_config = lora_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.speculative_config = speculative_config self.load_config = load_config self.decoding_config = decoding_config or DecodingConfig() self.prompt_adapter_config = prompt_adapter_config self.observability_config = observability_config or ObservabilityConfig() self.log_stats = log_stats self.use_cached_outputs = use_cached_outputs if not self.model_config.skip_tokenizer_init: self.tokenizer = self._init_tokenizer(tokenizer) self.detokenizer = Detokenizer(self.tokenizer) tokenizer_group = self.get_tokenizer_group() else: self.tokenizer = None self.detokenizer = None tokenizer_group = None # Ensure that the function doesn't contain a reference to self, # to avoid engine GC issues def get_tokenizer_for_seq(sequence: Sequence) -> AnyTokenizer: assert tokenizer_group, \"tokenizer_group cannot be None, \" \"make sure skip_tokenizer_init is False\" return tokenizer_group.get_lora_tokenizer(sequence.lora_request) self.seq_counter = Counter() self.generation_config_fields = _load_generation_config_dict(model_config) self.input_preprocessor = InputPreprocessor(model_config, self.tokenizer) self.input_registry = input_registry self.input_processor = input_registry.create_input_processor(model_config) self.model_executor = executor_class( model=model, # add for spmd_gpu_executor model_config=model_config, cache_config=cache_config, parallel_config=parallel_config, scheduler_config=scheduler_config, device_config=device_config, lora_config=lora_config, speculative_config=speculative_config, load_config=load_config, prompt_adapter_config=prompt_adapter_config, observability_config=self.observability_config, ) if not self.model_config.embedding_mode: self._initialize_kv_caches() # If usage stat is enabled, collect relevant info. if is_usage_stats_enabled(): from vllm.model_executor.model_loader import get_architecture_class_name usage_message.report_usage( get_architecture_class_name(model_config), usage_context, extra_kvs={ # Common configuration \"dtyp","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__.get_tokenizer_for_seq","name":"get_tokenizer_for_seq","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L200-L202","signature":"def get_tokenizer_for_seq(sequence: Sequence) -> AnyTokenizer","exported":true,"origin":"ast","body_hash":"a4fc5f7c45dc3744ff69612ccd9adeba61d7d1368636edc91c1a87741156c8c1","body_text":"def get_tokenizer_for_seq(sequence: Sequence) -> AnyTokenizer: assert tokenizer_group, \"tokenizer_group cannot be None, \" \"make sure skip_tokenizer_init is False\" return tokenizer_group.get_lora_tokenizer(sequence.lora_request)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine._init_tokenizer","name":"_init_tokenizer","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L337-L342","signature":"def _init_tokenizer(self, tokenizer, **tokenizer_init_kwargs)","exported":false,"origin":"ast","body_hash":"14caea76c616136989b92122bb1c73ca013e47bed311a2b9dbc88616d313fe4a","body_text":"def _init_tokenizer(self, tokenizer, **tokenizer_init_kwargs): init_kwargs = dict(enable_lora=bool(self.lora_config), max_num_seqs=self.scheduler_config.max_num_seqs, max_input_length=None) init_kwargs.update(tokenizer_init_kwargs) return TokenizerGroup(tokenizer, **init_kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L344-L347","signature":"def init_cache_engine(self): # TODO: check whether we should rebuild the CUDAGraph every iter when offload/load KVCache # Re-capture CUDAGraph would be time-consuming","exported":true,"origin":"ast","body_hash":"b8670db02a6509c498d21709df59989b89cdc29c54be45a2684c9a9514906677","body_text":"def init_cache_engine(self): # TODO: check whether we should rebuild the CUDAGraph every iter when offload/load KVCache # Re-capture CUDAGraph would be time-consuming self.model_executor.init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L349-L350","signature":"def free_cache_engine(self)","exported":true,"origin":"ast","body_hash":"61329183cfad12da6f1cfae074a6cf5526b7ffa477e5571a36f81ab4ce9e6e46","body_text":"def free_cache_engine(self): self.model_executor.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine._get_executor_cls","name":"_get_executor_cls","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L355-L369","signature":"def _get_executor_cls(cls, engine_config: EngineConfig) -> Type[ExecutorBase]","exported":false,"origin":"ast","body_hash":"70d7dde8a56224cb196bd8cd57edcf4240f871b3f5f846174b9ebd12cd794e45","body_text":"def _get_executor_cls(cls, engine_config: EngineConfig) -> Type[ExecutorBase]: distributed_executor_backend = engine_config.parallel_config.distributed_executor_backend # Initialize the cluster and specify the executor class.] assert (engine_config.device_config.device_type == \"cuda\" ), \"Currently, the vllm in verl only support running on GPU\" # print('Waiting for debugger'); import os,debugpy; debugpy.listen(('localhost', 5678 + int(os.getenv('RANK', '0')))); debugpy.wait_for_client() if engine_config.parallel_config.world_size == 1: engine_config.load_config.load_format = \"dummy_hf\" from .spmd_gpu_executor import SPMDGPUExecutor executor_class = SPMDGPUExecutor return executor_class","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.from_engine_args","name":"from_engine_args","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L372-L402","signature":"def from_engine_args( cls, model, tokenizer, engine_args: EngineArgs, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, stat_loggers: Optional[Dict[str, StatLoggerBase]] = None, ) -> \"LLMEngine\"","exported":true,"origin":"ast","body_hash":"6a688879126b4d109b8039708c4e41e16ac8ec29c485464d43333ec32556a0f5","body_text":"def from_engine_args( cls, model, tokenizer, engine_args: EngineArgs, usage_context: UsageContext = UsageContext.ENGINE_CONTEXT, stat_loggers: Optional[Dict[str, StatLoggerBase]] = None, ) -> \"LLMEngine\": \"\"\"Creates an LLM engine from the engine arguments.\"\"\" # Create the engine configs. engine_config = engine_args.create_engine_config() executor_class = cls._get_executor_cls(engine_config) # Initialize the cluster and specify the executor class. assert (engine_config.device_config.device_type == \"cuda\" ), \"Currently, the vllm in verl only support running on GPU\" from .spmd_gpu_executor import SPMDGPUExecutor executor_class = SPMDGPUExecutor # Create the LLM engine. engine = cls( model, tokenizer, **engine_config.to_dict(), executor_class=executor_class, log_stats=not engine_args.disable_log_stats, usage_context=usage_context, stat_loggers=stat_loggers, ) return engine","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L404-L405","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None","exported":true,"origin":"ast","body_hash":"3e783d98664a0f1d1e1d5620565f52cce66f192cc38c10123eeb27fe5c8d2776","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None: self.model_executor.sync_model_weights(actor_weights=actor_weights, load_format=load_format)","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","span":"L407-L408","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"e77bccf4128d8406c1aae651fad5c639bd60d4b251e0a060f0fdb8196cb26eaf","body_text":"def offload_model_weights(self) -> None: self.model_executor.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"LLMEngine"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"functools","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.core.scheduler","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.arg_utils","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.llm_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.metrics_types","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.output_processor.interfaces","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.output_processor.stop_checker","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.executor.executor_base","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.inputs","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.inputs.preprocess","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.tracing","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.transformers_utils.detokenizer","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.transformers_utils.tokenizer","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.usage.usage_lib","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.version","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":".arg_utils","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":".tokenizer","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","relation":"extends","name":"LLMEngine","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"DecodingConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"ObservabilityConfig","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_init_tokenizer","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Detokenizer","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"get_tokenizer_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__.get_tokenizer_for_seq","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__.get_tokenizer_for_seq","relation":"calls","name":"get_lora_tokenizer","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Counter","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_load_generation_config_dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"InputPreprocessor","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"create_input_processor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"executor_class","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"_initialize_kv_caches","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"is_usage_stats_enabled","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.model_executor.model_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"report_usage","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"get_architecture_class_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"str","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"str","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"bool","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"bool","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"ping","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"SchedulerOutputState","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"SchedulerContext","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"weak_bind","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"partial","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"Scheduler","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":"vllm.engine.metrics","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"LoggingStatLogger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"PrometheusStatLogger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"init_tracer","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"create_output_processor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.__init__","relation":"calls","name":"StopChecker","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine._init_tokenizer","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"bool","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"update","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine._init_tokenizer","relation":"calls","name":"TokenizerGroup","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.init_cache_engine","relation":"calls","name":"init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"executor_class"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"executor_class"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine._get_executor_cls","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":".spmd_gpu_executor","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.from_engine_args","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"create_engine_config","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"EngineArgs"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"_get_executor_cls","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"LLMEngine"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","relation":"imports","specifier":".spmd_gpu_executor","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"cls","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.from_engine_args","relation":"calls","name":"to_dict","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"executor_class"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py#LLMEngine.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/llm_engine_sp.py","recvType":"executor_class"}]},"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py":{"size":13909,"mtimeMs":1789275830318.203,"hash":"d5c3463d34dace74fbea556bd21655f9027572d72ccba407ce035be6074f668a","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","name":"megatron_weight_loaders.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L1-L309","signature":null,"exported":true,"origin":"ast","body_hash":"d5c3463d34dace74fbea556bd21655f9027572d72ccba407ce035be6074f668a","chars":13909,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/model_loader from typing import Dict import torch import torch.nn as nn from vllm.model_executor.layers.linear import * from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead, VocabParallelEmbedding from vllm.model_executor.models import ModelRegistry # NOTE(shengguangming): replace the origin weight loader function in the class __LAYER_WEIGHT_MEGATRON_LOADER_REGISTRY__ = { ColumnParallelLinear: parallel_weight_loader, MergedColumnParallelLinear: parallel_weight_loader, QKVParallelLinear: parallel_weight_loader, RowParallelLinear: parallel_weight_loader, VocabParallelEmbedding: parallel_weight_loader, ParallelLMHead: parallel_weight_loader, # \"ScaledActivation.weight_loader\": ScaledActivation, # TODO(shengguangming): latest commit in vllm fix awq for this function and add load_weights # \"default_weight_loader\": default_weight_loader } # for layer_class, weight_loader in __LAYER_WEIGHT_MEGATRON_LOADER_REGISTRY__.items(): # # setattr(layer_class, 'megatron_weight_loader', weight_loader) # layer_class.weight_loader = weight_loader __MODEL_MEGATRON_WEIGHT_LOADER_REGISTRY__ = { \"GPT2LMHeadModel\": gpt2_weight_loader, \"LlamaForCausalLM\": llama_megatron_weight_loader, # use te backend for open-source megatron \"LLaMAForCausalLM\": llama_megatron_weight_loader, \"MistralForCausalLM\": mistral_megatron_weight_loader, } # the actor model is .state_dict() # Load megatron weights"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#parallel_weight_loader","name":"parallel_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L26-L34","signature":"def parallel_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None","exported":true,"origin":"ast","body_hash":"ce69c9b5888a376326d5a2da38d482358b65ddb842496e072f11a4e5a2873513","body_text":"def parallel_weight_loader(self, param: torch.Tensor, loaded_weight: torch.Tensor) -> None: \"\"\"Parallel Linear weight loader.\"\"\" assert (param.size() == loaded_weight.size( )), \"the parameter size is not align with the loaded weight size, param size: {}, loaded_weight size: {}\".format( param.size(), loaded_weight.size()) assert (param.data.dtype == loaded_weight.data.dtype ), \"if we want to shared weights, the data type should also be the same\" param.data = loaded_weight.data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#default_weight_loader","name":"default_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L37-L43","signature":"def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None","exported":true,"origin":"ast","body_hash":"effe2162346a6b8bc54a24d7202bad8adca23384cda266718d338c3bd9f1ca68","body_text":"def default_weight_loader(param: torch.Tensor, loaded_weight: torch.Tensor) -> None: \"\"\"Default weight loader.\"\"\" assert param.size() == loaded_weight.size() assert (param.data.dtype == loaded_weight.data.dtype ), \"if we want to shared weights, the data type should also be the same\" param.data = loaded_weight.data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#gpt2_weight_loader","name":"gpt2_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L46-L71","signature":"def gpt2_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"d7af0ddffa9a06f7c8b64a06c96f5ce1cbbbd5f2c03114b137268dd5d69048eb","body_text":"def gpt2_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_dict = dict(vllm_model.named_parameters(remove_duplicate=False)) for name, loaded_weight in actor_weights.items(): if \"lm_head.weight\" in name: # GPT-2 ties the weights of the embedding layer and the final # linear layer. continue if \".attn.bias\" in name or \".attn.masked_bias\" in name: # Skip attention mask. # NOTE: \"c_attn.bias\" should not be skipped. continue if not name.startswith(\"transformer.\"): name = \"transformer.\" + name param = params_dict[name] # The HF's GPT-2 implementation uses Conv1D instead of Linear. # Because of this, we need to transpose the weights. # Note(zhuohan): the logic below might break quantized models. for conv1d_weight_name in [\"c_attn\", \"c_proj\", \"c_fc\"]: if conv1d_weight_name not in name: continue if not name.endswith(\".weight\"): continue # TODO: check megatron loaded_weight = loaded_weight.t() weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_weight_loader","name":"llama_megatron_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L74-L83","signature":"def llama_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # NOTE(shengguangming): the megatron llama may have this prefix","exported":true,"origin":"ast","body_hash":"e85f937f0d1b176d431ed32c7fdb4f44e023e98630b81ff16d89c4179292cf97","body_text":"def llama_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","name":"llama_megatron_core_te_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L86-L114","signature":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"2e81b18a30e01a282b9f71e107349b913cfe3801fa38a5ae7fe5e715d76e8d1f","body_text":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"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.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", \"self_attn.o_proj\"), (\"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\"), (\"decoder.final_layernorm\", \"model.norm\"), (\"output_layer\", \"lm_head\"), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith(\".bias\") and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader","name":"llama_megatron_core_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L117-L144","signature":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"2a435caf684ebc42f496c6218db1c087247a32d9412bd57ece7df4491413340f","body_text":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"self_attention.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", \"self_attn.o_proj\"), ( \"input_layernorm\", \"input_layernorm\", ), (\"pre_mlp_layernorm\", \"post_attention_layernorm\"), (\"mlp.linear_fc1\", \"mlp.gate_up_proj\"), (\"mlp.linear_fc2\", \"mlp.down_proj\"), (\"decoder.final_layernorm\", \"model.norm\"), (\"output_layer\", \"lm_head\"), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith(\".bias\") and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name","name":"_replace_name","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L147-L167","signature":"def _replace_name(megatron_name, name_mapping)","exported":false,"origin":"ast","body_hash":"0156f093825b0689948598884b8c719bed3bf7f43f482b40a143ca5058d9aa48","body_text":"def _replace_name(megatron_name, name_mapping): for m_name, v_name in name_mapping: if m_name not in megatron_name: continue if \"layers\" in megatron_name: # deal with decoder layers megatron_name = megatron_name.replace(\"decoder\", \"model\") megatron_name_list = megatron_name.split(\".\") if \"layer_norm_weight\" in megatron_name_list or \"layer_norm_bias\" in megatron_name_list: param_name_list = megatron_name_list[:3] param_name_list.append(v_name) param_name = \".\".join(param_name_list) else: param_name_list = megatron_name_list[:3] weight_or_bias = megatron_name_list[-1] param_name_list.append(v_name) param_name_list.append(weight_or_bias) param_name = \".\".join(param_name_list) return param_name else: param_name = megatron_name.replace(m_name, v_name) return param_name","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","name":"llama_megatron_core_te_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L170-L198","signature":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"2e81b18a30e01a282b9f71e107349b913cfe3801fa38a5ae7fe5e715d76e8d1f","body_text":"def llama_megatron_core_te_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"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.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", \"self_attn.o_proj\"), (\"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\"), (\"decoder.final_layernorm\", \"model.norm\"), (\"output_layer\", \"lm_head\"), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith(\".bias\") and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","name":"llama_megatron_core_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L201-L228","signature":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module","exported":true,"origin":"ast","body_hash":"2a435caf684ebc42f496c6218db1c087247a32d9412bd57ece7df4491413340f","body_text":"def llama_megatron_core_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: params_mapping = [ # (megatron core gpt model name, vllm model name) (\"embedding.word_embeddings\", \"model.embed_tokens\"), (\"self_attention.linear_qkv\", \"self_attn.qkv_proj\"), (\"self_attention.linear_proj\", \"self_attn.o_proj\"), ( \"input_layernorm\", \"input_layernorm\", ), (\"pre_mlp_layernorm\", \"post_attention_layernorm\"), (\"mlp.linear_fc1\", \"mlp.gate_up_proj\"), (\"mlp.linear_fc2\", \"mlp.down_proj\"), (\"decoder.final_layernorm\", \"model.norm\"), (\"output_layer\", \"lm_head\"), ] # NOTE(shengguangming): the megatron llama may have this prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): name = _replace_name(name, params_mapping) if name.endswith(\".bias\") and name not in params_dict: continue if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name~2","name":"_replace_name","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L231-L251","signature":"def _replace_name(megatron_name, name_mapping)","exported":false,"origin":"ast","body_hash":"0156f093825b0689948598884b8c719bed3bf7f43f482b40a143ca5058d9aa48","body_text":"def _replace_name(megatron_name, name_mapping): for m_name, v_name in name_mapping: if m_name not in megatron_name: continue if \"layers\" in megatron_name: # deal with decoder layers megatron_name = megatron_name.replace(\"decoder\", \"model\") megatron_name_list = megatron_name.split(\".\") if \"layer_norm_weight\" in megatron_name_list or \"layer_norm_bias\" in megatron_name_list: param_name_list = megatron_name_list[:3] param_name_list.append(v_name) param_name = \".\".join(param_name_list) else: param_name_list = megatron_name_list[:3] weight_or_bias = megatron_name_list[-1] param_name_list.append(v_name) param_name_list.append(weight_or_bias) param_name = \".\".join(param_name_list) return param_name else: param_name = megatron_name.replace(m_name, v_name) return param_name","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#mistral_megatron_weight_loader","name":"mistral_megatron_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L254-L263","signature":"def mistral_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # TODO: need to implement a general way to deal with prefix","exported":true,"origin":"ast","body_hash":"da2308149252bdda6caeac91a48a8428bb8f71adb703df0b69a2008805ed02cd","body_text":"def mistral_megatron_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module: # TODO: need to implement a general way to deal with prefix params_dict = dict(vllm_model.named_parameters()) for name, loaded_weight in actor_weights.items(): if \"rotary_emb.inv_freq\" in name: continue else: param = params_dict[name] weight_loader = getattr(param, \"weight_loader\", default_weight_loader) weight_loader(param, loaded_weight)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#load_megatron_weights","name":"load_megatron_weights","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L291-L296","signature":"def load_megatron_weights(actor_weights: Dict, vllm_model: nn.Module)","exported":true,"origin":"ast","body_hash":"308ec912ee64a38977bfc7739ea5f87730fbe66a16585ef28a7e106241f7e5b9","body_text":"def load_megatron_weights(actor_weights: Dict, vllm_model: nn.Module): weight_loader = _get_model_weight_loader(vllm_model.__class__.__name__) weight_loader(actor_weights, vllm_model) # NOTE(sgm) to reduce peak memory usage, we offload vllm model to cpu # after init, and we need this after sync model weights for in first iter. vllm_model = vllm_model.cuda()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_get_model_weight_loader","name":"_get_model_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L299-L303","signature":"def _get_model_weight_loader(arch: str)","exported":false,"origin":"ast","body_hash":"f208266263ddc85ba8ce8716c63c6dcb827fbdab60f6cde510bc807de3ed6ed3","body_text":"def _get_model_weight_loader(arch: str): if arch in __MODEL_MEGATRON_WEIGHT_LOADER_REGISTRY__: return __MODEL_MEGATRON_WEIGHT_LOADER_REGISTRY__[arch] raise ValueError(f\"Model architectures {arch} are not supported for now. \" f\"Supported architectures: {ModelRegistry.get_supported_archs()}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#update_megatron_weight_loader","name":"update_megatron_weight_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","span":"L306-L308","signature":"def update_megatron_weight_loader()","exported":true,"origin":"ast","body_hash":"795025216be7d96096150c3271a691f3f03257fc38134d8fe0313242e0025e67","body_text":"def update_megatron_weight_loader(): for layer_class, weight_loader in __LAYER_WEIGHT_MEGATRON_LOADER_REGISTRY__.items(): layer_class.weight_loader = weight_loader","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.linear","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.layers.vocab_parallel_embedding","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"imports","specifier":"vllm.model_executor.models","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#parallel_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#parallel_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#default_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#default_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#default_weight_loader","relation":"calls","name":"size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#gpt2_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"startswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"t","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#gpt2_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"split","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_te_weight_loader~2","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"_replace_name","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"endswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"_replace_name"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#llama_megatron_core_weight_loader~2","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name~2","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"split","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"join","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_replace_name~2","relation":"calls","name":"replace","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#mistral_megatron_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","recvType":"Dict"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#mistral_megatron_weight_loader","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#load_megatron_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#load_megatron_weights","relation":"calls","name":"_get_model_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#load_megatron_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#load_megatron_weights","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_get_model_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#_get_model_weight_loader","relation":"calls","name":"get_supported_archs","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#update_megatron_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py#update_megatron_weight_loader","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/megatron_weight_loaders.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py":{"size":14439,"mtimeMs":1789275830318.203,"hash":"3f033175550ad920a7398e36a7f0b4e16859b692c4098eeb3a5eca7b97219945","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","name":"model_loader.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L1-L339","signature":null,"exported":true,"origin":"ast","body_hash":"3f033175550ad920a7398e36a7f0b4e16859b692c4098eeb3a5eca7b97219945","chars":14439,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/tree/main/vllm/model_executor/models \"\"\"Utilities for selecting and loading models.\"\"\" from typing import Dict, Optional, Union import torch import torch.nn as nn from transformers import PreTrainedModel from vllm.config import CacheConfig, DeviceConfig, LoadConfig, LoRAConfig, ModelConfig, ParallelConfig, SchedulerConfig from vllm.distributed.communication_op import tensor_model_parallel_all_gather from vllm.model_executor.model_loader import BaseModelLoader from vllm.model_executor.model_loader.loader import _initialize_model from vllm.model_executor.model_loader.utils import set_default_torch_dtype from .config import LoadConfig, LoadFormat, ModelConfig from .dtensor_weight_loaders import load_dtensor_weights, update_dtensor_weight_loader from .hf_weight_loader import update_hf_weight_loader from .megatron_weight_loaders import load_megatron_weights, update_megatron_weight_loader # FIXME(sgm): hack the _get_logits function in vllm v0.4.2 # as they use ray, the _get_logits result will only need to return to the driver node, # therefore gather is enough. However, we use SPMD instead of a central scheduler, # all_gather is required (aligned with v0.2.6) from vllm.model_executor.layers.logits_processor import LogitsProcessor LogitsProcessor.__init__ = logitsprocessor_init # use all_gather"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model","name":"get_model","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L33-L62","signature":"def get_model( actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, load_config: LoadConfig, device_config: DeviceConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, lora_config: Optional[LoRAConfig], cache_config: CacheConfig = None, ) -> nn.Module","exported":true,"origin":"ast","body_hash":"636981f000029ee23bb9a673884b46f1605a362843b19acc93c54037a0f718b8","body_text":"def get_model( actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, load_config: LoadConfig, device_config: DeviceConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, lora_config: Optional[LoRAConfig], cache_config: CacheConfig = None, ) -> nn.Module: loader = get_model_loader(load_config) if load_config.load_format.startswith(\"dummy\"): return loader.load_model( model_config=model_config, device_config=device_config, lora_config=lora_config, parallel_config=parallel_config, scheduler_config=scheduler_config, cache_config=cache_config, ) else: return loader.load_model( actor_model=actor_model, model_config=model_config, device_config=device_config, lora_config=lora_config, parallel_config=parallel_config, scheduler_config=scheduler_config, cache_config=cache_config, )","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","name":"get_model_loader","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L65-L101","signature":"def get_model_loader(load_config: LoadConfig) -> BaseModelLoader","exported":true,"origin":"ast","body_hash":"d7c3c9db090b0afd14ff491405a1c8859eddba662b4a5b36c4474d2329ab8f47","body_text":"def get_model_loader(load_config: LoadConfig) -> BaseModelLoader: \"\"\"Get a model loader based on the load format.\"\"\" if isinstance(load_config.load_format, type): return load_config.load_format(load_config) if load_config.load_format == LoadFormat.AUTO: update_megatron_weight_loader() return MegatronLoader(load_config) # NOTE(sgm): change the weight_loader function in runtime if load_config.load_format == LoadFormat.MEGATRON: update_megatron_weight_loader() return MegatronLoader(load_config) if load_config.load_format == LoadFormat.HF: update_hf_weight_loader() return HFLoader(load_config) if load_config.load_format == LoadFormat.DTENSOR: update_dtensor_weight_loader() return DTensorLoader(load_config) if load_config.load_format == LoadFormat.DUMMY_HF: update_hf_weight_loader() return DummyModelLoader(load_config) if load_config.load_format == LoadFormat.DUMMY_MEGATRON: update_megatron_weight_loader() return DummyModelLoader(load_config) if load_config.load_format == LoadFormat.DUMMY_DTENSOR: update_dtensor_weight_loader() return DummyModelLoader(load_config) raise ValueError(\"load format not supported in verl: {}, only support {} and {}\".format( load_config.load_format, LoadFormat.MEGATRON, LoadFormat.HF))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader","name":"DummyModelLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L104-L132","signature":"class DummyModelLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"36af2f37e34b7f8b1d60d34890a68cb68fa64b9286c5c4f650623fe1cfbcd25d","body_text":"class DummyModelLoader(BaseModelLoader): \"\"\"Model loader that will set model weights to random values.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def download_model(self, model_config: ModelConfig) -> None: pass def load_model( self, *, model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, cache_config, scheduler_config) # NOTE(woosuk): For accurate performance evaluation, we assign # random values to the weights. # initialize_dummy_weights(model) return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L107-L111","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"DummyModelLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.download_model","name":"download_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L113-L114","signature":"def download_model(self, model_config: ModelConfig) -> None","exported":true,"origin":"ast","body_hash":"c8227b44f0dc0754d3a16c4793a9bdb91a827b33bad198ceb01715938aff53bb","body_text":"def download_model(self, model_config: ModelConfig) -> None: pass","summary_state":"pending","summary":null,"crux":null,"owner":"DummyModelLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L116-L132","signature":"def load_model( self, *, model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module","exported":true,"origin":"ast","body_hash":"9a20bd93efb5801b59f78a7e6fc726a00e8bf4b735aabca7991511d2368509c5","body_text":"def load_model( self, *, model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, cache_config, scheduler_config) # NOTE(woosuk): For accurate performance evaluation, we assign # random values to the weights. # initialize_dummy_weights(model) return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"DummyModelLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader","name":"MegatronLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L135-L187","signature":"class MegatronLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"fdf448b0c9592602eb767efd03d5851cc05964399aaf4b579f2fda4540fba6fb","body_text":"class MegatronLoader(BaseModelLoader): \"\"\"Model loader that can load the model weights from partitioned megatron model.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def download_model(self, model_config: ModelConfig) -> None: pass # Nothing to download def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model def load_model( self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, cache_config, scheduler_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_megatron_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_megatron_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L138-L142","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.download_model","name":"download_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L144-L145","signature":"def download_model(self, model_config: ModelConfig) -> None","exported":true,"origin":"ast","body_hash":"3ce141370ab2a06509ca223bfcbd712fdc81896b53ddbf1264e36c332a6e6ee6","body_text":"def download_model(self, model_config: ModelConfig) -> None: pass # Nothing to download","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader._get_weights_iterator","name":"_get_weights_iterator","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L147-L154","signature":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model","exported":false,"origin":"ast","body_hash":"694724aa8c3287baa407271e8b4f4dab178819219fb63707f139916d9a3e0340","body_text":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L156-L187","signature":"def load_model( self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module","exported":true,"origin":"ast","body_hash":"6163e626c477996ba100284f04c4a1bd51e122865851d860123580366ec8b55c","body_text":"def load_model( self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, cache_config, scheduler_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_megatron_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_megatron_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader","name":"HFLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L190-L235","signature":"class HFLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"ace6f271e24930dbcb16868dbb16decc3c6a5fc654c1982431a2de685038f14c","body_text":"class HFLoader(BaseModelLoader): \"\"\"Model loader that can load the model weights from model's full params.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def download_model(self, model_config: ModelConfig) -> None: pass # Nothing to download def _get_weights_iterator(self, actor_model: Union[PreTrainedModel, Dict]): if isinstance(actor_model, Dict): return actor_model.items() elif isinstance(actor_model, nn.Module): return dict(actor_model.named_parameters()).items() else: raise ValueError(f\"actor model should be Dict or nn.Module, but get {type(actor_model)}\") def load_model( self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module: with set_default_torch_dtype(model_config.dtype): # with torch.device(device_config.device): # NOTE(sgm): init the model in cpu model = _initialize_model(model_config, self.load_config, lora_config, cache_config, scheduler_config) model.load_weights(self._get_weights_iterator(actor_model)) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L193-L197","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"HFLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.download_model","name":"download_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L199-L200","signature":"def download_model(self, model_config: ModelConfig) -> None","exported":true,"origin":"ast","body_hash":"3ce141370ab2a06509ca223bfcbd712fdc81896b53ddbf1264e36c332a6e6ee6","body_text":"def download_model(self, model_config: ModelConfig) -> None: pass # Nothing to download","summary_state":"pending","summary":null,"crux":null,"owner":"HFLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader._get_weights_iterator","name":"_get_weights_iterator","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L202-L208","signature":"def _get_weights_iterator(self, actor_model: Union[PreTrainedModel, Dict])","exported":false,"origin":"ast","body_hash":"74f208cbfb11e08d772bf1f2634dfe0e81203e3582bc49aa5556b0b65f388546","body_text":"def _get_weights_iterator(self, actor_model: Union[PreTrainedModel, Dict]): if isinstance(actor_model, Dict): return actor_model.items() elif isinstance(actor_model, nn.Module): return dict(actor_model.named_parameters()).items() else: raise ValueError(f\"actor model should be Dict or nn.Module, but get {type(actor_model)}\")","summary_state":"pending","summary":null,"crux":null,"owner":"HFLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L210-L235","signature":"def load_model( self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module","exported":true,"origin":"ast","body_hash":"c624952d8de3c40ec0cc8ca8862d3959c1222e02cec42a44ac0655063820a03c","body_text":"def load_model( self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module: with set_default_torch_dtype(model_config.dtype): # with torch.device(device_config.device): # NOTE(sgm): init the model in cpu model = _initialize_model(model_config, self.load_config, lora_config, cache_config, scheduler_config) model.load_weights(self._get_weights_iterator(actor_model)) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"HFLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader","name":"DTensorLoader","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L238-L290","signature":"class DTensorLoader(BaseModelLoader)","exported":true,"origin":"ast","body_hash":"3c71b9f5ca8882ec337e16010ece6222e8fe09bd59552951b378f5dc1bb54af5","body_text":"class DTensorLoader(BaseModelLoader): \"\"\"Model loader that can load the model weights from partitioned megatron model.\"\"\" def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\") def download_model(self, model_config: ModelConfig) -> None: pass # Nothing to download def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model def load_model( self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, cache_config, scheduler_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_dtensor_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_dtensor_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L241-L245","signature":"def __init__(self, load_config: LoadConfig)","exported":false,"origin":"ast","body_hash":"99b71996e90876d493c934aa94f1876756a1c67275908b34337a0fd7ecb2f8eb","body_text":"def __init__(self, load_config: LoadConfig): super().__init__(load_config) if load_config.model_loader_extra_config: raise ValueError(f\"Model loader extra config is not supported for \" f\"load format {load_config.load_format}\")","summary_state":"pending","summary":null,"crux":null,"owner":"DTensorLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.download_model","name":"download_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L247-L248","signature":"def download_model(self, model_config: ModelConfig) -> None","exported":true,"origin":"ast","body_hash":"3ce141370ab2a06509ca223bfcbd712fdc81896b53ddbf1264e36c332a6e6ee6","body_text":"def download_model(self, model_config: ModelConfig) -> None: pass # Nothing to download","summary_state":"pending","summary":null,"crux":null,"owner":"DTensorLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader._get_weights_iterator","name":"_get_weights_iterator","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L250-L257","signature":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model","exported":false,"origin":"ast","body_hash":"694724aa8c3287baa407271e8b4f4dab178819219fb63707f139916d9a3e0340","body_text":"def _get_weights_iterator(actor_model: Union[PreTrainedModel, Dict]): # NOTE(shengguangming) Load the weights from the actor model pass # if isinstance(actor_model, nn.Module): # load_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) # else: # load_weights(actor_weights=actor_model, vllm_model=model) # return actor_model","summary_state":"pending","summary":null,"crux":null,"owner":"DTensorLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L259-L290","signature":"def load_model( self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module","exported":true,"origin":"ast","body_hash":"d290773d47d481c8040c359cc9a03cdccbf14ccb4137f8c68c2d45b214dfc48f","body_text":"def load_model( self, actor_model: Union[PreTrainedModel, Dict], model_config: ModelConfig, device_config: DeviceConfig, lora_config: Optional[LoRAConfig], parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, cache_config: CacheConfig, ) -> nn.Module: with set_default_torch_dtype(model_config.dtype): with torch.device(device_config.device): model = _initialize_model(model_config, self.load_config, lora_config, cache_config, scheduler_config) # TODO(sgm): This is a hack, we need to register the load_weight() func for each model in vllm if isinstance(actor_model, nn.Module): load_dtensor_weights(actor_weights=dict(actor_model.named_parameters(remove_duplicate=False)), vllm_model=model) else: load_dtensor_weights(actor_weights=actor_model, vllm_model=model) for _, module in model.named_modules(): quant_method = getattr(module, \"quant_method\", None) if quant_method is not None: quant_method.process_weights_after_loading(module) # FIXME: Remove this after Mixtral is updated # to use quant_method. if hasattr(module, \"process_weights_after_loading\"): module.process_weights_after_loading() # NOTE(sgm) Some weights are point to gpu, but still need this. model = model.cuda() # NOTE (zhangchi.usc1992) We need this for vllm to profile memory usage return model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"DTensorLoader"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#_get_logits","name":"_get_logits","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L297-L307","signature":"def _get_logits(self, hidden_states: torch.Tensor, embedding: torch.Tensor, embedding_bias: Optional[torch.Tensor]) -> torch.Tensor: # Get the logits for the next tokens.","exported":false,"origin":"ast","body_hash":"d294e846526a09f100321505cbacf5666d45ef972200a51c4e4c25557ec585e0","body_text":"def _get_logits(self, hidden_states: torch.Tensor, embedding: torch.Tensor, embedding_bias: Optional[torch.Tensor]) -> torch.Tensor: # Get the logits for the next tokens. logits = torch.matmul(hidden_states, embedding.t()) if embedding_bias is not None: logits += embedding_bias logits = tensor_model_parallel_all_gather(logits) # Remove paddings in vocab (if any). if logits is not None: logits = logits[:, :self.org_vocab_size] return logits","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#logitsprocessor_init","name":"logitsprocessor_init","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","span":"L313-L335","signature":"def logitsprocessor_init( self, vocab_size: int, org_vocab_size: Optional[int] = None, scale: float = 1.0, logits_as_input: bool = False, soft_cap: Optional[float] = None, ) -> None","exported":true,"origin":"ast","body_hash":"fe9faa876a70c10cef4f4aa6e6feb4e403c1e66e0a9ca05bdf07329a0fa676a2","body_text":"def logitsprocessor_init( self, vocab_size: int, org_vocab_size: Optional[int] = None, scale: float = 1.0, logits_as_input: bool = False, soft_cap: Optional[float] = None, ) -> None: \"\"\" Args: scale: A scaling factor to apply to the logits. \"\"\" super(LogitsProcessor, self).__init__() self.scale = scale self.vocab_size = vocab_size # Whether the input is logits (default is hidden states). self.logits_as_input = logits_as_input # original vocabulary size (without LoRA). self.org_vocab_size = org_vocab_size or vocab_size # Soft cap the logits. Used in Gemma 2. self.soft_cap = soft_cap # Whether to use gather or all-gather to gather the logits. self.use_gather = False","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":"vllm.distributed.communication_op","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.loader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":"vllm.model_executor.model_loader.utils","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":".dtensor_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":".hf_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":".megatron_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model","relation":"calls","name":"get_model_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model","relation":"calls","name":"startswith","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"get_model_loader"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"get_model_loader"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"load_format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"LoadConfig"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"update_megatron_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"MegatronLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"update_megatron_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"MegatronLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"update_hf_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"HFLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"update_dtensor_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"DTensorLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"update_hf_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"DummyModelLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"update_megatron_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"DummyModelLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"update_dtensor_weight_loader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"DummyModelLoader","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#get_model_loader","relation":"calls","name":"format","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.download_model","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DummyModelLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.download_model","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader._get_weights_iterator","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"load_megatron_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"load_megatron_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#MegatronLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.download_model","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader._get_weights_iterator","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"items","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader._get_weights_iterator","relation":"calls","name":"type","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"load_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"_get_weights_iterator","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"HFLoader"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#HFLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader","relation":"extends","name":"BaseModelLoader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.download_model","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader._get_weights_iterator","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"set_default_torch_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"_initialize_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"load_dtensor_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"dict","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"load_dtensor_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"named_modules","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"getattr"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"process_weights_after_loading","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#DTensorLoader.load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","recvType":"_initialize_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#_get_logits","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#_get_logits","relation":"calls","name":"matmul","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#_get_logits","relation":"calls","name":"t","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#_get_logits","relation":"calls","name":"tensor_model_parallel_all_gather","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"imports","specifier":"vllm.model_executor.layers.logits_processor","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#logitsprocessor_init","file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#logitsprocessor_init","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py#logitsprocessor_init","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_loader.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py":{"size":7587,"mtimeMs":1789275830319.203,"hash":"a30b7696071da0d9a08f30a0b42335cce3acd5331abbfdc0bc64fa3ef56a21ac","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","name":"model_runner.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","span":"L1-L183","signature":null,"exported":true,"origin":"ast","body_hash":"a30b7696071da0d9a08f30a0b42335cce3acd5331abbfdc0bc64fa3ef56a21ac","chars":7587,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/worker/model_runner.py import warnings from enum import IntEnum from typing import Dict, Optional, Union import torch import torch.nn as nn import vllm.envs as envs from vllm.compilation.levels import CompilationLevel from vllm.config import ( CacheConfig, DeviceConfig, LoadConfig, LoRAConfig, ModelConfig, ObservabilityConfig, ParallelConfig, PromptAdapterConfig, SchedulerConfig, ) from vllm.inputs import INPUT_REGISTRY, InputRegistry from vllm.logger import init_logger from vllm.lora.worker_manager import LRUCacheWorkerLoRAManager from vllm.model_executor.models.interfaces import supports_lora from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry from vllm.prompt_adapter.worker_manager import LRUCacheWorkerPromptAdapterManager from vllm.utils import DeviceMemoryProfiler, is_hip, supports_dynamo from vllm.worker.model_runner import ModelRunner from .config import LoadConfig, ModelConfig from .model_loader import get_model logger = init_logger(__name__) # How batches are constructed."},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#BatchType","name":"BatchType","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","span":"L51-L57","signature":"class BatchType(IntEnum): # Every batch is prefill.","exported":true,"origin":"ast","body_hash":"2936fda8ae00fc0b54ec6c67c640b8e048c86af85365d8366a5dd10212f1a15d","body_text":"class BatchType(IntEnum): # Every batch is prefill. PREFILL = 0 # Every batch is decode. DECODE = 1 # Batch is a mixture of prefill and decode. MIXED = 2","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner","name":"ModelRunner","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","span":"L60-L182","signature":"class ModelRunner(ModelRunner)","exported":true,"origin":"ast","body_hash":"e99f643068c1c920778ce6a9d34ae67c951de16b16dd15a7bb6c79bfe10d6d2d","body_text":"class ModelRunner(ModelRunner): def __init__( self, model: Union[nn.Module, Dict], # [verl] model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", is_driver_worker: bool = False, prompt_adapter_config: Optional[PromptAdapterConfig] = None, return_hidden_states: bool = False, observability_config: Optional[ObservabilityConfig] = None, input_registry: InputRegistry = INPUT_REGISTRY, mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, ): super().__init__( model_config, parallel_config, scheduler_config, device_config, cache_config, load_config, lora_config, kv_cache_dtype, is_driver_worker=True, # a hack prompt_adapter_config=prompt_adapter_config, return_hidden_states=return_hidden_states, observability_config=observability_config, input_registry=input_registry, mm_registry=mm_registry, ) # NOTE(sgm): add for verl self.model = model # this will be replaced by get_model() def load_model(self) -> None: logger.info(\"Starting to load model %s...\", self.model_config.model) with DeviceMemoryProfiler() as m: self.model = get_model( self.model, model_config=self.model_config, device_config=self.device_config, load_config=self.load_config, lora_config=self.lora_config, parallel_config=self.parallel_config, scheduler_config=self.scheduler_config, cache_config=self.cache_config, ) self.model_memory_usage = m.consumed_memory logger.info(\"Loading model weights took %.4f GB\", self.model_memory_usage / float(2**30)) if self.lora_config: assert supports_lora(self.model), f\"{self.model.__class__.__name__} does not support LoRA yet.\" if supports_multimodal(self.model): logger.warning(\"Regarding multimodal models, vLLM currently \" \"only supports adding LoRA to language model.\") # It's necessary to distinguish between the max_position_embeddings # of VLMs and LLMs. if hasattr(self.model.config, \"max_position_embeddings\"): max_pos_embeddings = self.model.config.max_position_embeddings else: max_pos_embeddings = self.model.config.text_config.max_position_embeddings self.lora_manager = LRUCacheWorkerLoRAManager( self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens, self.vocab_size, self.lora_config, self.device, self.model.embedding_modules, self.model.embedding_padding_modules, max_position_embeddings=max_pos_embeddings, ) self.model = self.lora_manager.create_lora_manager(self.model) if self.prompt_adapter_config: self.prompt_adapter_manager = LRUCacheWorkerPromptAdapterManager( self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens, self.device, self.prompt_adapter_config, ) self.model = self.prompt_adapter_manager.create_prompt_adapter_manager(self.model) if self.kv_cache_dtype == \"fp8\" and is_hip(): # Currently only ROCm accepts kv-cache scaling factors # via quantization_param_path and this will be deprecated # in the future. if self.model_config.quantization_param_path is not None: if callable(getattr(self.model, \"load_kv_cache_scales\", None)): warnings.warn( \"Loading kv cache scaling factor from JSON is \" \"deprecated and will be removed. Please include \" \"kv cache scaling factors in the model checkpoint.\", FutureWarning, stacklevel=2, ) self.model.load_kv_cache_scales(self.model_config.quantization_param_path) logger.info(\"Loaded KV cache scaling factors from %s\", self.model_config.quantization_param_path) else: raise RuntimeError( \"Using FP8 KV cache and scaling factors provided but \" \"model %s does not support loading scaling factors.\", self.model.__class__, ) else: logger.warning(\"Using FP8 KV cache but no scaling factors \" \"provided. Defaulting to scaling factors of 1.0. \" \"This may lead to less accurate results!\") if envs.VLLM_TORCH_COMPILE_LEVEL == CompilationLevel.DYNAMO_AS_IS and supports_dynamo(): from vllm.plugins import get_torch_compile_backend backend = get_torch_compile_backend() or \"eager\" self.model = torch.compile(self.model, fullgraph=envs.VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE, backend=backend)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","span":"L62-L99","signature":"def __init__( self, model: Union[nn.Module, Dict], # [verl] model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", is_driver_worker: bool = False, prompt_adapter_config: Optional[PromptAdapterConfig] = None, return_hidden_states: bool = False, observability_config: Optional[ObservabilityConfig] = None, input_registry: InputRegistry = INPUT_REGISTRY, mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, )","exported":false,"origin":"ast","body_hash":"17d8712c075dbc867827175eb79eef7dacd50acc150a54bbb78441f7d65221ca","body_text":"def __init__( self, model: Union[nn.Module, Dict], # [verl] model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], kv_cache_dtype: Optional[str] = \"auto\", is_driver_worker: bool = False, prompt_adapter_config: Optional[PromptAdapterConfig] = None, return_hidden_states: bool = False, observability_config: Optional[ObservabilityConfig] = None, input_registry: InputRegistry = INPUT_REGISTRY, mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, ): super().__init__( model_config, parallel_config, scheduler_config, device_config, cache_config, load_config, lora_config, kv_cache_dtype, is_driver_worker=True, # a hack prompt_adapter_config=prompt_adapter_config, return_hidden_states=return_hidden_states, observability_config=observability_config, input_registry=input_registry, mm_registry=mm_registry, ) # NOTE(sgm): add for verl self.model = model # this will be replaced by get_model()","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","name":"load_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","span":"L101-L182","signature":"def load_model(self) -> None","exported":true,"origin":"ast","body_hash":"cb64b7e0e31a10d21b90e9a741c1960ee1265ff64c3dce82b6c0467b837f6faa","body_text":"def load_model(self) -> None: logger.info(\"Starting to load model %s...\", self.model_config.model) with DeviceMemoryProfiler() as m: self.model = get_model( self.model, model_config=self.model_config, device_config=self.device_config, load_config=self.load_config, lora_config=self.lora_config, parallel_config=self.parallel_config, scheduler_config=self.scheduler_config, cache_config=self.cache_config, ) self.model_memory_usage = m.consumed_memory logger.info(\"Loading model weights took %.4f GB\", self.model_memory_usage / float(2**30)) if self.lora_config: assert supports_lora(self.model), f\"{self.model.__class__.__name__} does not support LoRA yet.\" if supports_multimodal(self.model): logger.warning(\"Regarding multimodal models, vLLM currently \" \"only supports adding LoRA to language model.\") # It's necessary to distinguish between the max_position_embeddings # of VLMs and LLMs. if hasattr(self.model.config, \"max_position_embeddings\"): max_pos_embeddings = self.model.config.max_position_embeddings else: max_pos_embeddings = self.model.config.text_config.max_position_embeddings self.lora_manager = LRUCacheWorkerLoRAManager( self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens, self.vocab_size, self.lora_config, self.device, self.model.embedding_modules, self.model.embedding_padding_modules, max_position_embeddings=max_pos_embeddings, ) self.model = self.lora_manager.create_lora_manager(self.model) if self.prompt_adapter_config: self.prompt_adapter_manager = LRUCacheWorkerPromptAdapterManager( self.scheduler_config.max_num_seqs, self.scheduler_config.max_num_batched_tokens, self.device, self.prompt_adapter_config, ) self.model = self.prompt_adapter_manager.create_prompt_adapter_manager(self.model) if self.kv_cache_dtype == \"fp8\" and is_hip(): # Currently only ROCm accepts kv-cache scaling factors # via quantization_param_path and this will be deprecated # in the future. if self.model_config.quantization_param_path is not None: if callable(getattr(self.model, \"load_kv_cache_scales\", None)): warnings.warn( \"Loading kv cache scaling factor from JSON is \" \"deprecated and will be removed. Please include \" \"kv cache scaling factors in the model checkpoint.\", FutureWarning, stacklevel=2, ) self.model.load_kv_cache_scales(self.model_config.quantization_param_path) logger.info(\"Loaded KV cache scaling factors from %s\", self.model_config.quantization_param_path) else: raise RuntimeError( \"Using FP8 KV cache and scaling factors provided but \" \"model %s does not support loading scaling factors.\", self.model.__class__, ) else: logger.warning(\"Using FP8 KV cache but no scaling factors \" \"provided. Defaulting to scaling factors of 1.0. \" \"This may lead to less accurate results!\") if envs.VLLM_TORCH_COMPILE_LEVEL == CompilationLevel.DYNAMO_AS_IS and supports_dynamo(): from vllm.plugins import get_torch_compile_backend backend = get_torch_compile_backend() or \"eager\" self.model = torch.compile(self.model, fullgraph=envs.VLLM_TEST_DYNAMO_FULLGRAPH_CAPTURE, backend=backend)","summary_state":"pending","summary":null,"crux":null,"owner":"ModelRunner"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"warnings","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"enum","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.compilation.levels","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.inputs","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.lora.worker_manager","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.model_executor.models.interfaces","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.multimodal","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.prompt_adapter.worker_manager","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.worker.model_runner","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":".model_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#BatchType","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#BatchType","relation":"extends","name":"IntEnum","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner","relation":"extends","name":"ModelRunner","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"DeviceMemoryProfiler","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"get_model","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"float","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"supports_lora","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"supports_multimodal","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"LRUCacheWorkerLoRAManager","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"create_lora_manager","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","recvType":"LRUCacheWorkerLoRAManager"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"LRUCacheWorkerPromptAdapterManager","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"create_prompt_adapter_manager","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","recvType":"LRUCacheWorkerPromptAdapterManager"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"is_hip","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"callable","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"getattr","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"warn","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"load_kv_cache_scales","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","recvType":"get_model"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"warning","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"supports_dynamo","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py","relation":"imports","specifier":"vllm.plugins","file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"get_torch_compile_backend","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py#ModelRunner.load_model","relation":"calls","name":"compile","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/model_runner.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py":{"size":12915,"mtimeMs":1789275830319.203,"hash":"c63e53effe7be48e3291d790bd02a7c95795692594655c39922a49381d93b56b","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","name":"parallel_state.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L1-L313","signature":null,"exported":true,"origin":"ast","body_hash":"c63e53effe7be48e3291d790bd02a7c95795692594655c39922a49381d93b56b","chars":12915,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM team. # Adapted from # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py # Copyright (c) 2022, NVIDIA CORPORATION. All rights reserved. \"\"\"Model and data parallel groups.\"\"\" import os from typing import Optional import torch import torch.distributed import vllm.distributed.parallel_state as ps from vllm.distributed.parallel_state import ( get_pp_group, get_world_group, init_distributed_environment, init_model_parallel_group, ) from vllm.logger import init_logger logger = init_logger(__name__) \"\"\" This version is strongly tied with Megatron to implement HybridEngine and weight sharing between vllm and Megatron. - We assume the Megatron tp+dp+pp world is already established before calling this function. \"\"\" # Device mesh for using DTensor _DEVICE_MESH = None # Tensor model parallel group that the current rank belongs to. _TP = None # Pipeline model parallel group that the current rank belongs to. _PP = None # This method is for initializing the ParallelGroup when using HybridEngine # TODO(sgm): deviate from the v0.5.4, not pp now \"\"\" Device mesh utilities \"\"\" \"\"\" Tensor model parallel utilities \"\"\""},{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","name":"initialize_parallel_state","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L38-L68","signature":"def initialize_parallel_state( distributed_init_method: str = \"env://\", backend: str = \"nccl\", tensor_model_parallel_size: int = 1, num_tp_per_train_tp: int = 1, pipeline_model_parallel_size: int = 1, ): # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573","exported":true,"origin":"ast","body_hash":"931a53e92b7fcfa5210aff6a6b9470c21ecc757baa4e45b7056799cbe1b0c502","body_text":"def initialize_parallel_state( distributed_init_method: str = \"env://\", backend: str = \"nccl\", tensor_model_parallel_size: int = 1, num_tp_per_train_tp: int = 1, pipeline_model_parallel_size: int = 1, ): # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # NOTE(sgm): Modify for verl, Env vars will be set by TORCHRUN. rank = int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) # Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" init_distributed_environment(world_size, rank, distributed_init_method, local_rank, backend) if torch.distributed.get_world_size() > 1: # NOTE: build a sepearate inference group with infer tp & micro dp initialize_model_parallel_for_vllm( tensor_model_parallel_size=tensor_model_parallel_size, num_tensor_model_parallel_groups_per_train_tp=num_tp_per_train_tp, ) else: initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size, backend)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#ensure_model_parallel_initialized","name":"ensure_model_parallel_initialized","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L71-L94","signature":"def ensure_model_parallel_initialized( tensor_model_parallel_size: int, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None","exported":true,"origin":"ast","body_hash":"ccd13efd3a054fd4dd656d5701ea59fba92a9c01e29f36cc2adfc075bb26ec76","body_text":"def ensure_model_parallel_initialized( tensor_model_parallel_size: int, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None: \"\"\"Helper to initialize model parallel groups if they are not initialized, or ensure tensor-parallel and pipeline-parallel sizes are equal to expected values if the model parallel groups are initialized. \"\"\" # get the backend of _DEVICE_WORLD_GROUP backend = backend or torch.distributed.get_backend(get_world_group().device_group) if not model_parallel_is_initialized(): initialize_model_parallel(tensor_model_parallel_size, pipeline_model_parallel_size, backend) return assert get_tensor_model_parallel_world_size() == tensor_model_parallel_size, ( \"tensor parallel group already initialized, but of unexpected size: \" f\"{get_tensor_model_parallel_world_size()=} vs. \" f\"{tensor_model_parallel_size=}\") pp_world_size = get_pp_group().world_size assert pp_world_size == pipeline_model_parallel_size, ( \"pipeline parallel group already initialized, but of unexpected size: \" f\"{pp_world_size=} vs. \" f\"{pipeline_model_parallel_size=}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#model_parallel_is_initialized","name":"model_parallel_is_initialized","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L98-L101","signature":"def model_parallel_is_initialized()","exported":true,"origin":"ast","body_hash":"e1109e6b0d712f0b1948d5233e4311a4be242c0bd12191eda3739925b8926af7","body_text":"def model_parallel_is_initialized(): \"\"\"Check if tensor and pipeline parallel groups are initialized.\"\"\" return ps._TP is not None # and _PIPELINE_MODEL_PARALLEL_GROUP is not None)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","name":"initialize_model_parallel_for_vllm","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L104-L196","signature":"def initialize_model_parallel_for_vllm( tensor_model_parallel_size: int, num_tensor_model_parallel_groups_per_train_tp: int = 1, pipeline_model_parallel_size: int = 1, ) -> None","exported":true,"origin":"ast","body_hash":"0331c354407dd00b3872cce8b0e0f03e10d28d65c67270d12f0926c34f1523a9","body_text":"def initialize_model_parallel_for_vllm( tensor_model_parallel_size: int, num_tensor_model_parallel_groups_per_train_tp: int = 1, pipeline_model_parallel_size: int = 1, ) -> None: pass # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() assert isinstance(tensor_model_parallel_size, int) # assert num_tensor_model_parallel_groups_per_train_tp == 1 and not different_tp_group # assert num_tensor_model_parallel_groups_per_train_tp > 1 and different_tp_group # Build the tensor model-parallel groups. assert ps._TP is None, \"tensor model parallel group is already initialized\" global _TP world_size: int = torch.distributed.get_world_size() rank = torch.distributed.get_rank() backend = torch.distributed.get_backend() num_tensor_model_parallel_groups = world_size // tensor_model_parallel_size if num_tensor_model_parallel_groups_per_train_tp == 1: # if tensor_model_parallel_size == train_tensor_parallel_size: # using the same tp group as Megatron/vllm assert _TP is None, \"tensor model parallel group is already initialized\" group_ranks = [] for i in range(num_tensor_model_parallel_groups): ranks = range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size) group_ranks.append(ranks) _TP = init_model_parallel_group( group_ranks=group_ranks, local_rank=get_world_group().local_rank, backend=backend, use_custom_allreduce=False, # TODO: check why True is not work in Ray trainer use_message_queue_broadcaster=True, ) ps._TP = _TP # _MICRO_DATA_PARALLEL_GROUP is move to hybrid engine else: # initialize a micro_dp group and a tp group # assume training tp=4, infer tp=2, then, weight is partitioned as # [1], [2], [3], [4] for training and [1,2], [1,2], [3,4], [3,4] for inference # Build the inference tp groups # train_tp = train_tensor_parallel_size train_tp = num_tensor_model_parallel_groups_per_train_tp * tensor_model_parallel_size # num_tensor_model_parallel_groups_per_train_tp = train_tp // tensor_model_parallel_size assert _TP is None, \"tensor model parallel group is already initialized\" group_ranks = [] for i in range(num_tensor_model_parallel_groups // num_tensor_model_parallel_groups_per_train_tp): start = train_tp * i end = train_tp * (i + 1) for j in range(num_tensor_model_parallel_groups_per_train_tp): ranks = list(range(start, end, num_tensor_model_parallel_groups_per_train_tp)) for i in range(len(ranks)): ranks[i] += j group_ranks.append(ranks) _TP = init_model_parallel_group( group_ranks=group_ranks, local_rank=get_world_group().local_rank, backend=backend, use_custom_allreduce=False, # TODO: check why True is not work in Ray trainer use_message_queue_broadcaster=True, ) ps._TP = _TP # Build the pipeline model-parallel groups. # global _PIPELINE_MODEL_PARALLEL_GROUP # global _PIPELINE_GLOBAL_RANKS # assert ps._PIPELINE_MODEL_PARALLEL_GROUP is None, (\"pipeline model parallel group is already initialized\") # ps._PIPELINE_MODEL_PARALLEL_GROUP = mpu.get_pipeline_model_parallel_group() # ps._PIPELINE_GLOBAL_RANKS = mpu.get_pipeline_model_parallel_ranks() # TODO: init using device mesh (not support hybrid engine now) # Build the pipeline model-parallel groups. num_pipeline_model_parallel_groups: int = world_size // pipeline_model_parallel_size global _PP assert _PP is None, \"pipeline model parallel group is already initialized\" group_ranks = [] for i in range(num_pipeline_model_parallel_groups): ranks = list(range(i, world_size, num_pipeline_model_parallel_groups)) group_ranks.append(ranks) # pipeline parallel does not need custom allreduce _PP = init_model_parallel_group(group_ranks, get_world_group().local_rank, backend, use_custom_allreduce=False) ps._PP = _PP # for verl","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","name":"initialize_model_parallel","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L199-L273","signature":"def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None","exported":true,"origin":"ast","body_hash":"a93b8df2ab15f06ca7af22b4b4de30519fa1897a6b6ef91466b617014dcc1251","body_text":"def initialize_model_parallel( tensor_model_parallel_size: int = 1, pipeline_model_parallel_size: int = 1, backend: Optional[str] = None, ) -> None: \"\"\" NOTE: This method is a hack from the open-sourced version without asertion of world_size = tp * pp Initialize model parallel groups. Arguments: tensor_model_parallel_size: number of GPUs used for tensor model parallelism. pipeline_model_parallel_size: number of GPUs used for pipeline model parallelism. Let's say we have a total of 8 GPUs denoted by g0 ... g7 and we use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize the model pipeline. The present function will create 4 tensor model-parallel groups and 2 pipeline model-parallel groups: 4 tensor model-parallel groups: [g0, g1], [g2, g3], [g4, g5], [g6, g7] 2 pipeline model-parallel groups: [g0, g2, g4, g6], [g1, g3, g5, g7] Note that for efficiency, the caller should make sure adjacent ranks are on the same DGX box. For example if we are using 2 DGX-1 boxes with a total of 16 GPUs, rank 0 to 7 belong to the first box and ranks 8 to 15 belong to the second box. \"\"\" # Get world size and rank. Ensure some consistencies. assert torch.distributed.is_initialized() world_size: int = torch.distributed.get_world_size() backend = backend or torch.distributed.get_backend(ps.get_world_group().device_group) # NOTE(sgm) we don't assert world_size == tp * pp # DP is not managed by vllm but by the VeRL WorkerGroup # 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})\") num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size rank = torch.distributed.get_rank() global _TP assert _TP is None, \"tensor model parallel group is already initialized\" group_ranks = [] for i in range(num_tensor_model_parallel_groups): ranks = list(range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size)) group_ranks.append(ranks) # message queue broadcaster is only used in tensor model parallel group _TP = init_model_parallel_group( group_ranks, get_world_group().local_rank, backend, use_custom_allreduce=False, # TODO: check why True is not work in Ray trainer use_message_queue_broadcaster=True, ) ps._TP = _TP # TODO: init using device mesh (not support hybrid engine now) # Build the pipeline model-parallel groups. num_pipeline_model_parallel_groups: int = world_size // pipeline_model_parallel_size global _PP assert _PP is None, \"pipeline model parallel group is already initialized\" group_ranks = [] for i in range(num_pipeline_model_parallel_groups): ranks = list(range(i, world_size, num_pipeline_model_parallel_groups)) group_ranks.append(ranks) # pipeline parallel does not need custom allreduce _PP = init_model_parallel_group(group_ranks, get_world_group().local_rank, backend, use_custom_allreduce=False) ps._PP = _PP # for verl","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_device_mesh","name":"get_device_mesh","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L281-L283","signature":"def get_device_mesh()","exported":true,"origin":"ast","body_hash":"7c4796ddbe8c71e9689239de03a0eb9923276f6491bedee4fada607cfa54978f","body_text":"def get_device_mesh(): assert _DEVICE_MESH is not None, \"device mesh is not initialized\" return _DEVICE_MESH","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_group","name":"get_tensor_model_parallel_group","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L291-L294","signature":"def get_tensor_model_parallel_group()","exported":true,"origin":"ast","body_hash":"a44b5afa314810b63fe76acfdee850ba47f0f89d74049f895156b89c5f8905d4","body_text":"def get_tensor_model_parallel_group(): \"\"\"Get the tensor model parallel group the caller rank belongs to.\"\"\" assert _TP is not None, \"tensor model parallel group is not initialized\" return _TP.device_group","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_world_size","name":"get_tensor_model_parallel_world_size","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L297-L299","signature":"def get_tensor_model_parallel_world_size()","exported":true,"origin":"ast","body_hash":"c1af0d2aaf0c6f73d3dc5df62d1ec18f42133f1075e5262abc2daf509ff9dc98","body_text":"def get_tensor_model_parallel_world_size(): \"\"\"Return world size for the tensor model parallel group.\"\"\" return torch.distributed.get_world_size(group=get_tensor_model_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_rank","name":"get_tensor_model_parallel_rank","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L302-L304","signature":"def get_tensor_model_parallel_rank()","exported":true,"origin":"ast","body_hash":"7d7e3527eb0fbc1e8f6ca37c5449d0a44e76ded85a0b824cf4f92bda58558364","body_text":"def get_tensor_model_parallel_rank(): \"\"\"Return my rank for the tensor model parallel group.\"\"\" return torch.distributed.get_rank(group=get_tensor_model_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_src_rank","name":"get_tensor_model_parallel_src_rank","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","span":"L307-L312","signature":"def get_tensor_model_parallel_src_rank()","exported":true,"origin":"ast","body_hash":"cfa50af80e4b8151b634ff31e0fe23018462a6f55018e05a8cc9c4918ef4e48e","body_text":"def get_tensor_model_parallel_src_rank(): \"\"\"Calculate the global rank corresponding to the first local rank in the tensor model parallel group.\"\"\" global_rank = torch.distributed.get_rank() local_world_size = get_tensor_model_parallel_world_size() return (global_rank // local_world_size) * local_world_size","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"imports","specifier":"torch.distributed","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"imports","specifier":"vllm.distributed.parallel_state","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","relation":"calls","name":"init_distributed_environment","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","relation":"calls","name":"initialize_model_parallel_for_vllm","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_parallel_state","relation":"calls","name":"initialize_model_parallel","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#ensure_model_parallel_initialized","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_backend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"model_parallel_is_initialized","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"initialize_model_parallel","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#ensure_model_parallel_initialized","relation":"calls","name":"get_pp_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#model_parallel_is_initialized","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_backend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"init_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"len","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"init_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"init_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel_for_vllm","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_backend","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_world_group","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"init_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"list","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"range","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"append","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"init_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#initialize_model_parallel","relation":"calls","name":"get_world_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_device_mesh","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_group","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_world_size","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_world_size","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_world_size","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_rank","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_rank","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_src_rank","file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_src_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py#get_tensor_model_parallel_src_rank","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/parallel_state.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py":{"size":10115,"mtimeMs":1789275830319.203,"hash":"0dd0c264240aeaac9f3bdd86054f1493334194143833de54a8636cfc200d27ae","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","name":"spmd_gpu_executor.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L1-L257","signature":null,"exported":true,"origin":"ast","body_hash":"0dd0c264240aeaac9f3bdd86054f1493334194143833de54a8636cfc200d27ae","chars":10115,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/executor/gpu_executor.py import os import socket from typing import Dict, List, Optional, Set, Tuple import torch from vllm.config import ( CacheConfig, DeviceConfig, LoRAConfig, ObservabilityConfig, ParallelConfig, PromptAdapterConfig, SchedulerConfig, SpeculativeConfig, ) from vllm.executor.executor_base import ExecutorAsyncBase, ExecutorBase from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.model_executor.layers.sampler import SamplerOutput from vllm.sequence import ExecuteModelRequest from .config import LoadConfig, ModelConfig logger = init_logger(__name__) # TODO(sgm): not implemented async executor yet"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","name":"SPMDGPUExecutor","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L42-L213","signature":"class SPMDGPUExecutor(ExecutorBase)","exported":true,"origin":"ast","body_hash":"64280b5732bd8065655e13bbf323d9a5d90fb9db7c0c79883115f8ab111c4ca5","body_text":"class SPMDGPUExecutor(ExecutorBase): \"\"\"SPMD-based multi-GPU executor implementations.\"\"\" def __init__( self, model, # pytorch model itself or its parameter dict model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], speculative_config: Optional[SpeculativeConfig], prompt_adapter_config: Optional[PromptAdapterConfig], observability_config: Optional[ObservabilityConfig], ) -> None: self.model_config = model_config self.cache_config = cache_config self.lora_config = lora_config self.load_config = load_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.speculative_config = speculative_config self.prompt_adapter_config = prompt_adapter_config self.observability_config = observability_config distributed_init_method = initialize_cluster(parallel_config) self._init_executor(model, distributed_init_method) # TODO(sgm): verl not support speculative decode now def _init_executor(self, model, distributed_init_method) -> None: assert not self.speculative_config, \"Speculative decoding not yet supported for multi-GPU backend.\" # Create the parallel worker for each GPU. self._init_workers_sp(model, distributed_init_method) def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker from .worker import Worker # pylint: disable=import-outside-toplevel rank = int(os.getenv(\"RANK\")) local_rank = int(os.getenv(\"LOCAL_RANK\")) print(f\"local rank {local_rank}\") # see https://github.com/NVIDIA/nccl/issues/1234 os.environ[\"NCCL_CUMEM_ENABLE\"] = \"0\" self.worker = Worker( model, self.model_config, self.parallel_config, self.scheduler_config, self.device_config, self.cache_config, self.load_config, local_rank, rank, distributed_init_method, lora_config=self.lora_config, speculative_config=None, prompt_adapter_config=self.speculative_config, is_driver_worker=True, model_runner_cls=None, # use the default one ) # NOTE(shengguangming): torch.distributed.init_process_group will be called inside the init_model() self.worker.init_device() self.worker.load_model() def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Determine the number of available KV blocks. This invokes `determine_num_available_blocks` on each worker and takes the min of the results, guaranteeing that the selected cache sizes are compatible with all workers. Returns: - tuple[num_gpu_blocks, num_cpu_blocks] \"\"\" # Get the maximum number of blocks that can be allocated on GPU and CPU. num_blocks = self.worker.determine_num_available_blocks() # NOTE(shengguangming): Now we don't use a shared centralized controler but each process will # have its own scheduler num_gpu_blocks = num_blocks[0] num_cpu_blocks = num_blocks[1] return num_gpu_blocks, num_cpu_blocks def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: \"\"\"Initialize the KV cache in all workers.\"\"\" # NOTE: We log here to avoid multiple logs when number of workers is # greater than one. We could log in the engine, but not all executors # have GPUs. logger.info(\"# GPU blocks: %d, # CPU blocks: %d\", num_gpu_blocks, num_cpu_blocks) self.cache_config.num_gpu_blocks = num_gpu_blocks self.cache_config.num_cpu_blocks = num_cpu_blocks if torch.distributed.get_rank() == 0: print( f\"before init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB\" ) self.worker.initialize_cache(num_gpu_blocks=num_gpu_blocks, num_cpu_blocks=num_cpu_blocks) if torch.distributed.get_rank() == 0: print( f\"after init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB\" ) # NOTE(sgm): This will not profile & capture the model(CUDAGraph) when rebuilding KVCache def init_cache_engine(self) -> None: self.worker._init_cache_engine() def free_cache_engine(self) -> None: self.worker.free_cache_engine() def execute_model(self, execute_model_req) -> List[SamplerOutput]: all_outputs = self.worker.execute_model(execute_model_req=execute_model_req) # NOTE(sgm): # Each GPU in vllm under verl has its own spmd_gpu_executor, therefore all GPUs should return the outputs # In vllm with ray, only the driver worker returns the sampling results. return all_outputs def add_lora(self, lora_request: LoRARequest) -> bool: assert lora_request.lora_int_id > 0, \"lora_id must be greater than 0.\" return self.worker.add_lora(lora_request=lora_request) def remove_lora(self, lora_id: int) -> bool: assert lora_id > 0, \"lora_id must be greater than 0.\" return self.worker.remove_lora(lora_id=lora_id) def list_loras(self) -> Set[int]: return self.worker.list_loras() def check_health(self) -> None: # SPMDExecutor will always be healthy as long as # it's running. return # NOTE(sgm) add for","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L45-L71","signature":"def __init__( self, model, # pytorch model itself or its parameter dict model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], speculative_config: Optional[SpeculativeConfig], prompt_adapter_config: Optional[PromptAdapterConfig], observability_config: Optional[ObservabilityConfig], ) -> None","exported":false,"origin":"ast","body_hash":"020ec665ef93884c51ce878913aa475eaa75fb5dc19028e786b24727eb4f7436","body_text":"def __init__( self, model, # pytorch model itself or its parameter dict model_config: ModelConfig, cache_config: CacheConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, load_config: LoadConfig, lora_config: Optional[LoRAConfig], speculative_config: Optional[SpeculativeConfig], prompt_adapter_config: Optional[PromptAdapterConfig], observability_config: Optional[ObservabilityConfig], ) -> None: self.model_config = model_config self.cache_config = cache_config self.lora_config = lora_config self.load_config = load_config self.parallel_config = parallel_config self.scheduler_config = scheduler_config self.device_config = device_config self.speculative_config = speculative_config self.prompt_adapter_config = prompt_adapter_config self.observability_config = observability_config distributed_init_method = initialize_cluster(parallel_config) self._init_executor(model, distributed_init_method)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_executor","name":"_init_executor","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L74-L78","signature":"def _init_executor(self, model, distributed_init_method) -> None","exported":false,"origin":"ast","body_hash":"cebff6c646ef05776cad6000b8ef4e4d3cf53ab1a86f8b3c49a21683f9ab532a","body_text":"def _init_executor(self, model, distributed_init_method) -> None: assert not self.speculative_config, \"Speculative decoding not yet supported for multi-GPU backend.\" # Create the parallel worker for each GPU. self._init_workers_sp(model, distributed_init_method)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","name":"_init_workers_sp","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L80-L112","signature":"def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker","exported":false,"origin":"ast","body_hash":"e0c0fe84a8ae8f051a66d0e2b0889b8932bf7facda94a4aafd822f7e3c2b1d55","body_text":"def _init_workers_sp(self, model, distributed_init_method: str): # Lazy import the Worker to avoid importing torch.cuda/xformers # before CUDA_VISIBLE_DEVICES is set in the Worker from .worker import Worker # pylint: disable=import-outside-toplevel rank = int(os.getenv(\"RANK\")) local_rank = int(os.getenv(\"LOCAL_RANK\")) print(f\"local rank {local_rank}\") # see https://github.com/NVIDIA/nccl/issues/1234 os.environ[\"NCCL_CUMEM_ENABLE\"] = \"0\" self.worker = Worker( model, self.model_config, self.parallel_config, self.scheduler_config, self.device_config, self.cache_config, self.load_config, local_rank, rank, distributed_init_method, lora_config=self.lora_config, speculative_config=None, prompt_adapter_config=self.speculative_config, is_driver_worker=True, model_runner_cls=None, # use the default one ) # NOTE(shengguangming): torch.distributed.init_process_group will be called inside the init_model() self.worker.init_device() self.worker.load_model()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.determine_num_available_blocks","name":"determine_num_available_blocks","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L114-L132","signature":"def determine_num_available_blocks(self) -> Tuple[int, int]","exported":true,"origin":"ast","body_hash":"dbbbc1bd7aa827704263dacd159512c096e5693fe6fa928309c9ed5451aa5514","body_text":"def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Determine the number of available KV blocks. This invokes `determine_num_available_blocks` on each worker and takes the min of the results, guaranteeing that the selected cache sizes are compatible with all workers. Returns: - tuple[num_gpu_blocks, num_cpu_blocks] \"\"\" # Get the maximum number of blocks that can be allocated on GPU and CPU. num_blocks = self.worker.determine_num_available_blocks() # NOTE(shengguangming): Now we don't use a shared centralized controler but each process will # have its own scheduler num_gpu_blocks = num_blocks[0] num_cpu_blocks = num_blocks[1] return num_gpu_blocks, num_cpu_blocks","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","name":"initialize_cache","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L134-L153","signature":"def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None","exported":true,"origin":"ast","body_hash":"64192e15dd50371c61d44b2e5936e5a4f04383ca648ec023b79fc6823d44640e","body_text":"def initialize_cache(self, num_gpu_blocks: int, num_cpu_blocks: int) -> None: \"\"\"Initialize the KV cache in all workers.\"\"\" # NOTE: We log here to avoid multiple logs when number of workers is # greater than one. We could log in the engine, but not all executors # have GPUs. logger.info(\"# GPU blocks: %d, # CPU blocks: %d\", num_gpu_blocks, num_cpu_blocks) self.cache_config.num_gpu_blocks = num_gpu_blocks self.cache_config.num_cpu_blocks = num_cpu_blocks if torch.distributed.get_rank() == 0: print( f\"before init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB\" ) self.worker.initialize_cache(num_gpu_blocks=num_gpu_blocks, num_cpu_blocks=num_cpu_blocks) if torch.distributed.get_rank() == 0: print( f\"after init cache memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB\" )","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.init_cache_engine","name":"init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L156-L157","signature":"def init_cache_engine(self) -> None","exported":true,"origin":"ast","body_hash":"135998160b42b08111bb719b769019409c044dd0da5bd18e3cf1819bf579782e","body_text":"def init_cache_engine(self) -> None: self.worker._init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L159-L160","signature":"def free_cache_engine(self) -> None","exported":true,"origin":"ast","body_hash":"869ecde616def986458e35749106985c797ce877ae77ec9bab2ddf6490106436","body_text":"def free_cache_engine(self) -> None: self.worker.free_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.execute_model","name":"execute_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L162-L168","signature":"def execute_model(self, execute_model_req) -> List[SamplerOutput]","exported":true,"origin":"ast","body_hash":"aae188cc605ed5322f94d274292f73007610b372cfc40fda0aa9a0800811b7a5","body_text":"def execute_model(self, execute_model_req) -> List[SamplerOutput]: all_outputs = self.worker.execute_model(execute_model_req=execute_model_req) # NOTE(sgm): # Each GPU in vllm under verl has its own spmd_gpu_executor, therefore all GPUs should return the outputs # In vllm with ray, only the driver worker returns the sampling results. return all_outputs","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.add_lora","name":"add_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L170-L172","signature":"def add_lora(self, lora_request: LoRARequest) -> bool","exported":true,"origin":"ast","body_hash":"22a01766fd3ba5823ac7f31e7164043c28514e0346dceaf1705ede572c20a6c0","body_text":"def add_lora(self, lora_request: LoRARequest) -> bool: assert lora_request.lora_int_id > 0, \"lora_id must be greater than 0.\" return self.worker.add_lora(lora_request=lora_request)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.remove_lora","name":"remove_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L174-L176","signature":"def remove_lora(self, lora_id: int) -> bool","exported":true,"origin":"ast","body_hash":"13816aebed74124d7222df89c015c122a2dace11bbe74355d4b3e2b0d545e089","body_text":"def remove_lora(self, lora_id: int) -> bool: assert lora_id > 0, \"lora_id must be greater than 0.\" return self.worker.remove_lora(lora_id=lora_id)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.list_loras","name":"list_loras","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L178-L179","signature":"def list_loras(self) -> Set[int]","exported":true,"origin":"ast","body_hash":"84207e642ba5b1e08cab7f7c4d2b2014f03cbfebd74f378d55838bda1d0c31b7","body_text":"def list_loras(self) -> Set[int]: return self.worker.list_loras()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.check_health","name":"check_health","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L181-L184","signature":"def check_health(self) -> None: # SPMDExecutor will always be healthy as long as # it's running.","exported":true,"origin":"ast","body_hash":"48ca06f8f34b874ea336f29bb675c8616e07b51643c0ee9e4708d53fc9b3fa9a","body_text":"def check_health(self) -> None: # SPMDExecutor will always be healthy as long as # it's running. return","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.add_prompt_adapter","name":"add_prompt_adapter","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L189-L191","signature":"def add_prompt_adapter(self, prompt_adapter_request: PromptAdapterRequest) -> bool","exported":true,"origin":"ast","body_hash":"91234eab8c6c659b404512ee661b883ae99188ad842f44e7cc1bb4801caceb80","body_text":"def add_prompt_adapter(self, prompt_adapter_request: PromptAdapterRequest) -> bool: assert prompt_adapter_request.prompt_adapter_id > 0, \"prompt_adapter_id must be greater than 0.\" return self.worker.add_prompt_adapter(prompt_adapter_request)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.list_prompt_adapters","name":"list_prompt_adapters","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L193-L194","signature":"def list_prompt_adapters(self) -> Set[int]","exported":true,"origin":"ast","body_hash":"2f22e71786a397ac8276f938821afca6e4319a157e7b64045de915ac7d3c7c9f","body_text":"def list_prompt_adapters(self) -> Set[int]: return self.worker.list_prompt_adapters()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.pin_lora","name":"pin_lora","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L196-L198","signature":"def pin_lora(self, lora_id: int) -> bool","exported":true,"origin":"ast","body_hash":"c86bee657fbd9fd8a3dc326489d09c0027c46e8d655205b5ddc8af2122ea076d","body_text":"def pin_lora(self, lora_id: int) -> bool: assert lora_id > 0, \"lora_id must be greater than 0.\" return self.worker.pin_lora(lora_id)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.pin_prompt_adapter","name":"pin_prompt_adapter","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L200-L202","signature":"def pin_prompt_adapter(self, prompt_adapter_id: int) -> bool","exported":true,"origin":"ast","body_hash":"491c5cc6add1b18992d97139a0121c5c22bf84d55e5ed7741601a83a82517019","body_text":"def pin_prompt_adapter(self, prompt_adapter_id: int) -> bool: assert prompt_adapter_id > 0, \"prompt_adapter_id must be greater than 0.\" return self.worker.pin_prompt_adapter(prompt_adapter_id)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.remove_prompt_adapter","name":"remove_prompt_adapter","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L204-L206","signature":"def remove_prompt_adapter(self, prompt_adapter_id: int) -> bool","exported":true,"origin":"ast","body_hash":"1ff9255581e0d4b37c5430512839b034f166b94d21bf51efcdfff328d3ea48f9","body_text":"def remove_prompt_adapter(self, prompt_adapter_id: int) -> bool: assert prompt_adapter_id > 0, \"prompt_adapter_id must be greater than 0.\" return self.worker.remove_prompt_adapter(prompt_adapter_id)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L209-L210","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"c346b8cbee5a5fe63905573385e113fb7b59c46e58b26ecc132e331ba728ea49","body_text":"def offload_model_weights(self) -> None: self.worker.offload_model_weights()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L212-L213","signature":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None","exported":true,"origin":"ast","body_hash":"e7853eafaa9c2dd18668c019f3a34f0b873eb14ebc971798a00af8501b8de047","body_text":"def sync_model_weights(self, actor_weights: Dict[str, torch.Tensor], load_format: str) -> None: self.worker.sync_model_weights(actor_weights=actor_weights, load_format=load_format)","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutor"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#initialize_cluster","name":"initialize_cluster","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L216-L237","signature":"def initialize_cluster( parallel_config: ParallelConfig, engine_use_ray: bool = False, ray_address: Optional[str] = None, ) -> Tuple[str, Optional[None]]","exported":true,"origin":"ast","body_hash":"d8f3c9a0c68b9bd35ef3d5cbcaceb419d1b554f3f24311eda81f4f4d82564cbb","body_text":"def initialize_cluster( parallel_config: ParallelConfig, engine_use_ray: bool = False, ray_address: Optional[str] = None, ) -> Tuple[str, Optional[None]]: \"\"\"Initialize the distributed cluster probably with Ray. Args: parallel_config: The configurations for parallel execution. Returns: The `distributed_init_method` is the address for initializing the distributed backend. \"\"\" # Initialize cluster locally. port = get_open_port() # We need to setup the distributed init method to make sure # the distributed megatron code (e.g., get world size) works correctly. # distributed_init_method = f\"tcp://localhost:{port}\" distributed_init_method = \"env://\" return distributed_init_method","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#get_open_port","name":"get_open_port","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L240-L243","signature":"def get_open_port()","exported":true,"origin":"ast","body_hash":"d88039216a61e6ec2a0e1328331a7b83f0af92fe7306e2625e8d4f34b557c26e","body_text":"def get_open_port(): with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s: s.bind((\"\", 0)) return s.getsockname()[1]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync","name":"SPMDGPUExecutorAsync","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L247-L256","signature":"class SPMDGPUExecutorAsync(SPMDGPUExecutor, ExecutorAsyncBase)","exported":true,"origin":"ast","body_hash":"a4473dfeed05058503ca8f50986836a978e7445454d17a15c9a1a71dc1c1ae0c","body_text":"class SPMDGPUExecutorAsync(SPMDGPUExecutor, ExecutorAsyncBase): async def execute_model_async(self, execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]: \"\"\"Executes one model step on the given sequences.\"\"\" raise NotImplementedError async def check_health_async(self) -> None: \"\"\"Checks if the executor is healthy. If not, it should raise an exception.\"\"\" self.check_health()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync.execute_model_async","name":"execute_model_async","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L249-L251","signature":"async def execute_model_async(self, execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]","exported":true,"origin":"ast","body_hash":"832d82699df49a6c8509c3c7f505f16afd4d4b985ff4fd27e7e1f5bdf2d96c0c","body_text":"async def execute_model_async(self, execute_model_req: ExecuteModelRequest) -> List[SamplerOutput]: \"\"\"Executes one model step on the given sequences.\"\"\" raise NotImplementedError","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutorAsync"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync.check_health_async","name":"check_health_async","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","span":"L253-L256","signature":"async def check_health_async(self) -> None","exported":true,"origin":"ast","body_hash":"fe8c9a207c0e656691830f7feb9005f0015b9529086d8913552ff6134a65430b","body_text":"async def check_health_async(self) -> None: \"\"\"Checks if the executor is healthy. If not, it should raise an exception.\"\"\" self.check_health()","summary_state":"pending","summary":null,"crux":null,"owner":"SPMDGPUExecutorAsync"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"socket","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.executor.executor_base","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.logger","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.lora.request","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.model_executor.layers.sampler","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"calls","name":"init_logger","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"extends","name":"ExecutorBase","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","relation":"calls","name":"initialize_cluster","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.__init__","relation":"calls","name":"_init_executor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"SPMDGPUExecutor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_executor","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_executor","relation":"calls","name":"_init_workers_sp","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"SPMDGPUExecutor"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":".worker","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"Worker","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"init_device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor._init_workers_sp","relation":"calls","name":"load_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.determine_num_available_blocks","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.determine_num_available_blocks","relation":"calls","name":"determine_num_available_blocks","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"init_logger"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_allocated","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_reserved","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"initialize_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"print","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_allocated","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.initialize_cache","relation":"calls","name":"memory_reserved","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.init_cache_engine","relation":"calls","name":"_init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.free_cache_engine","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.execute_model","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.execute_model","relation":"calls","name":"execute_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.add_lora","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.add_lora","relation":"calls","name":"add_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.remove_lora","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.remove_lora","relation":"calls","name":"remove_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.list_loras","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.list_loras","relation":"calls","name":"list_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.check_health","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"imports","specifier":"vllm.prompt_adapter.request","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.add_prompt_adapter","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.add_prompt_adapter","relation":"calls","name":"add_prompt_adapter","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.list_prompt_adapters","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.list_prompt_adapters","relation":"calls","name":"list_prompt_adapters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.pin_lora","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.pin_lora","relation":"calls","name":"pin_lora","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.pin_prompt_adapter","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.pin_prompt_adapter","relation":"calls","name":"pin_prompt_adapter","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.remove_prompt_adapter","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.remove_prompt_adapter","relation":"calls","name":"remove_prompt_adapter","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.offload_model_weights","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutor.sync_model_weights","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#initialize_cluster","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#initialize_cluster","relation":"calls","name":"get_open_port","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#get_open_port","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#get_open_port","relation":"calls","name":"socket","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#get_open_port","relation":"calls","name":"bind","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#get_open_port","relation":"calls","name":"getsockname","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"extends","name":"SPMDGPUExecutor","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"extends","name":"ExecutorAsyncBase","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync.execute_model_async","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync.check_health_async","file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py#SPMDGPUExecutorAsync.check_health_async","relation":"calls","name":"check_health","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/spmd_gpu_executor.py","recvType":"SPMDGPUExecutorAsync"}]},"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py":{"size":1649,"mtimeMs":1789275830319.203,"hash":"f6f2517ebbac01f4cfb234ec71776d54e15a088cd883d48014f6c120d73a8da2","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","name":"tokenizer.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","span":"L1-L41","signature":null,"exported":true,"origin":"ast","body_hash":"f6f2517ebbac01f4cfb234ec71776d54e15a088cd883d48014f6c120d73a8da2","chars":1649,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/transformers_utils/tokenizer_group/tokenizer_group.py from typing import Optional from transformers import PreTrainedTokenizer from vllm.transformers_utils.tokenizer_group import TokenizerGroup from vllm.utils import LRUCache"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup","name":"TokenizerGroup","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","span":"L23-L40","signature":"class TokenizerGroup(TokenizerGroup)","exported":true,"origin":"ast","body_hash":"fa25ee4493fa6d8af45d315152da3b7e820cd8d2d934c90180cfc1ca30f38e3c","body_text":"class TokenizerGroup(TokenizerGroup): \"\"\"A group of tokenizers that can be used for LoRA adapters.\"\"\" def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int]): self.enable_lora = enable_lora self.max_input_length = max_input_length self.tokenizer = tokenizer self.lora_tokenizers = LRUCache[PreTrainedTokenizer](capacity=max_num_seqs) if enable_lora else None # FIXME(sgm): for simplicity, we assign the special token here @property def pad_token_id(self): return self.tokenizer.pad_token_id @property def eos_token_id(self): return self.tokenizer.eos_token_id","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","span":"L26-L31","signature":"def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int])","exported":false,"origin":"ast","body_hash":"ac7258b8cb647fbec870bbffdc1633018f85352db29d6a00fd04abe0a3a94186","body_text":"def __init__(self, tokenizer: PreTrainedTokenizer, enable_lora: bool, max_num_seqs: int, max_input_length: Optional[int]): self.enable_lora = enable_lora self.max_input_length = max_input_length self.tokenizer = tokenizer self.lora_tokenizers = LRUCache[PreTrainedTokenizer](capacity=max_num_seqs) if enable_lora else None","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup.pad_token_id","name":"pad_token_id","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","span":"L35-L36","signature":"def pad_token_id(self)","exported":true,"origin":"ast","body_hash":"09c75a3902b0bb0ca3d7cec97618ef037c84e09c3213f62cd5ee9aa99c2b6669","body_text":"def pad_token_id(self): return self.tokenizer.pad_token_id","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup.eos_token_id","name":"eos_token_id","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","span":"L39-L40","signature":"def eos_token_id(self)","exported":true,"origin":"ast","body_hash":"57455769f69da68b62a323b65045e6baf5ecb3dde6a703266578aeeb1bf096a2","body_text":"def eos_token_id(self): return self.tokenizer.eos_token_id","summary_state":"pending","summary":null,"crux":null,"owner":"TokenizerGroup"}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","relation":"imports","specifier":"transformers","file":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","relation":"imports","specifier":"vllm.transformers_utils.tokenizer_group","file":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup","file":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup","relation":"extends","name":"TokenizerGroup","file":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup.pad_token_id","file":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py#TokenizerGroup.eos_token_id","file":"verl/third_party/vllm/vllm_v_0_6_3/tokenizer.py"}]},"verl/third_party/vllm/vllm_v_0_6_3/worker.py":{"size":15138,"mtimeMs":1789275830319.203,"hash":"d63fc21f39c9a730c9ddd1e776a98f554a7adc7dc9de96ae895d63b476e0577a","nodes":[{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","name":"worker.py","kind":"file","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L1-L334","signature":null,"exported":true,"origin":"ast","body_hash":"d63fc21f39c9a730c9ddd1e776a98f554a7adc7dc9de96ae895d63b476e0577a","chars":15138,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2023 The vLLM 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. # Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/worker/worker.py \"\"\"A GPU worker class.\"\"\" import gc import os from typing import Dict, List, Optional, Tuple, Type, Union import torch import torch.distributed import torch.nn as nn from vllm.config import ( CacheConfig, DeviceConfig, LoRAConfig, ParallelConfig, PromptAdapterConfig, SchedulerConfig, SpeculativeConfig, ) # TODO(sgm): check why vllm has similar file in vllm.model_executor.parallel_utils.parallel_state from vllm.distributed import get_tensor_model_parallel_group, init_distributed_environment, set_custom_all_reduce from vllm.model_executor import set_random_seed from vllm.model_executor.layers.sampler import SamplerOutput from vllm.sequence import ExecuteModelRequest, IntermediateTensors from vllm.worker.cache_engine import CacheEngine from vllm.worker.embedding_model_runner import EmbeddingModelRunner from vllm.worker.model_runner import GPUModelRunnerBase from vllm.worker.model_runner_base import ModelRunnerInputBase from vllm.worker.worker import Worker, _check_if_gpu_supports_dtype from vllm.worker.worker_base import WorkerInput from .config import LoadConfig, LoadFormat, ModelConfig from .dtensor_weight_loaders import load_dtensor_weights from .hf_weight_loader import load_hf_weights from .megatron_weight_loaders import load_megatron_weights from .model_runner import ModelRunner from .parallel_state import ensure_model_parallel_initialized"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","name":"Worker","kind":"class","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L53-L291","signature":"class Worker(Worker)","exported":true,"origin":"ast","body_hash":"8be538ef4e846e957e318a08314645957acd916e39990b1a0c4c30c73a4329e1","body_text":"class Worker(Worker): \"\"\"A worker class that executes (a partition of) the model on a GPU. Each worker is associated with a single GPU. The worker is responsible for maintaining the KV cache and executing the model on the GPU. In case of distributed inference, each worker is assigned a partition of the model. \"\"\" def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, local_rank: int, rank: int, distributed_init_method: str, lora_config: Optional[LoRAConfig] = None, speculative_config: Optional[SpeculativeConfig] = None, prompt_adapter_config: Optional[PromptAdapterConfig] = None, is_driver_worker: bool = False, model_runner_cls: Optional[Type[GPUModelRunnerBase]] = None, ) -> None: # self.model = model # will be replaced in the init_model self.model_config = model_config self.parallel_config = parallel_config self.parallel_config.rank = rank self.scheduler_config = scheduler_config self.device_config = device_config self.cache_config = cache_config self.local_rank = local_rank self.rank = rank self.distributed_init_method = distributed_init_method self.lora_config = lora_config self.load_config = load_config self.prompt_adapter_config = prompt_adapter_config self.is_driver_worker = is_driver_worker # TODO: we don't need driver # if parallel_config and is_driver_worker: # assert rank % parallel_config.tensor_parallel_size == 0, \\ # \"Driver worker should be rank 0 of tensor parallel group.\" if self.model_config.trust_remote_code: # note: lazy import to avoid importing torch before initializing from vllm.utils import init_cached_hf_modules init_cached_hf_modules() # Return hidden states from target model if the draft model is an # mlp_speculator speculative_args = ( {} if speculative_config is None or (speculative_config.draft_model_config.model == model_config.model) or (speculative_config.draft_model_config.hf_config.model_type not in [\"medusa\", \"mlp_speculator\"]) else { \"return_hidden_states\": True }) # TODO(sgm): set correct model runner class ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner if model_runner_cls is not None: ModelRunnerClass = model_runner_cls elif self.model_config.embedding_mode: ModelRunnerClass = EmbeddingModelRunner self.model_runner: GPUModelRunnerBase = ModelRunnerClass( model, # [VERL]: add for verl model_config, parallel_config, scheduler_config, device_config, cache_config, load_config=load_config, lora_config=self.lora_config, kv_cache_dtype=self.cache_config.cache_dtype, is_driver_worker=is_driver_worker, prompt_adapter_config=prompt_adapter_config, **speculative_args, ) # Uninitialized cache engine. Will be initialized by # initialize_cache. self.cache_engine: List[CacheEngine] = None # Initialize gpu_cache as embedding models don't initialize kv_caches self.gpu_cache: Optional[List[List[torch.Tensor]]] = None # NOTE(sgm): [VERL] For offloading inference engine params self.cpu_model = None def init_device(self) -> None: if self.device_config.device.type == \"cuda\": # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # NOTE(sgm): Modify for verl, Env vars will be set by TORCHRUN. self.rank = self.rank if self.rank is not None else int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) self.device = torch.device(f\"cuda:{local_rank}\") if self.rank < 0: raise ValueError(\"Invalid or unspecified rank.\") torch.cuda.set_device(self.device) # Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" self.parallel_config.world_size = world_size _check_if_gpu_supports_dtype(self.model_config.dtype) torch.cuda.empty_cache() self.init_gpu_memory = torch.cuda.mem_get_info()[0] else: raise RuntimeError(f\"Not support device type: {self.device_config.device}\") # Initialize the distributed environment. init_worker_distributed_environment(self.parallel_config, self.rank, self.distributed_init_method, self.local_rank) # Set random seed. set_random_seed(self.model_config.seed) # self.model = get_model(actor_model=self.model, model_config=self.model_config) @torch.inference_mode() def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Profiles the peak memory usage of the model to determine how many KV blocks may be allocated without OOMs. The engine will first conduct a profiling of the existing memory usage. Then, it calculate the maximum possible number of GPU and CPU blocks that can be allocated with the remaining f","summary_state":"pending","summary":null,"crux":null},{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.__init__","name":"__init__","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L61-L138","signature":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, local_rank: int, rank: int, distributed_init_method: str, lora_config: Optional[LoRAConfig] = None, speculative_config: Optional[SpeculativeConfig] = None, prompt_adapter_config: Optional[PromptAdapterConfig] = None, is_driver_worker: bool = False, model_runner_cls: Optional[Type[GPUModelRunnerBase]] = None, ) -> None: # self.model = model # will be replaced in the init_model","exported":false,"origin":"ast","body_hash":"0073890f32a4cddd81d873ec901d6a480aa8253c8d823a1a7966868a4a48716e","body_text":"def __init__( self, model: Union[nn.Module, Dict], # model itself or its parameter dict model_config: ModelConfig, parallel_config: ParallelConfig, scheduler_config: SchedulerConfig, device_config: DeviceConfig, cache_config: CacheConfig, load_config: LoadConfig, local_rank: int, rank: int, distributed_init_method: str, lora_config: Optional[LoRAConfig] = None, speculative_config: Optional[SpeculativeConfig] = None, prompt_adapter_config: Optional[PromptAdapterConfig] = None, is_driver_worker: bool = False, model_runner_cls: Optional[Type[GPUModelRunnerBase]] = None, ) -> None: # self.model = model # will be replaced in the init_model self.model_config = model_config self.parallel_config = parallel_config self.parallel_config.rank = rank self.scheduler_config = scheduler_config self.device_config = device_config self.cache_config = cache_config self.local_rank = local_rank self.rank = rank self.distributed_init_method = distributed_init_method self.lora_config = lora_config self.load_config = load_config self.prompt_adapter_config = prompt_adapter_config self.is_driver_worker = is_driver_worker # TODO: we don't need driver # if parallel_config and is_driver_worker: # assert rank % parallel_config.tensor_parallel_size == 0, \\ # \"Driver worker should be rank 0 of tensor parallel group.\" if self.model_config.trust_remote_code: # note: lazy import to avoid importing torch before initializing from vllm.utils import init_cached_hf_modules init_cached_hf_modules() # Return hidden states from target model if the draft model is an # mlp_speculator speculative_args = ( {} if speculative_config is None or (speculative_config.draft_model_config.model == model_config.model) or (speculative_config.draft_model_config.hf_config.model_type not in [\"medusa\", \"mlp_speculator\"]) else { \"return_hidden_states\": True }) # TODO(sgm): set correct model runner class ModelRunnerClass: Type[GPUModelRunnerBase] = ModelRunner if model_runner_cls is not None: ModelRunnerClass = model_runner_cls elif self.model_config.embedding_mode: ModelRunnerClass = EmbeddingModelRunner self.model_runner: GPUModelRunnerBase = ModelRunnerClass( model, # [VERL]: add for verl model_config, parallel_config, scheduler_config, device_config, cache_config, load_config=load_config, lora_config=self.lora_config, kv_cache_dtype=self.cache_config.cache_dtype, is_driver_worker=is_driver_worker, prompt_adapter_config=prompt_adapter_config, **speculative_args, ) # Uninitialized cache engine. Will be initialized by # initialize_cache. self.cache_engine: List[CacheEngine] = None # Initialize gpu_cache as embedding models don't initialize kv_caches self.gpu_cache: Optional[List[List[torch.Tensor]]] = None # NOTE(sgm): [VERL] For offloading inference engine params self.cpu_model = None","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","name":"init_device","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L140-L174","signature":"def init_device(self) -> None","exported":true,"origin":"ast","body_hash":"762272295f8ebaa0d8a603527f1c7088967ebd4ff4562d0c715a5656dcddc70b","body_text":"def init_device(self) -> None: if self.device_config.device.type == \"cuda\": # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow # as the number of all_reduce calls increases. This env var disables # this behavior. # Related issue: # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\" # NOTE(sgm): Modify for verl, Env vars will be set by TORCHRUN. self.rank = self.rank if self.rank is not None else int(os.getenv(\"RANK\", \"-1\")) local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\")) self.device = torch.device(f\"cuda:{local_rank}\") if self.rank < 0: raise ValueError(\"Invalid or unspecified rank.\") torch.cuda.set_device(self.device) # Use the world_size set by TORCHRUN world_size = int(os.getenv(\"WORLD_SIZE\", \"-1\")) assert world_size != -1, \"The world_size is set to -1, not initialized by TORCHRUN\" self.parallel_config.world_size = world_size _check_if_gpu_supports_dtype(self.model_config.dtype) torch.cuda.empty_cache() self.init_gpu_memory = torch.cuda.mem_get_info()[0] else: raise RuntimeError(f\"Not support device type: {self.device_config.device}\") # Initialize the distributed environment. init_worker_distributed_environment(self.parallel_config, self.rank, self.distributed_init_method, self.local_rank) # Set random seed. set_random_seed(self.model_config.seed) # self.model = get_model(actor_model=self.model, model_config=self.model_config)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","name":"determine_num_available_blocks","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L177-L233","signature":"def determine_num_available_blocks(self) -> Tuple[int, int]","exported":true,"origin":"ast","body_hash":"a6a27c977327ccefe506e7c6b32378a4ee35a0534fc55dfa1e5300255c49241b","body_text":"def determine_num_available_blocks(self) -> Tuple[int, int]: \"\"\"Profiles the peak memory usage of the model to determine how many KV blocks may be allocated without OOMs. The engine will first conduct a profiling of the existing memory usage. Then, it calculate the maximum possible number of GPU and CPU blocks that can be allocated with the remaining free memory. .. tip:: You may limit the usage of GPU memory by adjusting the `gpu_memory_utilization` parameter. \"\"\" # Profile the memory usage of the model and get the maximum number of # cache blocks that can be allocated with the remaining free memory. torch.cuda.empty_cache() # torch.cuda.reset_peak_memory_stats() # Execute a forward pass with dummy inputs to profile the memory usage # of the model. self.model_runner.profile_run() # Calculate the number of blocks that can be allocated with the # profiled peak memory. torch.cuda.synchronize() free_gpu_memory, total_gpu_memory = torch.cuda.mem_get_info() peak_memory = total_gpu_memory - free_gpu_memory assert peak_memory > 0, (\"Error in memory profiling. This happens when the GPU memory was \" \"not properly cleaned up before initializing the vLLM instance.\") cache_block_size = self.get_cache_block_size_bytes() # NOTE(sgm) [VERL] use the remaining memory num_gpu_blocks = int((free_gpu_memory * self.cache_config.gpu_memory_utilization) // cache_block_size) # num_gpu_blocks = int((total_gpu_memory * self.cache_config.gpu_memory_utilization - peak_memory) // cache_block_size) num_cpu_blocks = int(self.cache_config.swap_space_bytes // cache_block_size) num_gpu_blocks = max(num_gpu_blocks, 0) num_cpu_blocks = max(num_cpu_blocks, 0) if self.model_runner.lora_manager: self.model_runner.remove_all_loras() # NOTE(sgm): Add for [VERL], synchronize number of blocks with all the rank num_gpu_blocks = torch.tensor([num_gpu_blocks], device=\"cuda\") num_cpu_blocks = torch.tensor([num_cpu_blocks], device=\"cuda\") torch.distributed.all_reduce(num_gpu_blocks, op=torch.distributed.ReduceOp.MIN, group=get_tensor_model_parallel_group().device_group) torch.distributed.all_reduce(num_cpu_blocks, op=torch.distributed.ReduceOp.MIN, group=get_tensor_model_parallel_group().device_group) num_gpu_blocks = num_gpu_blocks.item() num_cpu_blocks = num_cpu_blocks.item() gc.collect() torch.cuda.empty_cache() return num_gpu_blocks, num_cpu_blocks","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker._init_cache_engine","name":"_init_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L235-L237","signature":"def _init_cache_engine(self)","exported":false,"origin":"ast","body_hash":"ddbf730d543fa64f637055bf2b4047182f388f84f7375434e0765a7cd602b012","body_text":"def _init_cache_engine(self): if self.cache_engine is None and self.gpu_cache is None: super()._init_cache_engine()","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.free_cache_engine","name":"free_cache_engine","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L239-L242","signature":"def free_cache_engine(self): # ensure `enforce_eager=True`","exported":true,"origin":"ast","body_hash":"d384948b80b2c04ee44319bcd88c9b941e3967b68e2050a10073ade04af52223","body_text":"def free_cache_engine(self): # ensure `enforce_eager=True` self.cache_engine = None self.gpu_cache = None","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.execute_model","name":"execute_model","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L245-L271","signature":"def execute_model(self, execute_model_req: ExecuteModelRequest, intermediate_tensors: Optional[IntermediateTensors] = None) -> Optional[List[SamplerOutput]]","exported":true,"origin":"ast","body_hash":"a0e5942a1547d3f1ce9cbc20cc38ecc20af4d38daec3d3efc7ee476bcdafaced","body_text":"def execute_model(self, execute_model_req: ExecuteModelRequest, intermediate_tensors: Optional[IntermediateTensors] = None) -> Optional[List[SamplerOutput]]: \"\"\" Execute model in Single Program Multiple Data (SPMD) fashion. All workers take the same request, prepare the input and execute the model. \"\"\" assert execute_model_req is not None, (\"_execute_model_spmd() requires each worker to take in an \" \"ExecuteModelRequest\") worker_input: WorkerInput = self.prepare_worker_input(execute_model_req=execute_model_req) model_input: ModelRunnerInputBase = self.model_runner.prepare_model_input( execute_model_req.seq_group_metadata_list) # verl.worker.workerbase.WorkerBase # swap cache super().execute_worker(worker_input) # If there is no input, we don't need to execute the model. if worker_input.num_seq_groups == 0: return [] return self.model_runner.execute_model( model_input, self.kv_cache[worker_input.virtual_engine] if self.kv_cache is not None else None, intermediate_tensors, )","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.sync_model_weights","name":"sync_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L274-L281","signature":"def sync_model_weights(self, actor_weights: Dict, load_format: str)","exported":true,"origin":"ast","body_hash":"61f049debd75c5e97bb95de1fbaf9d769e14978f2c8fd9b2ab6143756010d1d1","body_text":"def sync_model_weights(self, actor_weights: Dict, load_format: str): if load_format in [LoadFormat.MEGATRON, LoadFormat.AUTO]: load_megatron_weights(actor_weights, self.model_runner.model) elif load_format == LoadFormat.HF: # full model state dict without no sharding load_hf_weights(actor_weights, self.model_runner.model) elif load_format == LoadFormat.DTENSOR: load_dtensor_weights(actor_weights, self.model_runner.model)","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.offload_model_weights","name":"offload_model_weights","kind":"method","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L283-L291","signature":"def offload_model_weights(self) -> None","exported":true,"origin":"ast","body_hash":"245d1bbb4107732af3041d61b174c1f7f8d3baa4b06c477277f49b53f93ec260","body_text":"def offload_model_weights(self) -> None: if self.cpu_model == None: self.cpu_model = {} for name, params in self.model_runner.model.named_parameters(): self.cpu_model[name] = torch.empty_like(params, device=\"cpu\") params.data = self.cpu_model[name] else: for name, params in self.model_runner.model.named_parameters(): params.data = self.cpu_model[name]","summary_state":"pending","summary":null,"crux":null,"owner":"Worker"},{"id":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#init_worker_distributed_environment","name":"init_worker_distributed_environment","kind":"function","path":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","span":"L294-L333","signature":"def init_worker_distributed_environment( parallel_config: ParallelConfig, rank: int, distributed_init_method: Optional[str] = \"env://\", local_rank: int = -1, ) -> None","exported":true,"origin":"ast","body_hash":"c69b02456fd24b41bc6e00a78cc2e73e76ff46f4a19b5942d83e358f902ccd85","body_text":"def init_worker_distributed_environment( parallel_config: ParallelConfig, rank: int, distributed_init_method: Optional[str] = \"env://\", local_rank: int = -1, ) -> None: \"\"\"Initialize the distributed environment.\"\"\" set_custom_all_reduce(not parallel_config.disable_custom_all_reduce) # NOTE(sgm) use tcp://localhost:xxxx will hang in HF setting without megatron init_distributed_environment(parallel_config.world_size, rank, distributed_init_method, local_rank) ensure_model_parallel_initialized( tensor_model_parallel_size=parallel_config.tensor_parallel_size, pipeline_model_parallel_size=parallel_config.pipeline_parallel_size, ) # TODO(sgm): check whether need this # if pynccl_utils.is_initialized(): # pynccl_world_size = pynccl_utils.get_world_size() # if pynccl_world_size != parallel_config.world_size: # raise RuntimeError( # \"pynccl is already initialized but the pynccl world \" # \"size does not match parallel_config.world_size \" # f\"({pynccl_world_size} vs. {parallel_config.world_size}).\") # elif parallel_config.world_size > 1: # # NOTE(woosuk): We don't initialize pynccl process group when world size # # is 1. # # NOTE(kaichao): By default, pynccl is initialized for tp group. # pynccl_utils.init_process_group( # group=get_tensor_model_parallel_cpu_group()) # # Initialize a custom fast all-reduce implementation. # if not parallel_config.disable_custom_all_reduce: # init_custom_ar() # A small all_reduce for warmup. torch.distributed.all_reduce(torch.zeros(1).cuda()) # if pynccl_utils.is_initialized(): # pynccl_utils.all_reduce(torch.zeros(1).cuda())","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"gc","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"os","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"typing","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"torch","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"torch.distributed","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.config","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.distributed","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.model_executor","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.model_executor.layers.sampler","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.sequence","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.worker.cache_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.worker.embedding_model_runner","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.worker.model_runner","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.worker.model_runner_base","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.worker.worker","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.worker.worker_base","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":".config","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":".dtensor_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":".hf_weight_loader","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":".megatron_weight_loaders","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":".model_runner","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":".parallel_state","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","relation":"extends","name":"Worker","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.__init__","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"imports","specifier":"vllm.utils","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.__init__","relation":"calls","name":"init_cached_hf_modules","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.__init__","relation":"calls","name":"ModelRunnerClass","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"set_device","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"getenv","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"_check_if_gpu_supports_dtype","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"mem_get_info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"init_worker_distributed_environment","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.init_device","relation":"calls","name":"set_random_seed","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","relation":"calls","name":"inference_mode","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"profile_run","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","recvType":"ModelRunnerClass"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"synchronize","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"mem_get_info","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"get_cache_block_size_bytes","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"int","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"max","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"remove_all_loras","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","recvType":"ModelRunnerClass"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"tensor","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"item","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","recvType":"max"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"item","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","recvType":"max"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"collect","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.determine_num_available_blocks","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker._init_cache_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker._init_cache_engine","relation":"calls","name":"_init_cache_engine","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker._init_cache_engine","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.free_cache_engine","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.execute_model","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.execute_model","relation":"calls","name":"prepare_worker_input","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","recvType":"Worker"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.execute_model","relation":"calls","name":"prepare_model_input","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","recvType":"ModelRunnerClass"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.execute_model","relation":"calls","name":"execute_worker","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.execute_model","relation":"calls","name":"super","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.execute_model","relation":"calls","name":"execute_model","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","recvType":"ModelRunnerClass"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.sync_model_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.sync_model_weights","relation":"calls","name":"load_megatron_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.sync_model_weights","relation":"calls","name":"load_hf_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.sync_model_weights","relation":"calls","name":"load_dtensor_weights","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.offload_model_weights","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.offload_model_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.offload_model_weights","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#Worker.offload_model_weights","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py","relation":"contains","targetId":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#init_worker_distributed_environment","file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#init_worker_distributed_environment","relation":"calls","name":"set_custom_all_reduce","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#init_worker_distributed_environment","relation":"calls","name":"init_distributed_environment","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#init_worker_distributed_environment","relation":"calls","name":"ensure_model_parallel_initialized","viaMember":false,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#init_worker_distributed_environment","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#init_worker_distributed_environment","relation":"calls","name":"cuda","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"},{"source":"verl/third_party/vllm/vllm_v_0_6_3/worker.py#init_worker_distributed_environment","relation":"calls","name":"zeros","viaMember":true,"file":"verl/third_party/vllm/vllm_v_0_6_3/worker.py"}]},"verl/trainer/__init__.py":{"size":600,"mtimeMs":1789275830319.203,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/trainer/__init__.py","name":"__init__.py","kind":"file","path":"verl/trainer/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/trainer/conf_adv_config.py":{"size":10663,"mtimeMs":1789275830319.203,"hash":"e383f793d22345f7368c0fe1f968fddf3d93f5db0efd6c98a25f0346274b4a31","nodes":[{"id":"verl/trainer/conf_adv_config.py","name":"conf_adv_config.py","kind":"file","path":"verl/trainer/conf_adv_config.py","span":"L1-L291","signature":null,"exported":true,"origin":"ast","body_hash":"e383f793d22345f7368c0fe1f968fddf3d93f5db0efd6c98a25f0346274b4a31","chars":9769,"summary_state":"pending","summary":null,"crux":null,"body_text":"\"\"\"DualStake confidence-calibration reward configuration. The default values reproduce the balanced configuration in the paper. \"\"\" import numpy as np # Evidence confidence is the final <confidence> emitted after retrieval; # answer confidence is the <final-confidence> emitted after the answer. LSC_MODE = \"stake\" FC_MODE = \"stake\" # Balanced DualStake setting: alpha = beta = 0.25. # Both coefficients are linearly warmed up from step 100 to step 300. LSC_ALPHA_MAX = 0.25 LSC_START_STEP = 100 LSC_FULL_STEP = 300 FC_ALPHA_MAX = 0.25 FC_START_STEP = 100 FC_FULL_STEP = 300 # The paper uses margin clipping [m-, m+] = [0.1, 0.9]. STAKE_MARGIN_ENABLE = True STAKE_CORRECT_MARGIN = 0.9 STAKE_WRONG_MARGIN = 0.1 # ============================================================================= # Curriculum 工具函数 # ============================================================================= # ============================================================================= # Calib Reward 计算函数 # ============================================================================= # ============================================================================= # 主入口：双路 calib 叠加 # ============================================================================= # ============================================================================= # 校准评估工具（供 ray_trainer 和 main_ppo 调用，不涉及奖励，保持不变） # ============================================================================="},{"id":"verl/trainer/conf_adv_config.py#get_lsc_alpha","name":"get_lsc_alpha","kind":"function","path":"verl/trainer/conf_adv_config.py","span":"L32-L39","signature":"def get_lsc_alpha(cur_step: int) -> float","exported":true,"origin":"ast","body_hash":"7f11dd07e4b10a046795452e248a0e2327faed4d27c1d7a2870b72ed0f58a4a3","body_text":"def get_lsc_alpha(cur_step: int) -> float: \"\"\"返回当前 step 对应的 lsc alpha 值（线性 warmup）\"\"\" if cur_step < LSC_START_STEP: return 0.0 if cur_step >= LSC_FULL_STEP: return LSC_ALPHA_MAX ratio = (cur_step - LSC_START_STEP) / (LSC_FULL_STEP - LSC_START_STEP) return LSC_ALPHA_MAX * ratio","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/conf_adv_config.py#get_fc_alpha","name":"get_fc_alpha","kind":"function","path":"verl/trainer/conf_adv_config.py","span":"L42-L49","signature":"def get_fc_alpha(cur_step: int) -> float","exported":true,"origin":"ast","body_hash":"b3be37c48e7003f1fe88da11074a9fa31536489d44cd4342caf93f183f934bc8","body_text":"def get_fc_alpha(cur_step: int) -> float: \"\"\"返回当前 step 对应的 fc alpha 值（线性 warmup）\"\"\" if cur_step < FC_START_STEP: return 0.0 if cur_step >= FC_FULL_STEP: return FC_ALPHA_MAX ratio = (cur_step - FC_START_STEP) / (FC_FULL_STEP - FC_START_STEP) return FC_ALPHA_MAX * ratio","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/conf_adv_config.py#compute_calib_reward_stake","name":"compute_calib_reward_stake","kind":"function","path":"verl/trainer/conf_adv_config.py","span":"L56-L81","signature":"def compute_calib_reward_stake(q: np.ndarray, c: np.ndarray) -> np.ndarray","exported":true,"origin":"ast","body_hash":"0878d66929f1529ad16f140888e51aaac26b369b28e855d082a513738add31b7","body_text":"def compute_calib_reward_stake(q: np.ndarray, c: np.ndarray) -> np.ndarray: \"\"\" Stake-Calib：R_i^calib = (2*q_i - 1) * c_eff lsc 和 fc 的 stake 监督共用本函数。 c = -1 的格式错误样本 calib 项为 0（中性）。 q, c: shape (G,)，已归一化到 [0,1]，c=-1 表示无效 当 STAKE_MARGIN_ENABLE=True 时启用松弛余量： 答对（q>=0.5）：c_eff = min(c, STAKE_CORRECT_MARGIN)，超过上限后奖励封顶 答错（q< 0.5）：c_eff = max(c, STAKE_WRONG_MARGIN)，低于下限后惩罚封底 \"\"\" G = len(q) r = np.zeros(G, dtype=np.float32) for i in range(G): if c[i] < 0: # 格式错误 r[i] = 0.0 else: if STAKE_MARGIN_ENABLE: if q[i] >= 0.5: # 答对：置信度奖励封顶 c_eff = min(c[i], STAKE_CORRECT_MARGIN) else: # 答错：置信度惩罚封底 c_eff = max(c[i], STAKE_WRONG_MARGIN) else: c_eff = c[i] r[i] = (2.0 * q[i] - 1.0) * c_eff return r","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/conf_adv_config.py#compute_fc_consistency_reward","name":"compute_fc_consistency_reward","kind":"function","path":"verl/trainer/conf_adv_config.py","span":"L84-L112","signature":"def compute_fc_consistency_reward( c_fc: np.ndarray, c_lsc: np.ndarray, threshold: float = FC_CONSISTENCY_THRESHOLD, ) -> np.ndarray","exported":true,"origin":"ast","body_hash":"a75b70eb4200a955f029573cf0406fab53988cd906036000780d50343c15920f","body_text":"def compute_fc_consistency_reward( c_fc: np.ndarray, c_lsc: np.ndarray, threshold: float = FC_CONSISTENCY_THRESHOLD, ) -> np.ndarray: \"\"\" FC per-sample 方向一致性惩罚： s_fc = sign(c_fc - threshold) s_lsc = sign(c_lsc - threshold) R_i = -1 如果 s_fc != s_lsc 且两者方向均非 0（方向明确且相反） R_i = 0 其他情况（任意一方无效 / tie / 方向一致） 范围：{-1, 0}，纯惩罚信号。 c_fc, c_lsc: shape (G,)，已归一化到 [0,1]，-1 表示无效 \"\"\" G = len(c_fc) r = np.zeros(G, dtype=np.float32) for i in range(G): if c_fc[i] < 0 or c_lsc[i] < 0: # 任意一方无效 r[i] = 0.0 continue s_fc = 1.0 if c_fc[i] > threshold else (-1.0 if c_fc[i] < threshold else 0.0) s_lsc = 1.0 if c_lsc[i] > threshold else (-1.0 if c_lsc[i] < threshold else 0.0) if s_fc == 0.0 or s_lsc == 0.0: # 任意一方恰好等于阈值，不惩罚 r[i] = 0.0 elif s_fc != s_lsc: # 方向相反 r[i] = -1.0 else: # 方向一致 r[i] = 0.0 return r","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","name":"compute_fc_rank_consistency_reward","kind":"function","path":"verl/trainer/conf_adv_config.py","span":"L115-L155","signature":"def compute_fc_rank_consistency_reward( c_fc: np.ndarray, c_lsc: np.ndarray, delta: float = FC_RANK_DELTA, ) -> np.ndarray","exported":true,"origin":"ast","body_hash":"5b5d96b02dc0b9c33caafc9757463314ba45eb49f444905d74501ff03958d84e","body_text":"def compute_fc_rank_consistency_reward( c_fc: np.ndarray, c_lsc: np.ndarray, delta: float = FC_RANK_DELTA, ) -> np.ndarray: \"\"\" FC group 内排序一致性（fc 与 lsc 的 Kendall τ，不涉及 q）： 对每对 (i,j)，定义带 margin 的 soft sign： sgn_d(a,b) = +1 if a>b+d, -1 if a<b-d, 0 otherwise kappa_ij = sgn_d(c_fc_i, c_fc_j) * sgn_d(c_lsc_i, c_lsc_j) = 0 如果任意一方（fc 或 lsc）在该样本上为 -1（无效） R_i = (1/(G-1)) * Σ_{j≠i} kappa_ij 范围：[-1, +1]，group 内零和（对称结构）。 分母固定为 G-1（无效 pair 贡献 0 但不减分母）。 c_fc, c_lsc: shape (G,)，已归一化到 [0,1]，-1 表示无效 \"\"\" G = len(c_fc) if G <= 1: return np.zeros(G, dtype=np.float32) def _sgn(a, b, d): if a > b + d: return 1.0 elif a < b - d: return -1.0 else: return 0.0 r = np.zeros(G, dtype=np.float32) for i in range(G): for j in range(G): if i == j: continue # 任意一方无效，该 pair 贡献 0 if c_fc[i] < 0 or c_lsc[i] < 0 or c_fc[j] < 0 or c_lsc[j] < 0: continue s_fc = _sgn(c_fc[i], c_fc[j], delta) s_lsc = _sgn(c_lsc[i], c_lsc[j], delta) r[i] += s_fc * s_lsc r = r / (G - 1) return r.astype(np.float32)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward._sgn","name":"_sgn","kind":"function","path":"verl/trainer/conf_adv_config.py","span":"L135-L141","signature":"def _sgn(a, b, d)","exported":false,"origin":"ast","body_hash":"4a45b84c3a6105a315ee95dc43b72c3db3c193f10eabb8c9c09207aa68748deb","body_text":"def _sgn(a, b, d): if a > b + d: return 1.0 elif a < b - d: return -1.0 else: return 0.0","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","name":"apply_dual_calib_to_scores","kind":"function","path":"verl/trainer/conf_adv_config.py","span":"L162-L228","signature":"def apply_dual_calib_to_scores( scores: \"torch.Tensor\", # (batch_size,) 原始 outcome score（q） lsc_confs: \"torch.Tensor\", # (batch_size,) last_step_confidences，-1 表示无效 fc_confs: \"torch.Tensor\", # (batch_size,) final_confidences，-1 表示无效 uid_list: list, cur_step: int, ) -> \"torch.Tensor\"","exported":true,"origin":"ast","body_hash":"a0cae77573d74def1b7a5e5949aac2c106f3639ae4cdba07c12e8e54cc81e6a4","body_text":"def apply_dual_calib_to_scores( scores: \"torch.Tensor\", # (batch_size,) 原始 outcome score（q） lsc_confs: \"torch.Tensor\", # (batch_size,) last_step_confidences，-1 表示无效 fc_confs: \"torch.Tensor\", # (batch_size,) final_confidences，-1 表示无效 uid_list: list, cur_step: int, ) -> \"torch.Tensor\": \"\"\" 双路 calib 叠加： R_i = q_i + lsc_alpha(t) * R_i^lsc + fc_alpha(t) * R_i^fc LSC_MODE / FC_MODE 均为 \"none\" 时直接返回原始 scores。 lsc 和 fc 各自按独立 alpha curriculum warmup。 \"\"\" import torch lsc_alpha = get_lsc_alpha(cur_step) if LSC_MODE != \"none\" else 0.0 fc_alpha = get_fc_alpha(cur_step) if FC_MODE != \"none\" else 0.0 if lsc_alpha == 0.0 and fc_alpha == 0.0: return scores new_scores = scores.clone() unique_uids = list(dict.fromkeys(uid_list)) # 保序去重 for u in unique_uids: idx = [i for i, x in enumerate(uid_list) if x == u] G = len(idx) if G <= 1: continue q_arr = np.array([scores[i].item() for i in idx], dtype=np.float32) lsc_arr = np.array([lsc_confs[i].item() for i in idx], dtype=np.float32) fc_arr = np.array([fc_confs[i].item() for i in idx], dtype=np.float32) # 归一化 confidence 到 [0,1]（原始范围 0-10）；-1 保持 -1 表示无效 c_lsc = np.where(lsc_arr >= 0, lsc_arr / 10.0, lsc_arr) c_fc = np.where(fc_arr >= 0, fc_arr / 10.0, fc_arr) delta = np.zeros(G, dtype=np.float32) # ------ LSC 分量 ------ if LSC_MODE == \"stake\" and lsc_alpha > 0.0: r_lsc = compute_calib_reward_stake(q_arr, c_lsc) delta += lsc_alpha * r_lsc # ------ FC 分量 ------ if FC_MODE == \"stake\" and fc_alpha > 0.0: r_fc = compute_calib_reward_stake(q_arr, c_fc) delta += fc_alpha * r_fc elif FC_MODE == \"consistency\" and fc_alpha > 0.0: r_fc = compute_fc_consistency_reward(c_fc, c_lsc) delta += fc_alpha * r_fc elif FC_MODE == \"rank_consistency\" and fc_alpha > 0.0: r_fc = compute_fc_rank_consistency_reward(c_fc, c_lsc) delta += fc_alpha * r_fc for local_i, global_i in enumerate(idx): new_scores[global_i] = new_scores[global_i] + float(delta[local_i]) score_delta = float((new_scores - scores).abs().mean()) margin_str = (f\"on(correct<={STAKE_CORRECT_MARGIN}, wrong>={STAKE_WRONG_MARGIN})\" if STAKE_MARGIN_ENABLE else \"off\") print(f\"[DUAL_CALIB] step={cur_step}, lsc_mode={LSC_MODE}, fc_mode={FC_MODE}, \" f\"lsc_alpha={lsc_alpha:.4f}, fc_alpha={fc_alpha:.4f}, \" f\"margin={margin_str}, score_delta_mean={score_delta:.4f}\") return new_scores","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/conf_adv_config.py#compute_ece","name":"compute_ece","kind":"function","path":"verl/trainer/conf_adv_config.py","span":"L235-L267","signature":"def compute_ece(confidences: list, is_correct: list, n_bins: int = 20)","exported":true,"origin":"ast","body_hash":"4e9ec959a1cf7addc73c5890485578244e29c192aafa0017d4ef0839141931cb","body_text":"def compute_ece(confidences: list, is_correct: list, n_bins: int = 20): \"\"\" 计算 Expected Calibration Error。 confidences: List[Optional[float]]，原始值域 [1, 10]（integer），None 表示无效。 is_correct: List[bool] 返回 float 或 None（有效样本不足时） \"\"\" valid_pairs = [ (c / 10.0, float(acc)) for c, acc in zip(confidences, is_correct) if c is not None ] if len(valid_pairs) < n_bins: return None bin_conf_sums = np.zeros(n_bins) bin_acc_sums = np.zeros(n_bins) bin_counts = np.zeros(n_bins) for conf, acc in valid_pairs: bin_idx = min(int(conf * n_bins), n_bins - 1) bin_conf_sums[bin_idx] += conf bin_acc_sums[bin_idx] += acc bin_counts[bin_idx] += 1 ece = 0.0 total = len(valid_pairs) for b in range(n_bins): if bin_counts[b] > 0: avg_conf = bin_conf_sums[b] / bin_counts[b] avg_acc = bin_acc_sums[b] / bin_counts[b] ece += (bin_counts[b] / total) * abs(avg_conf - avg_acc) return float(ece)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/conf_adv_config.py#compute_auroc","name":"compute_auroc","kind":"function","path":"verl/trainer/conf_adv_config.py","span":"L270-L290","signature":"def compute_auroc(confidences, is_correct)","exported":true,"origin":"ast","body_hash":"5042fb2b553552eb5090c85a8370e3fd8bd311f67d016c990e3248bcf3e3a44b","body_text":"def compute_auroc(confidences, is_correct): \"\"\" 纯 NumPy 计算 AUROC，考虑 None 和 ties。 \"\"\" valid_pairs = [(c, int(is_correct[i])) for i, c in enumerate(confidences) if c is not None] if len(valid_pairs) < 2: return None confs = np.array([p[0] for p in valid_pairs], dtype=float) labels = np.array([p[1] for p in valid_pairs], dtype=int) n_pos = labels.sum() n_neg = len(labels) - n_pos if n_pos == 0 or n_neg == 0: return None pos_confs = confs[labels == 1][:, None] neg_confs = confs[labels == 0][None, :] diff = pos_confs - neg_confs u_stat = (diff > 0).sum() + 0.5 * (diff == 0).sum() return float(u_stat / (n_pos * n_neg))","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/trainer/conf_adv_config.py","relation":"contains","targetId":"verl/trainer/conf_adv_config.py#get_lsc_alpha","file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py","relation":"contains","targetId":"verl/trainer/conf_adv_config.py#get_fc_alpha","file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py","relation":"contains","targetId":"verl/trainer/conf_adv_config.py#compute_calib_reward_stake","file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_calib_reward_stake","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_calib_reward_stake","relation":"calls","name":"zeros","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_calib_reward_stake","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_calib_reward_stake","relation":"calls","name":"min","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_calib_reward_stake","relation":"calls","name":"max","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py","relation":"contains","targetId":"verl/trainer/conf_adv_config.py#compute_fc_consistency_reward","file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_consistency_reward","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_consistency_reward","relation":"calls","name":"zeros","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_consistency_reward","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py","relation":"contains","targetId":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","relation":"calls","name":"zeros","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","relation":"contains","targetId":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward._sgn","file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","relation":"calls","name":"zeros","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","relation":"calls","name":"_sgn","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","relation":"calls","name":"_sgn","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_fc_rank_consistency_reward","relation":"calls","name":"astype","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py","relation":"contains","targetId":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py","relation":"imports","specifier":"torch","file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"get_lsc_alpha","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"get_fc_alpha","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"clone","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"fromkeys","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"where","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"where","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"zeros","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"compute_calib_reward_stake","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"compute_calib_reward_stake","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"compute_fc_consistency_reward","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"compute_fc_rank_consistency_reward","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"abs","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#apply_dual_calib_to_scores","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py","relation":"contains","targetId":"verl/trainer/conf_adv_config.py#compute_ece","file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"zip","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"zeros","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"zeros","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"zeros","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"min","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"abs","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_ece","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py","relation":"contains","targetId":"verl/trainer/conf_adv_config.py#compute_auroc","file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_auroc","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_auroc","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_auroc","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_auroc","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_auroc","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_auroc","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_auroc","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_auroc","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_auroc","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/conf_adv_config.py"},{"source":"verl/trainer/conf_adv_config.py#compute_auroc","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/conf_adv_config.py"}]},"verl/trainer/fsdp_sft_trainer.py":{"size":19679,"mtimeMs":1789275830320.203,"hash":"7d0fe2f63fdc2f2026a8d731ff71dbd5ddddfb3413bc9273336d8514d894ad08","nodes":[{"id":"verl/trainer/fsdp_sft_trainer.py","name":"fsdp_sft_trainer.py","kind":"file","path":"verl/trainer/fsdp_sft_trainer.py","span":"L1-L436","signature":null,"exported":true,"origin":"ast","body_hash":"7d0fe2f63fdc2f2026a8d731ff71dbd5ddddfb3413bc9273336d8514d894ad08","chars":19679,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 lightweight one-file FSDP SFT Trainer TODO(zhangchi.usc1992) - Add calculation of mfu - Add validation \"\"\" import os os.environ['NCCL_DEBUG'] = 'WARN' os.environ['TOKENIZERS_PARALLELISM'] = 'true' import logging import re import torch import torch.distributed from torch import nn, optim from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, MixedPrecision, ShardingStrategy, CPUOffload from transformers import AutoTokenizer, AutoModelForCausalLM, PreTrainedModel, AutoConfig from verl.utils.torch_functional import get_cosine_schedule_with_warmup from tensordict import TensorDict from torch.utils.data import DataLoader, DistributedSampler from verl.utils.fsdp_utils import get_fsdp_wrap_policy, init_fn, get_init_weight_context_manager from verl.utils.dataset import SFTDataset from verl.utils.fs import copy_local_path_from_hdfs from verl.utils.tracking import Tracking from torch.distributed.device_mesh import DeviceMesh import verl.utils.hdfs_io as hdfs_io from verl.utils.debug import log_gpu_memory_usage from peft import LoraConfig, TaskType, get_peft_model logger = logging.getLogger(__file__) logger.setLevel(os.getenv('VERL_SFT_LOGGING_LEVEL', 'WARN')) from verl.trainer.fsdp_sft_trainer import FSDPSFTTrainer import hydra from torch.distributed.device_mesh import init_device_mesh from verl.utils.distributed import initialize_global_process_group @hydra.main(config_path='config', config_name='sft_trainer', version_base=None) if __name__ == '__main__': main()"},{"id":"verl/trainer/fsdp_sft_trainer.py#extract_step","name":"extract_step","kind":"function","path":"verl/trainer/fsdp_sft_trainer.py","span":"L52-L56","signature":"def extract_step(path)","exported":true,"origin":"ast","body_hash":"a189baafe9082729382f09d60f8ce08337babd4498ecbdaed76669de8e2a7dc2","body_text":"def extract_step(path): match = re.search(r'global_step_(\\d+)', path) if match: return int(match.group(1)) return None","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","name":"convert_to_regular_types","kind":"function","path":"verl/trainer/fsdp_sft_trainer.py","span":"L59-L68","signature":"def convert_to_regular_types(obj)","exported":true,"origin":"ast","body_hash":"5da247be7705217799c2e261c951edf37a8a1053d338f40de468ecbedf46bdc8","body_text":"def convert_to_regular_types(obj): \"\"\"Convert Hydra configs and other special types to regular Python types.\"\"\" from omegaconf import ListConfig, DictConfig if isinstance(obj, (ListConfig, DictConfig)): return {k: convert_to_regular_types(v) for k, v in obj.items()} if isinstance(obj, DictConfig) else list(obj) elif isinstance(obj, (list, tuple)): return [convert_to_regular_types(x) for x in obj] elif isinstance(obj, dict): return {k: convert_to_regular_types(v) for k, v in obj.items()} return obj","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","name":"FSDPSFTTrainer","kind":"class","path":"verl/trainer/fsdp_sft_trainer.py","span":"L71-L414","signature":"class FSDPSFTTrainer(object)","exported":true,"origin":"ast","body_hash":"0eff2ce18a4fa61938ac0b06084191f4d2ff2e7d2356759aaa6bb46502623534","body_text":"class FSDPSFTTrainer(object): def __init__(self, config, device_mesh: DeviceMesh): self.config = config self.device_mesh = device_mesh # build tokenizer first local_model_path = copy_local_path_from_hdfs(src=self.config.model.partial_pretrain, verbose=True) from verl.utils import hf_tokenizer self.tokenizer = hf_tokenizer(local_model_path, trust_remote_code=self.config.model.trust_remote_code) if self.config.data.chat_template is not None: raise ValueError('Apply Chat template from config is not supported yet.') # normalize dp size self._normalize_config_bsz() self._build_dataloader() # build model self._build_model_optimizer() # TODO: add checkpoint manager if self.device_mesh.get_rank() == 0: print(self.config) def _normalize_config_bsz(self): dp_size = self.device_mesh.size() if self.device_mesh.get_rank() == 0: print(f'Normalize batch size by dp {dp_size}') assert self.config.data.train_batch_size % dp_size == 0 assert self.config.data.micro_batch_size % dp_size == 0 self.config.data.train_batch_size //= dp_size self.config.data.micro_batch_size //= dp_size def _build_dataloader(self): config = self.config # build dataset self.train_dataset = SFTDataset(parquet_files=config.data.train_files, tokenizer=self.tokenizer, prompt_key=config.data.prompt_key, prompt_dict_keys=config.data.get('prompt_dict_keys', None), response_key=config.data.response_key, response_dict_keys=config.data.get('response_dict_keys', None), max_length=config.data.max_length, truncation=config.data.truncation) self.val_dataset = SFTDataset(parquet_files=config.data.val_files, tokenizer=self.tokenizer, prompt_key=config.data.prompt_key, prompt_dict_keys=config.data.get('prompt_dict_keys', None), response_key=config.data.response_key, response_dict_keys=config.data.get('response_dict_keys', None), max_length=config.data.max_length, truncation=config.data.truncation) # build dataloader rank = self.device_mesh.get_rank() world_size = self.device_mesh.size() self.train_sampler = DistributedSampler(self.train_dataset, shuffle=True, num_replicas=world_size, rank=rank, drop_last=True) self.train_dataloader = DataLoader(dataset=self.train_dataset, batch_size=config.data.train_batch_size, sampler=self.train_sampler, num_workers=8, pin_memory=True, drop_last=True) self.val_sampler = DistributedSampler(self.val_dataset, shuffle=True, num_replicas=world_size, rank=rank, drop_last=True) self.val_dataloader = DataLoader(dataset=self.val_dataset, batch_size=config.data.micro_batch_size, sampler=self.val_sampler, num_workers=8, pin_memory=True, drop_last=True) def _build_model_optimizer(self): # TODO (zhangchi.usc1992): # 1. support pretrain from random weights # 2. support init directly from sharded weights local_model_path = copy_local_path_from_hdfs(src=self.config.model.partial_pretrain, verbose=True) if self.config.model.get('external_lib', None) is not None: # This is used to import external_lib into the huggingface systems import importlib importlib.import_module(self.config.model.external_lib) log_gpu_memory_usage('Before model allocation', logger=logger) trust_remote_code = self.config.model.trust_remote_code # load config first config = AutoConfig.from_pretrained(local_model_path, trust_remote_code=trust_remote_code) # This may be very large init_context = get_init_weight_context_manager(use_meta_tensor=not config.tie_word_embeddings) with init_context(): self.model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(local_model_path, config=config, torch_dtype=torch.float32, attn_implementation='flash_attention_2', trust_remote_code=trust_remote_code) if self.config.model.get('lora_rank', 0) > 0: self.model.enable_input_require_grads() # Convert config to regular Python types before creating PEFT model lora_config = { 'task_type': TaskType.CAUSAL_LM, 'r': self.config.model.lora_rank, 'lora_alpha': self.config.model.lora_alpha, 'target_modules': convert_to_regular_types(self.config.model.target_modules), 'bias': \"none\" } self.model = get_peft_model(self.model, LoraConfig(**lora_config)) if self.config.model.enable_gradient_checkpointing: self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) log_gpu_memory_usage('After model allocation', logger=logger) mixed_precision = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32) auto_wrap_policy = get_fsdp_wrap_policy(self.model, config=self.config.model.fsdp_config.wrap_policy, is_lora=self.config.model.get('lora_rank', 0) > 0) if self.device_mesh.get_rank() == 0: print(auto_wrap_policy) if not self.config.model.fsdp_config.cpu_offload: cpu_offload = None else: cpu_offload = CPUOffload(offload_params=self.config.model.fsdp_config.offload_params) self.fsdp_model = FSDP(module=self.model, auto_wrap_policy=auto_wrap_policy, param_init_fn=init_fn, sharding_strategy=ShardingStrategy.FULL_SHARD, mixed_precision=mixed_precision, device_mesh=self.device_mesh, sync_module_states=True, device_id=torch.cuda.current_d","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.__init__","name":"__init__","kind":"method","path":"verl/trainer/fsdp_sft_trainer.py","span":"L73-L92","signature":"def __init__(self, config, device_mesh: DeviceMesh)","exported":false,"origin":"ast","body_hash":"3a2286696594348effa4f0c831cec1bfbbe8edb1e8a985b7faac0944450002b7","body_text":"def __init__(self, config, device_mesh: DeviceMesh): self.config = config self.device_mesh = device_mesh # build tokenizer first local_model_path = copy_local_path_from_hdfs(src=self.config.model.partial_pretrain, verbose=True) from verl.utils import hf_tokenizer self.tokenizer = hf_tokenizer(local_model_path, trust_remote_code=self.config.model.trust_remote_code) if self.config.data.chat_template is not None: raise ValueError('Apply Chat template from config is not supported yet.') # normalize dp size self._normalize_config_bsz() self._build_dataloader() # build model self._build_model_optimizer() # TODO: add checkpoint manager if self.device_mesh.get_rank() == 0: print(self.config)","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPSFTTrainer"},{"id":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._normalize_config_bsz","name":"_normalize_config_bsz","kind":"method","path":"verl/trainer/fsdp_sft_trainer.py","span":"L94-L103","signature":"def _normalize_config_bsz(self)","exported":false,"origin":"ast","body_hash":"5a5ff21dd1fc645dcd2f84be7bf7ffae4804e2553d271590735ae78e95713cd1","body_text":"def _normalize_config_bsz(self): dp_size = self.device_mesh.size() if self.device_mesh.get_rank() == 0: print(f'Normalize batch size by dp {dp_size}') assert self.config.data.train_batch_size % dp_size == 0 assert self.config.data.micro_batch_size % dp_size == 0 self.config.data.train_batch_size //= dp_size self.config.data.micro_batch_size //= dp_size","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPSFTTrainer"},{"id":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","name":"_build_dataloader","kind":"method","path":"verl/trainer/fsdp_sft_trainer.py","span":"L105-L150","signature":"def _build_dataloader(self)","exported":false,"origin":"ast","body_hash":"aa2529e3fa4b423dec31a934019c5e711cfbeff266ff0267105f3116710fddfc","body_text":"def _build_dataloader(self): config = self.config # build dataset self.train_dataset = SFTDataset(parquet_files=config.data.train_files, tokenizer=self.tokenizer, prompt_key=config.data.prompt_key, prompt_dict_keys=config.data.get('prompt_dict_keys', None), response_key=config.data.response_key, response_dict_keys=config.data.get('response_dict_keys', None), max_length=config.data.max_length, truncation=config.data.truncation) self.val_dataset = SFTDataset(parquet_files=config.data.val_files, tokenizer=self.tokenizer, prompt_key=config.data.prompt_key, prompt_dict_keys=config.data.get('prompt_dict_keys', None), response_key=config.data.response_key, response_dict_keys=config.data.get('response_dict_keys', None), max_length=config.data.max_length, truncation=config.data.truncation) # build dataloader rank = self.device_mesh.get_rank() world_size = self.device_mesh.size() self.train_sampler = DistributedSampler(self.train_dataset, shuffle=True, num_replicas=world_size, rank=rank, drop_last=True) self.train_dataloader = DataLoader(dataset=self.train_dataset, batch_size=config.data.train_batch_size, sampler=self.train_sampler, num_workers=8, pin_memory=True, drop_last=True) self.val_sampler = DistributedSampler(self.val_dataset, shuffle=True, num_replicas=world_size, rank=rank, drop_last=True) self.val_dataloader = DataLoader(dataset=self.val_dataset, batch_size=config.data.micro_batch_size, sampler=self.val_sampler, num_workers=8, pin_memory=True, drop_last=True)","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPSFTTrainer"},{"id":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","name":"_build_model_optimizer","kind":"method","path":"verl/trainer/fsdp_sft_trainer.py","span":"L152-L242","signature":"def _build_model_optimizer(self): # TODO (zhangchi.usc1992): # 1. support pretrain from random weights # 2. support init directly from sharded weights","exported":false,"origin":"ast","body_hash":"7f1f1f1ef5d1a38fad5dcb88ba1f09fa06b3fd34f9d59336544c0dbe2467f77c","body_text":"def _build_model_optimizer(self): # TODO (zhangchi.usc1992): # 1. support pretrain from random weights # 2. support init directly from sharded weights local_model_path = copy_local_path_from_hdfs(src=self.config.model.partial_pretrain, verbose=True) if self.config.model.get('external_lib', None) is not None: # This is used to import external_lib into the huggingface systems import importlib importlib.import_module(self.config.model.external_lib) log_gpu_memory_usage('Before model allocation', logger=logger) trust_remote_code = self.config.model.trust_remote_code # load config first config = AutoConfig.from_pretrained(local_model_path, trust_remote_code=trust_remote_code) # This may be very large init_context = get_init_weight_context_manager(use_meta_tensor=not config.tie_word_embeddings) with init_context(): self.model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(local_model_path, config=config, torch_dtype=torch.float32, attn_implementation='flash_attention_2', trust_remote_code=trust_remote_code) if self.config.model.get('lora_rank', 0) > 0: self.model.enable_input_require_grads() # Convert config to regular Python types before creating PEFT model lora_config = { 'task_type': TaskType.CAUSAL_LM, 'r': self.config.model.lora_rank, 'lora_alpha': self.config.model.lora_alpha, 'target_modules': convert_to_regular_types(self.config.model.target_modules), 'bias': \"none\" } self.model = get_peft_model(self.model, LoraConfig(**lora_config)) if self.config.model.enable_gradient_checkpointing: self.model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) log_gpu_memory_usage('After model allocation', logger=logger) mixed_precision = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32) auto_wrap_policy = get_fsdp_wrap_policy(self.model, config=self.config.model.fsdp_config.wrap_policy, is_lora=self.config.model.get('lora_rank', 0) > 0) if self.device_mesh.get_rank() == 0: print(auto_wrap_policy) if not self.config.model.fsdp_config.cpu_offload: cpu_offload = None else: cpu_offload = CPUOffload(offload_params=self.config.model.fsdp_config.offload_params) self.fsdp_model = FSDP(module=self.model, auto_wrap_policy=auto_wrap_policy, param_init_fn=init_fn, sharding_strategy=ShardingStrategy.FULL_SHARD, mixed_precision=mixed_precision, device_mesh=self.device_mesh, sync_module_states=True, device_id=torch.cuda.current_device(), cpu_offload=cpu_offload, use_orig_params=False) log_gpu_memory_usage('After FSDP wrapping', logger=logger) self.optimizer = optim.AdamW(self.fsdp_model.parameters(), lr=self.config.optim.lr, betas=self.config.optim.betas, weight_decay=self.config.optim.weight_decay) log_gpu_memory_usage('After initialize optimizer', logger=logger) steps_per_epoch = len(self.train_dataloader) total_steps = steps_per_epoch * self.config.trainer.total_epochs if self.device_mesh.get_rank() == 0: print( f'Number of steps/epoch {steps_per_epoch}, number of epochs {self.config.trainer.total_epochs}, total number of steps {total_steps}' ) num_warmup_steps = int(total_steps * self.config.optim.warmup_steps_ratio) self.lr_scheduler = get_cosine_schedule_with_warmup(optimizer=self.optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=total_steps)","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPSFTTrainer"},{"id":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","name":"_compute_loss","kind":"method","path":"verl/trainer/fsdp_sft_trainer.py","span":"L244-L276","signature":"def _compute_loss(self, batch)","exported":false,"origin":"ast","body_hash":"dcc5b0d458c48197e1412073cfede4bd3e11f8e674dc9381e99106cdacaa4266","body_text":"def _compute_loss(self, batch): loss_mask = batch.pop('loss_mask')[:, :-1].reshape(-1).cuda() labels = batch['input_ids'][:, 1:].cuda() with torch.autocast(device_type='cuda', dtype=torch.bfloat16): output = self.fsdp_model(input_ids=batch['input_ids'], attention_mask=batch['attention_mask'], position_ids=batch['position_ids'], use_cache=False) # prevent model thinks it it generating logits = output.logits shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels.contiguous() # Flatten the tokens loss_fct = nn.CrossEntropyLoss(reduction='none') shift_logits = shift_logits.view(-1, self.model.config.vocab_size) shift_labels = shift_labels.view(-1) # Enable model parallelism shift_labels = shift_labels.to(shift_logits.device) loss = loss_fct(shift_logits, shift_labels) loss = loss * loss_mask valid_token_this_rank = torch.sum(loss_mask) if self.config.data.balance_dp_token: torch.distributed.all_reduce(valid_token_this_rank) # becomes total valid tokens in all ranks dp_size = torch.distributed.get_world_size() else: dp_size = 1 loss = torch.sum(loss) / valid_token_this_rank * dp_size # possible bugs here for dp return loss","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPSFTTrainer"},{"id":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","name":"training_step","kind":"method","path":"verl/trainer/fsdp_sft_trainer.py","span":"L278-L312","signature":"def training_step(self, batch: TensorDict)","exported":true,"origin":"ast","body_hash":"3a721f7d535241f06096cc02079fa37992fbe471992a9955941ba22b0df0d69e","body_text":"def training_step(self, batch: TensorDict): self.fsdp_model.train() log_gpu_memory_usage('Before optimizer zero_grad', logger=logger) self.optimizer.zero_grad() log_gpu_memory_usage('After optimizer zero_grad', logger=logger) micro_batches = batch.split(self.config.data.micro_batch_size) n_micro_batches = len(micro_batches) step_loss = 0 for micro_batch in micro_batches: loss = self._compute_loss(batch=micro_batch) / n_micro_batches loss.backward() step_loss += loss.item() self.fsdp_model.clip_grad_norm_(max_norm=self.config.optim.clip_grad) log_gpu_memory_usage('Before optimizer step', logger=logger) self.optimizer.step() log_gpu_memory_usage('After optimizer step', logger=logger) self.lr_scheduler.step() # reduce loss across dp ranks lr = self.lr_scheduler.get_last_lr()[0] log_gpu_memory_usage('After offload weights', logger=logger) step_loss = torch.tensor(step_loss).cuda() torch.distributed.all_reduce(step_loss, op=torch.distributed.ReduceOp.AVG) return {'train/loss': step_loss.detach().item(), 'train/lr(1e-3)': lr * 1e3}","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPSFTTrainer"},{"id":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.validation_step","name":"validation_step","kind":"method","path":"verl/trainer/fsdp_sft_trainer.py","span":"L314-L319","signature":"def validation_step(self, batch: TensorDict)","exported":true,"origin":"ast","body_hash":"ef06c641954c2dfdcc1accc4a4373d45ed835db28f551ae5f2f43414fdedf829","body_text":"def validation_step(self, batch: TensorDict): self.fsdp_model.eval() with torch.no_grad(): loss = self._compute_loss(batch) torch.distributed.all_reduce(loss, op=torch.distributed.ReduceOp.AVG) return loss","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPSFTTrainer"},{"id":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","name":"save_checkpoint","kind":"method","path":"verl/trainer/fsdp_sft_trainer.py","span":"L321-L337","signature":"def save_checkpoint(self, step): # save checkpoint","exported":true,"origin":"ast","body_hash":"f3676ad43ffa834a317825113ef812d916c3e69e13d4859c20db079a86c17d8d","body_text":"def save_checkpoint(self, step): # save checkpoint from torch.distributed.fsdp import FullStateDictConfig, StateDictType cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) with FSDP.state_dict_type(self.fsdp_model, StateDictType.FULL_STATE_DICT, cfg): state_dict = self.fsdp_model.state_dict() path = os.path.join(self.config.trainer.default_local_dir, f'global_step_{step}') # save huggingface model if self.device_mesh.get_rank() == 0: os.makedirs(path, exist_ok=True) self.model.save_pretrained(path, state_dict=state_dict) self.tokenizer.save_pretrained(path) if self.config.trainer.default_hdfs_dir: hdfs_io.makedirs(self.config.trainer.default_hdfs_dir, exist_ok=True) hdfs_io.copy(src=path, dst=self.config.trainer.default_hdfs_dir, dirs_exist_ok=True) torch.distributed.barrier()","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPSFTTrainer"},{"id":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","name":"fit","kind":"method","path":"verl/trainer/fsdp_sft_trainer.py","span":"L339-L414","signature":"def fit(self)","exported":true,"origin":"ast","body_hash":"7faa8cd781b95f9ee7aa75adc8a5b2c44500ec1f46489d4ee9df25c275726bb8","body_text":"def fit(self): rank = self.device_mesh.get_rank() # TODO: add a unified tracking if rank == 0: tracking = Tracking(project_name=self.config.trainer.project_name, experiment_name=self.config.trainer.experiment_name, default_backend=self.config.trainer.logger) global_step = 0 # compute the total training steps. # the total training steps in SFT is mainly for early exit total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs if self.config.trainer.total_training_steps is not None: total_training_steps = self.config.trainer.total_training_steps self.total_training_steps = total_training_steps print(f'Total training steps: {self.total_training_steps}') # TODO (zhangchi.usc1992) add back checkpoint manager. Currently, it blocks when uploading to hdfs. So very slow. if self.config.trainer.validate_before_training: # validate before training val_losses = [] for data in self.val_dataloader: data = TensorDict(data, batch_size=self.config.data.micro_batch_size).cuda() val_loss = self.validation_step(data) val_losses.append(val_loss) if rank == 0: val_loss = torch.mean(torch.stack(val_losses)) metric = {'val/loss': val_loss.detach().item()} tracking.log(data=metric, step=global_step) torch.distributed.barrier() for epoch in range(self.config.trainer.total_epochs): self.train_sampler.set_epoch(epoch=epoch) for data in self.train_dataloader: data = TensorDict(data, batch_size=self.config.data.train_batch_size).cuda() metric = self.training_step(data) if rank == 0: tracking.log(data=metric, step=global_step) global_step += 1 # for early exit validation if global_step >= self.total_training_steps: # Perform final validation val_losses = [] for val_data in self.val_dataloader: val_data = TensorDict(val_data, batch_size=self.config.data.micro_batch_size).cuda() val_loss = self.validation_step(val_data) val_losses.append(val_loss) if rank == 0: avg_val_loss = torch.mean(torch.stack(val_losses)) metric = {'val/loss': avg_val_loss.detach().item()} tracking.log(data=metric, step=global_step) torch.distributed.barrier() # Save final checkpoint self.save_checkpoint(step=global_step) return # validation val_losses = [] for data in self.val_dataloader: data = TensorDict(data, batch_size=self.config.data.micro_batch_size).cuda() val_loss = self.validation_step(data) val_losses.append(val_loss) if rank == 0: val_loss = torch.mean(torch.stack(val_losses)) metric = {'val/loss': val_loss.detach().item()} tracking.log(data=metric, step=global_step) torch.distributed.barrier() # save checkpoint self.save_checkpoint(step=global_step)","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPSFTTrainer"},{"id":"verl/trainer/fsdp_sft_trainer.py#main","name":"main","kind":"function","path":"verl/trainer/fsdp_sft_trainer.py","span":"L426-L431","signature":"def main(config)","exported":true,"origin":"ast","body_hash":"c7e39d232363e17922d7ae45b1b6c8f8a9b77eb58f605b713c0697ca24a654bc","body_text":"def main(config): local_rank, rank, world_size = initialize_global_process_group() device_mesh = init_device_mesh(device_type='cuda', mesh_shape=(world_size,), mesh_dim_names=('dp',)) trainer = FSDPSFTTrainer(config=config, device_mesh=device_mesh) trainer.fit()","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"os","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"logging","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"re","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"torch","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"torch.distributed","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"torch","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"torch.distributed.fsdp","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"transformers","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"tensordict","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"torch.utils.data","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"verl.utils.fsdp_utils","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"verl.utils.dataset","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"verl.utils.fs","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"verl.utils.tracking","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"torch.distributed.device_mesh","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"verl.utils.debug","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"peft","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"calls","name":"getLogger","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"calls","name":"setLevel","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"calls","name":"getenv","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#extract_step","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#extract_step","relation":"calls","name":"search","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#extract_step","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#extract_step","relation":"calls","name":"group","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"omegaconf","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","relation":"calls","name":"convert_to_regular_types","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","relation":"calls","name":"convert_to_regular_types","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","relation":"calls","name":"convert_to_regular_types","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#convert_to_regular_types","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","relation":"extends","name":"object","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.__init__","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.__init__","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"verl.utils","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.__init__","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.__init__","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.__init__","relation":"calls","name":"_normalize_config_bsz","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.__init__","relation":"calls","name":"_build_dataloader","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.__init__","relation":"calls","name":"_build_model_optimizer","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.__init__","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._normalize_config_bsz","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._normalize_config_bsz","relation":"calls","name":"size","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._normalize_config_bsz","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._normalize_config_bsz","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"SFTDataset","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"SFTDataset","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"size","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"DistributedSampler","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"DataLoader","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"DistributedSampler","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_dataloader","relation":"calls","name":"DataLoader","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"importlib","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"import_module","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"get_init_weight_context_manager","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"init_context","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"enable_input_require_grads","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"get_peft_model"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"convert_to_regular_types","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"get_peft_model","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"LoraConfig","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"gradient_checkpointing_enable","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"get_peft_model"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"MixedPrecision","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"get_fsdp_wrap_policy","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"CPUOffload","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"FSDP","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"current_device","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"AdamW","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"parameters","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FullyShardedDataParallel"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._build_model_optimizer","relation":"calls","name":"get_cosine_schedule_with_warmup","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"cuda","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"reshape","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"pop","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"cuda","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"autocast","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"fsdp_model","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"CrossEntropyLoss","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"view","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"view","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"to","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"loss_fct","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer._compute_loss","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"train","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FullyShardedDataParallel"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"zero_grad","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"split","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"TensorDict"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"_compute_loss","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"backward","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"clip_grad_norm_","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FullyShardedDataParallel"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"step","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"step","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"get_cosine_schedule_with_warmup"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"get_last_lr","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"get_cosine_schedule_with_warmup"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"cuda","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"tensor","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.training_step","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.validation_step","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.validation_step","relation":"calls","name":"eval","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FullyShardedDataParallel"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.validation_step","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.validation_step","relation":"calls","name":"_compute_loss","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.validation_step","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"torch.distributed.fsdp","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"FullStateDictConfig","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"state_dict_type","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"state_dict","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FullyShardedDataParallel"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"join","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"save_pretrained","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"get_peft_model"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"save_pretrained","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"hf_tokenizer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"copy","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.save_checkpoint","relation":"calls","name":"barrier","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"Tracking","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"cuda","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"validation_step","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"stack","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"log","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"Tracking"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"barrier","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"set_epoch","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"DistributedSampler"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"cuda","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"training_step","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"log","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"Tracking"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"cuda","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"validation_step","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"stack","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"log","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"Tracking"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"barrier","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"save_checkpoint","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"cuda","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"validation_step","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"stack","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"log","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"Tracking"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"barrier","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#FSDPSFTTrainer.fit","relation":"calls","name":"save_checkpoint","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"verl.trainer.fsdp_sft_trainer","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"hydra","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"torch.distributed.device_mesh","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"imports","specifier":"verl.utils.distributed","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"calls","name":"main","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"contains","targetId":"verl/trainer/fsdp_sft_trainer.py#main","file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#main","relation":"calls","name":"initialize_global_process_group","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#main","relation":"calls","name":"init_device_mesh","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#main","relation":"calls","name":"FSDPSFTTrainer","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"},{"source":"verl/trainer/fsdp_sft_trainer.py#main","relation":"calls","name":"fit","viaMember":true,"file":"verl/trainer/fsdp_sft_trainer.py","recvType":"FSDPSFTTrainer"},{"source":"verl/trainer/fsdp_sft_trainer.py","relation":"calls","name":"main","viaMember":false,"file":"verl/trainer/fsdp_sft_trainer.py"}]},"verl/trainer/main_eval.py":{"size":2204,"mtimeMs":1789275830320.203,"hash":"91d913c8280df50eca396ea2c6d1961d9569a29ed7c7f104a597610e24cfda50","nodes":[{"id":"verl/trainer/main_eval.py","name":"main_eval.py","kind":"file","path":"verl/trainer/main_eval.py","span":"L1-L70","signature":null,"exported":true,"origin":"ast","body_hash":"91d913c8280df50eca396ea2c6d1961d9569a29ed7c7f104a597610e24cfda50","chars":2204,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Offline evaluate the performance of a generated file using reward model and ground truth verifier. The input is a parquet file that contains N generated sequences and (optional) the ground truth. \"\"\" import hydra from verl.utils.fs import copy_local_path_from_hdfs from verl.utils.reward_score import math, gsm8k import pandas as pd import numpy as np @hydra.main(config_path='config', config_name='evaluation', version_base=None) if __name__ == '__main__': main()"},{"id":"verl/trainer/main_eval.py#select_reward_fn","name":"select_reward_fn","kind":"function","path":"verl/trainer/main_eval.py","span":"L27-L31","signature":"def select_reward_fn(data_source)","exported":true,"origin":"ast","body_hash":"0b07a0fe86a5bffdf8ca9d037f939b73b8a62127cd78326a71e11ce5d097c163","body_text":"def select_reward_fn(data_source): if data_source == 'lighteval/MATH': return math.compute_score else: raise NotImplementedError","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/main_eval.py#main","name":"main","kind":"function","path":"verl/trainer/main_eval.py","span":"L35-L65","signature":"def main(config)","exported":true,"origin":"ast","body_hash":"0d545cd630e88aee6643d7c3d6f1a132532650becc6d3d2c9673bd14d4903a61","body_text":"def main(config): local_path = copy_local_path_from_hdfs(config.data.path) dataset = pd.read_parquet(local_path) prompts = dataset[config.data.prompt_key] responses = dataset[config.data.response_key] data_sources = dataset[config.data.data_source_key] reward_model_data = dataset[config.data.reward_model_key] passes = 0 total = len(dataset) for i in range(total): response_lst = responses[i] data_source = data_sources[i] # select reward score based on data_source prompt = prompts[i] reward_data = reward_model_data[i] reward_fn = select_reward_fn(data_source) ground_truth = reward_data['ground_truth'] score_lst = [] for r in response_lst: score = reward_fn(r, ground_truth) score_lst.append(score) max_score = np.max(score_lst) if max_score == 1: passes += 1 print(f'pass@5: {passes / total}')","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/trainer/main_eval.py","relation":"imports","specifier":"hydra","file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py","relation":"imports","specifier":"verl.utils.fs","file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py","relation":"imports","specifier":"verl.utils.reward_score","file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py","relation":"contains","targetId":"verl/trainer/main_eval.py#select_reward_fn","file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py","relation":"calls","name":"main","viaMember":true,"file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py","relation":"contains","targetId":"verl/trainer/main_eval.py#main","file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py#main","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py#main","relation":"calls","name":"read_parquet","viaMember":true,"file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py#main","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py#main","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py#main","relation":"calls","name":"select_reward_fn","viaMember":false,"file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py#main","relation":"calls","name":"reward_fn","viaMember":false,"file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py#main","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py#main","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py#main","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_eval.py"},{"source":"verl/trainer/main_eval.py","relation":"calls","name":"main","viaMember":false,"file":"verl/trainer/main_eval.py"}]},"verl/trainer/main_generation.py":{"size":5760,"mtimeMs":1789275830320.203,"hash":"19a40ee1d20bb06bff56d46b2c1bb9417e6e33539ff35eaae96fd3e55a4e5ae2","nodes":[{"id":"verl/trainer/main_generation.py","name":"main_generation.py","kind":"file","path":"verl/trainer/main_generation.py","span":"L1-L138","signature":null,"exported":true,"origin":"ast","body_hash":"19a40ee1d20bb06bff56d46b2c1bb9417e6e33539ff35eaae96fd3e55a4e5ae2","chars":5760,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Generate responses given a dataset of prompts \"\"\" import ray import numpy as np import hydra import os os.environ['NCCL_DEBUG'] = 'WARN' os.environ['TOKENIZERS_PARALLELISM'] = 'true' # os.environ['TORCH_COMPILE_DISABLE'] = '1' from verl.utils.model import compute_position_id_with_mask import pandas as pd from transformers import AutoTokenizer from verl import DataProto from verl.utils.fs import copy_local_path_from_hdfs from verl.workers.fsdp_workers import ActorRolloutRefWorker from verl.utils.hdfs_io import makedirs from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup @hydra.main(config_path='config', config_name='generation', version_base=None) if __name__ == '__main__': main()"},{"id":"verl/trainer/main_generation.py#main","name":"main","kind":"function","path":"verl/trainer/main_generation.py","span":"L40-L133","signature":"def main(config)","exported":true,"origin":"ast","body_hash":"472b9a656beb5eb434a9d5794df2c5f2cf646a0de5fed7bef31287e2ca57c695","body_text":"def main(config): from pprint import pprint from omegaconf import OmegaConf pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values OmegaConf.resolve(config) local_path = copy_local_path_from_hdfs(config.model.path) from verl.utils import hf_tokenizer tokenizer = hf_tokenizer(local_path) if config.rollout.temperature == 0.: assert config.data.n_samples == 1, 'When temperature=0, n_samples must be 1.' # read dataset. Note that the dataset should directly contain chat template format (e.g., a list of dictionary) dataset = pd.read_parquet(config.data.path) chat_lst = dataset[config.data.prompt_key].tolist() chat_lst = [chat.tolist() for chat in chat_lst] tokenizer.padding_side = 'left' if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token 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) wg.init_model() total_samples = len(dataset) # real_batch_size = data.batch['input_ids'].shape[0] config_batch_size = config.data.batch_size dp_size = wg.world_size // config.rollout.tensor_model_parallel_size num_batch = (total_samples // config_batch_size) + 1 output_lst = [[] for _ in range(config.data.n_samples)] for batch_idx in range(num_batch): print(f'[{batch_idx+1}/{num_batch}] Start to process.') batch_chat_lst = chat_lst[batch_idx * config_batch_size:(batch_idx + 1) * config_batch_size] inputs = tokenizer.apply_chat_template(batch_chat_lst, add_generation_prompt=True, padding=True, truncation=True, max_length=config.rollout.prompt_length, return_tensors='pt', return_dict=True, tokenize=True) input_ids = inputs['input_ids'] attention_mask = inputs['attention_mask'] position_ids = compute_position_id_with_mask(attention_mask) batch_dict = {'input_ids': input_ids, 'attention_mask': attention_mask, 'position_ids': position_ids} data = DataProto.from_dict(batch_dict) real_batch_size = data.batch['input_ids'].shape[0] if real_batch_size % dp_size != 0: dummy_data_size = dp_size - real_batch_size % dp_size dummy_data = data[:dummy_data_size] data = DataProto.concat([data, dummy_data]) print( f'dp_size {dp_size} is not divisible by real_batch_size {real_batch_size}, add {dummy_data_size} dummy data' ) batch_size = data.batch['input_ids'].shape[0] assert batch_size % dp_size == 0, f'batch_size {batch_size} is not divisible by dp_size {dp_size}' print(f'[{batch_idx+1}/{num_batch}] Start to generate.') # START TO GENERATE FOR n_samples TIMES for i in range(config.data.n_samples): output = wg.generate_sequences(data) # remove dummy data output = output[:real_batch_size] output_text = tokenizer.batch_decode(output.batch['input_ids'][:, -config.rollout.response_length:], skip_special_tokens=False) # remove the padding pad_token = tokenizer.pad_token output_text_unpad = [] for text in output_text: output_text_unpad.append(text.replace(pad_token, '')) output_lst[i].extend(output_text_unpad) # convert output_lst from (n_samples, n_data) to (n_data, n_sampels) output_lst = np.array(output_lst, dtype=object) output_lst = np.transpose(output_lst, axes=(1, 0)).tolist() # add to the data frame dataset[f'responses'] = output_lst # write to a new parquet output_dir = os.path.dirname(config.data.output_path) makedirs(output_dir, exist_ok=True) dataset.to_parquet(config.data.output_path) return output_text","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"ray","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"hydra","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"os","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"verl.utils.model","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"transformers","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"verl","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"verl.utils.fs","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"verl.workers.fsdp_workers","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"verl.utils.hdfs_io","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"verl.single_controller.ray","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"calls","name":"main","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"contains","targetId":"verl/trainer/main_generation.py#main","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"pprint","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"omegaconf","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"pprint","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"to_container","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"resolve","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"imports","specifier":"verl.utils","file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"read_parquet","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"tolist","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"tolist","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"RayClassWithInitArgs","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"remote","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"RayResourcePool","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"RayWorkerGroup","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"init_model","viaMember":true,"file":"verl/trainer/main_generation.py","recvType":"RayWorkerGroup"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"apply_chat_template","viaMember":true,"file":"verl/trainer/main_generation.py","recvType":"hf_tokenizer"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"compute_position_id_with_mask","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"from_dict","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"concat","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"generate_sequences","viaMember":true,"file":"verl/trainer/main_generation.py","recvType":"RayWorkerGroup"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"batch_decode","viaMember":true,"file":"verl/trainer/main_generation.py","recvType":"hf_tokenizer"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"replace","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"extend","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"tolist","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"transpose","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"dirname","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"makedirs","viaMember":false,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py#main","relation":"calls","name":"to_parquet","viaMember":true,"file":"verl/trainer/main_generation.py"},{"source":"verl/trainer/main_generation.py","relation":"calls","name":"main","viaMember":false,"file":"verl/trainer/main_generation.py"}]},"verl/trainer/main_ppo.py":{"size":30364,"mtimeMs":1789275830320.203,"hash":"8af5f6ec11f73b46ee9669080d1983283c4d43d93fc977b7a72384c1f07c1f0b","nodes":[{"id":"verl/trainer/main_ppo.py","name":"main_ppo.py","kind":"file","path":"verl/trainer/main_ppo.py","span":"L1-L677","signature":null,"exported":true,"origin":"ast","body_hash":"8af5f6ec11f73b46ee9669080d1983283c4d43d93fc977b7a72384c1f07c1f0b","chars":28836,"summary_state":"pending","summary":null,"crux":null,"body_text":"# main_ppo-try2.py # 修改：confidence作为系数，具体实现包括softmax归一化处理和温度 # 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. \"\"\" Note that we don't combine the main with ray_trainer as ray_trainer is used by other main. \"\"\" from verl import DataProto import torch from verl.utils.reward_score import qa_em, qa_f1 from verl.trainer.ppo.ray_trainer import RayPPOTrainer import re import numpy as np from verl.trainer.conf_adv_config import ( compute_ece, compute_auroc, LSC_MODE, FC_MODE, ) # ================== Reward 默认配置 ================== REWARD_TYPE = \"f1\" # reward_mode=\"base\" 表示只用 r_answer + r_format + r_calib，去掉搜索/回合/无答案惩罚 REWARD_MODE = \"base\" # 以下三项在新 reward_mode=\"base\" 下不参与计算，保留参数供旧模式兼容 NO_SEARCH_PENALTY = 0.05 TURN_FREE = 1 TURN_PENALTY_PER_TURN = 0 NO_ANSWER_PENALTY = 0.05 CURRICULUM_STEP = 0 # confidence 统计（用于 logging，不影响 reward 计算） # calib reward 由 ray_trainer.py 中 apply_dual_calib_to_scores 在组内计算 # 通过 conf_adv_config.py 的 LSC_MODE / FC_MODE 参数控制 CONFIDENCE_MODE = \"observe_only\" # RewardManager 只做统计，不计算 r_calib CONFIDENCE_MODE_VAL = \"observe_only\" # 验证时同样只统计 # F1 >= CORRECT_THRESHOLD 算\"答对\"，用于 logging 中 is_correct 判定 CORRECT_THRESHOLD = 0.8 # ---- r_format 参数 ---- # <answer> 格式正确时给的小正向奖励 FORMAT_SCORE = 0.1 # ================== Per-step Confidence 配置 ================== ENABLE_PERSTEP_REWARD = False # 细粒度 per-step reward 开关（默认关闭） PERSTEP_CONF_MISSING_PENALTY = 0.02 # per-step conf 缺失时的格式惩罚（粗粒度模式） import ray import hydra @hydra.main(config_path='config', config_name='ppo_trainer', version_base=None) @ray.remote if __name__ == '__main__': main()"},{"id":"verl/trainer/main_ppo.py#_select_rm_score_fn","name":"_select_rm_score_fn","kind":"function","path":"verl/trainer/main_ppo.py","span":"L67-L79","signature":"def _select_rm_score_fn(data_source, reward_type: str = \"em\")","exported":false,"origin":"ast","body_hash":"08a662a2eeb8d1fafb664422727c71949c811b0f639bb752281f5a08829853b2","body_text":"def _select_rm_score_fn(data_source, reward_type: str = \"em\"): \"\"\" 修改：新增正确性检测选择em/f1 \"\"\" if data_source.lower() in ['nq', 'triviaqa', 'popqa', 'hotpotqa', '2wikimultihopqa', '2wiki', 'musique', 'bamboogle', 'nq_hotpotqa', 'nqhotpot', 'simpleqa']: if reward_type == \"em\": return qa_em.compute_score_em elif reward_type == \"f1\": return qa_f1.compute_score_f1 else: raise ValueError(f\"Unknown reward_type: {reward_type}\") else: raise NotImplementedError(f\"Unknown data_source: '{data_source}'. Supported: ['nq', 'triviaqa', 'popqa', 'hotpotqa', '2wikimultihopqa', '2wiki', 'musique', 'bamboogle', 'nq_hotpotqa', 'nqhotpot', 'simpleqa']\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/main_ppo.py#RewardManager","name":"RewardManager","kind":"class","path":"verl/trainer/main_ppo.py","span":"L82-L547","signature":"class RewardManager()","exported":true,"origin":"ast","body_hash":"7ae7cdf3e580aa67b7cdf370058e6b7390cbebc31321c9fa6cb608b7d6de7fa3","body_text":"class RewardManager(): \"\"\"The reward manager. \"\"\" def __init__( self, tokenizer, num_examine, format_score: float = FORMAT_SCORE, reward_type: str = REWARD_TYPE, reward_mode: str = REWARD_MODE, no_search_penalty: float = NO_SEARCH_PENALTY, turn_free: int = TURN_FREE, turn_penalty_per_turn: float = TURN_PENALTY_PER_TURN, no_answer_penalty: float = NO_ANSWER_PENALTY, curriculum_step: int = CURRICULUM_STEP, # confidence 统计相关参数 confidence_mode: str = CONFIDENCE_MODE, correct_threshold: float = CORRECT_THRESHOLD, # 评估模式：final_score 只用 EM（base_score），不叠加任何惩罚 eval_only_em: bool = False, ) -> None: self.tokenizer = tokenizer self.num_examine = num_examine self.format_score = format_score self.reward_type = reward_type self.reward_mode = reward_mode self.no_search_penalty = float(no_search_penalty) self.turn_free = int(turn_free) self.turn_penalty_per_turn = float(turn_penalty_per_turn) self.no_answer_penalty = float(no_answer_penalty) self.curriculum_step = int(curriculum_step) self.confidence_mode = confidence_mode self.correct_threshold = float(correct_threshold) self.eval_only_em = eval_only_em print(\"=\" * 80) print(\"[REWARD_MANAGER] Initialization:\") print(f\" reward_type: {self.reward_type}\") print(f\" reward_mode: {self.reward_mode}\") print(f\" format_score: {self.format_score}\") print(f\" num_examine: {self.num_examine}\") print(f\" no_search_penalty: {self.no_search_penalty}\") print(f\" turn_free: {self.turn_free}\") print(f\" turn_penalty_per_turn: {self.turn_penalty_per_turn}\") print(f\" no_answer_penalty: {self.no_answer_penalty}\") print(f\" curriculum_step: {self.curriculum_step}\") print(f\" confidence_mode: {self.confidence_mode}\") print(f\" correct_threshold: {self.correct_threshold}\") print(f\" [calib reward] LSC_MODE={LSC_MODE}, FC_MODE={FC_MODE}\") print(f\" eval_only_em: {self.eval_only_em}\") print(\"=\" * 80) def __call__(self, data: DataProto): \"\"\"We will expand this function gradually based on the available datasets\"\"\" # ####debug打印确认confidence tensor信息#### # # 验证 confidence tensor # if 'confidences' in data.batch: # confidence_tensor = data.batch['confidences'] # print(f\"[DEBUG] confidence tensor shape: {confidence_tensor.shape}, dtype: {confidence_tensor.dtype}\") # print(f\"[DEBUG] confidence values: min={confidence_tensor.min():.2f}, max={confidence_tensor.max():.2f}\") # print(f\"[DEBUG] invalid count (-1): {(confidence_tensor == -1).sum().item()}\") # # 验证顺序一致性（可选） # # 比较第一个样本的 confidence 和 responses 内容是否对应 # if self.num_examine > 0: # first_conf = confidence_tensor[0].item() # first_response = self.tokenizer.decode(data.batch['responses'][0]) # print(f\"[DEBUG] Sample 0: confidence={first_conf}, response contains <confidence>: {'<confidence>' in first_response}\") # ###至此### # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn if 'rm_scores' in data.batch.keys(): return data.batch['rm_scores'] reward_tensor = torch.zeros_like(data.batch['responses'], dtype=torch.float32) print(f\"[DEBUG] Reward type in use: {self.reward_type}\") # ========== 修改：从batch取confidence tensor而不是meta_info ========== # confidence存在batch里，shape=(batch_size,)，dtype=long # -1表示格式错误或未输出，1-10表示合法值 if 'confidences' in data.batch: confidence_tensor = data.batch['confidences'] has_confidence = True else: confidence_tensor = None has_confidence = False print(\"[DEBUG] No confidence tensor found in batch\") # ===================================================================== # ========== 新增 ========== reward_statistics = { 'is_correct': [], 'predicted_answers': [], 'ground_truths': [], 'scores': [], 'data_sources': [], 'answer_tag_count': [], 'has_valid_format': [], 'n_search': [], # ========== 新增：confidence统计 ========== 'confidences': [], # 原始confidence值（None表示格式错误） 'confidence_valid': [], # bool，confidence格式是否正确 'confidence_reward': [], # 每个样本的confidence reward分量 'is_overconfident': [], # bool，答错但conf>=7 'is_underconfident': [], # bool，答对但conf<=3 # per-step confidence统计 'per_step_confidences': [], # 每个样本的 per_step_confidences list 'per_step_conf_count': [], # 每个样本有效 per-step conf 数量 'per_step_conf_mean': [], # 每个样本 per-step conf 均值（None 表示无有效值） # ========================================== } # ========== 新增结束 ========== already_print_data_sources = {} for i in range(len(data)): data_item = data[i] 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: 利用info-mask 分开 decode，后续只用 response 做匹配 ===== prompt_str = self.tokenizer.decode(valid_prompt_ids) # response token ids response_ids = data_item.batch['responses'] valid_response_length = int(data_item.batch['attention_mask'][prompt_length:].sum()) valid_response_ids = response_ids[:valid_response_length] # ---- action","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/main_ppo.py#RewardManager.__init__","name":"__init__","kind":"method","path":"verl/trainer/main_ppo.py","span":"L86-L135","signature":"def __init__( self, tokenizer, num_examine, format_score: float = FORMAT_SCORE, reward_type: str = REWARD_TYPE, reward_mode: str = REWARD_MODE, no_search_penalty: float = NO_SEARCH_PENALTY, turn_free: int = TURN_FREE, turn_penalty_per_turn: float = TURN_PENALTY_PER_TURN, no_answer_penalty: float = NO_ANSWER_PENALTY, curriculum_step: int = CURRICULUM_STEP, # confidence 统计相关参数 confidence_mode: str = CONFIDENCE_MODE, correct_threshold: float = CORRECT_THRESHOLD, # 评估模式：final_score 只用 EM（base_score），不叠加任何惩罚 eval_only_em: bool = False, ) -> None","exported":false,"origin":"ast","body_hash":"9ecff2a9d8e1ef031c16b8f6bf30c0602f7d328fd65c43e7533ff60ad7588b4f","body_text":"def __init__( self, tokenizer, num_examine, format_score: float = FORMAT_SCORE, reward_type: str = REWARD_TYPE, reward_mode: str = REWARD_MODE, no_search_penalty: float = NO_SEARCH_PENALTY, turn_free: int = TURN_FREE, turn_penalty_per_turn: float = TURN_PENALTY_PER_TURN, no_answer_penalty: float = NO_ANSWER_PENALTY, curriculum_step: int = CURRICULUM_STEP, # confidence 统计相关参数 confidence_mode: str = CONFIDENCE_MODE, correct_threshold: float = CORRECT_THRESHOLD, # 评估模式：final_score 只用 EM（base_score），不叠加任何惩罚 eval_only_em: bool = False, ) -> None: self.tokenizer = tokenizer self.num_examine = num_examine self.format_score = format_score self.reward_type = reward_type self.reward_mode = reward_mode self.no_search_penalty = float(no_search_penalty) self.turn_free = int(turn_free) self.turn_penalty_per_turn = float(turn_penalty_per_turn) self.no_answer_penalty = float(no_answer_penalty) self.curriculum_step = int(curriculum_step) self.confidence_mode = confidence_mode self.correct_threshold = float(correct_threshold) self.eval_only_em = eval_only_em print(\"=\" * 80) print(\"[REWARD_MANAGER] Initialization:\") print(f\" reward_type: {self.reward_type}\") print(f\" reward_mode: {self.reward_mode}\") print(f\" format_score: {self.format_score}\") print(f\" num_examine: {self.num_examine}\") print(f\" no_search_penalty: {self.no_search_penalty}\") print(f\" turn_free: {self.turn_free}\") print(f\" turn_penalty_per_turn: {self.turn_penalty_per_turn}\") print(f\" no_answer_penalty: {self.no_answer_penalty}\") print(f\" curriculum_step: {self.curriculum_step}\") print(f\" confidence_mode: {self.confidence_mode}\") print(f\" correct_threshold: {self.correct_threshold}\") print(f\" [calib reward] LSC_MODE={LSC_MODE}, FC_MODE={FC_MODE}\") print(f\" eval_only_em: {self.eval_only_em}\") print(\"=\" * 80)","summary_state":"pending","summary":null,"crux":null,"owner":"RewardManager"},{"id":"verl/trainer/main_ppo.py#RewardManager.__call__","name":"__call__","kind":"method","path":"verl/trainer/main_ppo.py","span":"L137-L511","signature":"def __call__(self, data: DataProto)","exported":false,"origin":"ast","body_hash":"785d9d36d1a5bcd34e93e92667059e91fe15d27956258093fadc1a29408dea60","body_text":"def __call__(self, data: DataProto): \"\"\"We will expand this function gradually based on the available datasets\"\"\" # ####debug打印确认confidence tensor信息#### # # 验证 confidence tensor # if 'confidences' in data.batch: # confidence_tensor = data.batch['confidences'] # print(f\"[DEBUG] confidence tensor shape: {confidence_tensor.shape}, dtype: {confidence_tensor.dtype}\") # print(f\"[DEBUG] confidence values: min={confidence_tensor.min():.2f}, max={confidence_tensor.max():.2f}\") # print(f\"[DEBUG] invalid count (-1): {(confidence_tensor == -1).sum().item()}\") # # 验证顺序一致性（可选） # # 比较第一个样本的 confidence 和 responses 内容是否对应 # if self.num_examine > 0: # first_conf = confidence_tensor[0].item() # first_response = self.tokenizer.decode(data.batch['responses'][0]) # print(f\"[DEBUG] Sample 0: confidence={first_conf}, response contains <confidence>: {'<confidence>' in first_response}\") # ###至此### # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn if 'rm_scores' in data.batch.keys(): return data.batch['rm_scores'] reward_tensor = torch.zeros_like(data.batch['responses'], dtype=torch.float32) print(f\"[DEBUG] Reward type in use: {self.reward_type}\") # ========== 修改：从batch取confidence tensor而不是meta_info ========== # confidence存在batch里，shape=(batch_size,)，dtype=long # -1表示格式错误或未输出，1-10表示合法值 if 'confidences' in data.batch: confidence_tensor = data.batch['confidences'] has_confidence = True else: confidence_tensor = None has_confidence = False print(\"[DEBUG] No confidence tensor found in batch\") # ===================================================================== # ========== 新增 ========== reward_statistics = { 'is_correct': [], 'predicted_answers': [], 'ground_truths': [], 'scores': [], 'data_sources': [], 'answer_tag_count': [], 'has_valid_format': [], 'n_search': [], # ========== 新增：confidence统计 ========== 'confidences': [], # 原始confidence值（None表示格式错误） 'confidence_valid': [], # bool，confidence格式是否正确 'confidence_reward': [], # 每个样本的confidence reward分量 'is_overconfident': [], # bool，答错但conf>=7 'is_underconfident': [], # bool，答对但conf<=3 # per-step confidence统计 'per_step_confidences': [], # 每个样本的 per_step_confidences list 'per_step_conf_count': [], # 每个样本有效 per-step conf 数量 'per_step_conf_mean': [], # 每个样本 per-step conf 均值（None 表示无有效值） # ========================================== } # ========== 新增结束 ========== already_print_data_sources = {} for i in range(len(data)): data_item = data[i] 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: 利用info-mask 分开 decode，后续只用 response 做匹配 ===== prompt_str = self.tokenizer.decode(valid_prompt_ids) # response token ids response_ids = data_item.batch['responses'] valid_response_length = int(data_item.batch['attention_mask'][prompt_length:].sum()) valid_response_ids = response_ids[:valid_response_length] # ---- action-only filtering (key!) ---- action_only_ids = valid_response_ids if 'info_mask' in data_item.batch: # info_mask 是整段 [prompt + response] 的 mask valid_info_mask = data_item.batch['info_mask'][prompt_length:prompt_length + valid_response_length] # 经验上：1 = action(token from model), 0 = obs/info action_only_ids = valid_response_ids[valid_info_mask.bool()] response_str = self.tokenizer.decode(action_only_ids) ground_truth = data_item.non_tensor_batch['reward_model']['ground_truth'] data_source = data_item.non_tensor_batch['data_source'] compute_score_fn = _select_rm_score_fn(data_source, reward_type=self.reward_type) # ===== base correctness reward: 只用 response_str ===== base_score = compute_score_fn( solution_str=response_str, ground_truth=ground_truth, format_score=self.format_score ) # ===== 取 global_step（用于 curriculum），ray_trainer 会补；拿不到就 -1 ===== cur_step = -1 if hasattr(data, \"meta_info\") and isinstance(data.meta_info, dict): cur_step = int(data.meta_info.get(\"global_step\", -1)) # ===== n_search（不搜惩罚用） ===== n_search = 0 if hasattr(data, \"meta_info\") and 'valid_search_stats' in data.meta_info: try: n_search = int(data.meta_info['valid_search_stats'][i]) except Exception: n_search = 0 # ===== turn_count（多回合惩罚用） ===== turn_count = 0 if hasattr(data, \"meta_info\") and \"turns_stats\" in data.meta_info: try: turn_count = int(data.meta_info[\"turns_stats\"][i]) except Exception: turn_count = 0 # ===== answer 提取：只从 response_str 里提取 <answer> ===== predicted_answer, answer_tag_count = self._extract_answer_with_format_check(response_str) # ===== GT 文本提取（你原来的安全转换逻辑保留）===== if isinstance(ground_truth, dict) and 'target' in ground_truth: raw_gt = ground_truth['target'] else: raw_gt = ground_truth if isinstance(raw_gt, np.ndarray): gt_text = str(raw_gt.reshape(-1)[0]) if raw_gt.size > 0 else \"\" elif isinstance(raw_gt, (l","summary_state":"pending","summary":null,"crux":null,"owner":"RewardManager"},{"id":"verl/trainer/main_ppo.py#RewardManager._extract_answer_with_format_check","name":"_extract_answer_with_format_check","kind":"method","path":"verl/trainer/main_ppo.py","span":"L514-L522","signature":"def _extract_answer_with_format_check(self, text: str) -> tuple","exported":false,"origin":"ast","body_hash":"48e44f31f8d7a89776acd68b3453509c00c016983d80106008ef0c98e72ab578","body_text":"def _extract_answer_with_format_check(self, text: str) -> tuple: \"\"\"只在 response 文本中提取 <answer>，返回(last_answer, tag_count)\"\"\" answer_pattern = r'<answer>(.*?)</answer>' matches = list(re.finditer(answer_pattern, text or \"\", re.DOTALL)) tag_count = len(matches) if tag_count == 0: return \"\", 0 answer = matches[-1].group(1).strip() return answer, tag_count","summary_state":"pending","summary":null,"crux":null,"owner":"RewardManager"},{"id":"verl/trainer/main_ppo.py#RewardManager._is_no_answer","name":"_is_no_answer","kind":"method","path":"verl/trainer/main_ppo.py","span":"L524-L545","signature":"def _is_no_answer(self, ans: str) -> bool","exported":false,"origin":"ast","body_hash":"ea4e027c7dfeac874ea55ce208d3ead31309b9e331ef50b5452dbbfe005ec2db","body_text":"def _is_no_answer(self, ans: str) -> bool: \"\"\"判定 answer 是否为空/占位（防止骗 answer 惩罚）\"\"\" if ans is None: return True x = str(ans).strip() if x == \"\": return True # 用和 F1 相同的 normalize（更稳） try: x_norm = qa_f1.normalize_answer(x) except Exception: x_norm = x.lower().strip() # 你提到的 and，以及一些常见占位 if x_norm in {\"and\", \"or\", \"the\", \"a\", \"an\", \"none\", \"null\"}: return True # 太短且全是停用占位（可选更严格） toks = x_norm.split() if len(toks) <= 2 and all(t in {\"and\", \"or\", \"the\", \"a\", \"an\"} for t in toks): return True return False","summary_state":"pending","summary":null,"crux":null,"owner":"RewardManager"},{"id":"verl/trainer/main_ppo.py#main","name":"main","kind":"function","path":"verl/trainer/main_ppo.py","span":"L555-L560","signature":"def main(config)","exported":true,"origin":"ast","body_hash":"52aac7b72331acfeae1b5b565b691c27831f20cbad0e528f8c6b31bcbd7d4fa2","body_text":"def main(config): if not ray.is_initialized(): # this is for local ray cluster ray.init(runtime_env={'env_vars': {'TOKENIZERS_PARALLELISM': 'true', 'NCCL_DEBUG': 'WARN'}}) ray.get(main_task.remote(config))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/main_ppo.py#main_task","name":"main_task","kind":"function","path":"verl/trainer/main_ppo.py","span":"L564-L672","signature":"def main_task(config)","exported":true,"origin":"ast","body_hash":"551fd0c0276cdbe1b44022b353d3e3630e8b480f0689dcac430b49911151ef7d","body_text":"def main_task(config): from verl.utils.fs import copy_local_path_from_hdfs from transformers import AutoTokenizer # print initial config from pprint import pprint from omegaconf import OmegaConf pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values OmegaConf.resolve(config) # env_class = ENV_CLASS_MAPPING[config.env.name] # download the checkpoint from hdfs local_path = copy_local_path_from_hdfs(config.actor_rollout_ref.model.path) # instantiate tokenizer from verl.utils import hf_tokenizer tokenizer = hf_tokenizer(local_path) # define worker classes if config.actor_rollout_ref.actor.strategy == 'fsdp': assert config.actor_rollout_ref.actor.strategy == config.critic.strategy from verl.workers.fsdp_workers import ActorRolloutRefWorker, CriticWorker from verl.single_controller.ray import RayWorkerGroup ray_worker_group_cls = RayWorkerGroup elif config.actor_rollout_ref.actor.strategy == 'megatron': assert config.actor_rollout_ref.actor.strategy == config.critic.strategy from verl.workers.megatron_workers import ActorRolloutRefWorker, CriticWorker from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup ray_worker_group_cls = NVMegatronRayWorkerGroup else: raise NotImplementedError from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role role_worker_mapping = { Role.ActorRollout: ray.remote(ActorRolloutRefWorker), Role.Critic: ray.remote(CriticWorker), Role.RefPolicy: ray.remote(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, } # 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: if config.reward_model.strategy == 'fsdp': from verl.workers.fsdp_workers import RewardModelWorker elif config.reward_model.strategy == 'megatron': from verl.workers.megatron_workers import RewardModelWorker else: raise NotImplementedError role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker) mapping[Role.RewardModel] = global_pool_id reward_fn = RewardManager( tokenizer=tokenizer, num_examine=0, reward_type=REWARD_TYPE, reward_mode=REWARD_MODE, no_search_penalty=NO_SEARCH_PENALTY, turn_free=TURN_FREE, turn_penalty_per_turn=TURN_PENALTY_PER_TURN, no_answer_penalty=NO_ANSWER_PENALTY, curriculum_step=CURRICULUM_STEP, confidence_mode=CONFIDENCE_MODE, correct_threshold=CORRECT_THRESHOLD, ) # val 建议只做纯 correctness（便于比较）em得分，方便多项目比较，也可以保持同一 mode 看稳定性 val_reward_fn = RewardManager( tokenizer=tokenizer, num_examine=1, reward_type='em', reward_mode=\"base\", # 推荐：验证只看 base correctness # observe_only：统计confidence但不影响EM分数 confidence_mode=CONFIDENCE_MODE_VAL, correct_threshold=CORRECT_THRESHOLD, # 评估时只用纯 EM 得分，不叠加任何惩罚 eval_only_em=True, # ==================================================== ) resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) 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()","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"torch","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.utils.reward_score","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.trainer.ppo.ray_trainer","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"re","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.trainer.conf_adv_config","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"contains","targetId":"verl/trainer/main_ppo.py#_select_rm_score_fn","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#_select_rm_score_fn","relation":"calls","name":"lower","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#_select_rm_score_fn","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#_select_rm_score_fn","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"contains","targetId":"verl/trainer/main_ppo.py#RewardManager","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager","relation":"contains","targetId":"verl/trainer/main_ppo.py#RewardManager.__init__","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager","relation":"contains","targetId":"verl/trainer/main_ppo.py#RewardManager.__call__","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"keys","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"zeros_like","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"decode","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"bool","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"decode","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"_select_rm_score_fn","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"compute_score_fn","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"_extract_answer_with_format_check","viaMember":true,"file":"verl/trainer/main_ppo.py","recvType":"RewardManager"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"reshape","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"_is_no_answer","viaMember":true,"file":"verl/trainer/main_ppo.py","recvType":"RewardManager"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"max","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"bool","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"bool","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"replace","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"replace","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"tolist","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"histogram","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"compute_ece","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"compute_auroc","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager.__call__","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager","relation":"contains","targetId":"verl/trainer/main_ppo.py#RewardManager._extract_answer_with_format_check","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._extract_answer_with_format_check","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._extract_answer_with_format_check","relation":"calls","name":"finditer","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._extract_answer_with_format_check","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._extract_answer_with_format_check","relation":"calls","name":"strip","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._extract_answer_with_format_check","relation":"calls","name":"group","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager","relation":"contains","targetId":"verl/trainer/main_ppo.py#RewardManager._is_no_answer","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._is_no_answer","relation":"calls","name":"strip","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._is_no_answer","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._is_no_answer","relation":"calls","name":"normalize_answer","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._is_no_answer","relation":"calls","name":"strip","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._is_no_answer","relation":"calls","name":"lower","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._is_no_answer","relation":"calls","name":"split","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._is_no_answer","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#RewardManager._is_no_answer","relation":"calls","name":"all","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"ray","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"hydra","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"calls","name":"main","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"contains","targetId":"verl/trainer/main_ppo.py#main","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main","relation":"calls","name":"init","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main","relation":"calls","name":"remote","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"contains","targetId":"verl/trainer/main_ppo.py#main_task","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.utils.fs","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"transformers","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"pprint","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"omegaconf","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"pprint","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"to_container","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"resolve","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.utils","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.workers.fsdp_workers","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.single_controller.ray","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.workers.megatron_workers","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.single_controller.ray.megatron","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.trainer.ppo.ray_trainer","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"remote","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"remote","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"remote","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.workers.fsdp_workers","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py","relation":"imports","specifier":"verl.workers.megatron_workers","file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"remote","viaMember":true,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"RewardManager","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"RewardManager","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"ResourcePoolManager","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"RayPPOTrainer","viaMember":false,"file":"verl/trainer/main_ppo.py"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"init_workers","viaMember":true,"file":"verl/trainer/main_ppo.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/main_ppo.py#main_task","relation":"calls","name":"fit","viaMember":true,"file":"verl/trainer/main_ppo.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/main_ppo.py","relation":"calls","name":"main","viaMember":false,"file":"verl/trainer/main_ppo.py"}]},"verl/trainer/ppo/__init__.py":{"size":600,"mtimeMs":1789275830320.203,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/trainer/ppo/__init__.py","name":"__init__.py","kind":"file","path":"verl/trainer/ppo/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/trainer/ppo/core_algos.py":{"size":19950,"mtimeMs":1789275830320.203,"hash":"168edc02eaf41b3726f38d0888532448eb4f53463ac79c125544a99208e7d0dd","nodes":[{"id":"verl/trainer/ppo/core_algos.py","name":"core_algos.py","kind":"file","path":"verl/trainer/ppo/core_algos.py","span":"L1-L558","signature":null,"exported":true,"origin":"ast","body_hash":"168edc02eaf41b3726f38d0888532448eb4f53463ac79c125544a99208e7d0dd","chars":19362,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2022 The HuggingFace Team. 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. \"\"\" Core functions to implement PPO algorithms. The function implemented in this file should be used by trainer with different distributed strategies to implement PPO \"\"\" import numpy as np import torch from collections import defaultdict, Counter import verl.utils.torch_functional as verl_F # NOTE(sgm): this implementation only consider outcome supervision, where the reward is a scalar."},{"id":"verl/trainer/ppo/core_algos.py#AdaptiveKLController","name":"AdaptiveKLController","kind":"class","path":"verl/trainer/ppo/core_algos.py","span":"L28-L43","signature":"class AdaptiveKLController","exported":true,"origin":"ast","body_hash":"1d19d44958312782745ca43b4b9353e5a293db2116237a81e7407cbd4f495359","body_text":"class AdaptiveKLController: \"\"\" Adaptive KL controller described in the paper: https://arxiv.org/pdf/1909.08593.pdf \"\"\" def __init__(self, init_kl_coef, target_kl, horizon): self.value = init_kl_coef self.target = target_kl self.horizon = horizon def update(self, current_kl, n_steps): target = self.target proportional_error = np.clip(current_kl / target - 1, -0.2, 0.2) mult = 1 + proportional_error * n_steps / self.horizon self.value *= mult","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#AdaptiveKLController.__init__","name":"__init__","kind":"method","path":"verl/trainer/ppo/core_algos.py","span":"L34-L37","signature":"def __init__(self, init_kl_coef, target_kl, horizon)","exported":false,"origin":"ast","body_hash":"d420ba46ecd2936ce45bb3cc59d8bfe3847fcb2f83a7c167b60c22ccba17daf6","body_text":"def __init__(self, init_kl_coef, target_kl, horizon): self.value = init_kl_coef self.target = target_kl self.horizon = horizon","summary_state":"pending","summary":null,"crux":null,"owner":"AdaptiveKLController"},{"id":"verl/trainer/ppo/core_algos.py#AdaptiveKLController.update","name":"update","kind":"method","path":"verl/trainer/ppo/core_algos.py","span":"L39-L43","signature":"def update(self, current_kl, n_steps)","exported":true,"origin":"ast","body_hash":"878aa3715b7b9b7afa6d98db90f231c43897ac58fc59942cd94931ce6ac07ba4","body_text":"def update(self, current_kl, n_steps): target = self.target proportional_error = np.clip(current_kl / target - 1, -0.2, 0.2) mult = 1 + proportional_error * n_steps / self.horizon self.value *= mult","summary_state":"pending","summary":null,"crux":null,"owner":"AdaptiveKLController"},{"id":"verl/trainer/ppo/core_algos.py#FixedKLController","name":"FixedKLController","kind":"class","path":"verl/trainer/ppo/core_algos.py","span":"L46-L53","signature":"class FixedKLController","exported":true,"origin":"ast","body_hash":"0ddd42812b3266dffd12aa64f5be8d1a1aab907a9489db2dd1d0f29a0fdbd145","body_text":"class FixedKLController: \"\"\"Fixed KL controller.\"\"\" def __init__(self, kl_coef): self.value = kl_coef def update(self, current_kl, n_steps): pass","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#FixedKLController.__init__","name":"__init__","kind":"method","path":"verl/trainer/ppo/core_algos.py","span":"L49-L50","signature":"def __init__(self, kl_coef)","exported":false,"origin":"ast","body_hash":"166e7784943101bd8cabe062e4054e91e20ee6a3e5f47b58708a7a7a34380f14","body_text":"def __init__(self, kl_coef): self.value = kl_coef","summary_state":"pending","summary":null,"crux":null,"owner":"FixedKLController"},{"id":"verl/trainer/ppo/core_algos.py#FixedKLController.update","name":"update","kind":"method","path":"verl/trainer/ppo/core_algos.py","span":"L52-L53","signature":"def update(self, current_kl, n_steps)","exported":true,"origin":"ast","body_hash":"95600a371b1f1262d9daeab28dee95c7eeb976983d84eb8b6d9310ced3b881dc","body_text":"def update(self, current_kl, n_steps): pass","summary_state":"pending","summary":null,"crux":null,"owner":"FixedKLController"},{"id":"verl/trainer/ppo/core_algos.py#get_kl_controller","name":"get_kl_controller","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L56-L67","signature":"def get_kl_controller(config): # seems never used?","exported":true,"origin":"ast","body_hash":"74c8212f5c454d575daf3aea435fc36d5f09f88b6618ae47ffd9f4cdac87a179","body_text":"def get_kl_controller(config): # seems never used? if config.critic.kl_ctrl.type == 'fixed': kl_ctrl = FixedKLController(kl_coef=config.critic.kl_ctrl.kl_coef) elif config.critic.kl_ctrl.type == 'adaptive': assert config.kl_ctrl.horizon > 0, f'horizon must be larger than 0. Got {config.critic.kl_ctrl.horizon}' kl_ctrl = AdaptiveKLController(init_kl_coef=config.critic.kl_ctrl.kl_coef, target_kl=config.critic.kl_ctrl.target_kl, horizon=config.critic.kl_ctrl.horizon) else: raise ValueError('Unknown kl_ctrl type') return kl_ctrl","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_gae_advantage_return","name":"compute_gae_advantage_return","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L70-L107","signature":"def compute_gae_advantage_return(token_level_rewards: torch.Tensor, values: torch.Tensor, eos_mask: torch.Tensor, gamma: torch.Tensor, lam: torch.Tensor)","exported":true,"origin":"ast","body_hash":"c9360b45b39908224fe4e87b1c3e4c19aa3274751f4d01749af71824fe2a25ce","body_text":"def compute_gae_advantage_return(token_level_rewards: torch.Tensor, values: torch.Tensor, eos_mask: torch.Tensor, gamma: torch.Tensor, lam: torch.Tensor): \"\"\"Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py Args: token_level_rewards: `(torch.Tensor)` shape: (bs, response_length) values: `(torch.Tensor)` shape: (bs, response_length) eos_mask: `(torch.Tensor)` shape: (bs, response_length). [EOS] mask. The token after [EOS] have mask zero. gamma: `(float)` discounted factor used in RL lam: `(float)` lambda value when computing Generalized Advantage Estimation (https://arxiv.org/abs/1506.02438) Returns: advantages: `(torch.Tensor)` shape: (bs, response_length) Returns: `(torch.Tensor)` shape: (bs, response_length) \"\"\" with torch.no_grad(): lastgaelam = 0 advantages_reversed = [] gen_len = token_level_rewards.shape[-1] for t in reversed(range(gen_len)): nextvalues = values[:, t + 1] if t < gen_len - 1 else 0.0 delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t] lastgaelam = delta + gamma * lam * lastgaelam advantages_reversed.append(lastgaelam) advantages = torch.stack(advantages_reversed[::-1], dim=1) returns = advantages + values advantages = verl_F.masked_whiten(advantages, eos_mask) return advantages, returns","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","name":"compute_grpo_outcome_advantage","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L111-L155","signature":"def compute_grpo_outcome_advantage(token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, index: torch.Tensor, epsilon: float = 1e-6)","exported":true,"origin":"ast","body_hash":"8a370d111cc20f18c0eb1799ae8e2c7aa616015d1f4da94ab689fe799e8b4729","body_text":"def compute_grpo_outcome_advantage(token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, index: torch.Tensor, epsilon: float = 1e-6): \"\"\" Compute advantage for GRPO, operating only on Outcome reward (with only one scalar reward for each response). Args: token_level_rewards: `(torch.Tensor)` shape: (bs, response_length) eos_mask: `(torch.Tensor)` shape: (bs, response_length) Returns: advantages: `(torch.Tensor)` shape: (bs, response_length) Returns: `(torch.Tensor)` shape: (bs, response_length) \"\"\" response_length = token_level_rewards.shape[-1] non_zero_mask = (token_level_rewards != 0) scores = (token_level_rewards * non_zero_mask).sum(dim=-1) id2score = defaultdict(list) id2mean = {} id2std = {} with torch.no_grad(): bsz = scores.shape[0] for i in range(bsz): id2score[index[i]].append(scores[i]) for idx in id2score: if len(id2score[idx]) == 1: id2mean[idx] = torch.tensor(0.0) id2std[idx] = torch.tensor(1.0) elif len(id2score[idx]) > 1: id2mean[idx] = torch.mean(torch.tensor(id2score[idx])) id2std[idx] = torch.std(torch.tensor([id2score[idx]])) else: raise ValueError(f\"no score in prompt index: {idx}\") for i in range(bsz): scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon) scores = scores.unsqueeze(-1).tile([1, response_length]) * eos_mask return scores, scores","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_reinforce_outcome_advantage","name":"compute_reinforce_outcome_advantage","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L157-L164","signature":"def compute_reinforce_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, )","exported":true,"origin":"ast","body_hash":"b2cd20a8800dc1995aa052421d0feb16e7eda882f9fe68139e1da3c8da3ae5b7","body_text":"def compute_reinforce_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, ): scores = token_level_rewards.sum(dim=-1) advantages = scores.unsqueeze(-1) * response_mask returns = advantages return advantages, returns","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","name":"_uid_check_and_print","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L166-L195","signature":"def _uid_check_and_print(index, expect_n=None, strict=False, prefix=\"[reinforce_mean]\", max_show=10)","exported":false,"origin":"ast","body_hash":"686f84d9b7742320350549ab9f06a2490d56bcbb4e7bb8fc5fa4bf060113a674","body_text":"def _uid_check_and_print(index, expect_n=None, strict=False, prefix=\"[reinforce_mean]\", max_show=10): \"\"\" index: list/np.array/torch tensor of uid (len = batch_size) expect_n: 期望每个 uid 出现次数（通常 = n_agent） strict: True 时如果不等于 expect_n 直接报错 \"\"\" uid_list = list(index) c = Counter(uid_list) counts = list(c.values()) count_min = min(counts) if counts else 0 count_max = max(counts) if counts else 0 # 统计分布（前几个） dist = Counter(counts) top_counts = dist.most_common(5) print(f\"{prefix}[uid_check] unique_uids={len(c)} batch={len(uid_list)} \" f\"count_min={count_min} count_max={count_max} top_counts={top_counts}\") if expect_n is not None: bad = {k: v for k, v in c.items() if v != expect_n} if bad: sample_bad = list(bad.items())[:max_show] msg = (f\"{prefix}[uid_check] group size mismatch: expect {expect_n}, \" f\"bad_groups(first {max_show})={sample_bad}, unique_uids={len(c)}, batch={len(uid_list)}\") if strict: raise RuntimeError(msg) else: print(msg)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","name":"compute_reinforce_mean_outcome_advantage","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L197-L247","signature":"def compute_reinforce_mean_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, index, expect_n: int = None, strict: bool = False, print_check: bool = False, )","exported":true,"origin":"ast","body_hash":"251400d84deca43f8c60251617910078317fde74fd511f9b168b12676bd7674e","body_text":"def compute_reinforce_mean_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, index, expect_n: int = None, strict: bool = False, print_check: bool = False, ): \"\"\" REINFORCE + 组内 baseline（减去同 uid 的均值） token_level_rewards: (bs, response_len) response_mask: (bs, response_len) 0/1 mask（EOS 后为 0） index: uid 列表/数组（len = bs） 返回： advantages: (bs, response_len) 每条样本一个标量 advantage，复制到 token 维度并乘 mask returns: (bs, response_len) 这里不训练 critic，所以 returns = advantages（保持接口一致） \"\"\" if print_check or strict: _uid_check_and_print(index, expect_n=expect_n, strict=strict, prefix=\"[reinforce_mean]\") # 1) 每条样本的总回报 R_i（outcome reward） scores = token_level_rewards.sum(dim=-1) # (bs,) # 2) 组内 baseline: mu_uid # 用 python 分组写法，和你 grpo 的风格一致，最稳妥 uid_list = list(index) id2scores = defaultdict(list) for i, uid in enumerate(uid_list): id2scores[uid].append(scores[i]) id2mean = {} with torch.no_grad(): for uid, s_list in id2scores.items(): if len(s_list) == 0: id2mean[uid] = torch.tensor(0.0, device=scores.device, dtype=scores.dtype) else: # s_list 是一堆 0-d tensor，stack 更稳 id2mean[uid] = torch.stack(s_list).mean() # 3) A_i = R_i - mu_uid advantages_scalar = scores.clone() with torch.no_grad(): for i, uid in enumerate(uid_list): advantages_scalar[i] = advantages_scalar[i] - id2mean[uid] # 4) 复制到 token 维度，并 mask advantages = advantages_scalar.unsqueeze(-1) * response_mask # (bs, response_len) returns = advantages return advantages, returns","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","name":"_uid_group_mean_and_count","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L250-L278","signature":"def _uid_group_mean_and_count(scores: torch.Tensor, uid: torch.Tensor)","exported":false,"origin":"ast","body_hash":"15cf89d3d47f2f3df41c653551f41b8f4add065af83e4b60a633c29f4c829fd9","body_text":"def _uid_group_mean_and_count(scores: torch.Tensor, uid: torch.Tensor): \"\"\" scores: (bs,) float/0-1 uid: (bs,) int/long returns: group_mean_per_sample: (bs,) 每条样本所属uid的组均值 uid2count: dict(uid -> count) 用于debug \"\"\" uid2vals = defaultdict(list) bsz = scores.shape[0] # 注意：uid 可能是 tensor，也可能是 numpy/int，统一转成 python int 做 key 更稳 for i in range(bsz): key = int(uid[i]) uid2vals[key].append(scores[i]) uid2mean = {} uid2count = {} for key, vals in uid2vals.items(): # vals 是 list[tensor]，stack 后 mean v = torch.stack(vals) uid2mean[key] = v.mean() uid2count[key] = len(vals) group_mean = torch.empty_like(scores) for i in range(bsz): key = int(uid[i]) group_mean[i] = uid2mean[key] return group_mean, uid2count","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","name":"compute_reinforce_da_outcome_advantage","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L281-L320","signature":"def compute_reinforce_da_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, uid: torch.Tensor, w_min: float = 0.5, w_max: float = 1.0, )","exported":true,"origin":"ast","body_hash":"afba4e94d05af8d72483c8542f586cfb22260e7d7c80e3f5f0e25f85764e6c08","body_text":"def compute_reinforce_da_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, uid: torch.Tensor, w_min: float = 0.5, w_max: float = 1.0, ): \"\"\" REINFORCE-DA: scores = sum(token_level_rewards) (outcome reward: 0/1) group_mean = mean(scores within same uid) w = w_min + (w_max-w_min)*(1-group_mean) A_{t} = w * scores * mask_t Returns: advantages: (bs, response_len) returns: same as advantages (actor-only) debug: dict with uid2count / stats \"\"\" # outcome scalar scores = token_level_rewards.sum(dim=-1) # (bs,) with torch.no_grad(): group_mean, uid2count = _uid_group_mean_and_count(scores, uid) # w in [w_min, w_max] w = w_min + (w_max - w_min) * (1.0 - group_mean) # 数值安全：可选 clamp w = torch.clamp(w, min=min(w_min, w_max), max=max(w_min, w_max)) advantages = (w * scores).unsqueeze(-1) * response_mask returns = advantages debug = { \"uid2count\": uid2count, \"scores_mean\": float(scores.mean().item()), \"group_mean_mean\": float(group_mean.mean().item()), \"w_min\": float(w.min().item()), \"w_max\": float(w.max().item()), \"w_mean\": float(w.mean().item()), } return advantages, returns, debug","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","name":"compute_reinforce_mean_da_outcome_advantage","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L323-L355","signature":"def compute_reinforce_mean_da_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, uid: torch.Tensor, w_min: float = 0.5, w_max: float = 1.0, )","exported":true,"origin":"ast","body_hash":"4c706527ddc8ba07dd3a8d29be4252d0bbe152301553bc96d3355e83bbcb3f0d","body_text":"def compute_reinforce_mean_da_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, uid: torch.Tensor, w_min: float = 0.5, w_max: float = 1.0, ): \"\"\" REINFORCE-MEAN-DA: scores = outcome group_mean = mean(scores within uid) w = linear difficulty weight in [w_min,w_max] A_t = w * (scores - group_mean) * mask_t \"\"\" scores = token_level_rewards.sum(dim=-1) # (bs,) with torch.no_grad(): group_mean, uid2count = _uid_group_mean_and_count(scores, uid) w = w_min + (w_max - w_min) * (1.0 - group_mean) w = torch.clamp(w, min=min(w_min, w_max), max=max(w_min, w_max)) advantages = (w * (scores - group_mean)).unsqueeze(-1) * response_mask returns = advantages debug = { \"uid2count\": uid2count, \"scores_mean\": float(scores.mean().item()), \"group_mean_mean\": float(group_mean.mean().item()), \"w_min\": float(w.min().item()), \"w_max\": float(w.max().item()), \"w_mean\": float(w.mean().item()), } return advantages, returns, debug","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","name":"compute_grpo_da_outcome_advantage","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L357-L436","signature":"def compute_grpo_da_outcome_advantage( token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, uid, # 参数名改为 uid w_min: float = 0.5, w_max: float = 1.0, epsilon: float = 1e-6, expect_n: int = None, strict: bool = False, print_check: bool = False, )","exported":true,"origin":"ast","body_hash":"ed32fcfb42b120f1484d862f823cd7dda3911a9c6974bad452858bc7151f38e0","body_text":"def compute_grpo_da_outcome_advantage( token_level_rewards: torch.Tensor, eos_mask: torch.Tensor, uid, # 参数名改为 uid w_min: float = 0.5, w_max: float = 1.0, epsilon: float = 1e-6, expect_n: int = None, strict: bool = False, print_check: bool = False, ): \"\"\" GRPO-DA: 1) outcome score: R_i = sum(token_level_rewards_i) 2) GRPO组内标准化: z_i = (R_i - mean_uid) / (std_uid + eps) 3) 样本难度权重（按uid）：w_uid in [w_min, w_max]，用该uid组内平均回报(mean_uid)映射 4) A_i = w_uid * z_i （注意：必须在标准化之后乘，否则会被抵消） 5) 复制到 token 维度，并乘 eos_mask \"\"\" response_length = token_level_rewards.shape[-1] non_zero_mask = (token_level_rewards != 0) scores = (token_level_rewards * non_zero_mask).sum(dim=-1) # (bs,) # --- uid check（可选） --- if print_check or strict: _uid_check_and_print(uid, expect_n=expect_n, strict=strict, prefix=\"[grpo-da]\", max_show=10) uid_list = list(uid) # --- 分组收集每个uid的 scores --- id2scores = defaultdict(list) for i, uid_val in enumerate(uid_list): # 改：uid -> uid_val，避免覆盖参数 id2scores[uid_val].append(scores[i]) id2mean = {} id2std = {} with torch.no_grad(): for uid_key, s_list in id2scores.items(): # 改：uid -> uid_key s = torch.stack(s_list) mu = s.mean() std = s.std(unbiased=False) if torch.isnan(std) or std.item() < epsilon: std = torch.tensor(1.0, device=scores.device, dtype=scores.dtype) id2mean[uid_key] = mu id2std[uid_key] = std # 1) 先做 GRPO z-score z = scores.clone() for i, uid_val in enumerate(uid_list): # 改：uid -> uid_val z[i] = (scores[i] - id2mean[uid_val]) / (id2std[uid_val] + epsilon) # 2) 再做 difficulty weight（按 uid） # 你目前 reward 是 EM(0/1)，mean_uid 就是这个 prompt 的\"成功率\" # mean 越低 -> 越难 -> 权重越接近 w_max w = torch.empty_like(scores) for i, uid_val in enumerate(uid_list): # 改：uid -> uid_val m = id2mean[uid_val] wi = w_min + (w_max - w_min) * (1.0 - m) wi = torch.clamp(wi, min=min(w_min, w_max), max=max(w_min, w_max)) w[i] = wi advantages_scalar = w * z # 关键：在标准化之后乘 advantages = advantages_scalar.unsqueeze(-1).expand(-1, response_length) * eos_mask returns = advantages debug = { \"scores_mean\": float(scores.mean().item()), \"group_mean_mean\": float(torch.stack([id2mean[u] for u in id2mean]).mean().item()) if len(id2mean) > 0 else 0.0, \"w_min\": float(w.min().item()), \"w_max\": float(w.max().item()), \"w_mean\": float(w.mean().item()), } if print_check: print(f\"[grpo-da][debug] scores_mean={debug['scores_mean']:.4f} \" f\"group_mean_mean={debug['group_mean_mean']:.4f} \" f\"w_min={debug['w_min']:.4f} w_max={debug['w_max']:.4f} w_mean={debug['w_mean']:.4f}\") return advantages, returns, debug","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_rewards","name":"compute_rewards","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L441-L443","signature":"def compute_rewards(token_level_scores, old_log_prob, ref_log_prob, kl_ratio)","exported":true,"origin":"ast","body_hash":"637529062e50b00dd696b7b82127db44623909c1c91604dfe40a2fec4a1e09ae","body_text":"def compute_rewards(token_level_scores, old_log_prob, ref_log_prob, kl_ratio): kl = old_log_prob - ref_log_prob return token_level_scores - kl * kl_ratio","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_policy_loss","name":"compute_policy_loss","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L446-L477","signature":"def compute_policy_loss(old_log_prob, log_prob, advantages, eos_mask, cliprange)","exported":true,"origin":"ast","body_hash":"1a982f4cafdca2b2c497b45256647575a0da22a2fe11e703a18f4f01f783056e","body_text":"def compute_policy_loss(old_log_prob, log_prob, advantages, eos_mask, cliprange): \"\"\"Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1122 Args: old_log_prob: `(torch.Tensor)` shape: (bs, response_length) log_prob: `(torch.Tensor)` shape: (bs, response_length) advantages: `(torch.Tensor)` shape: (bs, response_length) eos_mask: `(torch.Tensor)` shape: (bs, response_length) cliprange: (float) The clip range used in PPO. See https://arxiv.org/abs/1707.06347 Returns: pg_loss: `a scalar torch.Tensor` policy gradient loss computed via PPO pg_clipfrac: (float) a float number indicating the fraction of policy gradient loss being clipped \"\"\" negative_approx_kl = log_prob - old_log_prob ratio = torch.exp(negative_approx_kl) ppo_kl = verl_F.masked_mean(-negative_approx_kl, eos_mask) pg_losses = -advantages * ratio pg_losses2 = -advantages * torch.clamp(ratio, 1.0 - cliprange, 1.0 + cliprange) pg_loss = verl_F.masked_mean(torch.max(pg_losses, pg_losses2), eos_mask) pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses).float(), eos_mask) return pg_loss, pg_clipfrac, ppo_kl","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_entropy_loss","name":"compute_entropy_loss","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L480-L496","signature":"def compute_entropy_loss(logits, eos_mask)","exported":true,"origin":"ast","body_hash":"02d5d5ce39bd43f5d64925a07924c59c91db99d975e2b7c78f61e6b86c4be908","body_text":"def compute_entropy_loss(logits, eos_mask): \"\"\"Compute Categorical entropy loss Args: logits: `(torch.Tensor)` shape: (bs, response_length, vocab_size) eos_mask: `(torch.Tensor)` shape: (bs, response_length) Returns: entropy: a scalar torch.Tensor \"\"\" # compute entropy entropy = verl_F.entropy_from_logits(logits) # (bs, response_len) entropy_loss = verl_F.masked_mean(entropy, mask=eos_mask) return entropy_loss","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#compute_value_loss","name":"compute_value_loss","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L499-L522","signature":"def compute_value_loss(vpreds, returns, values, eos_mask, cliprange_value)","exported":true,"origin":"ast","body_hash":"bf58d543108e51510c3a78e381599138d42768b4f2980e95bff79ddc53df644a","body_text":"def compute_value_loss(vpreds, returns, values, eos_mask, cliprange_value): \"\"\"Compute the value loss. Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1151 Args: vpreds (`torch.FloatTensor`): Predicted values of the value head, shape (`batch_size`, `response_length`) values (`torch.FloatTensor`): Old values of value head, shape (`batch_size`, `response_length`) returns: (`torch.FloatTensor`): Ground truth returns, shape (`batch_size`, `response_length`) Returns: vf_loss: a scalar (`torch.FloatTensor`): value function loss vf_clipfrac: a float The ratio of vf being clipped \"\"\" vpredclipped = verl_F.clip_by_value(vpreds, values - cliprange_value, values + cliprange_value) vf_losses1 = (vpreds - returns)**2 vf_losses2 = (vpredclipped - returns)**2 vf_loss = 0.5 * verl_F.masked_mean(torch.max(vf_losses1, vf_losses2), eos_mask) vf_clipfrac = verl_F.masked_mean(torch.gt(vf_losses2, vf_losses1).float(), eos_mask) return vf_loss, vf_clipfrac","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/core_algos.py#kl_penalty","name":"kl_penalty","kind":"function","path":"verl/trainer/ppo/core_algos.py","span":"L525-L557","signature":"def kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty) -> torch.FloatTensor","exported":true,"origin":"ast","body_hash":"494f2e37c5b0c30d41c7d4945e33eaaf01b1fa291400f38ed7d8b032638b68cf","body_text":"def kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty) -> torch.FloatTensor: \"\"\"Compute KL divergence given logprob and ref_logprob. Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1104 Args: logprob: ref_logprob: Returns: \"\"\" if kl_penalty == \"kl\": return logprob - ref_logprob if kl_penalty == \"abs\": return (logprob - ref_logprob).abs() if kl_penalty == \"mse\": return 0.5 * (logprob - ref_logprob).square() # J. Schulman. Approximating kl divergence, 2020. # # URL http://joschu.net/blog/kl-approx.html. if kl_penalty == 'low_var_kl': kl = ref_logprob - logprob ratio = torch.exp(kl) kld = (ratio - kl - 1).contiguous() return torch.clamp(kld, min=-10, max=10) if kl_penalty == \"full\": # so, here logprob and ref_logprob should contain the logits for every token in vocabulary raise NotImplementedError raise NotImplementedError","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/trainer/ppo/core_algos.py","relation":"imports","specifier":"torch","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"imports","specifier":"collections","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#AdaptiveKLController","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#AdaptiveKLController","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#AdaptiveKLController.__init__","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#AdaptiveKLController","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#AdaptiveKLController.update","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#AdaptiveKLController.update","relation":"calls","name":"clip","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#FixedKLController","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#FixedKLController","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#FixedKLController.__init__","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#FixedKLController","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#FixedKLController.update","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#get_kl_controller","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#get_kl_controller","relation":"calls","name":"FixedKLController","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#get_kl_controller","relation":"calls","name":"AdaptiveKLController","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#get_kl_controller","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_gae_advantage_return","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_gae_advantage_return","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_gae_advantage_return","relation":"calls","name":"reversed","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_gae_advantage_return","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_gae_advantage_return","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_gae_advantage_return","relation":"calls","name":"stack","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_gae_advantage_return","relation":"calls","name":"masked_whiten","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"defaultdict","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"tensor","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"tensor","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"tensor","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"std","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"tensor","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"tile","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_outcome_advantage","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_reinforce_outcome_advantage","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_outcome_advantage","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_outcome_advantage","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"Counter","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"values","viaMember":true,"file":"verl/trainer/ppo/core_algos.py","recvType":"Counter"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"min","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"max","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"Counter","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"most_common","viaMember":true,"file":"verl/trainer/ppo/core_algos.py","recvType":"Counter"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/core_algos.py","recvType":"Counter"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_check_and_print","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"_uid_check_and_print","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"defaultdict","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/core_algos.py","recvType":"defaultdict"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"tensor","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"stack","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"clone","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_outcome_advantage","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"defaultdict","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/core_algos.py","recvType":"defaultdict"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"stack","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#_uid_group_mean_and_count","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"_uid_group_mean_and_count","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"clamp","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"min","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"max","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_da_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"_uid_group_mean_and_count","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"clamp","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"min","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"max","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_reinforce_mean_da_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"_uid_check_and_print","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"defaultdict","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/core_algos.py","recvType":"defaultdict"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"stack","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"std","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"isnan","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"tensor","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"clone","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"clamp","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"min","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"max","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"expand","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"stack","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_grpo_da_outcome_advantage","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_rewards","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_policy_loss","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_policy_loss","relation":"calls","name":"exp","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_policy_loss","relation":"calls","name":"masked_mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_policy_loss","relation":"calls","name":"clamp","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_policy_loss","relation":"calls","name":"masked_mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_policy_loss","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_policy_loss","relation":"calls","name":"masked_mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_policy_loss","relation":"calls","name":"float","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_policy_loss","relation":"calls","name":"gt","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_entropy_loss","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_entropy_loss","relation":"calls","name":"entropy_from_logits","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_entropy_loss","relation":"calls","name":"masked_mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#compute_value_loss","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_value_loss","relation":"calls","name":"clip_by_value","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_value_loss","relation":"calls","name":"masked_mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_value_loss","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_value_loss","relation":"calls","name":"masked_mean","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_value_loss","relation":"calls","name":"float","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#compute_value_loss","relation":"calls","name":"gt","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py","relation":"contains","targetId":"verl/trainer/ppo/core_algos.py#kl_penalty","file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#kl_penalty","relation":"calls","name":"abs","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#kl_penalty","relation":"calls","name":"square","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#kl_penalty","relation":"calls","name":"exp","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#kl_penalty","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"},{"source":"verl/trainer/ppo/core_algos.py#kl_penalty","relation":"calls","name":"clamp","viaMember":true,"file":"verl/trainer/ppo/core_algos.py"}]},"verl/trainer/ppo/ray_trainer.py":{"size":70845,"mtimeMs":1789275830321.203,"hash":"551cb0f6ba1f27649a83d838b4664adec88e3335685128a617511ac23c816842","nodes":[{"id":"verl/trainer/ppo/ray_trainer.py","name":"ray_trainer.py","kind":"file","path":"verl/trainer/ppo/ray_trainer.py","span":"L1-L1408","signature":null,"exported":true,"origin":"ast","body_hash":"551cb0f6ba1f27649a83d838b4664adec88e3335685128a617511ac23c816842","chars":69473,"summary_state":"pending","summary":null,"crux":null,"body_text":"# ray_trainer-try2.py #修改：参数温度等的传递；ece从main中导入 # 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. \"\"\" FSDP PPO Trainer with Ray-based single controller. This trainer supports model-agonistic model initialization with huggingface \"\"\" import os import uuid from contextlib import contextmanager from dataclasses import dataclass, field from enum import Enum from pprint import pprint from typing import Type, Dict, Optional import re import json from collections import defaultdict import numpy as np from codetiming import Timer from omegaconf import OmegaConf, open_dict from verl import DataProto from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto from verl.single_controller.base import Worker from verl.single_controller.ray import RayResourcePool, RayWorkerGroup, RayClassWithInitArgs from verl.single_controller.ray.base import create_colocated_worker_cls from verl.trainer.ppo import core_algos from verl.utils.seqlen_balancing import get_seqlen_balanced_partitions, log_seqlen_unbalance import re from search_r1.llm_agent.generation import LLMGenerationManager, GenerationConfig from verl.trainer.conf_adv_config import ( LSC_MODE, FC_MODE, compute_ece, compute_auroc, apply_dual_calib_to_scores, ) WorkerType = Type[Worker] @dataclass import torch from verl.utils.torch_functional import masked_mean @contextmanager"},{"id":"verl/trainer/ppo/ray_trainer.py#Role","name":"Role","kind":"class","path":"verl/trainer/ppo/ray_trainer.py","span":"L57-L67","signature":"class Role(Enum)","exported":true,"origin":"ast","body_hash":"e8ff555baf3aba9f706b87b759f90301b9cf44b14004d4ec68d85f86a1f62170","body_text":"class Role(Enum): \"\"\" To create more roles dynamically, you can subclass Role and add new members \"\"\" Actor = 0 Rollout = 1 ActorRollout = 2 Critic = 3 RefPolicy = 4 RewardModel = 5 ActorRolloutRef = 6","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager","name":"ResourcePoolManager","kind":"class","path":"verl/trainer/ppo/ray_trainer.py","span":"L71-L93","signature":"class ResourcePoolManager","exported":true,"origin":"ast","body_hash":"9795550a68750ab78c2b8438ec1549c7b670aa48a540712d7264c4110fe74f63","body_text":"class ResourcePoolManager: \"\"\" Define a resource pool specification. Resource pool will be initialized first. Mapping \"\"\" resource_pool_spec: dict[str, list[int]] mapping: dict[Role, str] resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict) def create_resource_pool(self): for resource_pool_name, process_on_nodes in self.resource_pool_spec.items(): # 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=process_on_nodes, use_gpu=True, max_colocate_count=1, name_prefix=resource_pool_name) self.resource_pool_dict[resource_pool_name] = resource_pool def get_resource_pool(self, role: Role) -> RayResourcePool: \"\"\"Get the resource pool of the worker_cls\"\"\" return self.resource_pool_dict[self.mapping[role]]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager.create_resource_pool","name":"create_resource_pool","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L80-L89","signature":"def create_resource_pool(self)","exported":true,"origin":"ast","body_hash":"f91f8f0b67b02c470bcbb708caf1764ae9a6685aeffd2469f1f8467bf4629f8c","body_text":"def create_resource_pool(self): for resource_pool_name, process_on_nodes in self.resource_pool_spec.items(): # 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=process_on_nodes, use_gpu=True, max_colocate_count=1, name_prefix=resource_pool_name) self.resource_pool_dict[resource_pool_name] = resource_pool","summary_state":"pending","summary":null,"crux":null,"owner":"ResourcePoolManager"},{"id":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager.get_resource_pool","name":"get_resource_pool","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L91-L93","signature":"def get_resource_pool(self, role: Role) -> RayResourcePool","exported":true,"origin":"ast","body_hash":"0085abda986abb99b666ed3d66a8a791b7ba30887f0369a9ed1508ae86a5d8c4","body_text":"def get_resource_pool(self, role: Role) -> RayResourcePool: \"\"\"Get the resource pool of the worker_cls\"\"\" return self.resource_pool_dict[self.mapping[role]]","summary_state":"pending","summary":null,"crux":null,"owner":"ResourcePoolManager"},{"id":"verl/trainer/ppo/ray_trainer.py#apply_kl_penalty","name":"apply_kl_penalty","kind":"function","path":"verl/trainer/ppo/ray_trainer.py","span":"L100-L129","signature":"def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty='kl')","exported":true,"origin":"ast","body_hash":"40461bfc8e49de82f35dc404142035e972ccfeec69468208efcd7ba8b7a2a1f6","body_text":"def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty='kl'): responses = data.batch['responses'] response_length = responses.size(1) token_level_scores = data.batch['token_level_scores'] batch_size = data.batch.batch_size[0] attention_mask = data.batch['info_mask'] response_mask = attention_mask[:, -response_length:] # compute kl between ref_policy and current policy if 'ref_log_prob' in data.batch.keys(): kld = core_algos.kl_penalty(data.batch['old_log_probs'], data.batch['ref_log_prob'], kl_penalty=kl_penalty) # (batch_size, response_length) kld = kld * response_mask beta = kl_ctrl.value else: beta = 0 kld = torch.zeros_like(response_mask, dtype=torch.float32) token_level_rewards = token_level_scores - beta * kld current_kl = masked_mean(kld, mask=response_mask, axis=-1) # average over sequence current_kl = torch.mean(current_kl, dim=0).item() # according to https://github.com/huggingface/trl/blob/951ca1841f29114b969b57b26c7d3e80a39f75a0/trl/trainer/ppo_trainer.py#L837 kl_ctrl.update(current_kl=current_kl, n_steps=batch_size) data.batch['token_level_rewards'] = token_level_rewards metrics = {'critic/kl': current_kl, 'critic/kl_coeff': beta} return data, metrics","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#compute_advantage","name":"compute_advantage","kind":"function","path":"verl/trainer/ppo/ray_trainer.py","span":"L131-L174","signature":"def compute_advantage(data: DataProto, adv_estimator, gamma=1.0, lam=1.0, num_repeat=1): # prepare response group # TODO: add other ways to estimate advantages","exported":true,"origin":"ast","body_hash":"6db04d78dc4fe50f179f057e8fc22f844ec5524897eb9a2ac8891337877b79e5","body_text":"def compute_advantage(data: DataProto, adv_estimator, gamma=1.0, lam=1.0, num_repeat=1): # prepare response group # TODO: add other ways to estimate advantages if adv_estimator == 'gae': values = data.batch['values'] responses = data.batch['responses'] response_length = responses.size(-1) attention_mask = data.batch['attention_mask'] response_mask = attention_mask[:, -response_length:] token_level_rewards = data.batch['token_level_rewards'] advantages, returns = core_algos.compute_gae_advantage_return(token_level_rewards=token_level_rewards, values=values, eos_mask=response_mask, gamma=gamma, lam=lam) data.batch['advantages'] = advantages data.batch['returns'] = returns elif adv_estimator == 'grpo': token_level_rewards = data.batch['token_level_rewards'] index = data.non_tensor_batch['uid'] responses = data.batch['responses'] response_length = responses.size(-1) attention_mask = data.batch['attention_mask'] response_mask = attention_mask[:, -response_length:] advantages, returns = core_algos.compute_grpo_outcome_advantage(token_level_rewards=token_level_rewards, eos_mask=response_mask, index=index) data.batch['advantages'] = advantages data.batch['returns'] = returns elif adv_estimator == 'reinforce': token_level_rewards = data.batch['token_level_rewards'] responses = data.batch['responses'] response_length = responses.size(-1) attention_mask = data.batch['attention_mask'] response_mask = attention_mask[:, -response_length:] advantages, returns = core_algos.compute_reinforce_outcome_advantage( token_level_rewards=token_level_rewards, response_mask=response_mask ) data.batch['advantages'] = advantages data.batch['returns'] = returns else: raise NotImplementedError return data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#reduce_metrics","name":"reduce_metrics","kind":"function","path":"verl/trainer/ppo/ray_trainer.py","span":"L177-L180","signature":"def reduce_metrics(metrics: dict)","exported":true,"origin":"ast","body_hash":"103bdd01174623c17b869623021a6a851950493ae7c2fa26a9bacb3e48324143","body_text":"def reduce_metrics(metrics: dict): for key, val in metrics.items(): metrics[key] = np.mean(val) return metrics","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#_compute_response_info","name":"_compute_response_info","kind":"function","path":"verl/trainer/ppo/ray_trainer.py","span":"L183-L196","signature":"def _compute_response_info(batch)","exported":false,"origin":"ast","body_hash":"def389af1b6f649123e8f4790c99231bb7ec63e328f6a1dc91449c7ed3607733","body_text":"def _compute_response_info(batch): response_length = batch.batch['responses'].shape[-1] prompt_mask = batch.batch['attention_mask'][:, :-response_length] response_mask = batch.batch['attention_mask'][:, -response_length:] prompt_length = prompt_mask.sum(-1).float() response_length = response_mask.sum(-1).float() # (batch_size,) return dict( response_mask=response_mask, prompt_length=prompt_length, response_length=response_length, )","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","name":"compute_data_metrics","kind":"function","path":"verl/trainer/ppo/ray_trainer.py","span":"L199-L301","signature":"def compute_data_metrics(batch, use_critic=True): # TODO: add response length","exported":true,"origin":"ast","body_hash":"509aa7115980ad02efed281a79dcf21218a866a6bf8f93f576824a0d57bfd4a8","body_text":"def compute_data_metrics(batch, use_critic=True): # TODO: add response length sequence_score = batch.batch['token_level_scores'].sum(-1) sequence_reward = batch.batch['token_level_rewards'].sum(-1) advantages = batch.batch['advantages'] returns = batch.batch['returns'] max_response_length = batch.batch['responses'].shape[-1] prompt_mask = batch.batch['attention_mask'][:, :-max_response_length].bool() response_mask = batch.batch['attention_mask'][:, -max_response_length:].bool() max_prompt_length = prompt_mask.size(-1) response_info = _compute_response_info(batch) prompt_length = response_info['prompt_length'] response_length = response_info['response_length'] valid_adv = torch.masked_select(advantages, response_mask) valid_returns = torch.masked_select(returns, response_mask) if use_critic: values = batch.batch['values'] valid_values = torch.masked_select(values, response_mask) return_diff_var = torch.var(valid_returns - valid_values) return_var = torch.var(valid_returns) metrics = { # score 'critic/score/mean': torch.mean(sequence_score).detach().item(), 'critic/score/max': torch.max(sequence_score).detach().item(), 'critic/score/min': torch.min(sequence_score).detach().item(), # reward 'critic/rewards/mean': torch.mean(sequence_reward).detach().item(), 'critic/rewards/max': torch.max(sequence_reward).detach().item(), 'critic/rewards/min': torch.min(sequence_reward).detach().item(), # adv 'critic/advantages/mean': torch.mean(valid_adv).detach().item(), 'critic/advantages/max': torch.max(valid_adv).detach().item(), 'critic/advantages/min': torch.min(valid_adv).detach().item(), # returns 'critic/returns/mean': torch.mean(valid_returns).detach().item(), 'critic/returns/max': torch.max(valid_returns).detach().item(), 'critic/returns/min': torch.min(valid_returns).detach().item(), **({ # values 'critic/values/mean': torch.mean(valid_values).detach().item(), 'critic/values/max': torch.max(valid_values).detach().item(), 'critic/values/min': torch.min(valid_values).detach().item(), # vf explained var 'critic/vf_explained_var': (1.0 - return_diff_var / (return_var + 1e-5)).detach().item(), } if use_critic else {}), # response length 'response_length/mean': torch.mean(response_length).detach().item(), 'response_length/max': torch.max(response_length).detach().item(), 'response_length/min': torch.min(response_length).detach().item(), 'response_length/clip_ratio': torch.mean(torch.eq(response_length, max_response_length).float()).detach().item(), # prompt length 'prompt_length/mean': torch.mean(prompt_length).detach().item(), 'prompt_length/max': torch.max(prompt_length).detach().item(), 'prompt_length/min': torch.min(prompt_length).detach().item(), 'prompt_length/clip_ratio': torch.mean(torch.eq(prompt_length, max_prompt_length).float()).detach().item(), # metrics for actions 'env/number_of_actions/mean': float(np.array(batch.meta_info['turns_stats'], dtype=np.int16).mean()), 'env/number_of_actions/max': float(np.array(batch.meta_info['turns_stats'], dtype=np.int16).max()), 'env/number_of_actions/min': float(np.array(batch.meta_info['turns_stats'], dtype=np.int16).min()), 'env/finish_ratio': 1 - float(np.array(batch.meta_info['active_mask'], dtype=np.int16).mean()), 'env/number_of_valid_action': float(np.array(batch.meta_info['valid_action_stats'], dtype=np.int16).mean()), 'env/ratio_of_valid_action': float((np.array(batch.meta_info['valid_action_stats'], dtype=np.int16) / np.array(batch.meta_info['turns_stats'], dtype=np.int16)).mean()), 'env/number_of_valid_search': float(np.array(batch.meta_info['valid_search_stats'], dtype=np.int16).mean()), } return metrics","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","name":"compute_timing_metrics","kind":"function","path":"verl/trainer/ppo/ray_trainer.py","span":"L304-L325","signature":"def compute_timing_metrics(batch, timing_raw)","exported":true,"origin":"ast","body_hash":"7c4434dbd4b68228c4a6b9abd74a7d20c421434f85ef5cfaf05f154ea493167a","body_text":"def compute_timing_metrics(batch, timing_raw): response_info = _compute_response_info(batch) num_prompt_tokens = torch.sum(response_info['prompt_length']).item() num_response_tokens = torch.sum(response_info['response_length']).item() num_overall_tokens = num_prompt_tokens + num_response_tokens num_tokens_of_section = { 'gen': num_response_tokens, **{ name: num_overall_tokens for name in ['ref', 'values', 'adv', 'update_critic', 'update_actor', 'rollout'] }, } return { **{ f'timing_s/{name}': value for name, value in timing_raw.items() }, **{ f'timing_per_token_ms/{name}': timing_raw[name] * 1000 / num_tokens_of_section[name] for name in set(num_tokens_of_section.keys( )) & set(timing_raw.keys()) }, }","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#_timer","name":"_timer","kind":"function","path":"verl/trainer/ppo/ray_trainer.py","span":"L329-L332","signature":"def _timer(name: str, timing_raw: Dict[str, float])","exported":false,"origin":"ast","body_hash":"e7ea6489da509530c3bfcb7c3dd87245cee9ef38da64b94bff4ae0d76b98b875","body_text":"def _timer(name: str, timing_raw: Dict[str, float]): with Timer(name=name, logger=None) as timer: yield timing_raw[name] = timer.last","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","name":"RayPPOTrainer","kind":"class","path":"verl/trainer/ppo/ray_trainer.py","span":"L335-L1407","signature":"class RayPPOTrainer(object)","exported":true,"origin":"ast","body_hash":"5d79c85ce77fbc345e634b1bb812de7cc7689ae7e6a0c4d02da66237a37c92db","body_text":"class RayPPOTrainer(object): \"\"\" Note that this trainer runs on the driver process on a single CPU/GPU node. \"\"\" # TODO: support each role have individual ray_worker_group_cls, # i.e., support different backend of different role def __init__(self, config, tokenizer, role_worker_mapping: dict[Role, WorkerType], resource_pool_manager: ResourcePoolManager, ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup, reward_fn=None, val_reward_fn=None): # assert torch.cuda.is_available(), 'cuda must be available on driver' self.tokenizer = tokenizer self.config = config self.reward_fn = reward_fn self.val_reward_fn = val_reward_fn self.hybrid_engine = config.actor_rollout_ref.hybrid_engine assert self.hybrid_engine, 'Currently, only support hybrid engine' if self.hybrid_engine: assert Role.ActorRollout in role_worker_mapping, f'{role_worker_mapping.keys()=}' self.role_worker_mapping = role_worker_mapping self.resource_pool_manager = resource_pool_manager self.use_reference_policy = Role.RefPolicy in role_worker_mapping self.use_rm = Role.RewardModel in role_worker_mapping self.ray_worker_group_cls = ray_worker_group_cls # define KL control if self.use_reference_policy: if config.algorithm.kl_ctrl.type == 'fixed': self.kl_ctrl = core_algos.FixedKLController(kl_coef=config.algorithm.kl_ctrl.kl_coef) elif config.algorithm.kl_ctrl.type == 'adaptive': assert config.algorithm.kl_ctrl.horizon > 0, f'horizon must be larger than 0. Got {config.critic.kl_ctrl.horizon}' self.kl_ctrl = core_algos.AdaptiveKLController(init_kl_coef=config.algorithm.kl_ctrl.kl_coef, target_kl=config.algorithm.kl_ctrl.target_kl, horizon=config.algorithm.kl_ctrl.horizon) else: raise NotImplementedError else: self.kl_ctrl = core_algos.FixedKLController(kl_coef=0.) self._create_dataloader() self._init_logger() def _init_logger(self): from verl.utils.tracking import Tracking self.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)) def _create_dataloader(self): from torch.utils.data import DataLoader # TODO: we have to make sure the batch size is divisible by the dp size from verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn self.train_dataset = RLHFDataset(parquet_files=self.config.data.train_files, tokenizer=self.tokenizer, prompt_key=self.config.data.prompt_key, max_prompt_length=self.config.data.max_prompt_length, filter_prompts=True, return_raw_chat=self.config.data.get('return_raw_chat', False), truncation='error') if self.config.data.train_data_num is not None: if self.config.data.train_data_num > len(self.train_dataset.dataframe): print(f\"[WARNING] training dataset size is smaller than desired size. Using the dataset as the original size {len(self.train_dataset.dataframe)}\") else: self.train_dataset.dataframe = self.train_dataset.dataframe.sample(self.config.data.train_data_num, random_state=42) print(f\"filtered training dataset size: {len(self.train_dataset.dataframe)}\") self.train_dataloader = DataLoader(dataset=self.train_dataset, batch_size=self.config.data.train_batch_size, shuffle=self.config.data.shuffle_train_dataloader, drop_last=True, collate_fn=collate_fn) self.val_dataset = RLHFDataset(parquet_files=self.config.data.val_files, tokenizer=self.tokenizer, prompt_key=self.config.data.prompt_key, max_prompt_length=self.config.data.max_prompt_length, filter_prompts=True, return_raw_chat=self.config.data.get('return_raw_chat', False), truncation='error') if self.config.data.val_data_num is not None: if self.config.data.val_data_num > len(self.val_dataset.dataframe): print(f\"[WARNING] validation dataset size is smaller than desired size. Using the dataset as the original size {len(self.val_dataset.dataframe)}\") else: self.val_dataset.dataframe = self.val_dataset.dataframe.sample(self.config.data.val_data_num, random_state=42) print(f\"filtered validation dataset size: {len(self.val_dataset.dataframe)}\") self.val_dataloader = DataLoader(dataset=self.val_dataset, batch_size=self.config.data.val_batch_size, shuffle=True, drop_last=True, collate_fn=collate_fn) print(f'Size of train dataloader: {len(self.train_dataloader)}') print(f'Size of val dataloader: {len(self.val_dataloader)}') assert len(self.train_dataloader) >= 1 assert len(self.val_dataloader) >= 1 # inject total_training_steps to actor/critic optim_config. This is hacky. total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs if self.config.trainer.total_training_steps is not None: total_training_steps = self.config.trainer.total_training_steps self.total_training_steps = total_training_steps print(f'Total training steps: {self.total_training_steps}') OmegaConf.set_struct(self.config, True) with open_dict(self.config): self.config.actor_rollout_ref.actor.optim.total_training_steps = total_training_steps self.config.critic.optim.total_training_steps = total_training_steps def _valida","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.__init__","name":"__init__","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L342-L385","signature":"def __init__(self, config, tokenizer, role_worker_mapping: dict[Role, WorkerType], resource_pool_manager: ResourcePoolManager, ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup, reward_fn=None, val_reward_fn=None): # assert torch.cuda.is_available(), 'cuda must be available on driver'","exported":false,"origin":"ast","body_hash":"1eb952da50762a356ca29bcef5a31e841eeb195d99af9d941c7b761066ef8f6c","body_text":"def __init__(self, config, tokenizer, role_worker_mapping: dict[Role, WorkerType], resource_pool_manager: ResourcePoolManager, ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup, reward_fn=None, val_reward_fn=None): # assert torch.cuda.is_available(), 'cuda must be available on driver' self.tokenizer = tokenizer self.config = config self.reward_fn = reward_fn self.val_reward_fn = val_reward_fn self.hybrid_engine = config.actor_rollout_ref.hybrid_engine assert self.hybrid_engine, 'Currently, only support hybrid engine' if self.hybrid_engine: assert Role.ActorRollout in role_worker_mapping, f'{role_worker_mapping.keys()=}' self.role_worker_mapping = role_worker_mapping self.resource_pool_manager = resource_pool_manager self.use_reference_policy = Role.RefPolicy in role_worker_mapping self.use_rm = Role.RewardModel in role_worker_mapping self.ray_worker_group_cls = ray_worker_group_cls # define KL control if self.use_reference_policy: if config.algorithm.kl_ctrl.type == 'fixed': self.kl_ctrl = core_algos.FixedKLController(kl_coef=config.algorithm.kl_ctrl.kl_coef) elif config.algorithm.kl_ctrl.type == 'adaptive': assert config.algorithm.kl_ctrl.horizon > 0, f'horizon must be larger than 0. Got {config.critic.kl_ctrl.horizon}' self.kl_ctrl = core_algos.AdaptiveKLController(init_kl_coef=config.algorithm.kl_ctrl.kl_coef, target_kl=config.algorithm.kl_ctrl.target_kl, horizon=config.algorithm.kl_ctrl.horizon) else: raise NotImplementedError else: self.kl_ctrl = core_algos.FixedKLController(kl_coef=0.) self._create_dataloader() self._init_logger()","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._init_logger","name":"_init_logger","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L387-L392","signature":"def _init_logger(self)","exported":false,"origin":"ast","body_hash":"81a2d810eb80b20385a3e6718a23e98cb02b52eaf16e4c1f055b11182a0d34af","body_text":"def _init_logger(self): from verl.utils.tracking import Tracking self.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))","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","name":"_create_dataloader","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L394-L456","signature":"def _create_dataloader(self)","exported":false,"origin":"ast","body_hash":"6ffc72e6d07f92347988cb6b11f07f08ebca9ecffdb65b3c865f6c310b8b7bd6","body_text":"def _create_dataloader(self): from torch.utils.data import DataLoader # TODO: we have to make sure the batch size is divisible by the dp size from verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn self.train_dataset = RLHFDataset(parquet_files=self.config.data.train_files, tokenizer=self.tokenizer, prompt_key=self.config.data.prompt_key, max_prompt_length=self.config.data.max_prompt_length, filter_prompts=True, return_raw_chat=self.config.data.get('return_raw_chat', False), truncation='error') if self.config.data.train_data_num is not None: if self.config.data.train_data_num > len(self.train_dataset.dataframe): print(f\"[WARNING] training dataset size is smaller than desired size. Using the dataset as the original size {len(self.train_dataset.dataframe)}\") else: self.train_dataset.dataframe = self.train_dataset.dataframe.sample(self.config.data.train_data_num, random_state=42) print(f\"filtered training dataset size: {len(self.train_dataset.dataframe)}\") self.train_dataloader = DataLoader(dataset=self.train_dataset, batch_size=self.config.data.train_batch_size, shuffle=self.config.data.shuffle_train_dataloader, drop_last=True, collate_fn=collate_fn) self.val_dataset = RLHFDataset(parquet_files=self.config.data.val_files, tokenizer=self.tokenizer, prompt_key=self.config.data.prompt_key, max_prompt_length=self.config.data.max_prompt_length, filter_prompts=True, return_raw_chat=self.config.data.get('return_raw_chat', False), truncation='error') if self.config.data.val_data_num is not None: if self.config.data.val_data_num > len(self.val_dataset.dataframe): print(f\"[WARNING] validation dataset size is smaller than desired size. Using the dataset as the original size {len(self.val_dataset.dataframe)}\") else: self.val_dataset.dataframe = self.val_dataset.dataframe.sample(self.config.data.val_data_num, random_state=42) print(f\"filtered validation dataset size: {len(self.val_dataset.dataframe)}\") self.val_dataloader = DataLoader(dataset=self.val_dataset, batch_size=self.config.data.val_batch_size, shuffle=True, drop_last=True, collate_fn=collate_fn) print(f'Size of train dataloader: {len(self.train_dataloader)}') print(f'Size of val dataloader: {len(self.val_dataloader)}') assert len(self.train_dataloader) >= 1 assert len(self.val_dataloader) >= 1 # inject total_training_steps to actor/critic optim_config. This is hacky. total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs if self.config.trainer.total_training_steps is not None: total_training_steps = self.config.trainer.total_training_steps self.total_training_steps = total_training_steps print(f'Total training steps: {self.total_training_steps}') OmegaConf.set_struct(self.config, True) with open_dict(self.config): self.config.actor_rollout_ref.actor.optim.total_training_steps = total_training_steps self.config.critic.optim.total_training_steps = total_training_steps","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","name":"_validate","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L458-L625","signature":"def _validate(self)","exported":false,"origin":"ast","body_hash":"f46743be42980babcecc8e1e6acaa5a33865617ba9e82cb46c2731d67a409445","body_text":"def _validate(self): \"\"\" The training loop of PPO with global metric computation. Accumulates metrics across all batches before computing final statistics. \"\"\" import torch import types reward_tensor_lst = [] data_source_lst = [] # ========== 跨 batch 累积统计 ========== # 所有字段均为 list，跨 batch 直接 extend all_reward_stats = None all_gen_stats = None def _merge_stats(accum, new_stats): \"\"\"将 new_stats 各字段 extend 进 accum，返回合并后的 accum。\"\"\" if accum is None: # 第一批：深拷贝 import copy return copy.deepcopy(new_stats) for k, v in new_stats.items(): if isinstance(v, list) and k in accum and isinstance(accum[k], list): accum[k].extend(v) # 非 list 字段（如不存在）忽略 return accum # ========================================== gen_config = GenerationConfig( max_turns=self.config.max_turns, max_start_length=self.config.data.max_start_length, max_prompt_length=self.config.data.max_prompt_length, max_response_length=self.config.data.max_response_length, max_obs_length=self.config.data.max_obs_length, num_gpus=self.config.trainer.n_gpus_per_node, no_think_rl=self.config.algorithm.no_think_rl, search_url = self.config.retriever.url, topk = self.config.retriever.topk, ) # Agent config preparation generation_manager = LLMGenerationManager( tokenizer=self.tokenizer, actor_rollout_wg=self.actor_rollout_wg, config=gen_config, is_validation = True, ) if not self.config.do_search: for test_data in self.val_dataloader: test_batch = DataProto.from_single_dict(test_data) # we only do validation on rule-based rm if self.config.reward_model.enable and test_batch[0].non_tensor_batch['reward_model']['style'] == 'model': return {} test_gen_batch = test_batch.pop(['input_ids', 'attention_mask', 'position_ids']) test_gen_batch.meta_info = { 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, 'recompute_log_prob': False, 'do_sample': False, 'validate': True, } # pad to be divisible by dp_size test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, self.actor_rollout_wg.world_size) test_output_gen_batch_padded = self.actor_rollout_wg.generate_sequences(test_gen_batch_padded) # unpad test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size) print('validation generation end') test_batch = test_batch.union(test_output_gen_batch) # evaluate using reward_function # for certain reward function (e.g. sandbox), the generation can overlap with reward reward_tensor = self.val_reward_fn(test_batch) # ========== 跨 batch 累积统计（不在 loop 内计算，仅 extend）========== all_reward_stats = _merge_stats(all_reward_stats, test_batch.meta_info.get('reward_statistics', {})) all_gen_stats = _merge_stats(all_gen_stats, test_batch.meta_info.get('generation_statistics', {})) # ======================================================== reward_tensor_lst.append(reward_tensor) data_source_lst.append(test_batch.non_tensor_batch.get('data_source', ['unknown'] * reward_tensor.shape[0])) else: for batch_dict in self.val_dataloader: timing_raw = {} test_batch: DataProto = DataProto.from_single_dict(batch_dict) # test_batch = test_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n_agent, interleave=True) test_gen_batch = test_batch.pop(batch_keys=['input_ids', 'attention_mask', 'position_ids']) test_gen_batch.meta_info = { 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, 'recompute_log_prob': False, 'do_sample': False, 'validate': True, } with _timer('step', timing_raw): first_input_ids = test_gen_batch.batch['input_ids'][:, -gen_config.max_start_length:].clone() with _timer('gen', timing_raw): generation_manager.timing_raw = timing_raw final_gen_batch_output = generation_manager.run_llm_loop( gen_batch=test_gen_batch, initial_input_ids=first_input_ids, ) test_batch = test_batch.union(final_gen_batch_output) # TODO-10 (val): transfer per_step_confidences from meta_info to non_tensor_batch if 'per_step_confidences' in final_gen_batch_output.meta_info: psc = final_gen_batch_output.meta_info['per_step_confidences'] test_batch.non_tensor_batch['per_step_confidences'] = np.array(psc, dtype=object) for key in test_batch.batch.keys(): test_batch.batch[key] = test_batch.batch[key].long() # evaluate using reward_function # for certain reward function (e.g. sandbox), the generation can overlap with reward reward_tensor = self.val_reward_fn(test_batch) # ========== 跨 batch 累积统计（不在 loop 内计算，仅 extend）========== all_reward_stats = _merge_stats(all_reward_stats, test_batch.meta_info.get('reward_statistics', {})) all_gen_stats = _merge_stats(all_gen_stats, test_batch.meta_info.get('generation_statistics', {})) # ======================================================== # ========== 每个 batch 立即保存样本级 CSV / JSONL ========== _batch_gen_stats = test_batch.meta_info.get('generation_statistics', None) _batch_reward_stats = test_batch.meta_info.get('reward_statistics', None) if _batch_gen_stats and _batch_reward_stats: self._save_sample_level_data(test_batch, _batch_gen_stats, _batch_reward_sta","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate._merge_stats","name":"_merge_stats","kind":"function","path":"verl/trainer/ppo/ray_trainer.py","span":"L473-L483","signature":"def _merge_stats(accum, new_stats)","exported":false,"origin":"ast","body_hash":"0d685857117882f956b84d6faa2fd99590f62966050245b7e6c8c01f23bcc1b8","body_text":"def _merge_stats(accum, new_stats): \"\"\"将 new_stats 各字段 extend 进 accum，返回合并后的 accum。\"\"\" if accum is None: # 第一批：深拷贝 import copy return copy.deepcopy(new_stats) for k, v in new_stats.items(): if isinstance(v, list) and k in accum and isinstance(accum[k], list): accum[k].extend(v) # 非 list 字段（如不存在）忽略 return accum","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","name":"init_workers","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L628-L699","signature":"def init_workers(self)","exported":true,"origin":"ast","body_hash":"dbd28649655cd78f736f66b5d85d9e3bd05912f679a85f062c6dd9496230bd1a","body_text":"def init_workers(self): \"\"\"Init resource pool and worker group\"\"\" self.resource_pool_manager.create_resource_pool() self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()} # create actor and rollout if self.hybrid_engine: resource_pool = self.resource_pool_manager.get_resource_pool(Role.ActorRollout) actor_rollout_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.ActorRollout], config=self.config.actor_rollout_ref, role='actor_rollout') self.resource_pool_to_cls[resource_pool]['actor_rollout'] = actor_rollout_cls else: raise NotImplementedError # create critic if self.config.algorithm.adv_estimator == 'gae': resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic) critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=self.config.critic) self.resource_pool_to_cls[resource_pool]['critic'] = critic_cls self.use_critic = True elif self.config.algorithm.adv_estimator in ['grpo', 'reinforce']: self.use_critic = False else: self.use_critic = False # create reference policy if needed if self.use_reference_policy: resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy) ref_policy_cls = RayClassWithInitArgs(self.role_worker_mapping[Role.RefPolicy], config=self.config.actor_rollout_ref, role='ref') self.resource_pool_to_cls[resource_pool]['ref'] = ref_policy_cls # create a reward model if reward_fn is None if self.use_rm: # we create a RM here resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) rm_cls = RayClassWithInitArgs(self.role_worker_mapping[Role.RewardModel], config=self.config.reward_model) self.resource_pool_to_cls[resource_pool]['rm'] = rm_cls # 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 https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information. all_wg = {} self.wg_dicts = [] 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) # keep the referece of WorkerDict to support ray >= 2.31. Ref: https://github.com/ray-project/ray/pull/45699 self.wg_dicts.append(wg_dict) 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()","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_checkpoint","name":"_save_checkpoint","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L701-L713","signature":"def _save_checkpoint(self)","exported":false,"origin":"ast","body_hash":"b4c1b3c583a30d454703aa7bdf36f9ff941cb31c71969cae2badde0bee034f53","body_text":"def _save_checkpoint(self): actor_local_path = os.path.join(self.config.trainer.default_local_dir, 'actor', f'global_step_{self.global_steps}') actor_remote_path = None if self.config.trainer.default_hdfs_dir is None else 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_{self.global_steps}') critic_remote_path = None if self.config.trainer.default_hdfs_dir is None else os.path.join( self.config.trainer.default_hdfs_dir, 'critic') self.critic_wg.save_checkpoint(critic_local_path, critic_remote_path)","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._balance_batch","name":"_balance_batch","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L715-L730","signature":"def _balance_batch(self, batch: DataProto, metrics, logging_prefix='global_seqlen')","exported":false,"origin":"ast","body_hash":"01ccfebb14dcb1923a44ed2f929818489b50c95747e89603eadd5ef29b338fec","body_text":"def _balance_batch(self, batch: DataProto, metrics, logging_prefix='global_seqlen'): \"\"\"Reorder the data on single controller such that each dp rank gets similar total tokens\"\"\" attention_mask = batch.batch['attention_mask'] batch_size = attention_mask.shape[0] global_seqlen_lst = attention_mask.view(batch_size, -1).sum(-1).tolist() # (train_batch_size,) world_size = self.actor_rollout_wg.world_size global_partition_lst = get_seqlen_balanced_partitions(global_seqlen_lst, k_partitions=world_size, equal_size=True) # reorder based on index. The data will be automatically equally partitioned by dispatch function global_idx = torch.tensor([j for partition in global_partition_lst for j in partition]) batch.reorder(global_idx) global_balance_stats = log_seqlen_unbalance(seqlen_list=global_seqlen_lst, partitions=global_partition_lst, prefix=logging_prefix) metrics.update(global_balance_stats)","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","name":"fit","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L732-L972","signature":"def fit(self)","exported":true,"origin":"ast","body_hash":"550f8f1b28f443a0608d118a052e303489480e8bae81900246596334f72a8d61","body_text":"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. \"\"\" logger = self.logger self.global_steps = 0 # perform validation before training # currently, we only support validation using the reward_function. if self.val_reward_fn is not None and self.config.trainer.get('val_before_train', True): val_metrics = self._validate() pprint(f'Initial validation metrics: {val_metrics}') logger.log(data=val_metrics, step=self.global_steps) if self.config.trainer.get('val_only', False): return # we start from step 1 self.global_steps += 1 # Agent config preparation gen_config = GenerationConfig( max_turns=self.config.max_turns, max_start_length=self.config.data.max_start_length, max_prompt_length=self.config.data.max_prompt_length, max_response_length=self.config.data.max_response_length, max_obs_length=self.config.data.max_obs_length, num_gpus=self.config.trainer.n_gpus_per_node, no_think_rl=self.config.algorithm.no_think_rl, search_url = self.config.retriever.url, topk = self.config.retriever.topk, ) generation_manager = LLMGenerationManager( tokenizer=self.tokenizer, actor_rollout_wg=self.actor_rollout_wg, config=gen_config, ) # start training loop for epoch in range(self.config.trainer.total_epochs): for batch_dict in self.train_dataloader: print(f'epoch {epoch}, step {self.global_steps}') metrics = {} timing_raw = {} batch: DataProto = DataProto.from_single_dict(batch_dict) batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n_agent, interleave=True) # pop those keys for generation gen_batch = batch.pop(batch_keys=['input_ids', 'attention_mask', 'position_ids']) #################### # original code here with _timer('step', timing_raw): if not self.config.do_search: gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) batch.non_tensor_batch['uid'] = np.array([str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object) # repeat to align with repeated responses in rollout batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) batch = batch.union(gen_batch_output) #################### # Below is aLL about agents - the \"LLM + forloop\" #################### # with _timer('step', timing_raw): else: first_input_ids = gen_batch.batch['input_ids'][:, -gen_config.max_start_length:].clone().long() with _timer('gen', timing_raw): generation_manager.timing_raw = timing_raw final_gen_batch_output = generation_manager.run_llm_loop( gen_batch=gen_batch, initial_input_ids=first_input_ids, ) # final_gen_batch_output.batch.apply(lambda x: x.long(), inplace=True) for key in final_gen_batch_output.batch.keys(): final_gen_batch_output.batch[key] = final_gen_batch_output.batch[key].long() with torch.no_grad(): output = self.actor_rollout_wg.compute_log_prob(final_gen_batch_output) final_gen_batch_output = final_gen_batch_output.union(output) # batch.non_tensor_batch['uid'] = np.array([str(uuid.uuid4()) for _ in range(len(batch.batch))], # dtype=object) batch.non_tensor_batch['uid'] = batch.non_tensor_batch['index'].copy() # repeat to align with repeated responses in rollout batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) batch = batch.union(final_gen_batch_output) # TODO-10 (train): transfer per_step_confidences from meta_info to non_tensor_batch if 'per_step_confidences' in final_gen_batch_output.meta_info: psc = final_gen_batch_output.meta_info['per_step_confidences'] batch.non_tensor_batch['per_step_confidences'] = np.array(psc, dtype=object) #################### #################### # balance the number of valid tokens on each dp rank. # Note that this breaks the order of data inside the batch. # Please take care when you implement group based adv computation such as GRPO and rloo self._balance_batch(batch, metrics=metrics) # compute global_valid tokens batch.meta_info['global_token_num'] = torch.sum(batch.batch['attention_mask'], dim=-1).tolist() batch.meta_info['global_step'] = self.global_steps # batch.batch.apply(lambda x, key: x.long() if key != \"old_log_probs\" else x, inplace=True, key=True) for key in batch.batch.keys(): if key not in ('old_log_probs', 'token_level_scores'): batch.batch[key] = batch.batch[key].long() if self.use_reference_policy: # compute reference log_prob with _timer('ref', timing_raw): ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch) batch = batch.union(ref_log_prob) # compute values if self.use_critic: with _timer('values', timing_raw): values = self.critic_wg.compute_values(batch) batch = batch.union(values) with _timer('adv', timing_raw): # 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 rewa","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","name":"_save_sample_level_data","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L976-L1190","signature":"def _save_sample_level_data(self, batch: DataProto, gen_stats: dict, reward_stats: dict, split='train')","exported":false,"origin":"ast","body_hash":"ba07805fd046a2ea39b16e000b24b19930d97ee567b3d0a11ab6c57820bffb01","body_text":"def _save_sample_level_data(self, batch: DataProto, gen_stats: dict, reward_stats: dict, split='train'): \"\"\" 保存每个样本的详细数据到 CSV 文件 Args: batch: DataProto 对象 gen_stats: generation 统计 reward_stats: reward 统计 \"\"\" import csv import os uid = batch.non_tensor_batch.get(\"uid\", None) if uid is None: uid = batch.non_tensor_batch.get(\"index\", None) # normalize uid into a python list for indexing if isinstance(uid, np.ndarray): uid = uid.tolist() elif isinstance(uid, (list, tuple)): uid = list(uid) else: uid = None # ======== [FIX END] ======== # 创建输出目录 output_dir = os.path.join(self.config.trainer.default_local_dir, 'sample_statistics') os.makedirs(output_dir, exist_ok=True) # 🔴 根据配置自动生成文件名 if split == 'val' and hasattr(self.config.trainer, 'csv_filename') and self.config.trainer.csv_filename: # 用户手动指定了文件名 csv_filename = self.config.trainer.csv_filename elif split == 'val': # 自动生成验证集文件名 # 1. 从 actor model path 提取信息 model_path = self.config.actor_rollout_ref.model.path # 提取最后两级目录作为模型标识 (例如: actor/global_step_100) model_parts = model_path.rstrip('/').split('/')[-2:] model_id = '_'.join(model_parts) # 例如: actor_global_step_100 # 2. 从验证数据集路径提取数据集名称 val_files = self.config.data.val_files if isinstance(val_files, (list, tuple)): val_files = val_files[0] # 提取数据集目录名 (例如: nq_hotpot-only) dataset_name = os.path.basename(os.path.dirname(val_files)) # 3. 组合文件名 csv_filename = f'{model_id}_{dataset_name}_val.csv' else: # 训练模式使用默认文件名 csv_filename = 'sample_level_data.csv' csv_path = os.path.join(output_dir, csv_filename) # 检查文件是否存在,决定是否写入表头 file_exists = os.path.exists(csv_path) # 打开文件并追加数据 with open(csv_path, 'a', newline='', encoding='utf-8') as f: fieldnames = [ 'global_step', 'split', 'sample_id', 'data_source', 'is_correct', 'score', 'total_output_tokens', 'reasoning_tokens', 'information_tokens', 'search_count', 'turn_count', 'reasoning_ratio', # reasoning_tokens / (reasoning_tokens + information_tokens) 'has_valid_format', 'answer_tag_count', 'predicted_answer', 'ground_truth', 'think_open', 'think_close', 'think_malformed', 'think_total', 'think_bad', 'search_open', 'search_close', 'search_malformed', 'search_total', 'search_bad', 'answer_open', 'answer_close', 'answer_malformed', 'answer_total', 'answer_bad', # optional: all tags json (if you saved) 'all_tag_open_json', 'all_tag_close_json', 'all_tag_malformed_json', 'confidence', 'confidence_valid', 'confidence_reward', 'is_overconfident', 'is_underconfident', # per-step confidence (TODO-11) 'per_step_conf_count', 'per_step_conf_mean', 'per_step_conf_values', 'per_step_action_types', ] writer = csv.DictWriter(f, fieldnames=fieldnames) # 如果是新文件,写入表头 if not file_exists: writer.writeheader() # 写入每个样本的数据 num_samples = len(reward_stats['is_correct']) for i in range(num_samples): total_tokens = gen_stats['total_output_tokens'][i] reasoning_tokens = gen_stats['reasoning_tokens'][i] info_tokens = gen_stats['information_tokens'][i] # 计算 reasoning ratio if reasoning_tokens + info_tokens > 0: reasoning_ratio = reasoning_tokens / (reasoning_tokens + info_tokens) else: reasoning_ratio = 0.0 row = { 'global_step': self.global_steps, 'split': split, 'sample_id': str(uid[i]) if (uid is not None and i < len(uid)) else str(i), 'data_source': reward_stats['data_sources'][i], 'is_correct': int(reward_stats['is_correct'][i]), 'score': reward_stats['scores'][i], 'total_output_tokens': total_tokens, 'reasoning_tokens': reasoning_tokens, 'information_tokens': info_tokens, 'search_count': gen_stats['search_count'][i], 'turn_count': gen_stats['turn_count'][i], 'reasoning_ratio': reasoning_ratio, 'has_valid_format': int(reward_stats['has_valid_format'][i]), 'answer_tag_count': reward_stats['answer_tag_count'][i], 'predicted_answer': reward_stats['predicted_answers'][i][:200], # 截断 'ground_truth': reward_stats['ground_truths'][i][:200], 'think_open': int(reward_stats.get('think_open', [0]*num_samples)[i]), 'think_close': int(reward_stats.get('think_close', [0]*num_samples)[i]), 'think_malformed': int(reward_stats.get('think_malformed', [0]*num_samples)[i]), 'think_total': int(reward_stats.get('think_total', [0]*num_samples)[i]), 'think_bad': int(reward_stats.get('think_bad', [0]*num_samples)[i]), 'search_open': int(reward_stats.get('search_open', [0]*num_samples)[i]), 'search_close': int(reward_stats.get('search_close', [0]*num_samples)[i]), 'search_malformed': int(reward_stats.get('search_malformed', [0]*num_samples)[i]), 'search_total': int(reward_stats.get('search_total', [0]*num_samples)[i]), 'search_bad': int(reward_stats.get('search_bad', [0]*num_samples)[i]), 'answer_open': int(reward_stats.get('answer_open', [0]*num_samples)[i]), 'answer_close': int(reward_stats.get('answer_close', [0]*num_samples)[i]), 'answer_malformed': int(reward_stats.get('answer_malformed', [0]*num_samples)[i]), 'answer_total': int(reward_stats.get('answer_total', [0]*num_samples)[i]), 'answer_bad': int(reward_stats.get('answer_bad', [0]*num_samples)[i]), 'all_tag_open_json': reward_stats.get('all_tag_open_json', ['{}']","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","name":"_collect_and_log_statistics","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L1192-L1390","signature":"def _collect_and_log_statistics(self, batch, metrics: dict, split='train'): # 检查是否有统计信息（支持 DataProto 和 SimpleNamespace）","exported":false,"origin":"ast","body_hash":"40a7ec531dcc6233d6a129601a18a6138b6f5b1fada058d58c7001730875ed32","body_text":"def _collect_and_log_statistics(self, batch, metrics: dict, split='train'): # 检查是否有统计信息（支持 DataProto 和 SimpleNamespace） gen_stats = batch.meta_info.get('generation_statistics', None) reward_stats = batch.meta_info.get('reward_statistics', None) has_gen_stats = gen_stats is not None and len(gen_stats) > 0 has_reward_stats = reward_stats is not None and len(reward_stats) > 0 if not (has_gen_stats or has_reward_stats): return print(\"\\n\" + \"=\" * 80) print(\"[COMBINED STATISTICS]\") print(\"=\" * 80) # 1) generation if has_gen_stats: metrics.update({ 'generation/avg_total_tokens': np.mean(gen_stats['total_output_tokens']), 'generation/max_total_tokens': np.max(gen_stats['total_output_tokens']), 'generation/min_total_tokens': np.min(gen_stats['total_output_tokens']), 'generation/avg_reasoning_tokens': np.mean(gen_stats['reasoning_tokens']), 'generation/avg_information_tokens': np.mean(gen_stats['information_tokens']), 'generation/avg_search_count': np.mean(gen_stats['search_count']), 'generation/max_search_count': np.max(gen_stats['search_count']), 'generation/avg_turn_count': np.mean(gen_stats['turn_count']), }) total_tokens = np.sum(gen_stats['total_output_tokens']) info_tokens = np.sum(gen_stats['information_tokens']) if total_tokens > 0: metrics['generation/info_token_ratio'] = info_tokens / total_tokens print(f\"Generation Stats:\") print(f\" Avg total tokens: {metrics['generation/avg_total_tokens']:.1f}\") print(f\" Avg reasoning tokens: {metrics['generation/avg_reasoning_tokens']:.1f}\") print(f\" Avg information tokens: {metrics['generation/avg_information_tokens']:.1f}\") print(f\" Avg search count: {metrics['generation/avg_search_count']:.1f}\") print(f\" Info token ratio: {metrics.get('generation/info_token_ratio', 0):.2%}\") # 2) reward if has_reward_stats: total = len(reward_stats['is_correct']) correct = sum(reward_stats['is_correct']) accuracy = correct / total if total > 0 else 0 metrics.update({ 'reward/accuracy': accuracy, 'reward/correct_count': correct, 'reward/total_count': total, 'reward/avg_score': np.mean(reward_stats['scores']), 'reward/valid_format_ratio': sum(reward_stats['has_valid_format']) / total if total > 0 else 0, 'reward/avg_answer_tags': np.mean(reward_stats['answer_tag_count']), }) # ========== 新增：confidence metrics ========== if has_reward_stats and 'confidences' in reward_stats: total = len(reward_stats['is_correct']) valid_confs = [c for c in reward_stats['confidences'] if c is not None] conf_valid_count = sum(reward_stats['confidence_valid']) overconf_count = sum(reward_stats['is_overconfident']) underconf_count = sum(reward_stats['is_underconfident']) metrics.update({ # 格式合规率 f'{split}/confidence/valid_rate': conf_valid_count / total if total > 0 else 0, # 平均confidence（只统计格式合规的） f'{split}/confidence/mean': float(np.mean(valid_confs)) if valid_confs else 0.0, f'{split}/confidence/std': float(np.std(valid_confs)) if valid_confs else 0.0, # overconfidence / underconfidence比例 f'{split}/confidence/overconfident_rate': overconf_count / total if total > 0 else 0, f'{split}/confidence/underconfident_rate': underconf_count / total if total > 0 else 0, # confidence reward分量均值 f'{split}/confidence/avg_conf_reward': float(np.mean(reward_stats['confidence_reward'])) if reward_stats['confidence_reward'] else 0.0, }) # 答对/答错时的平均confidence（核心校准指标） correct_confs = [ reward_stats['confidences'][i] for i in range(total) if reward_stats['is_correct'][i] and reward_stats['confidences'][i] is not None ] incorrect_confs = [ reward_stats['confidences'][i] for i in range(total) if not reward_stats['is_correct'][i] and reward_stats['confidences'][i] is not None ] if correct_confs: metrics[f'{split}/confidence/mean_when_correct'] = float(np.mean(correct_confs)) if incorrect_confs: metrics[f'{split}/confidence/mean_when_incorrect'] = float(np.mean(incorrect_confs)) # 关键校准指标：答对时conf是否高于答错时conf # 这个gap越大说明calibration越好 if correct_confs and incorrect_confs: metrics[f'{split}/confidence/calibration_gap'] = ( float(np.mean(correct_confs)) - float(np.mean(incorrect_confs)) ) # ECE近似（分桶计算） ece = compute_ece(reward_stats['confidences'], reward_stats['is_correct']) if ece is not None: metrics[f'{split}/confidence/ece'] = ece print(f\" ECE: {ece:.4f} (越小越好，0=完美校准)\") else: print(f\" ECE: N/A (有效样本不足5个)\"), auroc = compute_auroc(reward_stats['confidences'], reward_stats['is_correct']) if auroc is not None: metrics[f'{split}/confidence/auroc'] = auroc print(f\" AUROC: {auroc:.4f} (越大越好，1=完美区分)\") else: print(f\" AUROC: N/A (有效样本不足或只有一类标签)\") # ========== 当前奖励方案配置信息 ========== if split == 'train': metrics['conf/lsc_mode'] = hash(LSC_MODE) % 1000 # 字符串转数字便于记录 metrics['conf/fc_mode'] = hash(FC_MODE) % 1000 # ================================================ # ======== [FIX] tag totals only for TRAIN (wandb needs these) ======== if split == \"train\": def _mean_list(key, default=0.0): v = reward_stats.get(key, None) if v is None or len(v) == 0: return float(default) return float(np.mean(v)) metrics.update({ \"train/tags/think_total_mean\": _me","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics._mean_list","name":"_mean_list","kind":"function","path":"verl/trainer/ppo/ray_trainer.py","span":"L1330-L1334","signature":"def _mean_list(key, default=0.0)","exported":false,"origin":"ast","body_hash":"68eb2790c884b8186e47e17143acd8e657a0d0ab378bdce5d810dbe770f61f88","body_text":"def _mean_list(key, default=0.0): v = reward_stats.get(key, None) if v is None or len(v) == 0: return float(default) return float(np.mean(v))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_loss_mask","name":"_create_loss_mask","kind":"method","path":"verl/trainer/ppo/ray_trainer.py","span":"L1394-L1407","signature":"def _create_loss_mask(self, batch, metrics)","exported":false,"origin":"ast","body_hash":"6ffaf0798d9988ac351ae9bb5fa1d1d459bd9114ac56093f6b29161c02f08d32","body_text":"def _create_loss_mask(self, batch, metrics): \"\"\"Create loss mask for state tokens.\"\"\" response_length = batch.batch['responses'].shape[-1] response_mask = batch.batch['attention_mask'][:, -response_length:] loss_mask = batch.batch['info_mask'][:, -response_length:] batch.batch['loss_mask'] = loss_mask metrics.update({ 'state_tokens/total': loss_mask.sum().item(), 'state_tokens/coverage': (loss_mask.sum() / response_mask.sum()).item(), }) return batch, metrics","summary_state":"pending","summary":null,"crux":null,"owner":"RayPPOTrainer"}],"rawEdges":[{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"os","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"uuid","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"contextlib","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"dataclasses","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"enum","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"pprint","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"typing","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"re","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"json","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"collections","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"codetiming","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"omegaconf","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl.protocol","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl.single_controller.base","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl.single_controller.ray","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl.single_controller.ray.base","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl.trainer.ppo","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl.utils.seqlen_balancing","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"re","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"search_r1.llm_agent.generation","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl.trainer.conf_adv_config","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#Role","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#Role","relation":"extends","name":"Enum","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager","relation":"calls","name":"field","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager.create_resource_pool","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager.create_resource_pool","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager.create_resource_pool","relation":"calls","name":"RayResourcePool","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#ResourcePoolManager.get_resource_pool","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"torch","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#apply_kl_penalty","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#apply_kl_penalty","relation":"calls","name":"size","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#apply_kl_penalty","relation":"calls","name":"keys","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#apply_kl_penalty","relation":"calls","name":"kl_penalty","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#apply_kl_penalty","relation":"calls","name":"zeros_like","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#apply_kl_penalty","relation":"calls","name":"masked_mean","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#apply_kl_penalty","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#apply_kl_penalty","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#apply_kl_penalty","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#compute_advantage","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_advantage","relation":"calls","name":"size","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_advantage","relation":"calls","name":"compute_gae_advantage_return","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_advantage","relation":"calls","name":"size","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_advantage","relation":"calls","name":"compute_grpo_outcome_advantage","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_advantage","relation":"calls","name":"size","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_advantage","relation":"calls","name":"compute_reinforce_outcome_advantage","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#reduce_metrics","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#reduce_metrics","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#reduce_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#_compute_response_info","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#_compute_response_info","relation":"calls","name":"float","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#_compute_response_info","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#_compute_response_info","relation":"calls","name":"float","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#_compute_response_info","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#_compute_response_info","relation":"calls","name":"dict","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"bool","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"bool","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"size","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"_compute_response_info","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"masked_select","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"masked_select","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"masked_select","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"var","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"var","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"float","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"eq","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"detach","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"float","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"eq","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_data_metrics","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","relation":"calls","name":"_compute_response_info","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","relation":"calls","name":"set","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","relation":"calls","name":"keys","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","relation":"calls","name":"set","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#compute_timing_metrics","relation":"calls","name":"keys","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#_timer","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#_timer","relation":"calls","name":"Timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"extends","name":"object","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.__init__","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.__init__","relation":"calls","name":"keys","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.__init__","relation":"calls","name":"FixedKLController","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.__init__","relation":"calls","name":"AdaptiveKLController","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.__init__","relation":"calls","name":"FixedKLController","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.__init__","relation":"calls","name":"_create_dataloader","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.__init__","relation":"calls","name":"_init_logger","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._init_logger","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl.utils.tracking","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._init_logger","relation":"calls","name":"Tracking","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._init_logger","relation":"calls","name":"to_container","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"torch.utils.data","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"verl.utils.dataset.rl_dataset","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"RLHFDataset","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"sample","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"DataLoader","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"RLHFDataset","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"sample","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"DataLoader","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"set_struct","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_dataloader","relation":"calls","name":"open_dict","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"torch","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"types","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate._merge_stats","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"copy","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate._merge_stats","relation":"calls","name":"deepcopy","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate._merge_stats","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate._merge_stats","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate._merge_stats","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate._merge_stats","relation":"calls","name":"extend","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"GenerationConfig","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"LLMGenerationManager","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"from_single_dict","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"pop","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"DataProto"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"pad_dataproto_to_divisor","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"generate_sequences","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"unpad_dataproto","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"union","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"DataProto"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"val_reward_fn","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"_merge_stats","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"_merge_stats","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"from_single_dict","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"pop","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"DataProto"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"clone","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"run_llm_loop","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"LLMGenerationManager"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"union","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"DataProto"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"keys","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"long","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"val_reward_fn","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"_merge_stats","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"_merge_stats","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"_save_sample_level_data","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"cpu","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"cat","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"concatenate","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"SimpleNamespace","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"_collect_and_log_statistics","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._validate","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"create_resource_pool","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"values","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"get_resource_pool","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"RayClassWithInitArgs","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"get_resource_pool","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"RayClassWithInitArgs","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"get_resource_pool","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"RayClassWithInitArgs","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"get_resource_pool","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"RayClassWithInitArgs","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"create_colocated_worker_cls","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"ray_worker_group_cls","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"spawn","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"keys","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"append","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"init_model","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"init_model","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"init_model","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.init_workers","relation":"calls","name":"init_model","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_checkpoint","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_checkpoint","relation":"calls","name":"join","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_checkpoint","relation":"calls","name":"join","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_checkpoint","relation":"calls","name":"save_checkpoint","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_checkpoint","relation":"calls","name":"join","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_checkpoint","relation":"calls","name":"join","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_checkpoint","relation":"calls","name":"save_checkpoint","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._balance_batch","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._balance_batch","relation":"calls","name":"tolist","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._balance_batch","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._balance_batch","relation":"calls","name":"view","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._balance_batch","relation":"calls","name":"get_seqlen_balanced_partitions","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._balance_batch","relation":"calls","name":"tensor","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._balance_batch","relation":"calls","name":"reorder","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"DataProto"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._balance_batch","relation":"calls","name":"log_seqlen_unbalance","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._balance_batch","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_validate","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"pprint","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"log","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"GenerationConfig","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"LLMGenerationManager","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"from_single_dict","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"repeat","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"compute_advantage"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"pop","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"compute_advantage"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"generate_sequences","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"uuid4","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"repeat","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"compute_advantage"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"union","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"compute_advantage"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"long","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"clone","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"run_llm_loop","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"LLMGenerationManager"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"keys","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"long","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"compute_log_prob","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"union","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"copy","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"repeat","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"compute_advantage"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"union","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"compute_advantage"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"array","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_balance_batch","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"tolist","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"keys","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"long","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"compute_ref_log_prob","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"union","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"compute_advantage"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"compute_values","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"union","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"compute_advantage"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"compute_rm_score","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"union","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"compute_advantage"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"reward_fn","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"tolist","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"apply_dual_calib_to_scores","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"argmax","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"flip","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"long","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_collect_and_log_statistics","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"apply_kl_penalty","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"compute_advantage","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"update_critic","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"reduce_metrics","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_create_loss_mask","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"update_actor","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"reduce_metrics","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_validate","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_timer","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_save_checkpoint","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"compute_data_metrics","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"compute_timing_metrics","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"log","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"_validate","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"pprint","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer.fit","relation":"calls","name":"log","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"csv","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"os","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"tolist","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"list"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"join","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"split","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"rstrip","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"join","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"basename","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"dirname","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"join","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"exists","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"open","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"DictWriter","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"writeheader","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"writerow","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py","relation":"imports","specifier":"json","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"replace","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"join","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"open","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"bool","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"decode","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"tolist","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"str","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"int","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"write","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._save_sample_level_data","relation":"calls","name":"dumps","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"min","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"max","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"sum","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"std","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"range","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"compute_ece","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"compute_auroc","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"hash","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"hash","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics._mean_list","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics._mean_list","relation":"calls","name":"get","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics._mean_list","relation":"calls","name":"len","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics._mean_list","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics._mean_list","relation":"calls","name":"float","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics._mean_list","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"dict"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"_mean_list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"_mean_list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"_mean_list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"_mean_list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"_mean_list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"_mean_list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"_mean_list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"_mean_list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"_mean_list","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"setdefault","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"items","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"mean","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"_save_sample_level_data","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py","recvType":"RayPPOTrainer"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._collect_and_log_statistics","relation":"calls","name":"print","viaMember":false,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer","relation":"contains","targetId":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_loss_mask","file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_loss_mask","relation":"calls","name":"update","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_loss_mask","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_loss_mask","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_loss_mask","relation":"calls","name":"item","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_loss_mask","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"},{"source":"verl/trainer/ppo/ray_trainer.py#RayPPOTrainer._create_loss_mask","relation":"calls","name":"sum","viaMember":true,"file":"verl/trainer/ppo/ray_trainer.py"}]},"verl/utils/__init__.py":{"size":678,"mtimeMs":1789275830321.203,"hash":"2a347bac8c8dd4741d639d32cd86e58a081d272b129d0c1125df9398e8661105","nodes":[{"id":"verl/utils/__init__.py","name":"__init__.py","kind":"file","path":"verl/utils/__init__.py","span":"L1-L18","signature":null,"exported":true,"origin":"ast","body_hash":"2a347bac8c8dd4741d639d32cd86e58a081d272b129d0c1125df9398e8661105","chars":678,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 . import tokenizer from .tokenizer import * __all__ = tokenizer.__all__"}],"rawEdges":[{"source":"verl/utils/__init__.py","relation":"imports","specifier":".","file":"verl/utils/__init__.py"},{"source":"verl/utils/__init__.py","relation":"imports","specifier":".tokenizer","file":"verl/utils/__init__.py"}]},"verl/utils/config.py":{"size":839,"mtimeMs":1789275830321.203,"hash":"b59763b8653abf6a011abf5304c2679b9e68690d2acc33a9a05fd6cfe0aab542","nodes":[{"id":"verl/utils/config.py","name":"config.py","kind":"file","path":"verl/utils/config.py","span":"L1-L24","signature":null,"exported":true,"origin":"ast","body_hash":"b59763b8653abf6a011abf5304c2679b9e68690d2acc33a9a05fd6cfe0aab542","chars":839,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 typing import Dict from omegaconf import DictConfig"},{"id":"verl/utils/config.py#update_dict_with_config","name":"update_dict_with_config","kind":"function","path":"verl/utils/config.py","span":"L20-L23","signature":"def update_dict_with_config(dictionary: Dict, config: DictConfig)","exported":true,"origin":"ast","body_hash":"f8e57104aa8190084737ee0325eed955ae30c37335b4e2583878089081b733ac","body_text":"def update_dict_with_config(dictionary: Dict, config: DictConfig): for key in dictionary: if hasattr(config, key): dictionary[key] = getattr(config, key)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/config.py","relation":"imports","specifier":"typing","file":"verl/utils/config.py"},{"source":"verl/utils/config.py","relation":"imports","specifier":"omegaconf","file":"verl/utils/config.py"},{"source":"verl/utils/config.py","relation":"contains","targetId":"verl/utils/config.py#update_dict_with_config","file":"verl/utils/config.py"},{"source":"verl/utils/config.py#update_dict_with_config","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/utils/config.py"},{"source":"verl/utils/config.py#update_dict_with_config","relation":"calls","name":"getattr","viaMember":false,"file":"verl/utils/config.py"}]},"verl/utils/dataset/__init__.py":{"size":670,"mtimeMs":1789275830321.203,"hash":"85b71f93709dab550ee5705fe8fd88e9cc77d4b9465249eaeb92bc7646bda7e3","nodes":[{"id":"verl/utils/dataset/__init__.py","name":"__init__.py","kind":"file","path":"verl/utils/dataset/__init__.py","span":"L1-L16","signature":null,"exported":true,"origin":"ast","body_hash":"85b71f93709dab550ee5705fe8fd88e9cc77d4b9465249eaeb92bc7646bda7e3","chars":670,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 .rl_dataset import RLHFDataset from .rm_dataset import RMDataset"}],"rawEdges":[{"source":"verl/utils/dataset/__init__.py","relation":"imports","specifier":".rl_dataset","file":"verl/utils/dataset/__init__.py"},{"source":"verl/utils/dataset/__init__.py","relation":"imports","specifier":".rm_dataset","file":"verl/utils/dataset/__init__.py"}]},"verl/utils/dataset/rl_dataset.py":{"size":5690,"mtimeMs":1789275830322.203,"hash":"e95b34a590956ed37e69631c87712c9dacbddc0fae632ca1e5403ed19e409815","nodes":[{"id":"verl/utils/dataset/rl_dataset.py","name":"rl_dataset.py","kind":"file","path":"verl/utils/dataset/rl_dataset.py","span":"L1-L156","signature":null,"exported":true,"origin":"ast","body_hash":"e95b34a590956ed37e69631c87712c9dacbddc0fae632ca1e5403ed19e409815","chars":5690,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 omegaconf import ListConfig import os from typing import List, Union import pandas as pd import torch import numpy as np from torch.utils.data import Dataset, DataLoader from transformers import AutoTokenizer, PreTrainedTokenizer from verl.utils.fs import copy_local_path_from_hdfs from verl.utils.model import compute_position_id_with_mask import verl.utils.torch_functional as verl_F"},{"id":"verl/utils/dataset/rl_dataset.py#collate_fn","name":"collate_fn","kind":"function","path":"verl/utils/dataset/rl_dataset.py","span":"L31-L55","signature":"def collate_fn(data_list: list[dict]) -> dict","exported":true,"origin":"ast","body_hash":"d6815837963dafb2d82a51baf34080308805aa9363770164d6a268f4f29bc0b8","body_text":"def collate_fn(data_list: list[dict]) -> dict: tensors = {} non_tensors = {} for data in data_list: for key, val in data.items(): if isinstance(val, torch.Tensor): if key not in tensors: tensors[key] = [] tensors[key].append(val) else: if key not in non_tensors: non_tensors[key] = [] non_tensors[key].append(val) for key, val in tensors.items(): tensors[key] = torch.stack(val, dim=0) for key, val in non_tensors.items(): non_tensors[key] = np.array(val, dtype=object) output = {} output.update(tensors) output.update(non_tensors) return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/dataset/rl_dataset.py#RLHFDataset","name":"RLHFDataset","kind":"class","path":"verl/utils/dataset/rl_dataset.py","span":"L58-L155","signature":"class RLHFDataset(Dataset)","exported":true,"origin":"ast","body_hash":"1bc8d838f7c67d2ef5dd1ce42ceee6d40d1d070f7ea59778fb41892fcdac36aa","body_text":"class RLHFDataset(Dataset): \"\"\" We assume the dataset contains a column that contains prompts and other information \"\"\" def __init__(self, parquet_files: Union[str, List[str]], tokenizer: PreTrainedTokenizer, prompt_key='prompt', max_prompt_length=1024, filter_prompts=True, cache_dir='~/.cache/verl/rlhf', chat_template_func=None, return_raw_chat=False, truncation='error'): if not isinstance(parquet_files, (List, ListConfig)): parquet_files = [parquet_files] self.parquet_files = parquet_files self.cache_dir = os.path.expanduser(cache_dir) self.tokenizer = tokenizer self.prompt_key = prompt_key self.max_prompt_length = max_prompt_length self.filter_prompts = filter_prompts self.return_raw_chat = return_raw_chat self.chat_template_func = chat_template_func self.truncation = truncation self._download() self._read_files_and_tokenize() def _download(self): from verl.utils.fs import copy_local_path_from_hdfs for i, parquet_file in enumerate(self.parquet_files): self.parquet_files[i] = copy_local_path_from_hdfs(src=parquet_file, cache_dir=self.cache_dir) def _read_files_and_tokenize(self): dataframes = [] for parquet_file in self.parquet_files: # read parquet files and cache dataframe = pd.read_parquet(parquet_file) dataframes.append(dataframe) self.dataframe = pd.concat(dataframes) print(f'original dataset len: {len(self.dataframe)}') # filter out too long prompts tokenizer = self.tokenizer prompt_key = self.prompt_key # nvm if prompt is too long # self.dataframe = self.dataframe[self.dataframe.apply(lambda doc: len( # tokenizer.apply_chat_template(doc[prompt_key], add_generation_prompt=True)) <= self.max_prompt_length, # axis=1)] print(f'filter dataset len: {len(self.dataframe)}') def __len__(self): return len(self.dataframe) def __getitem__(self, item): \"\"\" Note that we also return the raw_input_ids so that it can be combined with other chat template \"\"\" row_dict = self.dataframe.iloc[item].to_dict() chat = row_dict.pop(self.prompt_key) if self.tokenizer.chat_template: prompt_with_chat_template = self.tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=False) else: prompt_with_chat_template = chat[0]['content'] # prompt_with_chat_template = chat input_ids, attention_mask = verl_F.tokenize_and_postprocess_data(prompt=prompt_with_chat_template, tokenizer=self.tokenizer, max_length=self.max_prompt_length, pad_token_id=self.tokenizer.pad_token_id, left_pad=True, truncation=self.truncation) position_ids = compute_position_id_with_mask(attention_mask) row_dict['input_ids'] = input_ids[0] row_dict['attention_mask'] = attention_mask[0] row_dict['position_ids'] = position_ids[0] # encode prompts without chat template if self.return_raw_chat: row_dict['raw_prompt'] = chat.tolist() # add index for each prompt index = row_dict.get(\"extra_info\", {}).get(\"index\", 0) row_dict[\"index\"] = index return row_dict","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__init__","name":"__init__","kind":"method","path":"verl/utils/dataset/rl_dataset.py","span":"L63-L89","signature":"def __init__(self, parquet_files: Union[str, List[str]], tokenizer: PreTrainedTokenizer, prompt_key='prompt', max_prompt_length=1024, filter_prompts=True, cache_dir='~/.cache/verl/rlhf', chat_template_func=None, return_raw_chat=False, truncation='error')","exported":false,"origin":"ast","body_hash":"fc107d7cc11f30304e5caef8eaf39d94f749196f1956ee0eaaf3bdb85fd32556","body_text":"def __init__(self, parquet_files: Union[str, List[str]], tokenizer: PreTrainedTokenizer, prompt_key='prompt', max_prompt_length=1024, filter_prompts=True, cache_dir='~/.cache/verl/rlhf', chat_template_func=None, return_raw_chat=False, truncation='error'): if not isinstance(parquet_files, (List, ListConfig)): parquet_files = [parquet_files] self.parquet_files = parquet_files self.cache_dir = os.path.expanduser(cache_dir) self.tokenizer = tokenizer self.prompt_key = prompt_key self.max_prompt_length = max_prompt_length self.filter_prompts = filter_prompts self.return_raw_chat = return_raw_chat self.chat_template_func = chat_template_func self.truncation = truncation self._download() self._read_files_and_tokenize()","summary_state":"pending","summary":null,"crux":null,"owner":"RLHFDataset"},{"id":"verl/utils/dataset/rl_dataset.py#RLHFDataset._download","name":"_download","kind":"method","path":"verl/utils/dataset/rl_dataset.py","span":"L91-L94","signature":"def _download(self)","exported":false,"origin":"ast","body_hash":"dd3a2fe270c93c78e8c8f4ee96d331619eade93cda7943b2a334f5f3bbfab55e","body_text":"def _download(self): from verl.utils.fs import copy_local_path_from_hdfs for i, parquet_file in enumerate(self.parquet_files): self.parquet_files[i] = copy_local_path_from_hdfs(src=parquet_file, cache_dir=self.cache_dir)","summary_state":"pending","summary":null,"crux":null,"owner":"RLHFDataset"},{"id":"verl/utils/dataset/rl_dataset.py#RLHFDataset._read_files_and_tokenize","name":"_read_files_and_tokenize","kind":"method","path":"verl/utils/dataset/rl_dataset.py","span":"L96-L115","signature":"def _read_files_and_tokenize(self)","exported":false,"origin":"ast","body_hash":"bba69c686a046a9990644236b3da72415dba75e32c5b9e106e28ab5b06bbb977","body_text":"def _read_files_and_tokenize(self): dataframes = [] for parquet_file in self.parquet_files: # read parquet files and cache dataframe = pd.read_parquet(parquet_file) dataframes.append(dataframe) self.dataframe = pd.concat(dataframes) print(f'original dataset len: {len(self.dataframe)}') # filter out too long prompts tokenizer = self.tokenizer prompt_key = self.prompt_key # nvm if prompt is too long # self.dataframe = self.dataframe[self.dataframe.apply(lambda doc: len( # tokenizer.apply_chat_template(doc[prompt_key], add_generation_prompt=True)) <= self.max_prompt_length, # axis=1)] print(f'filter dataset len: {len(self.dataframe)}')","summary_state":"pending","summary":null,"crux":null,"owner":"RLHFDataset"},{"id":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__len__","name":"__len__","kind":"method","path":"verl/utils/dataset/rl_dataset.py","span":"L117-L118","signature":"def __len__(self)","exported":false,"origin":"ast","body_hash":"bdc6551e4036bc82795afadb8c779143efb9cdab162a30f71f84f75003692a05","body_text":"def __len__(self): return len(self.dataframe)","summary_state":"pending","summary":null,"crux":null,"owner":"RLHFDataset"},{"id":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__getitem__","name":"__getitem__","kind":"method","path":"verl/utils/dataset/rl_dataset.py","span":"L120-L155","signature":"def __getitem__(self, item)","exported":false,"origin":"ast","body_hash":"7c17a0f75e95e4c8e58e3f4bc51c11041d136707900f72c059d6d756efd472dc","body_text":"def __getitem__(self, item): \"\"\" Note that we also return the raw_input_ids so that it can be combined with other chat template \"\"\" row_dict = self.dataframe.iloc[item].to_dict() chat = row_dict.pop(self.prompt_key) if self.tokenizer.chat_template: prompt_with_chat_template = self.tokenizer.apply_chat_template(chat, add_generation_prompt=True, tokenize=False) else: prompt_with_chat_template = chat[0]['content'] # prompt_with_chat_template = chat input_ids, attention_mask = verl_F.tokenize_and_postprocess_data(prompt=prompt_with_chat_template, tokenizer=self.tokenizer, max_length=self.max_prompt_length, pad_token_id=self.tokenizer.pad_token_id, left_pad=True, truncation=self.truncation) position_ids = compute_position_id_with_mask(attention_mask) row_dict['input_ids'] = input_ids[0] row_dict['attention_mask'] = attention_mask[0] row_dict['position_ids'] = position_ids[0] # encode prompts without chat template if self.return_raw_chat: row_dict['raw_prompt'] = chat.tolist() # add index for each prompt index = row_dict.get(\"extra_info\", {}).get(\"index\", 0) row_dict[\"index\"] = index return row_dict","summary_state":"pending","summary":null,"crux":null,"owner":"RLHFDataset"}],"rawEdges":[{"source":"verl/utils/dataset/rl_dataset.py","relation":"imports","specifier":"omegaconf","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py","relation":"imports","specifier":"os","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py","relation":"imports","specifier":"typing","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py","relation":"imports","specifier":"torch","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py","relation":"imports","specifier":"torch.utils.data","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py","relation":"imports","specifier":"transformers","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py","relation":"imports","specifier":"verl.utils.fs","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py","relation":"imports","specifier":"verl.utils.model","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py","relation":"contains","targetId":"verl/utils/dataset/rl_dataset.py#collate_fn","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#collate_fn","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#collate_fn","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#collate_fn","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#collate_fn","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#collate_fn","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#collate_fn","relation":"calls","name":"stack","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#collate_fn","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#collate_fn","relation":"calls","name":"array","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#collate_fn","relation":"calls","name":"update","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#collate_fn","relation":"calls","name":"update","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py","relation":"contains","targetId":"verl/utils/dataset/rl_dataset.py#RLHFDataset","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset","relation":"extends","name":"Dataset","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset","relation":"contains","targetId":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__init__","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__init__","relation":"calls","name":"expanduser","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__init__","relation":"calls","name":"_download","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py","recvType":"RLHFDataset"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__init__","relation":"calls","name":"_read_files_and_tokenize","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py","recvType":"RLHFDataset"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset","relation":"contains","targetId":"verl/utils/dataset/rl_dataset.py#RLHFDataset._download","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py","relation":"imports","specifier":"verl.utils.fs","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset._download","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset._download","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset","relation":"contains","targetId":"verl/utils/dataset/rl_dataset.py#RLHFDataset._read_files_and_tokenize","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset._read_files_and_tokenize","relation":"calls","name":"read_parquet","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset._read_files_and_tokenize","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset._read_files_and_tokenize","relation":"calls","name":"concat","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset._read_files_and_tokenize","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset._read_files_and_tokenize","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset._read_files_and_tokenize","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset._read_files_and_tokenize","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset","relation":"contains","targetId":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__len__","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__len__","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset","relation":"contains","targetId":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__getitem__","file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__getitem__","relation":"calls","name":"to_dict","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__getitem__","relation":"calls","name":"pop","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__getitem__","relation":"calls","name":"apply_chat_template","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__getitem__","relation":"calls","name":"tokenize_and_postprocess_data","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__getitem__","relation":"calls","name":"compute_position_id_with_mask","viaMember":false,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__getitem__","relation":"calls","name":"tolist","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__getitem__","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"},{"source":"verl/utils/dataset/rl_dataset.py#RLHFDataset.__getitem__","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/dataset/rl_dataset.py"}]},"verl/utils/dataset/rm_dataset.py":{"size":5494,"mtimeMs":1789275830322.203,"hash":"927ff2808f7c0928d4399cc1cd213fab2d922365c055a67ad6d4228b162c5d1d","nodes":[{"id":"verl/utils/dataset/rm_dataset.py","name":"rm_dataset.py","kind":"file","path":"verl/utils/dataset/rm_dataset.py","span":"L1-L143","signature":null,"exported":true,"origin":"ast","body_hash":"927ff2808f7c0928d4399cc1cd213fab2d922365c055a67ad6d4228b162c5d1d","chars":5494,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 typing import List, Union import pandas as pd import torch from torch.utils.data import Dataset from transformers import AutoTokenizer from verl.utils import hf_tokenizer"},{"id":"verl/utils/dataset/rm_dataset.py#download_files_distributed","name":"download_files_distributed","kind":"function","path":"verl/utils/dataset/rm_dataset.py","span":"L27-L37","signature":"def download_files_distributed(download_fn)","exported":true,"origin":"ast","body_hash":"25783b33b80a86aaf4a82440c7ada187d853b5997655efaf6eef02b74ef0e044","body_text":"def download_files_distributed(download_fn): import torch.distributed if torch.distributed.is_initialized(): if torch.distributed.get_rank() == 0: # download files download_fn() torch.distributed.barrier() else: # download anyway download_fn()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/dataset/rm_dataset.py#RMDataset","name":"RMDataset","kind":"class","path":"verl/utils/dataset/rm_dataset.py","span":"L40-L143","signature":"class RMDataset(Dataset)","exported":true,"origin":"ast","body_hash":"e65f05aa08f265eb2c6cd2b5f0968cd35ebceb16ed848968cc29be582a111d3f","body_text":"class RMDataset(Dataset): def __init__(self, parquet_files: Union[str, List[str]], tokenizer, prompt_key='prompt', chosen_key='chosen', rejected_key='rejected', max_length=1024, add_eos=True, cache_dir='~/.cache/verl/rm'): if not isinstance(parquet_files, List): parquet_files = [parquet_files] self.parquet_files = parquet_files self.cache_dir = os.path.expanduser(cache_dir) if isinstance(tokenizer, str): tokenizer = hf_tokenizer(tokenizer) self.tokenizer = tokenizer self.prompt_key = prompt_key self.chosen_key = chosen_key self.rejected_key = rejected_key self.add_eos = add_eos self.max_length = max_length self._download() self._read_files_and_tokenize() def _download(self): def _download_files(): from verl.utils.fs import copy, _is_non_local os.makedirs(self.cache_dir, exist_ok=True) assert os.path.exists(self.cache_dir) for i, parquet_file in enumerate(self.parquet_files): if _is_non_local(parquet_file): dst = os.path.join(self.cache_dir, os.path.basename(parquet_file)) if not os.path.exists(dst): copy(src=parquet_file, dst=dst) self.parquet_files[i] = dst download_files_distributed(_download_files) def _read_files_and_tokenize(self): dataframes = [] for parquet_file in self.parquet_files: # read parquet files and cache dataframe = pd.read_parquet(parquet_file) dataframes.append(dataframe) self.dataframe = pd.concat(dataframes) self.prompts = self.dataframe[self.prompt_key].tolist() self.chosen_responses = self.dataframe[self.chosen_key].tolist() self.rejected_responses = self.dataframe[self.rejected_key].tolist() def __len__(self): return len(self.prompts) def _pad_to_length(self, input_ids, attention_mask): curr_length = input_ids.shape[-1] if curr_length < self.max_length: input_ids = torch.cat( (input_ids, torch.zeros(size=(self.max_length - curr_length,), dtype=input_ids.dtype)), dim=-1) attention_mask = torch.cat( (attention_mask, torch.zeros(size=(self.max_length - curr_length,), dtype=attention_mask.dtype)), dim=-1) elif curr_length > self.max_length: input_ids = input_ids[:self.max_length] attention_mask = attention_mask[:self.max_length] return input_ids, attention_mask def __getitem__(self, item): prompt = self.prompts[item] chosen_response = self.chosen_responses[item] rejected_response = self.rejected_responses[item] prompt_ids = self.tokenizer(prompt, return_tensors='pt')['input_ids'][0] chosen_response_ids = self.tokenizer(chosen_response, return_tensors='pt')['input_ids'][0] rejected_response_ids = self.tokenizer(rejected_response, return_tensors='pt')['input_ids'][0] if self.add_eos: chosen_response_ids = torch.cat((chosen_response_ids, torch.tensor([self.tokenizer.eos_token_id])), dim=-1) rejected_response_ids = torch.cat((rejected_response_ids, torch.tensor([self.tokenizer.eos_token_id])), dim=-1) chosen_input_ids = torch.cat((prompt_ids, chosen_response_ids), dim=-1) chosen_attention_mask = torch.ones_like(chosen_input_ids) rejected_input_ids = torch.cat((prompt_ids, rejected_response_ids), dim=-1) rejected_attention_mask = torch.ones_like(rejected_input_ids) chosen_input_ids, chosen_attention_mask = self._pad_to_length(chosen_input_ids, chosen_attention_mask) rejected_input_ids, rejected_attention_mask = self._pad_to_length(rejected_input_ids, rejected_attention_mask) input_ids = torch.stack((chosen_input_ids, rejected_input_ids), dim=0) attention_mask = torch.stack((rejected_input_ids, rejected_attention_mask), dim=0) return { 'input_ids': input_ids, 'attention_mask': attention_mask, }","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/dataset/rm_dataset.py#RMDataset.__init__","name":"__init__","kind":"method","path":"verl/utils/dataset/rm_dataset.py","span":"L42-L68","signature":"def __init__(self, parquet_files: Union[str, List[str]], tokenizer, prompt_key='prompt', chosen_key='chosen', rejected_key='rejected', max_length=1024, add_eos=True, cache_dir='~/.cache/verl/rm')","exported":false,"origin":"ast","body_hash":"080b43c67df9d7f26009e63d2fb7d2f5c55271bd9318cd1d4cfcb9ef40fd18a3","body_text":"def __init__(self, parquet_files: Union[str, List[str]], tokenizer, prompt_key='prompt', chosen_key='chosen', rejected_key='rejected', max_length=1024, add_eos=True, cache_dir='~/.cache/verl/rm'): if not isinstance(parquet_files, List): parquet_files = [parquet_files] self.parquet_files = parquet_files self.cache_dir = os.path.expanduser(cache_dir) if isinstance(tokenizer, str): tokenizer = hf_tokenizer(tokenizer) self.tokenizer = tokenizer self.prompt_key = prompt_key self.chosen_key = chosen_key self.rejected_key = rejected_key self.add_eos = add_eos self.max_length = max_length self._download() self._read_files_and_tokenize()","summary_state":"pending","summary":null,"crux":null,"owner":"RMDataset"},{"id":"verl/utils/dataset/rm_dataset.py#RMDataset._download","name":"_download","kind":"method","path":"verl/utils/dataset/rm_dataset.py","span":"L70-L83","signature":"def _download(self)","exported":false,"origin":"ast","body_hash":"58b23eeed746d2f26020330d5f15d3684b0ad85e950eaa96a744c90fa216aac9","body_text":"def _download(self): def _download_files(): from verl.utils.fs import copy, _is_non_local os.makedirs(self.cache_dir, exist_ok=True) assert os.path.exists(self.cache_dir) for i, parquet_file in enumerate(self.parquet_files): if _is_non_local(parquet_file): dst = os.path.join(self.cache_dir, os.path.basename(parquet_file)) if not os.path.exists(dst): copy(src=parquet_file, dst=dst) self.parquet_files[i] = dst download_files_distributed(_download_files)","summary_state":"pending","summary":null,"crux":null,"owner":"RMDataset"},{"id":"verl/utils/dataset/rm_dataset.py#RMDataset._download._download_files","name":"_download_files","kind":"function","path":"verl/utils/dataset/rm_dataset.py","span":"L72-L81","signature":"def _download_files()","exported":false,"origin":"ast","body_hash":"37ec85c99188b886f6ab05bd86bd5e1839bd466f1dbdf6c59a585deae16388c0","body_text":"def _download_files(): from verl.utils.fs import copy, _is_non_local os.makedirs(self.cache_dir, exist_ok=True) assert os.path.exists(self.cache_dir) for i, parquet_file in enumerate(self.parquet_files): if _is_non_local(parquet_file): dst = os.path.join(self.cache_dir, os.path.basename(parquet_file)) if not os.path.exists(dst): copy(src=parquet_file, dst=dst) self.parquet_files[i] = dst","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/dataset/rm_dataset.py#RMDataset._read_files_and_tokenize","name":"_read_files_and_tokenize","kind":"method","path":"verl/utils/dataset/rm_dataset.py","span":"L85-L94","signature":"def _read_files_and_tokenize(self)","exported":false,"origin":"ast","body_hash":"334cf80541c2fb27349fbacc2ea08d134a165c7363e2eeba4875ac7508d811ae","body_text":"def _read_files_and_tokenize(self): dataframes = [] for parquet_file in self.parquet_files: # read parquet files and cache dataframe = pd.read_parquet(parquet_file) dataframes.append(dataframe) self.dataframe = pd.concat(dataframes) self.prompts = self.dataframe[self.prompt_key].tolist() self.chosen_responses = self.dataframe[self.chosen_key].tolist() self.rejected_responses = self.dataframe[self.rejected_key].tolist()","summary_state":"pending","summary":null,"crux":null,"owner":"RMDataset"},{"id":"verl/utils/dataset/rm_dataset.py#RMDataset.__len__","name":"__len__","kind":"method","path":"verl/utils/dataset/rm_dataset.py","span":"L96-L97","signature":"def __len__(self)","exported":false,"origin":"ast","body_hash":"d3ec705f558c053786c913764f65f0afa378fba5294781471e2b56efc012fa97","body_text":"def __len__(self): return len(self.prompts)","summary_state":"pending","summary":null,"crux":null,"owner":"RMDataset"},{"id":"verl/utils/dataset/rm_dataset.py#RMDataset._pad_to_length","name":"_pad_to_length","kind":"method","path":"verl/utils/dataset/rm_dataset.py","span":"L99-L112","signature":"def _pad_to_length(self, input_ids, attention_mask)","exported":false,"origin":"ast","body_hash":"3768980b7c68059e90bff4654e7af3830a815d882412afe5ed4c4f586e5ff58f","body_text":"def _pad_to_length(self, input_ids, attention_mask): curr_length = input_ids.shape[-1] if curr_length < self.max_length: input_ids = torch.cat( (input_ids, torch.zeros(size=(self.max_length - curr_length,), dtype=input_ids.dtype)), dim=-1) attention_mask = torch.cat( (attention_mask, torch.zeros(size=(self.max_length - curr_length,), dtype=attention_mask.dtype)), dim=-1) elif curr_length > self.max_length: input_ids = input_ids[:self.max_length] attention_mask = attention_mask[:self.max_length] return input_ids, attention_mask","summary_state":"pending","summary":null,"crux":null,"owner":"RMDataset"},{"id":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","name":"__getitem__","kind":"method","path":"verl/utils/dataset/rm_dataset.py","span":"L114-L143","signature":"def __getitem__(self, item)","exported":false,"origin":"ast","body_hash":"2320ceaf44c487dd5b1baae3682076963200a70ddc88c27752555bccc4e65eaf","body_text":"def __getitem__(self, item): prompt = self.prompts[item] chosen_response = self.chosen_responses[item] rejected_response = self.rejected_responses[item] prompt_ids = self.tokenizer(prompt, return_tensors='pt')['input_ids'][0] chosen_response_ids = self.tokenizer(chosen_response, return_tensors='pt')['input_ids'][0] rejected_response_ids = self.tokenizer(rejected_response, return_tensors='pt')['input_ids'][0] if self.add_eos: chosen_response_ids = torch.cat((chosen_response_ids, torch.tensor([self.tokenizer.eos_token_id])), dim=-1) rejected_response_ids = torch.cat((rejected_response_ids, torch.tensor([self.tokenizer.eos_token_id])), dim=-1) chosen_input_ids = torch.cat((prompt_ids, chosen_response_ids), dim=-1) chosen_attention_mask = torch.ones_like(chosen_input_ids) rejected_input_ids = torch.cat((prompt_ids, rejected_response_ids), dim=-1) rejected_attention_mask = torch.ones_like(rejected_input_ids) chosen_input_ids, chosen_attention_mask = self._pad_to_length(chosen_input_ids, chosen_attention_mask) rejected_input_ids, rejected_attention_mask = self._pad_to_length(rejected_input_ids, rejected_attention_mask) input_ids = torch.stack((chosen_input_ids, rejected_input_ids), dim=0) attention_mask = torch.stack((rejected_input_ids, rejected_attention_mask), dim=0) return { 'input_ids': input_ids, 'attention_mask': attention_mask, }","summary_state":"pending","summary":null,"crux":null,"owner":"RMDataset"}],"rawEdges":[{"source":"verl/utils/dataset/rm_dataset.py","relation":"imports","specifier":"os","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py","relation":"imports","specifier":"typing","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py","relation":"imports","specifier":"torch","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py","relation":"imports","specifier":"torch.utils.data","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py","relation":"imports","specifier":"transformers","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py","relation":"imports","specifier":"verl.utils","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py","relation":"contains","targetId":"verl/utils/dataset/rm_dataset.py#download_files_distributed","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py","relation":"imports","specifier":"torch.distributed","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#download_files_distributed","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#download_files_distributed","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#download_files_distributed","relation":"calls","name":"download_fn","viaMember":false,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#download_files_distributed","relation":"calls","name":"barrier","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#download_files_distributed","relation":"calls","name":"download_fn","viaMember":false,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py","relation":"contains","targetId":"verl/utils/dataset/rm_dataset.py#RMDataset","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset","relation":"extends","name":"Dataset","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset","relation":"contains","targetId":"verl/utils/dataset/rm_dataset.py#RMDataset.__init__","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__init__","relation":"calls","name":"expanduser","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__init__","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__init__","relation":"calls","name":"_download","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py","recvType":"RMDataset"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__init__","relation":"calls","name":"_read_files_and_tokenize","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py","recvType":"RMDataset"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset","relation":"contains","targetId":"verl/utils/dataset/rm_dataset.py#RMDataset._download","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._download","relation":"contains","targetId":"verl/utils/dataset/rm_dataset.py#RMDataset._download._download_files","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py","relation":"imports","specifier":"verl.utils.fs","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._download._download_files","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._download._download_files","relation":"calls","name":"exists","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._download._download_files","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._download._download_files","relation":"calls","name":"_is_non_local","viaMember":false,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._download._download_files","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._download._download_files","relation":"calls","name":"basename","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._download._download_files","relation":"calls","name":"exists","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._download._download_files","relation":"calls","name":"copy","viaMember":false,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._download","relation":"calls","name":"download_files_distributed","viaMember":false,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset","relation":"contains","targetId":"verl/utils/dataset/rm_dataset.py#RMDataset._read_files_and_tokenize","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._read_files_and_tokenize","relation":"calls","name":"read_parquet","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._read_files_and_tokenize","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._read_files_and_tokenize","relation":"calls","name":"concat","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._read_files_and_tokenize","relation":"calls","name":"tolist","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._read_files_and_tokenize","relation":"calls","name":"tolist","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._read_files_and_tokenize","relation":"calls","name":"tolist","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset","relation":"contains","targetId":"verl/utils/dataset/rm_dataset.py#RMDataset.__len__","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__len__","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset","relation":"contains","targetId":"verl/utils/dataset/rm_dataset.py#RMDataset._pad_to_length","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._pad_to_length","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._pad_to_length","relation":"calls","name":"zeros","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._pad_to_length","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset._pad_to_length","relation":"calls","name":"zeros","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset","relation":"contains","targetId":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"tokenizer","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py","recvType":"RMDataset"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"tokenizer","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py","recvType":"RMDataset"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"tokenizer","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py","recvType":"RMDataset"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"tensor","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"tensor","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"ones_like","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"ones_like","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"_pad_to_length","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py","recvType":"RMDataset"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"_pad_to_length","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py","recvType":"RMDataset"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"stack","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"},{"source":"verl/utils/dataset/rm_dataset.py#RMDataset.__getitem__","relation":"calls","name":"stack","viaMember":true,"file":"verl/utils/dataset/rm_dataset.py"}]},"verl/utils/debug/__init__.py":{"size":646,"mtimeMs":1789275830322.203,"hash":"7c2e66dafef4ae2d091deec16c763299acc927bc7486268190239d9290c0b77c","nodes":[{"id":"verl/utils/debug/__init__.py","name":"__init__.py","kind":"file","path":"verl/utils/debug/__init__.py","span":"L1-L15","signature":null,"exported":true,"origin":"ast","body_hash":"7c2e66dafef4ae2d091deec16c763299acc927bc7486268190239d9290c0b77c","chars":646,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 .performance import log_gpu_memory_usage"}],"rawEdges":[{"source":"verl/utils/debug/__init__.py","relation":"imports","specifier":".performance","file":"verl/utils/debug/__init__.py"}]},"verl/utils/debug/performance.py":{"size":1214,"mtimeMs":1789275830322.203,"hash":"a27524298cce8e7c5b0dc060a9d58f4c79d5a63a2da36293dee491dbb927a211","nodes":[{"id":"verl/utils/debug/performance.py","name":"performance.py","kind":"file","path":"verl/utils/debug/performance.py","span":"L1-L31","signature":null,"exported":true,"origin":"ast","body_hash":"a27524298cce8e7c5b0dc060a9d58f4c79d5a63a2da36293dee491dbb927a211","chars":1214,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 import torch.distributed as dist import logging"},{"id":"verl/utils/debug/performance.py#log_gpu_memory_usage","name":"log_gpu_memory_usage","kind":"function","path":"verl/utils/debug/performance.py","span":"L20-L30","signature":"def log_gpu_memory_usage(head: str, logger: logging.Logger = None, level=logging.DEBUG, rank: int = 0)","exported":true,"origin":"ast","body_hash":"c8d5d2476fe5c411995c8d66ad883d58936a749c8688e52d1b6da806bbce7d5e","body_text":"def log_gpu_memory_usage(head: str, logger: logging.Logger = None, level=logging.DEBUG, rank: int = 0): if (not dist.is_initialized()) or (rank is None) or (dist.get_rank() == rank): memory_allocated = torch.cuda.memory_allocated() / 1024**3 memory_reserved = torch.cuda.memory_reserved() / 1024**3 message = f'{head}, memory allocated (GB): {memory_allocated}, memory reserved (GB): {memory_reserved}' if logger is None: print(message) else: logger.log(msg=message, level=level)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/debug/performance.py","relation":"imports","specifier":"torch","file":"verl/utils/debug/performance.py"},{"source":"verl/utils/debug/performance.py","relation":"imports","specifier":"logging","file":"verl/utils/debug/performance.py"},{"source":"verl/utils/debug/performance.py","relation":"contains","targetId":"verl/utils/debug/performance.py#log_gpu_memory_usage","file":"verl/utils/debug/performance.py"},{"source":"verl/utils/debug/performance.py#log_gpu_memory_usage","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/utils/debug/performance.py"},{"source":"verl/utils/debug/performance.py#log_gpu_memory_usage","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/utils/debug/performance.py"},{"source":"verl/utils/debug/performance.py#log_gpu_memory_usage","relation":"calls","name":"memory_allocated","viaMember":true,"file":"verl/utils/debug/performance.py"},{"source":"verl/utils/debug/performance.py#log_gpu_memory_usage","relation":"calls","name":"memory_reserved","viaMember":true,"file":"verl/utils/debug/performance.py"},{"source":"verl/utils/debug/performance.py#log_gpu_memory_usage","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/debug/performance.py"},{"source":"verl/utils/debug/performance.py#log_gpu_memory_usage","relation":"calls","name":"log","viaMember":true,"file":"verl/utils/debug/performance.py"}]},"verl/utils/debug/trajectory_tracker.py":{"size":3174,"mtimeMs":1789275830322.203,"hash":"f47db3824fe10b827c62110cb49151e075ae3e8bc9ea95905bb5dd915d60b897","nodes":[{"id":"verl/utils/debug/trajectory_tracker.py","name":"trajectory_tracker.py","kind":"file","path":"verl/utils/debug/trajectory_tracker.py","span":"L1-L109","signature":null,"exported":true,"origin":"ast","body_hash":"f47db3824fe10b827c62110cb49151e075ae3e8bc9ea95905bb5dd915d60b897","chars":3174,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Trajectory tracker can be inserted into code to save the intermediate results. The results will be dump to hdfs for offline comparison. Each process will have a client that first move all the tensors to CPU \"\"\" from verl.utils.hdfs_io import makedirs, copy import torch import os import ray import io import tempfile from collections import deque remote_copy = ray.remote(copy) @ray.remote @ray.remote if __name__ == '__main__': # testing os.environ['VERL_ENABLE_TRACKER'] = '1' os.environ['VERL_TRACKER_HDFS_DIR'] = '~/debug/test' @ray.remote ray.init() output_lst = [] for i in range(10): output_lst.append(process.remote(i)) out = ray.get(output_lst) tracker = get_trajectory_tracker() ray.get(tracker.wait_for_hdfs.remote())"},{"id":"verl/utils/debug/trajectory_tracker.py#save_to_hdfs","name":"save_to_hdfs","kind":"function","path":"verl/utils/debug/trajectory_tracker.py","span":"L33-L46","signature":"def save_to_hdfs(data: io.BytesIO, name, hdfs_dir, verbose)","exported":true,"origin":"ast","body_hash":"440259a8f5d20a69d6522429c53e52353c83f9c3768028f2b55e09311ebeddcb","body_text":"def save_to_hdfs(data: io.BytesIO, name, hdfs_dir, verbose): filename = name + '.pth' with tempfile.TemporaryDirectory() as tmpdirname: local_filepath = os.path.join(tmpdirname, filename) with open(local_filepath, 'wb') as f: f.write(data.getbuffer()) # upload to hdfs if verbose: print(f'Saving {local_filepath} to {hdfs_dir}') try: copy(local_filepath, hdfs_dir) except Exception as e: print(e)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker","name":"TrajectoryTracker","kind":"class","path":"verl/utils/debug/trajectory_tracker.py","span":"L50-L66","signature":"class TrajectoryTracker()","exported":true,"origin":"ast","body_hash":"03f90d59c1171832e33d172e9b8def1dc5f1b7c9204d5a92a1ec028552db796c","body_text":"class TrajectoryTracker(): def __init__(self, hdfs_dir, verbose) -> None: self.hdfs_dir = hdfs_dir makedirs(hdfs_dir) self.verbose = verbose self.handle = deque() def dump(self, data: io.BytesIO, name): # get a temp file and write to it self.handle.append(save_to_hdfs.remote(data, name, self.hdfs_dir, self.verbose)) def wait_for_hdfs(self): while len(self.handle) != 0: future = self.handle.popleft() ray.get(future)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.__init__","name":"__init__","kind":"method","path":"verl/utils/debug/trajectory_tracker.py","span":"L52-L57","signature":"def __init__(self, hdfs_dir, verbose) -> None","exported":false,"origin":"ast","body_hash":"a162a9a03b4c731aa3e0443f15bb5b6d0cd363ea39d5a4e90adaa7a4262c5e5a","body_text":"def __init__(self, hdfs_dir, verbose) -> None: self.hdfs_dir = hdfs_dir makedirs(hdfs_dir) self.verbose = verbose self.handle = deque()","summary_state":"pending","summary":null,"crux":null,"owner":"TrajectoryTracker"},{"id":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.dump","name":"dump","kind":"method","path":"verl/utils/debug/trajectory_tracker.py","span":"L59-L61","signature":"def dump(self, data: io.BytesIO, name): # get a temp file and write to it","exported":true,"origin":"ast","body_hash":"faf7cfbd714681301cf6b061294e4ede78edc32be2136fb8fb8098de19441214","body_text":"def dump(self, data: io.BytesIO, name): # get a temp file and write to it self.handle.append(save_to_hdfs.remote(data, name, self.hdfs_dir, self.verbose))","summary_state":"pending","summary":null,"crux":null,"owner":"TrajectoryTracker"},{"id":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.wait_for_hdfs","name":"wait_for_hdfs","kind":"method","path":"verl/utils/debug/trajectory_tracker.py","span":"L63-L66","signature":"def wait_for_hdfs(self)","exported":true,"origin":"ast","body_hash":"69d606fffac7c6c28c4e8ec7395a18772d6faefaa71ad42728a65dab4f2e7abd","body_text":"def wait_for_hdfs(self): while len(self.handle) != 0: future = self.handle.popleft() ray.get(future)","summary_state":"pending","summary":null,"crux":null,"owner":"TrajectoryTracker"},{"id":"verl/utils/debug/trajectory_tracker.py#dump_data","name":"dump_data","kind":"function","path":"verl/utils/debug/trajectory_tracker.py","span":"L69-L76","signature":"def dump_data(data, name)","exported":true,"origin":"ast","body_hash":"c108432ca791900d48f3cd2f032be63a7bd590b27b3aa20cf5710a6561b8b3f1","body_text":"def dump_data(data, name): enable = os.getenv('VERL_ENABLE_TRACKER', '0') == '1' if not enable: return buffer = io.BytesIO() torch.save(data, buffer) tracker = get_trajectory_tracker() ray.get(tracker.dump.remote(buffer, name))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/debug/trajectory_tracker.py#get_trajectory_tracker","name":"get_trajectory_tracker","kind":"function","path":"verl/utils/debug/trajectory_tracker.py","span":"L79-L85","signature":"def get_trajectory_tracker()","exported":true,"origin":"ast","body_hash":"e5df0d883a3f01f535d7102258bf81289734a66143fddf8ffc8cbde0c6cd1857","body_text":"def get_trajectory_tracker(): hdfs_dir = os.getenv('VERL_TRACKER_HDFS_DIR', default=None) verbose = os.getenv('VERL_TRACKER_VERBOSE', default='0') == '1' assert hdfs_dir is not None tracker = TrajectoryTracker.options(name=\"global_tracker\", get_if_exists=True, lifetime=\"detached\").remote(hdfs_dir, verbose) return tracker","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/debug/trajectory_tracker.py#process","name":"process","kind":"function","path":"verl/utils/debug/trajectory_tracker.py","span":"L94-L96","signature":"def process(iter)","exported":true,"origin":"ast","body_hash":"162022ef06fc9ad8e33f53c3aa7e3617089f495b60430468035c47fc65be37ba","body_text":"def process(iter): data = {'obs': torch.randn(10, 20)} dump_data(data, f'process_{iter}_obs')","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/debug/trajectory_tracker.py","relation":"imports","specifier":"verl.utils.hdfs_io","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"imports","specifier":"torch","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"imports","specifier":"os","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"imports","specifier":"ray","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"imports","specifier":"io","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"imports","specifier":"tempfile","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"imports","specifier":"collections","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"calls","name":"remote","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"contains","targetId":"verl/utils/debug/trajectory_tracker.py#save_to_hdfs","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#save_to_hdfs","relation":"calls","name":"TemporaryDirectory","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#save_to_hdfs","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#save_to_hdfs","relation":"calls","name":"open","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#save_to_hdfs","relation":"calls","name":"write","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#save_to_hdfs","relation":"calls","name":"getbuffer","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#save_to_hdfs","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#save_to_hdfs","relation":"calls","name":"copy","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#save_to_hdfs","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"contains","targetId":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker","relation":"contains","targetId":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.__init__","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.__init__","relation":"calls","name":"makedirs","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.__init__","relation":"calls","name":"deque","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker","relation":"contains","targetId":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.dump","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.dump","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py","recvType":"deque"},{"source":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.dump","relation":"calls","name":"remote","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker","relation":"contains","targetId":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.wait_for_hdfs","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.wait_for_hdfs","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.wait_for_hdfs","relation":"calls","name":"popleft","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py","recvType":"deque"},{"source":"verl/utils/debug/trajectory_tracker.py#TrajectoryTracker.wait_for_hdfs","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"contains","targetId":"verl/utils/debug/trajectory_tracker.py#dump_data","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#dump_data","relation":"calls","name":"getenv","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#dump_data","relation":"calls","name":"BytesIO","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#dump_data","relation":"calls","name":"save","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#dump_data","relation":"calls","name":"get_trajectory_tracker","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#dump_data","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#dump_data","relation":"calls","name":"remote","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"contains","targetId":"verl/utils/debug/trajectory_tracker.py#get_trajectory_tracker","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#get_trajectory_tracker","relation":"calls","name":"getenv","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#get_trajectory_tracker","relation":"calls","name":"getenv","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#get_trajectory_tracker","relation":"calls","name":"remote","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#get_trajectory_tracker","relation":"calls","name":"options","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"contains","targetId":"verl/utils/debug/trajectory_tracker.py#process","file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#process","relation":"calls","name":"randn","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py#process","relation":"calls","name":"dump_data","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"calls","name":"init","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"calls","name":"remote","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"calls","name":"get_trajectory_tracker","viaMember":false,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"},{"source":"verl/utils/debug/trajectory_tracker.py","relation":"calls","name":"remote","viaMember":true,"file":"verl/utils/debug/trajectory_tracker.py"}]},"verl/utils/distributed.py":{"size":1124,"mtimeMs":1789275830322.203,"hash":"ee90007c15948b647338f404dd865cc5a0e2586903b8790ce6331d1cf01d2244","nodes":[{"id":"verl/utils/distributed.py","name":"distributed.py","kind":"file","path":"verl/utils/distributed.py","span":"L1-L29","signature":null,"exported":true,"origin":"ast","body_hash":"ee90007c15948b647338f404dd865cc5a0e2586903b8790ce6331d1cf01d2244","chars":1124,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 distributed training.\"\"\" import os"},{"id":"verl/utils/distributed.py#initialize_global_process_group","name":"initialize_global_process_group","kind":"function","path":"verl/utils/distributed.py","span":"L18-L28","signature":"def initialize_global_process_group(timeout_second=36000)","exported":true,"origin":"ast","body_hash":"0b74afc3389013e179370434b6826752e906c6ecd596cc97734f20d3bd260874","body_text":"def initialize_global_process_group(timeout_second=36000): import torch.distributed from datetime import timedelta torch.distributed.init_process_group('nccl', timeout=timedelta(seconds=timeout_second)) local_rank = int(os.environ[\"LOCAL_RANK\"]) rank = int(os.environ[\"RANK\"]) world_size = int(os.environ[\"WORLD_SIZE\"]) if torch.distributed.is_initialized(): torch.cuda.set_device(local_rank) return local_rank, rank, world_size","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/distributed.py","relation":"imports","specifier":"os","file":"verl/utils/distributed.py"},{"source":"verl/utils/distributed.py","relation":"contains","targetId":"verl/utils/distributed.py#initialize_global_process_group","file":"verl/utils/distributed.py"},{"source":"verl/utils/distributed.py","relation":"imports","specifier":"torch.distributed","file":"verl/utils/distributed.py"},{"source":"verl/utils/distributed.py","relation":"imports","specifier":"datetime","file":"verl/utils/distributed.py"},{"source":"verl/utils/distributed.py#initialize_global_process_group","relation":"calls","name":"init_process_group","viaMember":true,"file":"verl/utils/distributed.py"},{"source":"verl/utils/distributed.py#initialize_global_process_group","relation":"calls","name":"timedelta","viaMember":false,"file":"verl/utils/distributed.py"},{"source":"verl/utils/distributed.py#initialize_global_process_group","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/distributed.py"},{"source":"verl/utils/distributed.py#initialize_global_process_group","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/distributed.py"},{"source":"verl/utils/distributed.py#initialize_global_process_group","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/distributed.py"},{"source":"verl/utils/distributed.py#initialize_global_process_group","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/utils/distributed.py"},{"source":"verl/utils/distributed.py#initialize_global_process_group","relation":"calls","name":"set_device","viaMember":true,"file":"verl/utils/distributed.py"}]},"verl/utils/flops_counter.py":{"size":4887,"mtimeMs":1789275830322.203,"hash":"753efda577398fe9378e2ce696eeb6288f153fc03086db63136633b4fc5fad1f","nodes":[{"id":"verl/utils/flops_counter.py","name":"flops_counter.py","kind":"file","path":"verl/utils/flops_counter.py","span":"L1-L124","signature":null,"exported":true,"origin":"ast","body_hash":"753efda577398fe9378e2ce696eeb6288f153fc03086db63136633b4fc5fad1f","chars":4887,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 PretrainedConfig, Qwen2Config, LlamaConfig VALID_CONFIG_TYPE = (Qwen2Config, LlamaConfig)"},{"id":"verl/utils/flops_counter.py#get_device_flops","name":"get_device_flops","kind":"function","path":"verl/utils/flops_counter.py","span":"L21-L48","signature":"def get_device_flops(unit=\"T\")","exported":true,"origin":"ast","body_hash":"47f9d2bb1eb463c92ed30c32716260ccd380a9bad73b2290c24a91af1da2cb4d","body_text":"def get_device_flops(unit=\"T\"): def unit_convert(number, level): units = [\"B\", \"K\", \"M\", \"G\", \"T\", \"P\"] if number <= 0: return number ptr = 0 while ptr < len(units) and units[ptr] != level: number /= 1000 ptr += 1 return number device_name = torch.cuda.get_device_name() flops = float(\"inf\") # INF flops for unkown gpu type if \"H100\" in device_name or \"H800\" in device_name: flops = 989e12 elif \"A100\" in device_name or \"A800\" in device_name: flops = 312e12 elif \"L40\" in device_name: flops = 181.05e12 elif \"L20\" in device_name: flops = 119.5e12 elif \"H20\" in device_name: flops = 148e12 elif \"910B\" in device_name: flops = 354e12 flops_unit = unit_convert(flops, unit) return flops_unit","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/flops_counter.py#get_device_flops.unit_convert","name":"unit_convert","kind":"function","path":"verl/utils/flops_counter.py","span":"L23-L31","signature":"def unit_convert(number, level)","exported":true,"origin":"ast","body_hash":"ef01aea74b83c691fbf2cc11798ada4e91b73c9baa30a20c147fac69ccdead2a","body_text":"def unit_convert(number, level): units = [\"B\", \"K\", \"M\", \"G\", \"T\", \"P\"] if number <= 0: return number ptr = 0 while ptr < len(units) and units[ptr] != level: number /= 1000 ptr += 1 return number","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/flops_counter.py#FlopsCounter","name":"FlopsCounter","kind":"class","path":"verl/utils/flops_counter.py","span":"L51-L123","signature":"class FlopsCounter","exported":true,"origin":"ast","body_hash":"09854625e88dd667cbfe51ad063c8881e8e07fe07155abd24fef80a3f23879be","body_text":"class FlopsCounter: \"\"\" Used to count mfu during training loop Example: flops_counter = FlopsCounter(config) flops_achieved, flops_promised = flops_counter.estimate_flops(tokens_list, delta_time) \"\"\" def __init__(self, config: PretrainedConfig): if not isinstance(config, VALID_CONFIG_TYPE): print(f\"Only support config type of {VALID_CONFIG_TYPE}, but got {type(config)}. \" f\"MFU will always be zero.\") self.estimate_func = {\"qwen2\": self._estimate_qwen2_flops, 'llama': self._estimate_qwen2_flops} self.config = config def _estimate_unknown_flops(self, tokens_sum, batch_seqlens, delta_time): return 0 def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): assert isinstance(self.config, (Qwen2Config, LlamaConfig)) hidden_size = self.config.hidden_size vocab_size = self.config.vocab_size num_hidden_layers = self.config.num_hidden_layers num_key_value_heads = self.config.num_key_value_heads num_attention_heads = self.config.num_attention_heads intermediate_size = self.config.intermediate_size head_dim = hidden_size // num_attention_heads q_size = num_attention_heads * head_dim k_size = num_key_value_heads * head_dim v_size = num_key_value_heads * head_dim # non-attn per layer parm # Qwen2/LLama use SwiGelu, gate, having up and down linear layer in mlp mlp_N = hidden_size * intermediate_size * 3 attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) emd_and_lm_head_N = vocab_size * hidden_size * 2 # non-attn all_layer parm dense_N = (mlp_N + attn_linear_N) * num_hidden_layers + emd_and_lm_head_N # non-attn all_layer & all_token fwd & bwd flops dense_N_flops = 6 * dense_N * tokens_sum # attn all_layer & all_token fwd & bwd flops seqlen_square_sum = 0 for seqlen in batch_seqlens: seqlen_square_sum += seqlen * seqlen attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers # all_layer & all_token fwd & bwd flops flops_all_token = dense_N_flops + attn_qkv_flops flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12 return flops_achieved def estimate_flops(self, batch_seqlens, delta_time): \"\"\" Estimate the FLOPS based on the number of valid tokens in the current batch and the time taken. Args: batch_seqlens (List[int]): A list where each element represents the number of valid tokens in the current batch. delta_time (float): The time taken to process the batch, in seconds. Returns: estimated_flops (float): The estimated FLOPS based on the input tokens and time. promised_flops (float): The expected FLOPS of the current device. \"\"\" tokens_sum = sum(batch_seqlens) func = self.estimate_func.get(self.config.model_type, self._estimate_unknown_flops) estimated_flops = func(tokens_sum, batch_seqlens, delta_time) promised_flops = get_device_flops() return estimated_flops, promised_flops","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/flops_counter.py#FlopsCounter.__init__","name":"__init__","kind":"method","path":"verl/utils/flops_counter.py","span":"L61-L67","signature":"def __init__(self, config: PretrainedConfig)","exported":false,"origin":"ast","body_hash":"3ab56d8c404632bece7d42f7dc374e83f370f393586612861fb11d72ea4f15c8","body_text":"def __init__(self, config: PretrainedConfig): if not isinstance(config, VALID_CONFIG_TYPE): print(f\"Only support config type of {VALID_CONFIG_TYPE}, but got {type(config)}. \" f\"MFU will always be zero.\") self.estimate_func = {\"qwen2\": self._estimate_qwen2_flops, 'llama': self._estimate_qwen2_flops} self.config = config","summary_state":"pending","summary":null,"crux":null,"owner":"FlopsCounter"},{"id":"verl/utils/flops_counter.py#FlopsCounter._estimate_unknown_flops","name":"_estimate_unknown_flops","kind":"method","path":"verl/utils/flops_counter.py","span":"L69-L70","signature":"def _estimate_unknown_flops(self, tokens_sum, batch_seqlens, delta_time)","exported":false,"origin":"ast","body_hash":"e3886e548b1aa49426ee51fa8de072b5e217922e67b106a7d33babb5458558ce","body_text":"def _estimate_unknown_flops(self, tokens_sum, batch_seqlens, delta_time): return 0","summary_state":"pending","summary":null,"crux":null,"owner":"FlopsCounter"},{"id":"verl/utils/flops_counter.py#FlopsCounter._estimate_qwen2_flops","name":"_estimate_qwen2_flops","kind":"method","path":"verl/utils/flops_counter.py","span":"L72-L105","signature":"def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time)","exported":false,"origin":"ast","body_hash":"024a3df5967f0523243a5c8ae101e09a93fe47bdb863e25132af21cf71c83e18","body_text":"def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): assert isinstance(self.config, (Qwen2Config, LlamaConfig)) hidden_size = self.config.hidden_size vocab_size = self.config.vocab_size num_hidden_layers = self.config.num_hidden_layers num_key_value_heads = self.config.num_key_value_heads num_attention_heads = self.config.num_attention_heads intermediate_size = self.config.intermediate_size head_dim = hidden_size // num_attention_heads q_size = num_attention_heads * head_dim k_size = num_key_value_heads * head_dim v_size = num_key_value_heads * head_dim # non-attn per layer parm # Qwen2/LLama use SwiGelu, gate, having up and down linear layer in mlp mlp_N = hidden_size * intermediate_size * 3 attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) emd_and_lm_head_N = vocab_size * hidden_size * 2 # non-attn all_layer parm dense_N = (mlp_N + attn_linear_N) * num_hidden_layers + emd_and_lm_head_N # non-attn all_layer & all_token fwd & bwd flops dense_N_flops = 6 * dense_N * tokens_sum # attn all_layer & all_token fwd & bwd flops seqlen_square_sum = 0 for seqlen in batch_seqlens: seqlen_square_sum += seqlen * seqlen attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers # all_layer & all_token fwd & bwd flops flops_all_token = dense_N_flops + attn_qkv_flops flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12 return flops_achieved","summary_state":"pending","summary":null,"crux":null,"owner":"FlopsCounter"},{"id":"verl/utils/flops_counter.py#FlopsCounter.estimate_flops","name":"estimate_flops","kind":"method","path":"verl/utils/flops_counter.py","span":"L107-L123","signature":"def estimate_flops(self, batch_seqlens, delta_time)","exported":true,"origin":"ast","body_hash":"edb52fec92c5bc3c75e063c39628976aa80ec1d4430630ea0b0acefaa1fefd45","body_text":"def estimate_flops(self, batch_seqlens, delta_time): \"\"\" Estimate the FLOPS based on the number of valid tokens in the current batch and the time taken. Args: batch_seqlens (List[int]): A list where each element represents the number of valid tokens in the current batch. delta_time (float): The time taken to process the batch, in seconds. Returns: estimated_flops (float): The estimated FLOPS based on the input tokens and time. promised_flops (float): The expected FLOPS of the current device. \"\"\" tokens_sum = sum(batch_seqlens) func = self.estimate_func.get(self.config.model_type, self._estimate_unknown_flops) estimated_flops = func(tokens_sum, batch_seqlens, delta_time) promised_flops = get_device_flops() return estimated_flops, promised_flops","summary_state":"pending","summary":null,"crux":null,"owner":"FlopsCounter"}],"rawEdges":[{"source":"verl/utils/flops_counter.py","relation":"imports","specifier":"torch","file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py","relation":"imports","specifier":"transformers","file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py","relation":"contains","targetId":"verl/utils/flops_counter.py#get_device_flops","file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#get_device_flops","relation":"contains","targetId":"verl/utils/flops_counter.py#get_device_flops.unit_convert","file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#get_device_flops.unit_convert","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#get_device_flops","relation":"calls","name":"get_device_name","viaMember":true,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#get_device_flops","relation":"calls","name":"float","viaMember":false,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#get_device_flops","relation":"calls","name":"unit_convert","viaMember":false,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py","relation":"contains","targetId":"verl/utils/flops_counter.py#FlopsCounter","file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter","relation":"contains","targetId":"verl/utils/flops_counter.py#FlopsCounter.__init__","file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter.__init__","relation":"calls","name":"type","viaMember":false,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter","relation":"contains","targetId":"verl/utils/flops_counter.py#FlopsCounter._estimate_unknown_flops","file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter","relation":"contains","targetId":"verl/utils/flops_counter.py#FlopsCounter._estimate_qwen2_flops","file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter._estimate_qwen2_flops","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter","relation":"contains","targetId":"verl/utils/flops_counter.py#FlopsCounter.estimate_flops","file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter.estimate_flops","relation":"calls","name":"sum","viaMember":false,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter.estimate_flops","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter.estimate_flops","relation":"calls","name":"func","viaMember":false,"file":"verl/utils/flops_counter.py"},{"source":"verl/utils/flops_counter.py#FlopsCounter.estimate_flops","relation":"calls","name":"get_device_flops","viaMember":false,"file":"verl/utils/flops_counter.py"}]},"verl/utils/fs.py":{"size":2821,"mtimeMs":1789275830322.203,"hash":"a922d349ce08eb4e23d878a33eba415e478ca467562b2082708fcc6f409ecb7e","nodes":[{"id":"verl/utils/fs.py","name":"fs.py","kind":"file","path":"verl/utils/fs.py","span":"L1-L89","signature":null,"exported":true,"origin":"ast","body_hash":"a922d349ce08eb4e23d878a33eba415e478ca467562b2082708fcc6f409ecb7e","chars":2821,"summary_state":"pending","summary":null,"crux":null,"body_text":"#!/usr/bin/env python # 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. # -*- coding: utf-8 -*- \"\"\"File-system agnostic IO APIs\"\"\" import os import tempfile import hashlib from .hdfs_io import copy, makedirs, exists __all__ = [\"copy\", \"exists\", \"makedirs\"] _HDFS_PREFIX = \"hdfs://\""},{"id":"verl/utils/fs.py#_is_non_local","name":"_is_non_local","kind":"function","path":"verl/utils/fs.py","span":"L29-L30","signature":"def _is_non_local(path)","exported":false,"origin":"ast","body_hash":"b5a7f0d8a3c33215e50ddc9432edf1aabadfa7a04b1cfa673d089c0a817cfce6","body_text":"def _is_non_local(path): return path.startswith(_HDFS_PREFIX)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fs.py#md5_encode","name":"md5_encode","kind":"function","path":"verl/utils/fs.py","span":"L33-L34","signature":"def md5_encode(path: str) -> str","exported":true,"origin":"ast","body_hash":"69dec58b30119dec212249be5da37262d6449c5c5e0a3cb0e45aaddb404c21bd","body_text":"def md5_encode(path: str) -> str: return hashlib.md5(path.encode()).hexdigest()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fs.py#get_local_temp_path","name":"get_local_temp_path","kind":"function","path":"verl/utils/fs.py","span":"L37-L52","signature":"def get_local_temp_path(hdfs_path: str, cache_dir: str) -> str","exported":true,"origin":"ast","body_hash":"d7540699e5b01fe5b869c3824d38e67a7ddaa39e6562847f3fa3d0c1c63e8e3a","body_text":"def get_local_temp_path(hdfs_path: str, cache_dir: str) -> str: \"\"\"Return a local temp path that joins cache_dir and basename of hdfs_path Args: hdfs_path: cache_dir: Returns: \"\"\" # make a base64 encoding of hdfs_path to avoid directory conflict encoded_hdfs_path = md5_encode(hdfs_path) temp_dir = os.path.join(cache_dir, encoded_hdfs_path) os.makedirs(temp_dir, exist_ok=True) dst = os.path.join(temp_dir, os.path.basename(hdfs_path)) return dst","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fs.py#copy_local_path_from_hdfs","name":"copy_local_path_from_hdfs","kind":"function","path":"verl/utils/fs.py","span":"L55-L88","signature":"def copy_local_path_from_hdfs(src: str, cache_dir=None, filelock='.file.lock', verbose=False) -> str","exported":true,"origin":"ast","body_hash":"659f35f5bb1b4e6ab8cab02f489d9d8f4d1bba024d8b15413378efb2f8558111","body_text":"def copy_local_path_from_hdfs(src: str, cache_dir=None, filelock='.file.lock', verbose=False) -> str: \"\"\"Copy src from hdfs to local if src is on hdfs or directly return src. If cache_dir is None, we will use the default cache dir of the system. Note that this may cause conflicts if the src name is the same between calls Args: src (str): a HDFS path of a local path Returns: a local path of the copied file \"\"\" from filelock import FileLock assert src[-1] != '/', f'Make sure the last char in src is not / because it will cause error. Got {src}' if _is_non_local(src): # download from hdfs to local if cache_dir is None: # get a temp folder cache_dir = tempfile.gettempdir() os.makedirs(cache_dir, exist_ok=True) assert os.path.exists(cache_dir) local_path = get_local_temp_path(src, cache_dir) # get a specific lock filelock = md5_encode(src) + '.lock' lock_file = os.path.join(cache_dir, filelock) with FileLock(lock_file=lock_file): if not os.path.exists(local_path): if verbose: print(f'Copy from {src} to {local_path}') copy(src, local_path) return local_path else: return src","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/fs.py","relation":"imports","specifier":"os","file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py","relation":"imports","specifier":"tempfile","file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py","relation":"imports","specifier":"hashlib","file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py","relation":"imports","specifier":".hdfs_io","file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py","relation":"contains","targetId":"verl/utils/fs.py#_is_non_local","file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#_is_non_local","relation":"calls","name":"startswith","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py","relation":"contains","targetId":"verl/utils/fs.py#md5_encode","file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#md5_encode","relation":"calls","name":"hexdigest","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#md5_encode","relation":"calls","name":"md5","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#md5_encode","relation":"calls","name":"encode","viaMember":true,"file":"verl/utils/fs.py","recvType":"str"},{"source":"verl/utils/fs.py","relation":"contains","targetId":"verl/utils/fs.py#get_local_temp_path","file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#get_local_temp_path","relation":"calls","name":"md5_encode","viaMember":false,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#get_local_temp_path","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#get_local_temp_path","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#get_local_temp_path","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#get_local_temp_path","relation":"calls","name":"basename","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py","relation":"contains","targetId":"verl/utils/fs.py#copy_local_path_from_hdfs","file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py","relation":"imports","specifier":"filelock","file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"_is_non_local","viaMember":false,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"gettempdir","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"exists","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"get_local_temp_path","viaMember":false,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"md5_encode","viaMember":false,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"FileLock","viaMember":false,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"exists","viaMember":true,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/fs.py"},{"source":"verl/utils/fs.py#copy_local_path_from_hdfs","relation":"calls","name":"copy","viaMember":false,"file":"verl/utils/fs.py"}]},"verl/utils/fsdp_utils.py":{"size":13201,"mtimeMs":1789275830322.203,"hash":"9b406351ace9f1a84b3d25ed3041395a10ea3508999cf4105de5fe0ee0dab2c0","nodes":[{"id":"verl/utils/fsdp_utils.py","name":"fsdp_utils.py","kind":"file","path":"verl/utils/fsdp_utils.py","span":"L1-L329","signature":null,"exported":true,"origin":"ast","body_hash":"9b406351ace9f1a84b3d25ed3041395a10ea3508999cf4105de5fe0ee0dab2c0","chars":13201,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 typing import Dict import functools import json import math import itertools import os from contextlib import contextmanager from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy from transformers.trainer_pt_utils import get_module_class_from_name import torch import torch.nn as nn import torch.distributed as dist # Copyright 2020-present the HuggingFace Inc. team. # Adapted from https://github.com/huggingface/transformers/src/transformers/trainer.py @contextmanager"},{"id":"verl/utils/fsdp_utils.py#init_fn","name":"init_fn","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L29-L33","signature":"def init_fn(x: torch.nn.Module)","exported":true,"origin":"ast","body_hash":"478453f418d8638b999a3dda1d37313a177987dd1a0341b5e8e0ac43227c048a","body_text":"def init_fn(x: torch.nn.Module): if not torch.distributed.get_rank() == 0: x = x.to_empty(device=torch.cuda.current_device(), recurse=False) torch.cuda.empty_cache() return x","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#get_init_weight_context_manager","name":"get_init_weight_context_manager","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L36-L43","signature":"def get_init_weight_context_manager(use_meta_tensor=True)","exported":true,"origin":"ast","body_hash":"809cf2e63287478dec78d0e9a62a4c933156400b131ae28f705cfc3f9a5b46e1","body_text":"def get_init_weight_context_manager(use_meta_tensor=True): from accelerate import init_empty_weights cpu_init_weights = lambda: torch.device('cpu') if use_meta_tensor: init_context = init_empty_weights if torch.distributed.get_rank() != 0 else cpu_init_weights else: init_context = cpu_init_weights return init_context","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","name":"get_fsdp_wrap_policy","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L48-L105","signature":"def get_fsdp_wrap_policy(module, config=None, is_lora=False)","exported":true,"origin":"ast","body_hash":"53558c0ba06821ad192cb754ca0c1d16940ca04ade94aec2a8af11af6c7e648b","body_text":"def get_fsdp_wrap_policy(module, config=None, is_lora=False): \"\"\"Get FSDP wrap policy for the module. Args: module: The module to get wrap policy for config: Configuration for wrap policy is_lora: Whether to enable lambda policy for LoRA modules \"\"\" if config is None: config = {} if config.get('disable', False): return None default_transformer_cls_names_to_wrap = getattr(module, \"_no_split_modules\", None) fsdp_transformer_layer_cls_to_wrap = config.get(\"transformer_layer_cls_to_wrap\", default_transformer_cls_names_to_wrap) min_num_params = config.get('min_num_params', 0) auto_wrap_policy = None policies = [] from torch.distributed.fsdp.wrap import _or_policy, lambda_auto_wrap_policy, transformer_auto_wrap_policy # Add lambda policy for LoRA modules if is_lora is True if is_lora: def lambda_policy_fn(module): if (len(list(module.named_children())) == 0 and getattr(module, \"weight\", None) is not None and module.weight.requires_grad): return True return False lambda_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=lambda_policy_fn) policies.append(lambda_policy) if min_num_params > 0: size_policy = functools.partial(size_based_auto_wrap_policy, min_num_params=min_num_params) policies.append(size_policy) elif fsdp_transformer_layer_cls_to_wrap is not None: transformer_cls_to_wrap = set() for layer_class in fsdp_transformer_layer_cls_to_wrap: transformer_cls = get_module_class_from_name(module, layer_class) if transformer_cls is None: raise Exception(\"Could not find the transformer layer class to wrap in the model.\") else: transformer_cls_to_wrap.add(transformer_cls) transformer_policy = functools.partial( transformer_auto_wrap_policy, transformer_layer_cls=transformer_cls_to_wrap, ) policies.append(transformer_policy) if len(policies) > 0: auto_wrap_policy = functools.partial(_or_policy, policies=policies) return auto_wrap_policy","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy.lambda_policy_fn","name":"lambda_policy_fn","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L75-L79","signature":"def lambda_policy_fn(module)","exported":true,"origin":"ast","body_hash":"c2b134de85a43e8a034ae1ccbc2567d215ee8a2fd8512615931390568b6906c6","body_text":"def lambda_policy_fn(module): if (len(list(module.named_children())) == 0 and getattr(module, \"weight\", None) is not None and module.weight.requires_grad): return True return False","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#offload_fsdp_grad","name":"offload_fsdp_grad","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L108-L112","signature":"def offload_fsdp_grad(module)","exported":true,"origin":"ast","body_hash":"6882e2967d6aed49e218adc4072a85c29a57bc6ceee22fd48d7e186406880af9","body_text":"def offload_fsdp_grad(module): for _, param in module.named_parameters(): if param.grad is not None: param.grad = param.grad.to(\"cpu\", non_blocking=True) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#load_fsdp_grad","name":"load_fsdp_grad","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L115-L119","signature":"def load_fsdp_grad(module, device_id)","exported":true,"origin":"ast","body_hash":"42e0431655e2a5aa161830a94e01c09a24fe61df256b53b6e48fd74777907184","body_text":"def load_fsdp_grad(module, device_id): for _, param in module.named_parameters(): if param.grad is not None: param.grad = param.grad.to(device_id, non_blocking=True) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#offload_fsdp_param_and_grad","name":"offload_fsdp_param_and_grad","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L122-L129","signature":"def offload_fsdp_param_and_grad(module, offload_grad=False)","exported":true,"origin":"ast","body_hash":"8c43f12f74769aece2279ae17ab6ae9d27c58143d8b916045a446361d73d905c","body_text":"def offload_fsdp_param_and_grad(module, offload_grad=False): for _, param in module.named_parameters(): if hasattr(param, \"_local_shard\"): param._local_shard = param._local_shard.to(\"cpu\", non_blocking=True) param.data = param.data.to('cpu', non_blocking=True) if offload_grad and param.grad is not None: param.grad = param.grad.to(\"cpu\", non_blocking=True) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#load_fsdp_param_and_grad","name":"load_fsdp_param_and_grad","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L132-L139","signature":"def load_fsdp_param_and_grad(module, device_id, load_grad=False)","exported":true,"origin":"ast","body_hash":"b356d01f261615f4136d96e6b164de34654fb81caa22d4279def1eb1055c5a04","body_text":"def load_fsdp_param_and_grad(module, device_id, load_grad=False): for _, param in module.named_parameters(): if hasattr(param, \"_local_shard\"): param._local_shard = param._local_shard.to(device_id, non_blocking=True) param.data = param.data.to(device_id, non_blocking=True) if load_grad and param.grad is not None: param.grad = param.grad.to(device_id, non_blocking=True) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#offload_fsdp_optimizer","name":"offload_fsdp_optimizer","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L142-L149","signature":"def offload_fsdp_optimizer(optimizer)","exported":true,"origin":"ast","body_hash":"fc449920c030ce94a9e26d261c3ad4ebf648ad6f4136a7f3f3a56642503fdbac","body_text":"def offload_fsdp_optimizer(optimizer): 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.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#load_fsdp_optimizer","name":"load_fsdp_optimizer","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L152-L159","signature":"def load_fsdp_optimizer(optimizer, device_id)","exported":true,"origin":"ast","body_hash":"d2f200950ed2a9af0ab17e171576cca2df60aa8138f7a7f11469f2eb623f1708","body_text":"def load_fsdp_optimizer(optimizer, device_id): 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) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#meta_device_init","name":"meta_device_init","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L163-L191","signature":"def meta_device_init()","exported":true,"origin":"ast","body_hash":"30859cacefc1bf39a145369651a8a2b26403766907c5a5961615913bc341610d","body_text":"def meta_device_init(): \"\"\" Create model parameters with meta device. Note buffers in model will still be initialized in default device (e.g., CPU), since the buffers can be non-persistent and filled with expected values that can NOT be captured in meta device. \"\"\" device = torch.device(\"meta\") old_register_parameter = nn.Module.register_parameter registered = set() def register_empty_parameter(module, name, param): old_register_parameter(module, name, param) # we will skip register shared parameters as it # is already registered previously if param is not None and param not in registered: param_cls = type(module._parameters[name]) kwargs = module._parameters[name].__dict__ kwargs[\"requires_grad\"] = param.requires_grad module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) registered.add(module._parameters[name]) try: nn.Module.register_parameter = register_empty_parameter yield finally: registered.clear() nn.Module.register_parameter = old_register_parameter","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#meta_device_init.register_empty_parameter","name":"register_empty_parameter","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L175-L184","signature":"def register_empty_parameter(module, name, param)","exported":true,"origin":"ast","body_hash":"8be9ad82d2aa30bbd4d3e2697844cb1f6bc0437270d49a0b3537d86b798faf33","body_text":"def register_empty_parameter(module, name, param): old_register_parameter(module, name, param) # we will skip register shared parameters as it # is already registered previously if param is not None and param not in registered: param_cls = type(module._parameters[name]) kwargs = module._parameters[name].__dict__ kwargs[\"requires_grad\"] = param.requires_grad module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs) registered.add(module._parameters[name])","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#parallel_load_safetensors","name":"parallel_load_safetensors","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L194-L247","signature":"def parallel_load_safetensors(filepath)","exported":true,"origin":"ast","body_hash":"65aa63eae67d62197f38e8936b2fe51c68e674a073d6fba6e040097d31b29b55","body_text":"def parallel_load_safetensors(filepath): \"\"\" Parallel load safetensors from huggingface checkpoint Huggingface checkpoint contains: - config.json: a json file for model configuration - model.safetensor.index.json: a json file for safetensors (parameters & buffers) index - model-000x-of-ooxx.safetensors: a binary file for safetensors (parameters & buffers) chunks Or (when model is small), - model.safetensors: a binary file for all parameters and buffers Each rank will own a part of model chunks and load them directly into GPU memory. \"\"\" from safetensors.torch import load_file safetensors2param = {} index_file = os.path.join(filepath, \"model.safetensors.index.json\") if os.path.exists(index_file): index = json.load(open(index_file, \"rb\")) for param_name, filename in index[\"weight_map\"].items(): safetensors2param.setdefault(filename, []).append(param_name) else: # in this case, the model is small and we can load it all at once param_file = os.path.join(filepath, \"model.safetensors\") assert os.path.exists(param_file), f\"Cannot find {param_file}\" states = load_file(param_file) for param_name in states: safetensors2param.setdefault(\"model.safetensors\", []).append(param_name) del states total_files = len(safetensors2param) ckpt_chunks = sorted(safetensors2param.keys()) world_size = dist.get_world_size() size = int(math.ceil(total_files / world_size)) ckpt_chunks = [ckpt_chunks[rank * size:rank * size + size] for rank in range(world_size)] shard_states = {} device = torch.cuda.current_device() for rank, files in enumerate(ckpt_chunks): if rank == dist.get_rank(): for file in files: file = os.path.join(filepath, file) states = load_file(file, device=device) # print(f\"rank {rank} loading {file}...\") shard_states.update(states) else: for file in files: for param_name in safetensors2param[file]: shard_states[param_name] = rank return shard_states","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#parallel_init_module_fn","name":"parallel_init_module_fn","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L250-L329","signature":"def parallel_init_module_fn(module: torch.nn.Module, shard_states: Dict[str, torch.nn.Parameter])","exported":true,"origin":"ast","body_hash":"7cba013268b43fa74f6c8744a023978e2256b773075fbaa50fc13ab7c0f3b94f","body_text":"def parallel_init_module_fn(module: torch.nn.Module, shard_states: Dict[str, torch.nn.Parameter]): \"\"\" Generate a function to initialize sub-modules in the `module` with `shard_states` from huggingface checkpoint. Args: module (torch.nn.Module): the global module to be initialized shard_states (Dict[str, torch.nn.Parameter]): the shard states from huggingface checkpoint Returns: init_fn (Callable): a function to initialize sub-modules in the `module` with `shard_states` \"\"\" state2fqn = {} for name, state in itertools.chain(module.named_parameters(remove_duplicate=False), module.named_buffers(remove_duplicate=False)): state2fqn.setdefault(state, []).append(name) # remove standalone parameters and buffers shared = {s for s, names in state2fqn.items() if len(names) > 1} materialized_states = {} @torch.no_grad() def create_and_sync_state(param_name, state, is_param): assert param_name in shard_states, f\"{param_name} not loaded\" device = torch.cuda.current_device() if is_param: param = torch.nn.Parameter(torch.empty_like(state.data, device=device), requires_grad=state.requires_grad) else: # buffer param = torch.empty_like(state.data, device=device) loaded = shard_states[param_name] if isinstance(loaded, (torch.nn.Parameter, torch.Tensor)): # NOTE: loaded.dtype can be different with param.dtype param.data.copy_(loaded.data) dist.broadcast(param.data, src=dist.get_rank()) else: assert isinstance(loaded, int) # the rank that holds the state dist.broadcast(param.data, src=loaded) shard_states.pop(param_name) del loaded return param def init_fn(sub_mod: torch.nn.Module, recurse: bool = True): param_and_buffers = tuple(sub_mod.named_parameters(recurse=False)) + tuple(sub_mod.named_buffers(recurse=False)) # param_and_buffers = sorted(sub_mod.named_parameters(recurse=False), key=lambda x: x[0]) for name, state in param_and_buffers: if not state.is_meta: continue is_param = name in sub_mod._parameters fqn = state2fqn[state].pop(0) # non-persistent buffers will not be saved in state dict, we can safely skip it if (not is_param) and fqn not in shard_states: if state.is_meta: raise RuntimeError( f\"find a non-persistent buffer ({fqn}) initiated with device meta. \" \"Such buffer is not saved in checkpoint and user should guarantee to init in CPU / GPU device.\") continue # for shared parameter, we get it from the first time it is created if state in shared: if state not in materialized_states: materialized_states[state] = create_and_sync_state(fqn, state, is_param) else: if fqn in shard_states: shard_states.pop(fqn) materialize_state = materialized_states[state] # for not shared parameter, we create it directly else: materialize_state = create_and_sync_state(fqn, state, is_param) if is_param: sub_mod._parameters[name] = materialize_state else: sub_mod._buffers[name] = materialize_state if recurse: for module in sub_mod.children(): init_fn(module, recurse=True) # for debug # if len(shard_states) == 0: print(\"clear\") return sub_mod return init_fn","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","name":"create_and_sync_state","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L272-L289","signature":"def create_and_sync_state(param_name, state, is_param)","exported":true,"origin":"ast","body_hash":"acb85d65ea151697ac7600f09629e8344ff04b6a3aa2104ad50d246cb9f2f838","body_text":"def create_and_sync_state(param_name, state, is_param): assert param_name in shard_states, f\"{param_name} not loaded\" device = torch.cuda.current_device() if is_param: param = torch.nn.Parameter(torch.empty_like(state.data, device=device), requires_grad=state.requires_grad) else: # buffer param = torch.empty_like(state.data, device=device) loaded = shard_states[param_name] if isinstance(loaded, (torch.nn.Parameter, torch.Tensor)): # NOTE: loaded.dtype can be different with param.dtype param.data.copy_(loaded.data) dist.broadcast(param.data, src=dist.get_rank()) else: assert isinstance(loaded, int) # the rank that holds the state dist.broadcast(param.data, src=loaded) shard_states.pop(param_name) del loaded return param","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","name":"init_fn","kind":"function","path":"verl/utils/fsdp_utils.py","span":"L291-L327","signature":"def init_fn(sub_mod: torch.nn.Module, recurse: bool = True)","exported":true,"origin":"ast","body_hash":"359c0b0307a1fa6e7c149b5c5785c315c7db483f135a8f82b3434c848bf7e849","body_text":"def init_fn(sub_mod: torch.nn.Module, recurse: bool = True): param_and_buffers = tuple(sub_mod.named_parameters(recurse=False)) + tuple(sub_mod.named_buffers(recurse=False)) # param_and_buffers = sorted(sub_mod.named_parameters(recurse=False), key=lambda x: x[0]) for name, state in param_and_buffers: if not state.is_meta: continue is_param = name in sub_mod._parameters fqn = state2fqn[state].pop(0) # non-persistent buffers will not be saved in state dict, we can safely skip it if (not is_param) and fqn not in shard_states: if state.is_meta: raise RuntimeError( f\"find a non-persistent buffer ({fqn}) initiated with device meta. \" \"Such buffer is not saved in checkpoint and user should guarantee to init in CPU / GPU device.\") continue # for shared parameter, we get it from the first time it is created if state in shared: if state not in materialized_states: materialized_states[state] = create_and_sync_state(fqn, state, is_param) else: if fqn in shard_states: shard_states.pop(fqn) materialize_state = materialized_states[state] # for not shared parameter, we create it directly else: materialize_state = create_and_sync_state(fqn, state, is_param) if is_param: sub_mod._parameters[name] = materialize_state else: sub_mod._buffers[name] = materialize_state if recurse: for module in sub_mod.children(): init_fn(module, recurse=True) # for debug # if len(shard_states) == 0: print(\"clear\") return sub_mod","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"typing","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"functools","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"json","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"math","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"itertools","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"os","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"contextlib","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"torch.distributed.fsdp.wrap","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"transformers.trainer_pt_utils","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"torch","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#init_fn","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#init_fn","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#init_fn","relation":"calls","name":"to_empty","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#init_fn","relation":"calls","name":"current_device","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#init_fn","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#get_init_weight_context_manager","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"accelerate","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_init_weight_context_manager","relation":"calls","name":"device","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_init_weight_context_manager","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"getattr","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"torch.distributed.fsdp.wrap","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"contains","targetId":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy.lambda_policy_fn","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy.lambda_policy_fn","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy.lambda_policy_fn","relation":"calls","name":"list","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy.lambda_policy_fn","relation":"calls","name":"named_children","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy.lambda_policy_fn","relation":"calls","name":"getattr","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"partial","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"partial","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"set","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"get_module_class_from_name","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"Exception","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"add","viaMember":true,"file":"verl/utils/fsdp_utils.py","recvType":"set"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"partial","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#get_fsdp_wrap_policy","relation":"calls","name":"partial","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#offload_fsdp_grad","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_grad","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_grad","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#load_fsdp_grad","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_grad","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_grad","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#offload_fsdp_param_and_grad","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_param_and_grad","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_param_and_grad","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_param_and_grad","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#load_fsdp_param_and_grad","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_param_and_grad","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_param_and_grad","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_param_and_grad","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#offload_fsdp_optimizer","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_optimizer","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_optimizer","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_optimizer","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#offload_fsdp_optimizer","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#load_fsdp_optimizer","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_optimizer","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_optimizer","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_optimizer","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#load_fsdp_optimizer","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#meta_device_init","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#meta_device_init","relation":"calls","name":"device","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#meta_device_init","relation":"calls","name":"set","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#meta_device_init","relation":"contains","targetId":"verl/utils/fsdp_utils.py#meta_device_init.register_empty_parameter","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#meta_device_init.register_empty_parameter","relation":"calls","name":"old_register_parameter","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#meta_device_init.register_empty_parameter","relation":"calls","name":"type","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#meta_device_init.register_empty_parameter","relation":"calls","name":"param_cls","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#meta_device_init.register_empty_parameter","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#meta_device_init.register_empty_parameter","relation":"calls","name":"add","viaMember":true,"file":"verl/utils/fsdp_utils.py","recvType":"set"},{"source":"verl/utils/fsdp_utils.py#meta_device_init","relation":"calls","name":"clear","viaMember":true,"file":"verl/utils/fsdp_utils.py","recvType":"set"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#parallel_load_safetensors","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"imports","specifier":"safetensors.torch","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"exists","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"load","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"open","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"setdefault","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"exists","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"load_file","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"setdefault","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"keys","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"ceil","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"current_device","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"load_file","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_load_safetensors","relation":"calls","name":"update","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py","relation":"contains","targetId":"verl/utils/fsdp_utils.py#parallel_init_module_fn","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn","relation":"calls","name":"chain","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn","relation":"calls","name":"named_buffers","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn","relation":"calls","name":"setdefault","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn","relation":"contains","targetId":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"current_device","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"Parameter","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"copy_","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.create_and_sync_state","relation":"calls","name":"pop","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn","relation":"contains","targetId":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"tuple","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"tuple","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"named_buffers","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"pop","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"create_and_sync_state","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"pop","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"create_and_sync_state","viaMember":false,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"children","viaMember":true,"file":"verl/utils/fsdp_utils.py"},{"source":"verl/utils/fsdp_utils.py#parallel_init_module_fn.init_fn","relation":"calls","name":"init_fn","viaMember":false,"file":"verl/utils/fsdp_utils.py"}]},"verl/utils/hdfs_io.py":{"size":4562,"mtimeMs":1789275830323.203,"hash":"3f006815223e1798ee829b325aa942453cdb6131731f056b07fd8de957d0bcf7","nodes":[{"id":"verl/utils/hdfs_io.py","name":"hdfs_io.py","kind":"file","path":"verl/utils/hdfs_io.py","span":"L1-L145","signature":null,"exported":true,"origin":"ast","body_hash":"3f006815223e1798ee829b325aa942453cdb6131731f056b07fd8de957d0bcf7","chars":4562,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 logging logger = logging.getLogger(__file__) logger.setLevel(os.getenv('VERL_SFT_LOGGING_LEVEL', 'WARN')) _HDFS_PREFIX = \"hdfs://\" _HDFS_BIN_PATH = shutil.which('hdfs')"},{"id":"verl/utils/hdfs_io.py#exists","name":"exists","kind":"function","path":"verl/utils/hdfs_io.py","span":"L27-L40","signature":"def exists(path: str, **kwargs) -> bool","exported":true,"origin":"ast","body_hash":"0c174720c6536682794ce9937ee6141b9027bcfb53f9b16b4a17effdabdf244a","body_text":"def exists(path: str, **kwargs) -> bool: r\"\"\"Works like os.path.exists() but supports hdfs. Test whether a path exists. Returns False for broken symbolic links. Args: path (str): path to test Returns: bool: True if the path exists, False otherwise \"\"\" if _is_non_local(path): return _exists(path, **kwargs) return os.path.exists(path)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/hdfs_io.py#_exists","name":"_exists","kind":"function","path":"verl/utils/hdfs_io.py","span":"L43-L47","signature":"def _exists(file_path: str)","exported":false,"origin":"ast","body_hash":"cedd6664e9c0a20ccbd7cda3abe94d5cb5d3aeeb21e8687c00ceaa6d3c70ae68","body_text":"def _exists(file_path: str): \"\"\" hdfs capable to check whether a file_path is exists \"\"\" if file_path.startswith(\"hdfs\"): return _run_cmd(_hdfs_cmd(f\"-test -e {file_path}\")) == 0 return os.path.exists(file_path)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/hdfs_io.py#makedirs","name":"makedirs","kind":"function","path":"verl/utils/hdfs_io.py","span":"L50-L72","signature":"def makedirs(name, mode=0o777, exist_ok=False, **kwargs) -> None","exported":true,"origin":"ast","body_hash":"eb24ff7dd4a0ee58a7dbbb6fcd4916fd93f1ad2f14b786efc545612083a33775","body_text":"def makedirs(name, mode=0o777, exist_ok=False, **kwargs) -> None: r\"\"\"Works like os.makedirs() but supports hdfs. Super-mkdir; create a leaf directory and all intermediate ones. Works like mkdir, except that any intermediate path segment (not just the rightmost) will be created if it does not exist. If the target directory already exists, raise an OSError if exist_ok is False. Otherwise no exception is raised. This is recursive. Args: name (str): directory to create mode (int): file mode bits exist_ok (bool): if True, do not raise an exception if the directory already exists kwargs: keyword arguments for hdfs \"\"\" if _is_non_local(name): # TODO(haibin.lin): # - handle OSError for hdfs(?) # - support exist_ok for hdfs(?) _mkdir(name, **kwargs) else: os.makedirs(name, mode=mode, exist_ok=exist_ok)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/hdfs_io.py#_mkdir","name":"_mkdir","kind":"function","path":"verl/utils/hdfs_io.py","span":"L75-L81","signature":"def _mkdir(file_path: str) -> bool","exported":false,"origin":"ast","body_hash":"d09ed96d9e95335ae1040da274f1d5ca696148971ef0b5da1161b4c2b3d101bb","body_text":"def _mkdir(file_path: str) -> bool: \"\"\"hdfs mkdir\"\"\" if file_path.startswith(\"hdfs\"): _run_cmd(_hdfs_cmd(f\"-mkdir -p {file_path}\")) else: os.makedirs(file_path, exist_ok=True) return True","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/hdfs_io.py#copy","name":"copy","kind":"function","path":"verl/utils/hdfs_io.py","span":"L84-L110","signature":"def copy(src: str, dst: str, **kwargs) -> bool","exported":true,"origin":"ast","body_hash":"2d2963d940836b5c3065446c9a5519070bafbdcf69e10ad0a59fa5a53d652b58","body_text":"def copy(src: str, dst: str, **kwargs) -> bool: r\"\"\"Works like shutil.copy() for file, and shutil.copytree for dir, and supports hdfs. Copy data and mode bits (\"cp src dst\"). Return the file's destination. The destination may be a directory. If source and destination are the same file, a SameFileError will be raised. Arg: src (str): source file path dst (str): destination file path kwargs: keyword arguments for hdfs copy Returns: str: destination file path \"\"\" if _is_non_local(src) or _is_non_local(dst): # TODO(haibin.lin): # - handle SameFileError for hdfs files(?) # - return file destination for hdfs files return _copy(src, dst) else: if os.path.isdir(src): return shutil.copytree(src, dst, **kwargs) else: return shutil.copy(src, dst, **kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/hdfs_io.py#_copy","name":"_copy","kind":"function","path":"verl/utils/hdfs_io.py","span":"L113-L132","signature":"def _copy(from_path: str, to_path: str, timeout: int = None) -> bool","exported":false,"origin":"ast","body_hash":"182dcb70bd19e98dbd865a48e620cad12a1796e861be7bfef6487fa0093fd601","body_text":"def _copy(from_path: str, to_path: str, timeout: int = None) -> bool: if to_path.startswith(\"hdfs\"): if from_path.startswith(\"hdfs\"): returncode = _run_cmd(_hdfs_cmd(f\"-cp -f {from_path} {to_path}\"), timeout=timeout) else: returncode = _run_cmd(_hdfs_cmd(f\"-put -f {from_path} {to_path}\"), timeout=timeout) else: if from_path.startswith(\"hdfs\"): returncode = _run_cmd(_hdfs_cmd(f\"-get \\ {from_path} {to_path}\"), timeout=timeout) else: try: shutil.copy(from_path, to_path) returncode = 0 except shutil.SameFileError: returncode = 0 except Exception as e: logger.warning(f\"copy {from_path} {to_path} failed: {e}\") returncode = -1 return returncode == 0","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/hdfs_io.py#_run_cmd","name":"_run_cmd","kind":"function","path":"verl/utils/hdfs_io.py","span":"L135-L136","signature":"def _run_cmd(cmd: str, timeout=None)","exported":false,"origin":"ast","body_hash":"7554272c4569a8d9ea48e4734694d6f2ae01a4892b460e63cb3ac8e2e6aba4b9","body_text":"def _run_cmd(cmd: str, timeout=None): return os.system(cmd)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/hdfs_io.py#_hdfs_cmd","name":"_hdfs_cmd","kind":"function","path":"verl/utils/hdfs_io.py","span":"L139-L140","signature":"def _hdfs_cmd(cmd: str) -> str","exported":false,"origin":"ast","body_hash":"fa1905e6fbdf93a5b05090bec8fc56c8da31c2335784b9f35f33a1c96a7ff9cb","body_text":"def _hdfs_cmd(cmd: str) -> str: return f\"{_HDFS_BIN_PATH} dfs {cmd}\"","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/hdfs_io.py#_is_non_local","name":"_is_non_local","kind":"function","path":"verl/utils/hdfs_io.py","span":"L143-L144","signature":"def _is_non_local(path: str)","exported":false,"origin":"ast","body_hash":"73f7d9a7b5fe3af9ab93520babd6ca6e00251e6b34b25d198f5a85fcda757ab1","body_text":"def _is_non_local(path: str): return path.startswith(_HDFS_PREFIX)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/hdfs_io.py","relation":"imports","specifier":"os","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"imports","specifier":"shutil","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"imports","specifier":"logging","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"calls","name":"getLogger","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"calls","name":"setLevel","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"calls","name":"getenv","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"calls","name":"which","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"contains","targetId":"verl/utils/hdfs_io.py#exists","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#exists","relation":"calls","name":"_is_non_local","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#exists","relation":"calls","name":"_exists","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#exists","relation":"calls","name":"exists","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"contains","targetId":"verl/utils/hdfs_io.py#_exists","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_exists","relation":"calls","name":"startswith","viaMember":true,"file":"verl/utils/hdfs_io.py","recvType":"str"},{"source":"verl/utils/hdfs_io.py#_exists","relation":"calls","name":"_run_cmd","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_exists","relation":"calls","name":"_hdfs_cmd","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_exists","relation":"calls","name":"exists","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"contains","targetId":"verl/utils/hdfs_io.py#makedirs","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#makedirs","relation":"calls","name":"_is_non_local","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#makedirs","relation":"calls","name":"_mkdir","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#makedirs","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"contains","targetId":"verl/utils/hdfs_io.py#_mkdir","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_mkdir","relation":"calls","name":"startswith","viaMember":true,"file":"verl/utils/hdfs_io.py","recvType":"str"},{"source":"verl/utils/hdfs_io.py#_mkdir","relation":"calls","name":"_run_cmd","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_mkdir","relation":"calls","name":"_hdfs_cmd","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_mkdir","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"contains","targetId":"verl/utils/hdfs_io.py#copy","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#copy","relation":"calls","name":"_is_non_local","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#copy","relation":"calls","name":"_is_non_local","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#copy","relation":"calls","name":"_copy","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#copy","relation":"calls","name":"isdir","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#copy","relation":"calls","name":"copytree","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#copy","relation":"calls","name":"copy","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"contains","targetId":"verl/utils/hdfs_io.py#_copy","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"startswith","viaMember":true,"file":"verl/utils/hdfs_io.py","recvType":"str"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"startswith","viaMember":true,"file":"verl/utils/hdfs_io.py","recvType":"str"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"_run_cmd","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"_hdfs_cmd","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"_run_cmd","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"_hdfs_cmd","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"startswith","viaMember":true,"file":"verl/utils/hdfs_io.py","recvType":"str"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"_run_cmd","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"_hdfs_cmd","viaMember":false,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"copy","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_copy","relation":"calls","name":"warning","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"contains","targetId":"verl/utils/hdfs_io.py#_run_cmd","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_run_cmd","relation":"calls","name":"system","viaMember":true,"file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"contains","targetId":"verl/utils/hdfs_io.py#_hdfs_cmd","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py","relation":"contains","targetId":"verl/utils/hdfs_io.py#_is_non_local","file":"verl/utils/hdfs_io.py"},{"source":"verl/utils/hdfs_io.py#_is_non_local","relation":"calls","name":"startswith","viaMember":true,"file":"verl/utils/hdfs_io.py","recvType":"str"}]},"verl/utils/import_utils.py":{"size":1351,"mtimeMs":1789275830323.203,"hash":"96fbdc91902ef5b62dee939d0a1f5a5cec7c5d93b240966017ef08278e644722","nodes":[{"id":"verl/utils/import_utils.py","name":"import_utils.py","kind":"file","path":"verl/utils/import_utils.py","span":"L1-L49","signature":null,"exported":true,"origin":"ast","body_hash":"96fbdc91902ef5b62dee939d0a1f5a5cec7c5d93b240966017ef08278e644722","chars":1351,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 to check if packages are available. We assume package availability won't change during runtime. \"\"\" from functools import cache from typing import List @cache @cache"},{"id":"verl/utils/import_utils.py#is_megatron_core_available","name":"is_megatron_core_available","kind":"function","path":"verl/utils/import_utils.py","span":"L24-L29","signature":"def is_megatron_core_available()","exported":true,"origin":"ast","body_hash":"4b0a4b37beb8320c3a0d9f897c4de697e9f6bbbbef137bdadcf6e9a321c38adf","body_text":"def is_megatron_core_available(): try: from megatron.core import parallel_state as mpu return True except ImportError: return False","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/import_utils.py#is_vllm_available","name":"is_vllm_available","kind":"function","path":"verl/utils/import_utils.py","span":"L33-L38","signature":"def is_vllm_available()","exported":true,"origin":"ast","body_hash":"1f9c6e2726c4b720a593a857eea77f8501567fa254c904ea75ff4c2cc5bf27ef","body_text":"def is_vllm_available(): try: import vllm return True except ImportError: return False","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/import_utils.py#import_external_libs","name":"import_external_libs","kind":"function","path":"verl/utils/import_utils.py","span":"L41-L48","signature":"def import_external_libs(external_libs=None)","exported":true,"origin":"ast","body_hash":"42b6f89d5b3db21e0d55bbe42905fe5832d0726bc661d7c4c4886082f4603f15","body_text":"def import_external_libs(external_libs=None): if external_libs is None: return if not isinstance(external_libs, List): external_libs = [external_libs] import importlib for external_lib in external_libs: importlib.import_module(external_lib)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/import_utils.py","relation":"imports","specifier":"functools","file":"verl/utils/import_utils.py"},{"source":"verl/utils/import_utils.py","relation":"imports","specifier":"typing","file":"verl/utils/import_utils.py"},{"source":"verl/utils/import_utils.py","relation":"contains","targetId":"verl/utils/import_utils.py#is_megatron_core_available","file":"verl/utils/import_utils.py"},{"source":"verl/utils/import_utils.py","relation":"imports","specifier":"megatron.core","file":"verl/utils/import_utils.py"},{"source":"verl/utils/import_utils.py","relation":"contains","targetId":"verl/utils/import_utils.py#is_vllm_available","file":"verl/utils/import_utils.py"},{"source":"verl/utils/import_utils.py","relation":"imports","specifier":"vllm","file":"verl/utils/import_utils.py"},{"source":"verl/utils/import_utils.py","relation":"contains","targetId":"verl/utils/import_utils.py#import_external_libs","file":"verl/utils/import_utils.py"},{"source":"verl/utils/import_utils.py#import_external_libs","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/import_utils.py"},{"source":"verl/utils/import_utils.py","relation":"imports","specifier":"importlib","file":"verl/utils/import_utils.py"},{"source":"verl/utils/import_utils.py#import_external_libs","relation":"calls","name":"import_module","viaMember":true,"file":"verl/utils/import_utils.py"}]},"verl/utils/logger/__init__.py":{"size":600,"mtimeMs":1789275830323.203,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/utils/logger/__init__.py","name":"__init__.py","kind":"file","path":"verl/utils/logger/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/utils/logger/aggregate_logger.py":{"size":1400,"mtimeMs":1789275830323.203,"hash":"9fcaf99c0e6bacb838c7d1e40f1a51dcab6b4fe633b9320d97f5434f032fee9d","nodes":[{"id":"verl/utils/logger/aggregate_logger.py","name":"aggregate_logger.py","kind":"file","path":"verl/utils/logger/aggregate_logger.py","span":"L1-L42","signature":null,"exported":true,"origin":"ast","body_hash":"9fcaf99c0e6bacb838c7d1e40f1a51dcab6b4fe633b9320d97f5434f032fee9d","chars":1400,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 Ray logger will receive logging info from different processes. \"\"\" import numbers from typing import Dict"},{"id":"verl/utils/logger/aggregate_logger.py#concat_dict_to_str","name":"concat_dict_to_str","kind":"function","path":"verl/utils/logger/aggregate_logger.py","span":"L21-L27","signature":"def concat_dict_to_str(dict: Dict, step)","exported":true,"origin":"ast","body_hash":"2324bf1ef94f22ca671549d11037f3bbcf5353b6c01c13d3aea6d20b5139b088","body_text":"def concat_dict_to_str(dict: Dict, step): output = [f'step:{step}'] for k, v in dict.items(): if isinstance(v, numbers.Number): output.append(f'{k}:{v:.3f}') output_str = ' - '.join(output) return output_str","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/logger/aggregate_logger.py#LocalLogger","name":"LocalLogger","kind":"class","path":"verl/utils/logger/aggregate_logger.py","span":"L30-L42","signature":"class LocalLogger","exported":true,"origin":"ast","body_hash":"bad179b5c225869530a4e7005b920ed0d9eec66f7732a2de92afcb5c0341cd0d","body_text":"class LocalLogger: def __init__(self, remote_logger=None, enable_wandb=False, print_to_console=False): self.print_to_console = print_to_console if print_to_console: print('Using LocalLogger is deprecated. The constructor API will change ') def flush(self): pass def log(self, data, step): if self.print_to_console: print(concat_dict_to_str(data, step=step), flush=True)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/logger/aggregate_logger.py#LocalLogger.__init__","name":"__init__","kind":"method","path":"verl/utils/logger/aggregate_logger.py","span":"L32-L35","signature":"def __init__(self, remote_logger=None, enable_wandb=False, print_to_console=False)","exported":false,"origin":"ast","body_hash":"989454bb6949ae7c161a7b069f4dac466e1a1383f0546d72c2480d7e194d26be","body_text":"def __init__(self, remote_logger=None, enable_wandb=False, print_to_console=False): self.print_to_console = print_to_console if print_to_console: print('Using LocalLogger is deprecated. The constructor API will change ')","summary_state":"pending","summary":null,"crux":null,"owner":"LocalLogger"},{"id":"verl/utils/logger/aggregate_logger.py#LocalLogger.flush","name":"flush","kind":"method","path":"verl/utils/logger/aggregate_logger.py","span":"L37-L38","signature":"def flush(self)","exported":true,"origin":"ast","body_hash":"b634cd44039f0de7f5248aae6081ee5481e7e9f7bbc0d0aa41df3468b2e5c433","body_text":"def flush(self): pass","summary_state":"pending","summary":null,"crux":null,"owner":"LocalLogger"},{"id":"verl/utils/logger/aggregate_logger.py#LocalLogger.log","name":"log","kind":"method","path":"verl/utils/logger/aggregate_logger.py","span":"L40-L42","signature":"def log(self, data, step)","exported":true,"origin":"ast","body_hash":"ffc7b15abfa541dc3814476e4be8652b1aa2bee6da696ffebfbe2e9d389cacb0","body_text":"def log(self, data, step): if self.print_to_console: print(concat_dict_to_str(data, step=step), flush=True)","summary_state":"pending","summary":null,"crux":null,"owner":"LocalLogger"}],"rawEdges":[{"source":"verl/utils/logger/aggregate_logger.py","relation":"imports","specifier":"numbers","file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py","relation":"imports","specifier":"typing","file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py","relation":"contains","targetId":"verl/utils/logger/aggregate_logger.py#concat_dict_to_str","file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py#concat_dict_to_str","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/logger/aggregate_logger.py","recvType":"Dict"},{"source":"verl/utils/logger/aggregate_logger.py#concat_dict_to_str","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py#concat_dict_to_str","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py#concat_dict_to_str","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py","relation":"contains","targetId":"verl/utils/logger/aggregate_logger.py#LocalLogger","file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py#LocalLogger","relation":"contains","targetId":"verl/utils/logger/aggregate_logger.py#LocalLogger.__init__","file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py#LocalLogger.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py#LocalLogger","relation":"contains","targetId":"verl/utils/logger/aggregate_logger.py#LocalLogger.flush","file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py#LocalLogger","relation":"contains","targetId":"verl/utils/logger/aggregate_logger.py#LocalLogger.log","file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py#LocalLogger.log","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/logger/aggregate_logger.py"},{"source":"verl/utils/logger/aggregate_logger.py#LocalLogger.log","relation":"calls","name":"concat_dict_to_str","viaMember":false,"file":"verl/utils/logger/aggregate_logger.py"}]},"verl/utils/logging_utils.py":{"size":843,"mtimeMs":1789275830323.203,"hash":"ee547e49dd03ebf83d493ee11ee71c65ef1ae8a6b68d7d21c195e0d24abcea4b","nodes":[{"id":"verl/utils/logging_utils.py","name":"logging_utils.py","kind":"file","path":"verl/utils/logging_utils.py","span":"L1-L23","signature":null,"exported":true,"origin":"ast","body_hash":"ee547e49dd03ebf83d493ee11ee71c65ef1ae8a6b68d7d21c195e0d24abcea4b","chars":843,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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"},{"id":"verl/utils/logging_utils.py#set_basic_config","name":"set_basic_config","kind":"function","path":"verl/utils/logging_utils.py","span":"L18-L22","signature":"def set_basic_config(level)","exported":true,"origin":"ast","body_hash":"d584bbdf458448eb706bccf57353df1759d37ce74f786b42c99e6ca45971cd9f","body_text":"def set_basic_config(level): \"\"\" This function sets the global logging format and level. It will be called when import verl \"\"\" logging.basicConfig(format='%(levelname)s:%(asctime)s:%(message)s', level=level)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/logging_utils.py","relation":"imports","specifier":"logging","file":"verl/utils/logging_utils.py"},{"source":"verl/utils/logging_utils.py","relation":"contains","targetId":"verl/utils/logging_utils.py#set_basic_config","file":"verl/utils/logging_utils.py"},{"source":"verl/utils/logging_utils.py#set_basic_config","relation":"calls","name":"basicConfig","viaMember":true,"file":"verl/utils/logging_utils.py"}]},"verl/utils/megatron/__init__.py":{"size":600,"mtimeMs":1789275830323.203,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/utils/megatron/__init__.py","name":"__init__.py","kind":"file","path":"verl/utils/megatron/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/utils/megatron/memory.py":{"size":1534,"mtimeMs":1789275830323.203,"hash":"422a9340d8aae030664b5e425bb89bffc5d428b34057d091905b84a5033291e1","nodes":[{"id":"verl/utils/megatron/memory.py","name":"memory.py","kind":"file","path":"verl/utils/megatron/memory.py","span":"L1-L42","signature":null,"exported":true,"origin":"ast","body_hash":"422a9340d8aae030664b5e425bb89bffc5d428b34057d091905b84a5033291e1","chars":1534,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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"},{"id":"verl/utils/megatron/memory.py#MemoryBuffer","name":"MemoryBuffer","kind":"class","path":"verl/utils/megatron/memory.py","span":"L18-L41","signature":"class MemoryBuffer","exported":true,"origin":"ast","body_hash":"c2763f32cf268585050206d7575cc1a19bd359ad06a1c59e120625d3412535ed","body_text":"class MemoryBuffer: def __init__(self, numel, numel_padded, dtype): self.numel = numel self.numel_padded = numel_padded self.dtype = dtype self.data = torch.zeros(self.numel_padded, dtype=self.dtype, device=torch.cuda.current_device(), requires_grad=False) def zero(self): \"\"\"Reset the buffer to zero.\"\"\" self.data.zero_() def get(self, shape, start_index): \"\"\"Return a tensor with the input `shape` as a view into the 1-D data starting at `start_index`.\"\"\" end_index = start_index + shape.numel() assert end_index <= self.numel, \\ 'requested tensor is out of the buffer range.' buffer_tensor = self.data[start_index:end_index] buffer_tensor = buffer_tensor.view(shape) return buffer_tensor","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/memory.py#MemoryBuffer.__init__","name":"__init__","kind":"method","path":"verl/utils/megatron/memory.py","span":"L20-L27","signature":"def __init__(self, numel, numel_padded, dtype)","exported":false,"origin":"ast","body_hash":"297768b388d05ec49ee74628e411103b6dacbb4057f8a0137790a1d5710d44a8","body_text":"def __init__(self, numel, numel_padded, dtype): self.numel = numel self.numel_padded = numel_padded self.dtype = dtype self.data = torch.zeros(self.numel_padded, dtype=self.dtype, device=torch.cuda.current_device(), requires_grad=False)","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryBuffer"},{"id":"verl/utils/megatron/memory.py#MemoryBuffer.zero","name":"zero","kind":"method","path":"verl/utils/megatron/memory.py","span":"L29-L31","signature":"def zero(self)","exported":true,"origin":"ast","body_hash":"e70d8222f4290f5089086292ef4d72fd8ed271457a83ab124bfc80775794b347","body_text":"def zero(self): \"\"\"Reset the buffer to zero.\"\"\" self.data.zero_()","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryBuffer"},{"id":"verl/utils/megatron/memory.py#MemoryBuffer.get","name":"get","kind":"method","path":"verl/utils/megatron/memory.py","span":"L33-L41","signature":"def get(self, shape, start_index)","exported":true,"origin":"ast","body_hash":"052a2f24dd6a6ff83b6b80f67d2d3bb7f7d53cbb40001e4bfde09fb78964f55a","body_text":"def get(self, shape, start_index): \"\"\"Return a tensor with the input `shape` as a view into the 1-D data starting at `start_index`.\"\"\" end_index = start_index + shape.numel() assert end_index <= self.numel, \\ 'requested tensor is out of the buffer range.' buffer_tensor = self.data[start_index:end_index] buffer_tensor = buffer_tensor.view(shape) return buffer_tensor","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryBuffer"}],"rawEdges":[{"source":"verl/utils/megatron/memory.py","relation":"imports","specifier":"torch","file":"verl/utils/megatron/memory.py"},{"source":"verl/utils/megatron/memory.py","relation":"contains","targetId":"verl/utils/megatron/memory.py#MemoryBuffer","file":"verl/utils/megatron/memory.py"},{"source":"verl/utils/megatron/memory.py#MemoryBuffer","relation":"contains","targetId":"verl/utils/megatron/memory.py#MemoryBuffer.__init__","file":"verl/utils/megatron/memory.py"},{"source":"verl/utils/megatron/memory.py#MemoryBuffer.__init__","relation":"calls","name":"zeros","viaMember":true,"file":"verl/utils/megatron/memory.py"},{"source":"verl/utils/megatron/memory.py#MemoryBuffer.__init__","relation":"calls","name":"current_device","viaMember":true,"file":"verl/utils/megatron/memory.py"},{"source":"verl/utils/megatron/memory.py#MemoryBuffer","relation":"contains","targetId":"verl/utils/megatron/memory.py#MemoryBuffer.zero","file":"verl/utils/megatron/memory.py"},{"source":"verl/utils/megatron/memory.py#MemoryBuffer.zero","relation":"calls","name":"zero_","viaMember":true,"file":"verl/utils/megatron/memory.py"},{"source":"verl/utils/megatron/memory.py#MemoryBuffer","relation":"contains","targetId":"verl/utils/megatron/memory.py#MemoryBuffer.get","file":"verl/utils/megatron/memory.py"},{"source":"verl/utils/megatron/memory.py#MemoryBuffer.get","relation":"calls","name":"numel","viaMember":true,"file":"verl/utils/megatron/memory.py"},{"source":"verl/utils/megatron/memory.py#MemoryBuffer.get","relation":"calls","name":"view","viaMember":true,"file":"verl/utils/megatron/memory.py"}]},"verl/utils/megatron/optimizer.py":{"size":4349,"mtimeMs":1789275830323.203,"hash":"f99cc03963e38b1ae6178564adac0f9e731d875589908290dbd844bba1642d80","nodes":[{"id":"verl/utils/megatron/optimizer.py","name":"optimizer.py","kind":"file","path":"verl/utils/megatron/optimizer.py","span":"L1-L93","signature":null,"exported":true,"origin":"ast","body_hash":"f99cc03963e38b1ae6178564adac0f9e731d875589908290dbd844bba1642d80","chars":4349,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, 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 apex.optimizers import FusedAdam as Adam from apex.optimizers import FusedSGD as SGD from megatron.optimizer.distrib_optimizer import DistributedOptimizer from megatron.optimizer.grad_scaler import ConstantGradScaler, DynamicGradScaler from megatron.optimizer import Float16OptimizerWithFloat16Params, FP32Optimizer from megatron.optimizer import get_param_groups from verl.utils.megatron.optimizer_config import OptimizerConfig"},{"id":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","name":"get_megatron_optimizer","kind":"function","path":"verl/utils/megatron/optimizer.py","span":"L26-L92","signature":"def get_megatron_optimizer( model, config: OptimizerConfig, no_weight_decay_cond=None, scale_lr_cond=None, lr_mult=1.0, check_for_nan_in_loss_and_grad=False, overlap_param_gather=False # add for verl ): # Base optimizer.","exported":true,"origin":"ast","body_hash":"3a9c0aabaea523fcc5cd48e7a3d0bae22d7090ea73742b44678ed409d36b326f","body_text":"def get_megatron_optimizer( model, config: OptimizerConfig, no_weight_decay_cond=None, scale_lr_cond=None, lr_mult=1.0, check_for_nan_in_loss_and_grad=False, overlap_param_gather=False # add for verl ): # Base optimizer. param_groups = get_param_groups(model, no_weight_decay_cond, scale_lr_cond, lr_mult) if config.optimizer == 'adam': optimizer = Adam(param_groups, lr=config.lr, weight_decay=config.weight_decay, betas=(config.adam_beta1, config.adam_beta2), eps=config.adam_eps) elif config.optimizer == 'sgd': optimizer = SGD(param_groups, lr=config.lr, weight_decay=config.weight_decay, momentum=config.sgd_momentum) else: raise Exception('{} optimizer is not supported.'.format(config.optimizer)) # Determine whether the params have main-grad field. params_have_main_grad = True # Mixed precision optimizer. # - Note: both the Float16Optimizer and the DistributedOptimizer inherit # from the MixedPrecisionOptimizer, which manages any optimizer where # the model params and main params are distinct. if config.fp16 or config.bf16 or config.use_distributed_optimizer: # Grad scaler: # if loss-scale is provided, instantiate the constant scaler. # if we are using fp16 and loss-scale is not present, use a # dynamic scaler. # otherwise we are running in bf16 with no loss-scale so # leave it as None. grad_scaler = None # Constant loss scale. if config.loss_scale: grad_scaler = ConstantGradScaler(config.loss_scale) # Dynamic loss scale. else: if config.fp16: grad_scaler = DynamicGradScaler(initial_scale=config.initial_loss_scale, min_scale=config.min_loss_scale, growth_factor=2.0, backoff_factor=0.5, growth_interval=config.loss_scale_window, hysteresis=config.hysteresis) # Megatron optimizer. if config.use_distributed_optimizer: return DistributedOptimizer(optimizer, config.clip_grad, config.log_num_zeros_in_grad, check_for_nan_in_loss_and_grad, params_have_main_grad, config.fp16, config.bf16, config.params_dtype, grad_scaler, model, overlap_param_gather) else: return Float16OptimizerWithFloat16Params(optimizer, config.clip_grad, config.log_num_zeros_in_grad, check_for_nan_in_loss_and_grad, params_have_main_grad, config.fp16, config.bf16, config.params_dtype, grad_scaler, model) # FP32. return FP32Optimizer(optimizer, config.clip_grad, config.log_num_zeros_in_grad, check_for_nan_in_loss_and_grad, params_have_main_grad, model)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/megatron/optimizer.py","relation":"imports","specifier":"apex.optimizers","file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py","relation":"imports","specifier":"apex.optimizers","file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py","relation":"imports","specifier":"megatron.optimizer.distrib_optimizer","file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py","relation":"imports","specifier":"megatron.optimizer.grad_scaler","file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py","relation":"imports","specifier":"megatron.optimizer","file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py","relation":"imports","specifier":"megatron.optimizer","file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py","relation":"imports","specifier":"verl.utils.megatron.optimizer_config","file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py","relation":"contains","targetId":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","relation":"calls","name":"get_param_groups","viaMember":false,"file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","relation":"calls","name":"Adam","viaMember":false,"file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","relation":"calls","name":"SGD","viaMember":false,"file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","relation":"calls","name":"Exception","viaMember":false,"file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","relation":"calls","name":"format","viaMember":true,"file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","relation":"calls","name":"ConstantGradScaler","viaMember":false,"file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","relation":"calls","name":"DynamicGradScaler","viaMember":false,"file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","relation":"calls","name":"DistributedOptimizer","viaMember":false,"file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","relation":"calls","name":"Float16OptimizerWithFloat16Params","viaMember":false,"file":"verl/utils/megatron/optimizer.py"},{"source":"verl/utils/megatron/optimizer.py#get_megatron_optimizer","relation":"calls","name":"FP32Optimizer","viaMember":false,"file":"verl/utils/megatron/optimizer.py"}]},"verl/utils/megatron/optimizer_config.py":{"size":4051,"mtimeMs":1789275830324.203,"hash":"d978c3e93e3353af1dc30c36d0b14cbd7a2d97189fa861caa9b9591819da2f14","nodes":[{"id":"verl/utils/megatron/optimizer_config.py","name":"optimizer_config.py","kind":"file","path":"verl/utils/megatron/optimizer_config.py","span":"L1-L130","signature":null,"exported":true,"origin":"ast","body_hash":"d978c3e93e3353af1dc30c36d0b14cbd7a2d97189fa861caa9b9591819da2f14","chars":4051,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, 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 dataclasses import dataclass from typing import Callable, Optional import torch @dataclass"},{"id":"verl/utils/megatron/optimizer_config.py#OptimizerConfig","name":"OptimizerConfig","kind":"class","path":"verl/utils/megatron/optimizer_config.py","span":"L23-L129","signature":"class OptimizerConfig","exported":true,"origin":"ast","body_hash":"d9f29c42c132f4921f94b4982276829a1ef3ddb48611d2c72a500a505a9991dc","body_text":"class OptimizerConfig: \"\"\"Configuration for optimizer.\"\"\" ############## # General ############## optimizer: str = 'adam' \"\"\"Optimizer to use (one of Adam or SGD).\"\"\" lr: Optional[float] = None \"\"\"Initial learning rate. Depending on decay style and initial warmup, the learning rate at each iteration would be different. \"\"\" min_lr: Optional[float] = None \"\"\"Minumum value for learning rate. The scheduler clip values below this threshold.\"\"\" decoupled_lr: Optional[float] = None \"\"\"Separate learning rate for the input and output layer.\"\"\" decoupled_min_lr: Optional[float] = None \"\"\"Minimum value for learning rate for the input and output layer. The scheduler clip values below this threshold. \"\"\" weight_decay: float = 0.01 \"\"\"Weight decay coefficient for L2 regularization.\"\"\" ############## # Precision ############## fp16: bool = False \"\"\"If true, train with fp16 mixed precision training. Defaults to False.\"\"\" bf16: bool = False \"\"\"If true, train with bf16 mixed precision training. Defaults to False.\"\"\" params_dtype: torch.dtype = torch.float32 \"\"\"dtype used when intializing the weights. Defaults to torch.float32.\"\"\" ############### # Loss scaling ############### loss_scale: Optional[float] = None \"\"\"Static loss scaling, positive power of 2 values can improve fp16 convergence. If None, dynamic loss scaling is used. \"\"\" initial_loss_scale: float = 2**32 \"\"\"Initial loss-scale for dynamic loss scaling.\"\"\" min_loss_scale: float = 1.0 \"\"\"Minimum loss scale for dynamic loss scaling.\"\"\" loss_scale_window: float = 1000 \"\"\"Window over which to raise/lower dynamic scale.\"\"\" hysteresis: int = 2 \"\"\"Hysteresis for dynamic loss scaling.\"\"\" ############## # Optimizer ############## # Adam adam_beta1: float = 0.9 \"\"\"First coefficient for computing running averages of gradient and its square in Adam optimizer. \"\"\" adam_beta2: float = 0.999 \"\"\"Second coefficient for computing running averages of gradient and its square in Adam optimizer. \"\"\" adam_eps: float = 1e-08 \"\"\"Term added to the denominator to improve numerical stability in Adam optimizer.\"\"\" # SGD. sgd_momentum: float = 0.9 \"\"\"Momentum factor for SGD optimizer.\"\"\" ####################### # Distributed optimizer ####################### use_distributed_optimizer: bool = False \"\"\"Distribute optimizer state over data-parallel replicas.\"\"\" overlap_grad_reduce: bool = False \"\"\"If true, overlap grad reduce-scatter with backward compute in distributed optimizer.\"\"\" overlap_param_gather: bool = False \"\"\"If true, overlap param all-gather with forward compute in distributed optimizer.\"\"\" ################ # Miscellaneous ################ clip_grad: float = 1.0 \"\"\"Gradient clipping based on global L2 norm.\"\"\" log_num_zeros_in_grad: bool = False \"\"\"If true, calculate and log the number of zeros in gradient.\"\"\" barrier_with_L1_time: bool = False \"\"\"If true, use barrier with level 1 time measurements.\"\"\" timers: Callable = None \"\"\"Function to get timers.\"\"\"","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/megatron/optimizer_config.py","relation":"imports","specifier":"dataclasses","file":"verl/utils/megatron/optimizer_config.py"},{"source":"verl/utils/megatron/optimizer_config.py","relation":"imports","specifier":"typing","file":"verl/utils/megatron/optimizer_config.py"},{"source":"verl/utils/megatron/optimizer_config.py","relation":"imports","specifier":"torch","file":"verl/utils/megatron/optimizer_config.py"},{"source":"verl/utils/megatron/optimizer_config.py","relation":"contains","targetId":"verl/utils/megatron/optimizer_config.py#OptimizerConfig","file":"verl/utils/megatron/optimizer_config.py"}]},"verl/utils/megatron/pipeline_parallel.py":{"size":2119,"mtimeMs":1789275830324.203,"hash":"322d879dbfd3ae6219944ec328b2b6c0fb19091a45bb4a2257402b4934f1dd46","nodes":[{"id":"verl/utils/megatron/pipeline_parallel.py","name":"pipeline_parallel.py","kind":"file","path":"verl/utils/megatron/pipeline_parallel.py","span":"L1-L52","signature":null,"exported":true,"origin":"ast","body_hash":"322d879dbfd3ae6219944ec328b2b6c0fb19091a45bb4a2257402b4934f1dd46","chars":2119,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, 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 torch from megatron.core import parallel_state as mpu from .sequence_parallel import pad_to_sequence_parallel"},{"id":"verl/utils/megatron/pipeline_parallel.py#compute_transformers_input_shapes","name":"compute_transformers_input_shapes","kind":"function","path":"verl/utils/megatron/pipeline_parallel.py","span":"L22-L40","signature":"def compute_transformers_input_shapes(batches, meta_info)","exported":true,"origin":"ast","body_hash":"a9942586ed8e890c6a73effbaa0ef4588038655f87e9809f27f9d3493730ea30","body_text":"def compute_transformers_input_shapes(batches, meta_info): from flash_attn.bert_padding import unpad_input # flash 2 is a must for Megatron # pre-compute input shapes for each micro-batch at each pp stage input_shapes = [] for model_inputs in batches: input_ids = model_inputs['input_ids'] attention_mask = model_inputs['attention_mask'] input_ids_rmpad = unpad_input(input_ids.unsqueeze(dim=-1), attention_mask)[0] # (total_nnz, 1) if meta_info['sequence_parallel']: input_ids_rmpad = pad_to_sequence_parallel(input_ids_rmpad) # compute shapes for model_inputs input_shapes.append( torch.Size([ input_ids_rmpad.shape[0] // mpu.get_tensor_model_parallel_world_size(), 1, meta_info['hidden_size'] ])) else: # compute shapes for model_inputs input_shapes.append(torch.Size([input_ids_rmpad.shape[0], 1, meta_info['hidden_size']])) return input_shapes","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/pipeline_parallel.py#make_batch_generator","name":"make_batch_generator","kind":"function","path":"verl/utils/megatron/pipeline_parallel.py","span":"L43-L51","signature":"def make_batch_generator(batches, vpp_size)","exported":true,"origin":"ast","body_hash":"9ab0f6b61bc371f28d6b8e692dbb3dbec09a8945d0cb524d518df22de3f8001d","body_text":"def make_batch_generator(batches, vpp_size): if vpp_size > 1: # has vpp batch_generator = [batches] * vpp_size # number of vpp chunks batch_generator = [iter(b) for b in batch_generator] else: # no vpp batch_generator = iter(batches) return batch_generator","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/megatron/pipeline_parallel.py","relation":"imports","specifier":"torch","file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py","relation":"imports","specifier":"megatron.core","file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py","relation":"imports","specifier":".sequence_parallel","file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py","relation":"contains","targetId":"verl/utils/megatron/pipeline_parallel.py#compute_transformers_input_shapes","file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py","relation":"imports","specifier":"flash_attn.bert_padding","file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py#compute_transformers_input_shapes","relation":"calls","name":"unpad_input","viaMember":false,"file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py#compute_transformers_input_shapes","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py#compute_transformers_input_shapes","relation":"calls","name":"pad_to_sequence_parallel","viaMember":false,"file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py#compute_transformers_input_shapes","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py#compute_transformers_input_shapes","relation":"calls","name":"Size","viaMember":true,"file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py#compute_transformers_input_shapes","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py#compute_transformers_input_shapes","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py#compute_transformers_input_shapes","relation":"calls","name":"Size","viaMember":true,"file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py","relation":"contains","targetId":"verl/utils/megatron/pipeline_parallel.py#make_batch_generator","file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py#make_batch_generator","relation":"calls","name":"iter","viaMember":false,"file":"verl/utils/megatron/pipeline_parallel.py"},{"source":"verl/utils/megatron/pipeline_parallel.py#make_batch_generator","relation":"calls","name":"iter","viaMember":false,"file":"verl/utils/megatron/pipeline_parallel.py"}]},"verl/utils/megatron/sequence_parallel.py":{"size":1810,"mtimeMs":1789275830324.203,"hash":"5291153114babd5010d809c671c17fe14a7ac06c2ea4d997844bf3b6f1f4154e","nodes":[{"id":"verl/utils/megatron/sequence_parallel.py","name":"sequence_parallel.py","kind":"file","path":"verl/utils/megatron/sequence_parallel.py","span":"L1-L55","signature":null,"exported":true,"origin":"ast","body_hash":"5291153114babd5010d809c671c17fe14a7ac06c2ea4d997844bf3b6f1f4154e","chars":1810,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, 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 torch import torch.nn.functional as F from megatron.core import parallel_state as mpu"},{"id":"verl/utils/megatron/sequence_parallel.py#mark_parameter_as_sequence_parallel","name":"mark_parameter_as_sequence_parallel","kind":"function","path":"verl/utils/megatron/sequence_parallel.py","span":"L21-L22","signature":"def mark_parameter_as_sequence_parallel(parameter)","exported":true,"origin":"ast","body_hash":"a8c6247a9e880f728471257018885b714a64083af3b26e5b8934d45b4489dbdb","body_text":"def mark_parameter_as_sequence_parallel(parameter): setattr(parameter, 'sequence_parallel', True)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/sequence_parallel.py#is_sequence_parallel_param","name":"is_sequence_parallel_param","kind":"function","path":"verl/utils/megatron/sequence_parallel.py","span":"L25-L26","signature":"def is_sequence_parallel_param(param)","exported":true,"origin":"ast","body_hash":"712d8b0e6faf858b2a5e12606c7f6ab4fc4b139a73bffff786d71c91b30bd01d","body_text":"def is_sequence_parallel_param(param): return hasattr(param, 'sequence_parallel') and param.sequence_parallel","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/sequence_parallel.py#pad_to_sequence_parallel","name":"pad_to_sequence_parallel","kind":"function","path":"verl/utils/megatron/sequence_parallel.py","span":"L29-L54","signature":"def pad_to_sequence_parallel(unpad_tokens: torch.Tensor)","exported":true,"origin":"ast","body_hash":"ed57517b38ab2650cfe42d00636935a25cf0c5c0d226619d6566ac7bfe3a6d79","body_text":"def pad_to_sequence_parallel(unpad_tokens: torch.Tensor): \"\"\"pad the tokens such that the total length is a multiple of sp world size Args: unpad_tokens: (total_nnz, ...). Tokens after removing padding Returns: \"\"\" total_nnz = unpad_tokens.shape[0] sp_world_size = mpu.get_tensor_model_parallel_world_size() if total_nnz % sp_world_size == 0: pad_size = 0 else: pad_size = sp_world_size - total_nnz % sp_world_size if pad_size > 0: if unpad_tokens.ndim == 1: unpad_tokens = F.pad(unpad_tokens, (0, pad_size)) elif unpad_tokens.ndim == 2: unpad_tokens = F.pad(unpad_tokens, (0, 0, 0, pad_size)) else: raise NotImplementedError(f'Padding dim {unpad_tokens.ndim()} is not supported') return unpad_tokens","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/megatron/sequence_parallel.py","relation":"imports","specifier":"torch","file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py","relation":"imports","specifier":"megatron.core","file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py","relation":"contains","targetId":"verl/utils/megatron/sequence_parallel.py#mark_parameter_as_sequence_parallel","file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py#mark_parameter_as_sequence_parallel","relation":"calls","name":"setattr","viaMember":false,"file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py","relation":"contains","targetId":"verl/utils/megatron/sequence_parallel.py#is_sequence_parallel_param","file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py#is_sequence_parallel_param","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py","relation":"contains","targetId":"verl/utils/megatron/sequence_parallel.py#pad_to_sequence_parallel","file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py#pad_to_sequence_parallel","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py#pad_to_sequence_parallel","relation":"calls","name":"pad","viaMember":true,"file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py#pad_to_sequence_parallel","relation":"calls","name":"pad","viaMember":true,"file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py#pad_to_sequence_parallel","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/utils/megatron/sequence_parallel.py"},{"source":"verl/utils/megatron/sequence_parallel.py#pad_to_sequence_parallel","relation":"calls","name":"ndim","viaMember":true,"file":"verl/utils/megatron/sequence_parallel.py"}]},"verl/utils/megatron/tensor_parallel.py":{"size":7257,"mtimeMs":1789275830324.203,"hash":"a6a059f8304728ae390c96252666cf8bfc742454dbec9af2746ef46d2cc1e2ba","nodes":[{"id":"verl/utils/megatron/tensor_parallel.py","name":"tensor_parallel.py","kind":"file","path":"verl/utils/megatron/tensor_parallel.py","span":"L1-L185","signature":null,"exported":true,"origin":"ast","body_hash":"a6a059f8304728ae390c96252666cf8bfc742454dbec9af2746ef46d2cc1e2ba","chars":7257,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, 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. \"\"\" Utilities for using tensor_parallel in megatron \"\"\" from typing import Dict import torch from torch.nn import init import torch.distributed as dist from megatron.core import ModelParallelConfig from megatron.core import parallel_state as mpu, tensor_parallel import verl.utils.torch_functional as verl_F"},{"id":"verl/utils/megatron/tensor_parallel.py#update_kwargs_with_config","name":"update_kwargs_with_config","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L27-L29","signature":"def update_kwargs_with_config(dictionary: Dict, config: ModelParallelConfig)","exported":true,"origin":"ast","body_hash":"3dd8cb902228fc4a154228782e80e5ee910ccfe97aa530e2b55de06d6af1a284","body_text":"def update_kwargs_with_config(dictionary: Dict, config: ModelParallelConfig): dictionary['config'] = config return dictionary","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_model_parallel_config","name":"get_default_kwargs_for_model_parallel_config","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L32-L40","signature":"def get_default_kwargs_for_model_parallel_config()","exported":true,"origin":"ast","body_hash":"33415a4b368eb7fdab7a35b5cc99d1364d6cd42fb1d689ebc6582bb1286291ee","body_text":"def get_default_kwargs_for_model_parallel_config(): model_parallel_config_kwargs = { 'params_dtype': torch.float32, 'use_cpu_initialization': False, 'perform_initialization': True, 'gradient_accumulation_fusion': False, 'sequence_parallel': False, } return model_parallel_config_kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#get_default_model_parallel_config","name":"get_default_model_parallel_config","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L43-L44","signature":"def get_default_model_parallel_config()","exported":true,"origin":"ast","body_hash":"0d620f63a859d58ea45bf00a868b7582f2015e9ed2c1d5191b42c30250eeab11","body_text":"def get_default_model_parallel_config(): return ModelParallelConfig(**get_default_kwargs_for_model_parallel_config())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#get_common_default_kwargs_for_parallel_linear","name":"get_common_default_kwargs_for_parallel_linear","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L47-L55","signature":"def get_common_default_kwargs_for_parallel_linear()","exported":true,"origin":"ast","body_hash":"7f2f07bb5b3ac42cf7a037ea386fb59e8f167c47da9d2274ff380bb1c1aa00c5","body_text":"def get_common_default_kwargs_for_parallel_linear(): default_model_parallel_config = get_default_model_parallel_config() common_default_kwargs = { 'init_method': init.xavier_normal_, 'stride': 1, 'keep_master_weight_for_test': False, 'config': default_model_parallel_config, } return common_default_kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_column_parallel_linear","name":"get_default_kwargs_for_column_parallel_linear","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L58-L69","signature":"def get_default_kwargs_for_column_parallel_linear()","exported":true,"origin":"ast","body_hash":"e9f9395bfae185dffa8ac79360cb0c0d297b5a3461341f94d36d0efa8beeed68","body_text":"def get_default_kwargs_for_column_parallel_linear(): model_parallel_config_kwargs = get_default_kwargs_for_model_parallel_config() column_parallel_config_kwargs = { 'async_tensor_model_parallel_allreduce': False, } model_parallel_config_kwargs.update(column_parallel_config_kwargs) column_default_kwargs = { 'config': ModelParallelConfig(**model_parallel_config_kwargs), } common_default_kwargs = get_common_default_kwargs_for_parallel_linear() common_default_kwargs.update(column_default_kwargs) return common_default_kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_row_parallel_linear","name":"get_default_kwargs_for_row_parallel_linear","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L72-L74","signature":"def get_default_kwargs_for_row_parallel_linear()","exported":true,"origin":"ast","body_hash":"3015282ab53cf4db3edfa29c76938e11f5f5bc3bce00eb354ab65d217dd05e66","body_text":"def get_default_kwargs_for_row_parallel_linear(): common_default_kwargs = get_common_default_kwargs_for_parallel_linear() return common_default_kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_parallel_embedding","name":"get_default_kwargs_for_parallel_embedding","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L77-L83","signature":"def get_default_kwargs_for_parallel_embedding()","exported":true,"origin":"ast","body_hash":"3963b5e3ce23de9734f45125759b149007c48118087700224055d56a3f722538","body_text":"def get_default_kwargs_for_parallel_embedding(): model_parallel_config_kwargs = get_default_kwargs_for_model_parallel_config() embedding_default_kwargs = { 'init_method': init.xavier_normal_, 'config': ModelParallelConfig(**model_parallel_config_kwargs), } return embedding_default_kwargs","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#is_tensor_parallel_param","name":"is_tensor_parallel_param","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L86-L87","signature":"def is_tensor_parallel_param(param)","exported":true,"origin":"ast","body_hash":"5cab1728c7554d6b133422149566e3d80643f8f3be06687f5199457216e5506d","body_text":"def is_tensor_parallel_param(param): return (hasattr(param, 'tensor_model_parallel') and param.tensor_model_parallel)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#get_tensor_parallel_partition_dim","name":"get_tensor_parallel_partition_dim","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L90-L92","signature":"def get_tensor_parallel_partition_dim(param)","exported":true,"origin":"ast","body_hash":"86d85aa809fb93ab620b9c57055c3a1bd15e2bdd4be334646184a76889fde767","body_text":"def get_tensor_parallel_partition_dim(param): assert is_tensor_parallel_param(param) return param.partition_dim","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#get_tensor_parallel_partition_stride","name":"get_tensor_parallel_partition_stride","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L95-L97","signature":"def get_tensor_parallel_partition_stride(param)","exported":true,"origin":"ast","body_hash":"6ada82a39bc832a1b2f72640e6873673bb01f696855e7ad3ce4cbc07631fa658","body_text":"def get_tensor_parallel_partition_stride(param): assert is_tensor_parallel_param(param) return param.partition_stride","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy","name":"_VocabParallelEntropy","kind":"class","path":"verl/utils/megatron/tensor_parallel.py","span":"L100-L121","signature":"class _VocabParallelEntropy(torch.autograd.Function)","exported":false,"origin":"ast","body_hash":"7b37475551e7d9847675a65a965404bec95ba0971bd85e4b4bc2d287df2c9a57","body_text":"class _VocabParallelEntropy(torch.autograd.Function): @staticmethod def forward(ctx, vocab_parallel_logits: torch.Tensor) -> torch.Tensor: logits_max = vocab_parallel_logits.max(dim=-1, keepdim=True).values dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=mpu.get_tensor_model_parallel_group()) normalized_vocab_parallel_logits = vocab_parallel_logits - logits_max normalized_exp_logits = normalized_vocab_parallel_logits.exp() normalized_sum_exp_logits = normalized_exp_logits.sum(dim=-1, keepdim=True) dist.all_reduce(normalized_sum_exp_logits, group=mpu.get_tensor_model_parallel_group()) softmax_logits = normalized_exp_logits / normalized_sum_exp_logits sum_softmax_times_logits = (softmax_logits * vocab_parallel_logits).sum(dim=-1, keepdim=True) dist.all_reduce(sum_softmax_times_logits, group=mpu.get_tensor_model_parallel_group()) entropy = logits_max + normalized_sum_exp_logits.log() - sum_softmax_times_logits ctx.save_for_backward(vocab_parallel_logits, softmax_logits, sum_softmax_times_logits) return entropy.squeeze(dim=-1) @staticmethod def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: vocab_parallel_logits, softmax_logits, sum_softmax_times_logits = ctx.saved_tensors grad_input = grad_output.unsqueeze(dim=-1) * softmax_logits * (sum_softmax_times_logits - vocab_parallel_logits) return grad_input","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","name":"forward","kind":"method","path":"verl/utils/megatron/tensor_parallel.py","span":"L103-L115","signature":"def forward(ctx, vocab_parallel_logits: torch.Tensor) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"7fa152a3e81c05fc2e71bf2d969d4b5ebc1e59e5dfddeca628934c425cfaf8a7","body_text":"def forward(ctx, vocab_parallel_logits: torch.Tensor) -> torch.Tensor: logits_max = vocab_parallel_logits.max(dim=-1, keepdim=True).values dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=mpu.get_tensor_model_parallel_group()) normalized_vocab_parallel_logits = vocab_parallel_logits - logits_max normalized_exp_logits = normalized_vocab_parallel_logits.exp() normalized_sum_exp_logits = normalized_exp_logits.sum(dim=-1, keepdim=True) dist.all_reduce(normalized_sum_exp_logits, group=mpu.get_tensor_model_parallel_group()) softmax_logits = normalized_exp_logits / normalized_sum_exp_logits sum_softmax_times_logits = (softmax_logits * vocab_parallel_logits).sum(dim=-1, keepdim=True) dist.all_reduce(sum_softmax_times_logits, group=mpu.get_tensor_model_parallel_group()) entropy = logits_max + normalized_sum_exp_logits.log() - sum_softmax_times_logits ctx.save_for_backward(vocab_parallel_logits, softmax_logits, sum_softmax_times_logits) return entropy.squeeze(dim=-1)","summary_state":"pending","summary":null,"crux":null,"owner":"_VocabParallelEntropy"},{"id":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.backward","name":"backward","kind":"method","path":"verl/utils/megatron/tensor_parallel.py","span":"L118-L121","signature":"def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"844f47cc144371fe27def1dc3f09d0bbd0f532ae72aa8c91341545ecd5ffdc82","body_text":"def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: vocab_parallel_logits, softmax_logits, sum_softmax_times_logits = ctx.saved_tensors grad_input = grad_output.unsqueeze(dim=-1) * softmax_logits * (sum_softmax_times_logits - vocab_parallel_logits) return grad_input","summary_state":"pending","summary":null,"crux":null,"owner":"_VocabParallelEntropy"},{"id":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_entropy","name":"vocab_parallel_entropy","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L124-L133","signature":"def vocab_parallel_entropy(vocab_parallel_logits: torch.Tensor) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"b83713a86d9570b38537805b042e18386f8c54c406f2a433068354c8c57e19c7","body_text":"def vocab_parallel_entropy(vocab_parallel_logits: torch.Tensor) -> torch.Tensor: \"\"\"Compute entropy when the logits are sharded in tp ranks Args: vocab_parallel_logits: (total_nnz, vocab_size // tp_size) Returns: (total_nnz,) \"\"\" return _VocabParallelEntropy.apply(vocab_parallel_logits)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits","name":"vocab_parallel_log_probs_from_logits","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L136-L138","signature":"def vocab_parallel_log_probs_from_logits(logits, labels)","exported":true,"origin":"ast","body_hash":"5db09d406b8217bce4a18a4f79b67b4383849b2059c65db177318d9c213c2580","body_text":"def vocab_parallel_log_probs_from_logits(logits, labels): \"\"\"TODO(zhangchi.usc1992): We may change the implementation later\"\"\" return -tensor_parallel.vocab_parallel_cross_entropy(vocab_parallel_logits=logits, target=labels)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits_response_rmpad","name":"vocab_parallel_log_probs_from_logits_response_rmpad","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L141-L165","signature":"def vocab_parallel_log_probs_from_logits_response_rmpad(input_ids, attention_mask, logits_rmpad, response_length)","exported":true,"origin":"ast","body_hash":"bf4a1f5c4d259d67466eb579b797d88cf9591953621cac911373d6f66ae171f8","body_text":"def vocab_parallel_log_probs_from_logits_response_rmpad(input_ids, attention_mask, logits_rmpad, response_length): \"\"\"Similar to log_probs_from_logits_response_rmpad, but the logits_rmpad is now spliited across tensor parallel region. This will further reduce the peak memory usage during training Args: input_ids: [batch_size, seqlen] attention_mask: [batch_size, seqlen] logits_rmpad: [total_nnz, vocab_size // tp_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 = vocab_parallel_log_probs_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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_compute_entropy_loss","name":"vocab_parallel_compute_entropy_loss","kind":"function","path":"verl/utils/megatron/tensor_parallel.py","span":"L168-L184","signature":"def vocab_parallel_compute_entropy_loss(logits, eos_mask)","exported":true,"origin":"ast","body_hash":"e79360be5eed4fe3a5cf553b592ace6d6f968e51e090494b21e42b0da3ad8e6b","body_text":"def vocab_parallel_compute_entropy_loss(logits, eos_mask): \"\"\"Compute Categorical entropy loss Args: logits: `(torch.Tensor)` shape: (bs, response_length, vocab_size) eos_mask: `(torch.Tensor)` shape: (bs, response_length) Returns: entropy: a scalar torch.Tensor \"\"\" # compute entropy entropy = vocab_parallel_entropy(logits) entropy_loss = verl_F.masked_mean(entropy, mask=eos_mask) return entropy_loss","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/megatron/tensor_parallel.py","relation":"imports","specifier":"typing","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"imports","specifier":"torch","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"imports","specifier":"torch.nn","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"imports","specifier":"megatron.core","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"imports","specifier":"megatron.core","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#update_kwargs_with_config","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_model_parallel_config","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#get_default_model_parallel_config","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_default_model_parallel_config","relation":"calls","name":"ModelParallelConfig","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_default_model_parallel_config","relation":"calls","name":"get_default_kwargs_for_model_parallel_config","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#get_common_default_kwargs_for_parallel_linear","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_common_default_kwargs_for_parallel_linear","relation":"calls","name":"get_default_model_parallel_config","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_column_parallel_linear","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_column_parallel_linear","relation":"calls","name":"get_default_kwargs_for_model_parallel_config","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_column_parallel_linear","relation":"calls","name":"update","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py","recvType":"get_default_kwargs_for_model_parallel_config"},{"source":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_column_parallel_linear","relation":"calls","name":"ModelParallelConfig","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_column_parallel_linear","relation":"calls","name":"get_common_default_kwargs_for_parallel_linear","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_column_parallel_linear","relation":"calls","name":"update","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py","recvType":"get_common_default_kwargs_for_parallel_linear"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_row_parallel_linear","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_row_parallel_linear","relation":"calls","name":"get_common_default_kwargs_for_parallel_linear","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_parallel_embedding","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_parallel_embedding","relation":"calls","name":"get_default_kwargs_for_model_parallel_config","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_default_kwargs_for_parallel_embedding","relation":"calls","name":"ModelParallelConfig","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#is_tensor_parallel_param","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#is_tensor_parallel_param","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#get_tensor_parallel_partition_dim","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_tensor_parallel_partition_dim","relation":"calls","name":"is_tensor_parallel_param","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#get_tensor_parallel_partition_stride","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#get_tensor_parallel_partition_stride","relation":"calls","name":"is_tensor_parallel_param","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"max","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"exp","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"log","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"save_for_backward","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.forward","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.backward","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#_VocabParallelEntropy.backward","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_entropy","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_entropy","relation":"calls","name":"apply","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits","relation":"calls","name":"vocab_parallel_cross_entropy","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits_response_rmpad","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"imports","specifier":"flash_attn.bert_padding","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits_response_rmpad","relation":"calls","name":"unpad_input","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits_response_rmpad","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits_response_rmpad","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits_response_rmpad","relation":"calls","name":"roll","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits_response_rmpad","relation":"calls","name":"vocab_parallel_log_probs_from_logits","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits_response_rmpad","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits_response_rmpad","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py","recvType":"vocab_parallel_log_probs_from_logits"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_log_probs_from_logits_response_rmpad","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py","recvType":"pad_input"},{"source":"verl/utils/megatron/tensor_parallel.py","relation":"contains","targetId":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_compute_entropy_loss","file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_compute_entropy_loss","relation":"calls","name":"vocab_parallel_entropy","viaMember":false,"file":"verl/utils/megatron/tensor_parallel.py"},{"source":"verl/utils/megatron/tensor_parallel.py#vocab_parallel_compute_entropy_loss","relation":"calls","name":"masked_mean","viaMember":true,"file":"verl/utils/megatron/tensor_parallel.py"}]},"verl/utils/megatron_utils.py":{"size":11812,"mtimeMs":1789275830324.203,"hash":"905d61e476d8062c298db358f8490a67e4b836b755ab6fbc0d45559f62d8dd23","nodes":[{"id":"verl/utils/megatron_utils.py","name":"megatron_utils.py","kind":"file","path":"verl/utils/megatron_utils.py","span":"L1-L254","signature":null,"exported":true,"origin":"ast","body_hash":"905d61e476d8062c298db358f8490a67e4b836b755ab6fbc0d45559f62d8dd23","chars":11812,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, 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. \"\"\"Pretrain utilities.\"\"\" from typing import Any, Dict import time from omegaconf import DictConfig from verl.utils.torch_dtypes import PrecisionType from verl.utils.memory_buffer import build_memory_reference_from_module import torch import torch.nn as nn import torch.nn.functional as F from megatron.core import mpu, tensor_parallel from megatron.core.utils import get_model_config from megatron.core.transformer import TransformerConfig from megatron.core.transformer.module import Float16Module # from megatron.core.distributed import DistributedDataParallelConfig from megatron.core.distributed import DistributedDataParallel as DDP from megatron.core.enums import ModelType ALL_MODULE_WRAPPER_CLASSNAMES = (DDP, Float16Module) from transformers import PretrainedConfig # from megatron.core.optimizer import OptimizerConfig from verl.utils.megatron.optimizer_config import OptimizerConfig from megatron.core import ModelParallelConfig"},{"id":"verl/utils/megatron_utils.py#get_model","name":"get_model","kind":"function","path":"verl/utils/megatron_utils.py","span":"L34-L116","signature":"def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True)","exported":true,"origin":"ast","body_hash":"ade3c583b2baa8e799ef297043fe3abb8af3ab6e0fef0cd4e45215bb0bdf9220","body_text":"def get_model(model_provider_func, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True): \"\"\"Build the model.\"\"\" # Build model. if mpu.get_pipeline_model_parallel_world_size() > 1 and \\ mpu.get_virtual_pipeline_model_parallel_world_size() is not None: assert model_type != ModelType.encoder_and_decoder, \\ \"Interleaved schedule not supported for model with both encoder and decoder\" model = [] for i in range(mpu.get_virtual_pipeline_model_parallel_world_size()): mpu.set_virtual_pipeline_model_parallel_rank(i) # Set pre_process and post_process only after virtual rank is set. pre_process = mpu.is_pipeline_first_stage() post_process = mpu.is_pipeline_last_stage() this_model = model_provider_func(pre_process=pre_process, post_process=post_process) this_model.model_type = model_type model.append(this_model) else: pre_process = mpu.is_pipeline_first_stage() post_process = mpu.is_pipeline_last_stage() add_encoder = True add_decoder = True if model_type == ModelType.encoder_and_decoder: if mpu.get_pipeline_model_parallel_world_size() > 1: assert mpu.get_pipeline_model_parallel_split_rank() is not None, \\ \"Split rank needs to be specified for model with both encoder and decoder\" rank = mpu.get_pipeline_model_parallel_rank() split_rank = mpu.get_pipeline_model_parallel_split_rank() world_size = mpu.get_pipeline_model_parallel_world_size() pre_process = rank == 0 or rank == split_rank post_process = (rank == (split_rank - 1)) or (rank == (world_size - 1)) add_encoder = mpu.is_pipeline_stage_before_split() add_decoder = mpu.is_pipeline_stage_after_split() model = model_provider_func(pre_process=pre_process, post_process=post_process, add_encoder=add_encoder, add_decoder=add_decoder) else: model = model_provider_func(pre_process=pre_process, post_process=post_process) model.model_type = model_type if not isinstance(model, list): model = [model] # Set tensor model parallel attributes if not set. # Only parameters that are already tensor model parallel have these # attributes set for them. We should make sure the default attributes # are set for all params so the optimizer can use them. for model_module in model: for param in model_module.parameters(): tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param) # Print number of parameters. if mpu.get_data_parallel_rank() == 0: print(' > number of parameters on (tensor, pipeline) ' 'model parallel rank ({}, {}): {}'.format( mpu.get_tensor_model_parallel_rank(), mpu.get_pipeline_model_parallel_rank(), sum([sum([p.nelement() for p in model_module.parameters()]) for model_module in model])), flush=True) # GPU allocation. for model_module in model: model_module.cuda(torch.cuda.current_device()) # Fp16 conversion. config = get_model_config(model[0]) if config.fp16 or config.bf16: # the ModelParallelConfig in GPTModel model = [Float16Module(config, model_module) for model_module in model] if wrap_with_ddp: model = [ DDP(config=config, module=model_chunk, data_parallel_group=mpu.get_data_parallel_group(with_context_parallel=True), accumulate_allreduce_grads_in_fp32=True, overlap_grad_reduce=False, use_distributed_optimizer=True, disable_bucketing=(model_chunk_idx > 0)) for (model_chunk_idx, model_chunk) in enumerate(model) ] # # Broadcast params from data parallel src rank to other data parallel ranks. # if args.data_parallel_random_init: for model_module in model: model_module.broadcast_params() return model","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron_utils.py#unwrap_model","name":"unwrap_model","kind":"function","path":"verl/utils/megatron_utils.py","span":"L122-L134","signature":"def unwrap_model(model, module_instances=ALL_MODULE_WRAPPER_CLASSNAMES)","exported":true,"origin":"ast","body_hash":"a537a4570de1f758c96a9cff634ecfeff7b84b32640072e034f57769165c0542","body_text":"def unwrap_model(model, module_instances=ALL_MODULE_WRAPPER_CLASSNAMES): return_list = True if not isinstance(model, list): model = [model] return_list = False unwrapped_model = [] for model_module in model: while isinstance(model_module, module_instances): model_module = model_module.module unwrapped_model.append(model_module) if not return_list: return unwrapped_model[0] return unwrapped_model","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron_utils.py#convert_config","name":"convert_config","kind":"function","path":"verl/utils/megatron_utils.py","span":"L140-L177","signature":"def convert_config(hf_config: PretrainedConfig, megatron_config) -> TransformerConfig","exported":true,"origin":"ast","body_hash":"cb5bc8628d79ad82aac2d6fd5c76ae05204bfcf1e8c42536c87518245a99be3e","body_text":"def convert_config(hf_config: PretrainedConfig, megatron_config) -> TransformerConfig: print(f'megatron config {megatron_config}') dt = PrecisionType.to_dtype(megatron_config['param_dtype']) print(f'pipeline_dtype=megatron_config {dt}') transformer_config = TransformerConfig( 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, # max_position_embeddings=hf_config.max_position_embeddings, activation_func=F.silu, normalization='RMSNorm', # rotary_percent=False, # default, gated_linear_unit=True, # for llama use_cpu_initialization=True, apply_residual_connection_post_layernorm=False, # check what's this mean add_bias_linear=False, tensor_model_parallel_size=mpu.get_tensor_model_parallel_world_size(), pipeline_model_parallel_size=mpu.get_pipeline_model_parallel_world_size(), virtual_pipeline_model_parallel_size=mpu.get_virtual_pipeline_model_parallel_world_size(), pipeline_dtype=PrecisionType.to_dtype(megatron_config['param_dtype']), params_dtype=PrecisionType.to_dtype(megatron_config['param_dtype']), sequence_parallel=megatron_config['sequence_parallel_enabled'], variable_seq_lengths=True, masked_softmax_fusion=True, bf16=PrecisionType.to_dtype(megatron_config['param_dtype']) is torch.bfloat16) if torch.distributed.get_rank() == 0: print(f'tensor_parallel_size={transformer_config.tensor_model_parallel_size} \\n \\ pipeline_model_parallel_size={transformer_config.pipeline_model_parallel_size} \\n \\ virtual_pipeline_model_parallel_size={transformer_config.virtual_pipeline_model_parallel_size} \\n \\ pipeline_dtype={transformer_config.pipeline_dtype} \\n \\ params_dtype={transformer_config.params_dtype} \\n \\ sequence_parallel={transformer_config.sequence_parallel} \\n \\ variable_seq_lengths={transformer_config.variable_seq_lengths} \\n \\ masked_softmax_fusion={transformer_config.masked_softmax_fusion} \\n ') return transformer_config","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron_utils.py#init_megatron_optim_config","name":"init_megatron_optim_config","kind":"function","path":"verl/utils/megatron_utils.py","span":"L185-L195","signature":"def init_megatron_optim_config(optim_config: Dict) -> OptimizerConfig","exported":true,"origin":"ast","body_hash":"24df7d09862a57b68cbfe8d7ae7a9200a34df88e7e0bcedae89468e029cee0ea","body_text":"def init_megatron_optim_config(optim_config: Dict) -> OptimizerConfig: config = OptimizerConfig( optimizer='adam', lr=optim_config.get('lr'), clip_grad=optim_config.get('clip_grad'), weight_decay=1e-2, bf16=True, params_dtype=torch.bfloat16, use_distributed_optimizer=True, ) return config","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron_utils.py#init_model_parallel_config","name":"init_model_parallel_config","kind":"function","path":"verl/utils/megatron_utils.py","span":"L201-L212","signature":"def init_model_parallel_config(config: DictConfig) -> ModelParallelConfig: # TODO(sgm): check how to disable megatron timers","exported":true,"origin":"ast","body_hash":"59d7489f7d7998c6b59743bb8123ae41eb7a40c046ca03d67470d2f7ca2cf457","body_text":"def init_model_parallel_config(config: DictConfig) -> ModelParallelConfig: # TODO(sgm): check how to disable megatron timers timers = FakeTimers() return ModelParallelConfig(tensor_model_parallel_size=config.get('tensor_model_parallel_size'), pipeline_model_parallel_size=config.get('pipeline_model_parallel_size'), virtual_pipeline_model_parallel_size=config.get('virtual_pipeline_model_parallel_size'), sequence_parallel=config.get('sequence_parallel'), params_dtype=PrecisionType.to_dtype(config.get('param_dtype')), pipeline_dtype=PrecisionType.to_dtype(config.get('param_dtype')), bf16=True, fp16=False, timers=timers)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron_utils.py#FakeTimers","name":"FakeTimers","kind":"class","path":"verl/utils/megatron_utils.py","span":"L215-L223","signature":"class FakeTimers","exported":true,"origin":"ast","body_hash":"335700ed88928305b8babb924411a0c957855d10334d8fb351d5ed9f7d368569","body_text":"class FakeTimers: \"\"\"Disable All Megatron Timing with FakeTimers\"\"\" def __init__(self): from megatron.timers import DummyTimer self.dummy_timer = DummyTimer() def __call__(self, *args: Any, **kwds: Any) -> Any: return self.dummy_timer","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron_utils.py#FakeTimers.__init__","name":"__init__","kind":"method","path":"verl/utils/megatron_utils.py","span":"L218-L220","signature":"def __init__(self)","exported":false,"origin":"ast","body_hash":"6cc9c76f6e5afb81abad8986c8fcfa06bb634e916625f7e0b49896a41e59e152","body_text":"def __init__(self): from megatron.timers import DummyTimer self.dummy_timer = DummyTimer()","summary_state":"pending","summary":null,"crux":null,"owner":"FakeTimers"},{"id":"verl/utils/megatron_utils.py#FakeTimers.__call__","name":"__call__","kind":"method","path":"verl/utils/megatron_utils.py","span":"L222-L223","signature":"def __call__(self, *args: Any, **kwds: Any) -> Any","exported":false,"origin":"ast","body_hash":"42bf4959fb835b75530ebcc97813e6bf71e9c1f772d5c12440a05fdb117c4886","body_text":"def __call__(self, *args: Any, **kwds: Any) -> Any: return self.dummy_timer","summary_state":"pending","summary":null,"crux":null,"owner":"FakeTimers"},{"id":"verl/utils/megatron_utils.py#offload_megatron_param_and_grad","name":"offload_megatron_param_and_grad","kind":"function","path":"verl/utils/megatron_utils.py","span":"L226-L238","signature":"def offload_megatron_param_and_grad(module_list: nn.ModuleList, offload_grad=False, hybrid_engine=None)","exported":true,"origin":"ast","body_hash":"18f56e697ddd468f96d758493dc35852d9ac0cd6b544d4cd0fa34704959a418f","body_text":"def offload_megatron_param_and_grad(module_list: nn.ModuleList, offload_grad=False, hybrid_engine=None): if hybrid_engine is not None: pp_rank = mpu.get_pipeline_model_parallel_rank() for buffer in hybrid_engine.memory_buffers[pp_rank].values(): buffer.data = buffer.data.to('cpu', non_blocking=True) build_memory_reference_from_module(module_list, hybrid_engine.memory_buffers[pp_rank], maintain_weight=True) else: for module in module_list: for _, param in module.named_parameters(): param.data = param.data.to('cpu', non_blocking=True) if offload_grad and param.grad is not None: param.grad = param.grad.to(\"cpu\", non_blocking=True) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/megatron_utils.py#load_megatron_param_and_grad","name":"load_megatron_param_and_grad","kind":"function","path":"verl/utils/megatron_utils.py","span":"L241-L253","signature":"def load_megatron_param_and_grad(module_list: nn.ModuleList, device_id, load_grad=False, hybrid_engine=None)","exported":true,"origin":"ast","body_hash":"613d7762ed9d7be88dfb2c7e07d24009c630e6e4cb174b12c33b6ff687b1bbaa","body_text":"def load_megatron_param_and_grad(module_list: nn.ModuleList, device_id, load_grad=False, hybrid_engine=None): if hybrid_engine is not None: pp_rank = mpu.get_pipeline_model_parallel_rank() for buffer in hybrid_engine.memory_buffers[pp_rank].values(): buffer.data = buffer.data.to(device_id, non_blocking=True) build_memory_reference_from_module(module_list, hybrid_engine.memory_buffers[pp_rank], maintain_weight=True) else: for module in module_list: for _, param in module.named_parameters(): param.data = param.data.to(device_id, non_blocking=True) if load_grad and param.grad is not None: param.grad = param.grad.to(device_id, non_blocking=True) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"typing","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"time","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"omegaconf","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"verl.utils.torch_dtypes","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"verl.utils.memory_buffer","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"torch","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"megatron.core","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"megatron.core.utils","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"megatron.core.transformer","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"megatron.core.transformer.module","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"megatron.core.distributed","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"megatron.core.enums","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"contains","targetId":"verl/utils/megatron_utils.py#get_model","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"set_virtual_pipeline_model_parallel_rank","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"is_pipeline_first_stage","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"is_pipeline_last_stage","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"model_provider_func","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/megatron_utils.py","recvType":"model_provider_func"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"is_pipeline_first_stage","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"is_pipeline_last_stage","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_pipeline_model_parallel_split_rank","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_pipeline_model_parallel_split_rank","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"is_pipeline_stage_before_split","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"is_pipeline_stage_after_split","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"model_provider_func","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"model_provider_func","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"parameters","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"set_defaults_if_not_set_tensor_model_parallel_attributes","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_data_parallel_rank","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"format","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_tensor_model_parallel_rank","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"sum","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"sum","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"nelement","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"parameters","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"current_device","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_model_config","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"Float16Module","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"DDP","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"get_data_parallel_group","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#get_model","relation":"calls","name":"broadcast_params","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"contains","targetId":"verl/utils/megatron_utils.py#unwrap_model","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#unwrap_model","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#unwrap_model","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#unwrap_model","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"transformers","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"contains","targetId":"verl/utils/megatron_utils.py#convert_config","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"TransformerConfig","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#convert_config","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"verl.utils.megatron.optimizer_config","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"contains","targetId":"verl/utils/megatron_utils.py#init_megatron_optim_config","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#init_megatron_optim_config","relation":"calls","name":"OptimizerConfig","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#init_megatron_optim_config","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/megatron_utils.py","recvType":"Dict"},{"source":"verl/utils/megatron_utils.py#init_megatron_optim_config","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/megatron_utils.py","recvType":"Dict"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"megatron.core","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"contains","targetId":"verl/utils/megatron_utils.py#init_model_parallel_config","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#init_model_parallel_config","relation":"calls","name":"FakeTimers","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#init_model_parallel_config","relation":"calls","name":"ModelParallelConfig","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#init_model_parallel_config","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/megatron_utils.py","recvType":"DictConfig"},{"source":"verl/utils/megatron_utils.py#init_model_parallel_config","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/megatron_utils.py","recvType":"DictConfig"},{"source":"verl/utils/megatron_utils.py#init_model_parallel_config","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/megatron_utils.py","recvType":"DictConfig"},{"source":"verl/utils/megatron_utils.py#init_model_parallel_config","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/megatron_utils.py","recvType":"DictConfig"},{"source":"verl/utils/megatron_utils.py#init_model_parallel_config","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#init_model_parallel_config","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/megatron_utils.py","recvType":"DictConfig"},{"source":"verl/utils/megatron_utils.py#init_model_parallel_config","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#init_model_parallel_config","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/megatron_utils.py","recvType":"DictConfig"},{"source":"verl/utils/megatron_utils.py","relation":"contains","targetId":"verl/utils/megatron_utils.py#FakeTimers","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#FakeTimers","relation":"contains","targetId":"verl/utils/megatron_utils.py#FakeTimers.__init__","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"imports","specifier":"megatron.timers","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#FakeTimers.__init__","relation":"calls","name":"DummyTimer","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#FakeTimers","relation":"contains","targetId":"verl/utils/megatron_utils.py#FakeTimers.__call__","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"contains","targetId":"verl/utils/megatron_utils.py#offload_megatron_param_and_grad","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#offload_megatron_param_and_grad","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#offload_megatron_param_and_grad","relation":"calls","name":"values","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#offload_megatron_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#offload_megatron_param_and_grad","relation":"calls","name":"build_memory_reference_from_module","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#offload_megatron_param_and_grad","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#offload_megatron_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#offload_megatron_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#offload_megatron_param_and_grad","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py","relation":"contains","targetId":"verl/utils/megatron_utils.py#load_megatron_param_and_grad","file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#load_megatron_param_and_grad","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#load_megatron_param_and_grad","relation":"calls","name":"values","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#load_megatron_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#load_megatron_param_and_grad","relation":"calls","name":"build_memory_reference_from_module","viaMember":false,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#load_megatron_param_and_grad","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#load_megatron_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#load_megatron_param_and_grad","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/megatron_utils.py"},{"source":"verl/utils/megatron_utils.py#load_megatron_param_and_grad","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/utils/megatron_utils.py"}]},"verl/utils/memory_buffer.py":{"size":7964,"mtimeMs":1789275830324.203,"hash":"736f122e7ca97cbfb173aaa97eab0f55fae19d8c0e3fc2da2e1722adeca8c518","nodes":[{"id":"verl/utils/memory_buffer.py","name":"memory_buffer.py","kind":"file","path":"verl/utils/memory_buffer.py","span":"L1-L215","signature":null,"exported":true,"origin":"ast","body_hash":"736f122e7ca97cbfb173aaa97eab0f55fae19d8c0e3fc2da2e1722adeca8c518","chars":7964,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 file contains utilities to manipulate torch memory buffers \"\"\" from typing import Dict, List import torch from torch import nn"},{"id":"verl/utils/memory_buffer.py#MemoryBuffer","name":"MemoryBuffer","kind":"class","path":"verl/utils/memory_buffer.py","span":"L24-L48","signature":"class MemoryBuffer","exported":true,"origin":"ast","body_hash":"3ecd89af1ea2c481fca6e98a12eed5e85d7fc5b1800b1df25c6dfc9636d3ad49","body_text":"class MemoryBuffer: \"\"\" A memory buffer is a contiguous torch tensor that may combine multiple tensors sharing with the underlying memory. It must have a unique type to support this behavior. \"\"\" def __init__(self, numel: int, numel_padded: int, dtype: torch.dtype): self.numel = numel self.numel_padded = numel_padded self.dtype = dtype self.data = torch.zeros(self.numel_padded, dtype=self.dtype, device='cuda', requires_grad=False) def zero(self): \"\"\"Reset the buffer to zero.\"\"\" self.data.zero_() def get(self, shape, start_index): \"\"\"Return a tensor with the input `shape` as a view into the 1-D data starting at `start_index`.\"\"\" end_index = start_index + shape.numel() assert end_index <= self.numel, \\ 'requested tensor is out of the buffer range.' buffer_tensor = self.data[start_index:end_index] buffer_tensor = buffer_tensor.view(shape) return buffer_tensor","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/memory_buffer.py#MemoryBuffer.__init__","name":"__init__","kind":"method","path":"verl/utils/memory_buffer.py","span":"L30-L34","signature":"def __init__(self, numel: int, numel_padded: int, dtype: torch.dtype)","exported":false,"origin":"ast","body_hash":"45ed35d28f1f74afc588946ea27d57cdcb23b36eeea4a303182c749758a64d5c","body_text":"def __init__(self, numel: int, numel_padded: int, dtype: torch.dtype): self.numel = numel self.numel_padded = numel_padded self.dtype = dtype self.data = torch.zeros(self.numel_padded, dtype=self.dtype, device='cuda', requires_grad=False)","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryBuffer"},{"id":"verl/utils/memory_buffer.py#MemoryBuffer.zero","name":"zero","kind":"method","path":"verl/utils/memory_buffer.py","span":"L36-L38","signature":"def zero(self)","exported":true,"origin":"ast","body_hash":"e70d8222f4290f5089086292ef4d72fd8ed271457a83ab124bfc80775794b347","body_text":"def zero(self): \"\"\"Reset the buffer to zero.\"\"\" self.data.zero_()","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryBuffer"},{"id":"verl/utils/memory_buffer.py#MemoryBuffer.get","name":"get","kind":"method","path":"verl/utils/memory_buffer.py","span":"L40-L48","signature":"def get(self, shape, start_index)","exported":true,"origin":"ast","body_hash":"052a2f24dd6a6ff83b6b80f67d2d3bb7f7d53cbb40001e4bfde09fb78964f55a","body_text":"def get(self, shape, start_index): \"\"\"Return a tensor with the input `shape` as a view into the 1-D data starting at `start_index`.\"\"\" end_index = start_index + shape.numel() assert end_index <= self.numel, \\ 'requested tensor is out of the buffer range.' buffer_tensor = self.data[start_index:end_index] buffer_tensor = buffer_tensor.view(shape) return buffer_tensor","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryBuffer"},{"id":"verl/utils/memory_buffer.py#calc_padded_numel","name":"calc_padded_numel","kind":"function","path":"verl/utils/memory_buffer.py","span":"L51-L55","signature":"def calc_padded_numel(shape: torch.Size, dtype: torch.dtype)","exported":true,"origin":"ast","body_hash":"440d4bed756c66e3facebc1e8496bceba0ba6aeacd8a34c081ba50f037c23115","body_text":"def calc_padded_numel(shape: torch.Size, dtype: torch.dtype): \"\"\"for cuda memory alignment, make sure alignment by 128-bits\"\"\" align_numel = 128 // torch.finfo(dtype).bits numel = shape.numel() return (numel + align_numel - 1) // align_numel * align_numel","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/memory_buffer.py#get_weight_buffer_meta_from_module","name":"get_weight_buffer_meta_from_module","kind":"function","path":"verl/utils/memory_buffer.py","span":"L58-L65","signature":"def get_weight_buffer_meta_from_module(module: nn.Module) -> Dict[str, Dict]","exported":true,"origin":"ast","body_hash":"80353751e2359c379be24cb2aa586b4fb8b4ba185c17c3596b4e9c2329e9a17a","body_text":"def get_weight_buffer_meta_from_module(module: nn.Module) -> Dict[str, Dict]: \"\"\" Return a dictionary containing name to a shape and dtype. \"\"\" weight_buffer_meta = {} for name, param in sorted(module.named_parameters()): weight_buffer_meta[name] = {'shape': param.shape, 'dtype': param.dtype} return weight_buffer_meta","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/memory_buffer.py#build_memory_buffer","name":"build_memory_buffer","kind":"function","path":"verl/utils/memory_buffer.py","span":"L68-L94","signature":"def build_memory_buffer(weight_buffer_meta: Dict[str, Dict]) -> Dict[torch.dtype, MemoryBuffer]","exported":true,"origin":"ast","body_hash":"dbaec98f08aea3bbe6e8dc37b2ddf0af23c47af3265bc460bf6e8b55061e40ba","body_text":"def build_memory_buffer(weight_buffer_meta: Dict[str, Dict]) -> Dict[torch.dtype, MemoryBuffer]: \"\"\"Build the memory buffer given weight_buffer_meta Args: weight_buffer_meta: contains mapping from name to a dictionary containing shape and dtype of the tensors Returns: a large memory buffer for each dtype that can hold all the tensors \"\"\" memory_buffers = {} total_numel_map = {} # map from dtype to the total numel for name, meta_info in sorted(weight_buffer_meta.items()): shape = meta_info['shape'] dtype = meta_info['dtype'] assert isinstance(shape, torch.Size) assert isinstance(dtype, torch.dtype) if dtype not in total_numel_map: total_numel_map[dtype] = 0 total_numel_map[dtype] += calc_padded_numel(shape, dtype) for dtype, total_numel in total_numel_map.items(): memory_buffers[dtype] = MemoryBuffer(total_numel, total_numel, dtype) return memory_buffers","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/memory_buffer.py#build_memory_reference_from_module","name":"build_memory_reference_from_module","kind":"function","path":"verl/utils/memory_buffer.py","span":"L97-L110","signature":"def build_memory_reference_from_module(module: torch.nn.Module, memory_buffers: Dict[torch.dtype, MemoryBuffer], maintain_weight=True)","exported":true,"origin":"ast","body_hash":"ef186980d70dfd9e65c73ff8c331ca4ed213d5f2257d9bf967a7999c3f4345d2","body_text":"def build_memory_reference_from_module(module: torch.nn.Module, memory_buffers: Dict[torch.dtype, MemoryBuffer], maintain_weight=True): start_index = {} for dtype in memory_buffers.keys(): start_index[dtype] = 0 for name, param in sorted(module.named_parameters()): memory_buffer = memory_buffers[param.dtype] buffer = memory_buffer.get(shape=param.shape, start_index=start_index[param.dtype]) # need to increment start_index start_index[param.dtype] += calc_padded_numel(param.shape, dtype) if maintain_weight: buffer.copy_(param.data) param.data = buffer","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/memory_buffer.py#build_memory_reference","name":"build_memory_reference","kind":"function","path":"verl/utils/memory_buffer.py","span":"L113-L137","signature":"def build_memory_reference(weight_buffer_meta: Dict[str, Dict], memory_buffers: Dict[torch.dtype, MemoryBuffer])","exported":true,"origin":"ast","body_hash":"69e9dbaa82ef1d9ac12bc5d39c8468461502e062a84056e70a31ef9119af9ed7","body_text":"def build_memory_reference(weight_buffer_meta: Dict[str, Dict], memory_buffers: Dict[torch.dtype, MemoryBuffer]): \"\"\"Build the memory references. The memory buffers are built using the build_memory_buffer API. This API will allocate a weight buffer pointer to the memory buffer according to the weight_buffer_meta. Args: weight_buffer_meta: memory_buffers: Returns: \"\"\" start_idx = {} weight_buffers = {} for dtype in memory_buffers.keys(): start_idx[dtype] = 0 for name, meta_info in sorted(weight_buffer_meta.items()): shape = meta_info['shape'] dtype = meta_info['dtype'] buffer = memory_buffers[dtype].get(shape, start_index=start_idx[dtype]) start_idx[dtype] += calc_padded_numel(shape, dtype) weight_buffers[name] = buffer return weight_buffers","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper","name":"MemoryBufferModuleWrapper","kind":"class","path":"verl/utils/memory_buffer.py","span":"L140-L157","signature":"class MemoryBufferModuleWrapper","exported":true,"origin":"ast","body_hash":"489aad535e7cb61c4f0ded4f4ebdc1aef0fba5dd393f39d83a74016e0724fc7b","body_text":"class MemoryBufferModuleWrapper: \"\"\" Note that we do not design MemoryBufferModuleWrapper as an nn.Module due to - It will change the checkpoint name \"\"\" def __init__(self, module: nn.Module): super().__init__() self.module = module self.weight_buffer_meta = get_weight_buffer_meta_from_module(self.module) self.memory_buffers = build_memory_buffer(self.weight_buffer_meta) build_memory_reference_from_module(self.module, self.memory_buffers) def get_memory_buffers(self): return self.memory_buffers def get_weight_buffer_meta(self): return self.weight_buffer_meta","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.__init__","name":"__init__","kind":"method","path":"verl/utils/memory_buffer.py","span":"L146-L151","signature":"def __init__(self, module: nn.Module)","exported":false,"origin":"ast","body_hash":"ff9afd60b383d3ff37a92c4fdbd8a7ae28bb840642861604dfff82bcad77947c","body_text":"def __init__(self, module: nn.Module): super().__init__() self.module = module self.weight_buffer_meta = get_weight_buffer_meta_from_module(self.module) self.memory_buffers = build_memory_buffer(self.weight_buffer_meta) build_memory_reference_from_module(self.module, self.memory_buffers)","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryBufferModuleWrapper"},{"id":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.get_memory_buffers","name":"get_memory_buffers","kind":"method","path":"verl/utils/memory_buffer.py","span":"L153-L154","signature":"def get_memory_buffers(self)","exported":true,"origin":"ast","body_hash":"6f75a82c20f9be08eef829c6201e07cfe5cd91ffb2c4614c607d102eb73e23d3","body_text":"def get_memory_buffers(self): return self.memory_buffers","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryBufferModuleWrapper"},{"id":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.get_weight_buffer_meta","name":"get_weight_buffer_meta","kind":"method","path":"verl/utils/memory_buffer.py","span":"L156-L157","signature":"def get_weight_buffer_meta(self)","exported":true,"origin":"ast","body_hash":"a378445dab07af09150da5f9d4a320b06ca8f9eb0b0eb8fb8e8c0a79af981af2","body_text":"def get_weight_buffer_meta(self): return self.weight_buffer_meta","summary_state":"pending","summary":null,"crux":null,"owner":"MemoryBufferModuleWrapper"},{"id":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout","name":"MegatronMemoryBufferForRollout","kind":"class","path":"verl/utils/memory_buffer.py","span":"L160-L214","signature":"class MegatronMemoryBufferForRollout(object)","exported":true,"origin":"ast","body_hash":"6fc5d22d65abc7d8756c75953f3a04bcca3656cbdf8968701c01a2929e2aeb65","body_text":"class MegatronMemoryBufferForRollout(object): \"\"\" We assume that - inference engine has tp + dp - actor has tp + pp + dp - the tp between inference engine and actor should be the same - memory_buffers: contains a list of memory_buffers, each is a dict from dtype to MemoryBuffer - weight_buffers: contains a list of weight_buffers, each is a dict from name to param - named_parameters: a dict from name to parameter that normalizes the names from pp and vpp. Note that the named_parameters may not be directly compatible with inference engine. User has to take care of this part such as the layout mismatches. (e.g. qkv transpose) - Note that weight_buffer, named_parameters and memory_buffers share the same underlying GPU memory. - When doing weight sync, the data is transfer via memory buffers \"\"\" def __init__(self, transform_memory_param_fn): self._memory_buffers = [] self._weight_buffers = [] self._named_parameters = {} self.transform_memory_param_fn = transform_memory_param_fn def initialize_weight_buffer(self, weight_buffer_meta_pp: List[Dict[str, Dict]]): \"\"\" Initialize the weight buffer. The weight buffer is obtained according to the actor. We will construct a large buffer for each dtype in the weight_buffer. Args: weight_buffer_meta: contains pp models, each pp models contains a dictionary of mapping from Returns: None \"\"\" self.weight_buffer_meta_pp = weight_buffer_meta_pp for weight_buffer_meta in self.weight_buffer_meta_pp: memory_buffer = build_memory_buffer(weight_buffer_meta) self._memory_buffers.append(memory_buffer) self._weight_buffers.append(None) def build_memory_reference(self): for i, weight_buffer_meta in enumerate(self.weight_buffer_meta_pp): self._weight_buffers[i] = build_memory_reference(weight_buffer_meta, self._memory_buffers[i]) self._named_parameters = self.transform_memory_param_fn(self._weight_buffers) @property def named_parameters(self): return self._named_parameters @property def weight_buffers(self): return self._weight_buffers @property def memory_buffers(self): return self._memory_buffers","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.__init__","name":"__init__","kind":"method","path":"verl/utils/memory_buffer.py","span":"L175-L179","signature":"def __init__(self, transform_memory_param_fn)","exported":false,"origin":"ast","body_hash":"89fcfe14dd34aef37eeb3e11c23866fc081af16621148f9731d219d494453dcf","body_text":"def __init__(self, transform_memory_param_fn): self._memory_buffers = [] self._weight_buffers = [] self._named_parameters = {} self.transform_memory_param_fn = transform_memory_param_fn","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronMemoryBufferForRollout"},{"id":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.initialize_weight_buffer","name":"initialize_weight_buffer","kind":"method","path":"verl/utils/memory_buffer.py","span":"L181-L197","signature":"def initialize_weight_buffer(self, weight_buffer_meta_pp: List[Dict[str, Dict]])","exported":true,"origin":"ast","body_hash":"e255c68b5694127f090ec0330ccd3bed95cf2e5fda13847d900e518fc90574ec","body_text":"def initialize_weight_buffer(self, weight_buffer_meta_pp: List[Dict[str, Dict]]): \"\"\" Initialize the weight buffer. The weight buffer is obtained according to the actor. We will construct a large buffer for each dtype in the weight_buffer. Args: weight_buffer_meta: contains pp models, each pp models contains a dictionary of mapping from Returns: None \"\"\" self.weight_buffer_meta_pp = weight_buffer_meta_pp for weight_buffer_meta in self.weight_buffer_meta_pp: memory_buffer = build_memory_buffer(weight_buffer_meta) self._memory_buffers.append(memory_buffer) self._weight_buffers.append(None)","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronMemoryBufferForRollout"},{"id":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.build_memory_reference","name":"build_memory_reference","kind":"method","path":"verl/utils/memory_buffer.py","span":"L199-L202","signature":"def build_memory_reference(self)","exported":true,"origin":"ast","body_hash":"daad6b17bf799cf4184a0745ae3a12eb183178e3d2f9a4524f2ba8ad57a7ad79","body_text":"def build_memory_reference(self): for i, weight_buffer_meta in enumerate(self.weight_buffer_meta_pp): self._weight_buffers[i] = build_memory_reference(weight_buffer_meta, self._memory_buffers[i]) self._named_parameters = self.transform_memory_param_fn(self._weight_buffers)","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronMemoryBufferForRollout"},{"id":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.named_parameters","name":"named_parameters","kind":"method","path":"verl/utils/memory_buffer.py","span":"L205-L206","signature":"def named_parameters(self)","exported":true,"origin":"ast","body_hash":"8544c1ddeaef6dc0a05d7cd646bc8e1578c2fb513d0aac1f9215c5547e3ea41d","body_text":"def named_parameters(self): return self._named_parameters","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronMemoryBufferForRollout"},{"id":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.weight_buffers","name":"weight_buffers","kind":"method","path":"verl/utils/memory_buffer.py","span":"L209-L210","signature":"def weight_buffers(self)","exported":true,"origin":"ast","body_hash":"d67cc585cfc8f0bd84ce4a3771cdc3ca05832cd5afc12a482d8624a256729730","body_text":"def weight_buffers(self): return self._weight_buffers","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronMemoryBufferForRollout"},{"id":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.memory_buffers","name":"memory_buffers","kind":"method","path":"verl/utils/memory_buffer.py","span":"L213-L214","signature":"def memory_buffers(self)","exported":true,"origin":"ast","body_hash":"aeb38c7c460b398a920d95a74393cbc582e9910f7a953176267c6ca05d66d7ba","body_text":"def memory_buffers(self): return self._memory_buffers","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronMemoryBufferForRollout"}],"rawEdges":[{"source":"verl/utils/memory_buffer.py","relation":"imports","specifier":"typing","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py","relation":"imports","specifier":"torch","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py","relation":"imports","specifier":"torch","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py","relation":"contains","targetId":"verl/utils/memory_buffer.py#MemoryBuffer","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBuffer","relation":"contains","targetId":"verl/utils/memory_buffer.py#MemoryBuffer.__init__","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBuffer.__init__","relation":"calls","name":"zeros","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBuffer","relation":"contains","targetId":"verl/utils/memory_buffer.py#MemoryBuffer.zero","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBuffer.zero","relation":"calls","name":"zero_","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBuffer","relation":"contains","targetId":"verl/utils/memory_buffer.py#MemoryBuffer.get","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBuffer.get","relation":"calls","name":"numel","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBuffer.get","relation":"calls","name":"view","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py","relation":"contains","targetId":"verl/utils/memory_buffer.py#calc_padded_numel","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#calc_padded_numel","relation":"calls","name":"finfo","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#calc_padded_numel","relation":"calls","name":"numel","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py","relation":"contains","targetId":"verl/utils/memory_buffer.py#get_weight_buffer_meta_from_module","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#get_weight_buffer_meta_from_module","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#get_weight_buffer_meta_from_module","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py","relation":"contains","targetId":"verl/utils/memory_buffer.py#build_memory_buffer","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_buffer","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_buffer","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_buffer","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_buffer","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_buffer","relation":"calls","name":"calc_padded_numel","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_buffer","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_buffer","relation":"calls","name":"MemoryBuffer","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py","relation":"contains","targetId":"verl/utils/memory_buffer.py#build_memory_reference_from_module","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference_from_module","relation":"calls","name":"keys","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference_from_module","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference_from_module","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference_from_module","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference_from_module","relation":"calls","name":"calc_padded_numel","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference_from_module","relation":"calls","name":"copy_","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py","relation":"contains","targetId":"verl/utils/memory_buffer.py#build_memory_reference","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference","relation":"calls","name":"keys","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#build_memory_reference","relation":"calls","name":"calc_padded_numel","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py","relation":"contains","targetId":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper","relation":"contains","targetId":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.__init__","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.__init__","relation":"calls","name":"get_weight_buffer_meta_from_module","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.__init__","relation":"calls","name":"build_memory_buffer","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.__init__","relation":"calls","name":"build_memory_reference_from_module","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper","relation":"contains","targetId":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.get_memory_buffers","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper","relation":"contains","targetId":"verl/utils/memory_buffer.py#MemoryBufferModuleWrapper.get_weight_buffer_meta","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py","relation":"contains","targetId":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout","relation":"extends","name":"object","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout","relation":"contains","targetId":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.__init__","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout","relation":"contains","targetId":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.initialize_weight_buffer","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.initialize_weight_buffer","relation":"calls","name":"build_memory_buffer","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.initialize_weight_buffer","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.initialize_weight_buffer","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout","relation":"contains","targetId":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.build_memory_reference","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.build_memory_reference","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.build_memory_reference","relation":"calls","name":"build_memory_reference","viaMember":false,"file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.build_memory_reference","relation":"calls","name":"transform_memory_param_fn","viaMember":true,"file":"verl/utils/memory_buffer.py","recvType":"MegatronMemoryBufferForRollout"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout","relation":"contains","targetId":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.named_parameters","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout","relation":"contains","targetId":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.weight_buffers","file":"verl/utils/memory_buffer.py"},{"source":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout","relation":"contains","targetId":"verl/utils/memory_buffer.py#MegatronMemoryBufferForRollout.memory_buffers","file":"verl/utils/memory_buffer.py"}]},"verl/utils/model.py":{"size":12885,"mtimeMs":1789275830324.203,"hash":"6b6def87d02f60692f175d3a4f8e8c2a07e914b896881c1ab3bf5e83287e359a","nodes":[{"id":"verl/utils/model.py","name":"model.py","kind":"file","path":"verl/utils/model.py","span":"L1-L333","signature":null,"exported":true,"origin":"ast","body_hash":"6b6def87d02f60692f175d3a4f8e8c2a07e914b896881c1ab3bf5e83287e359a","chars":12885,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 to create common models from huggingface \"\"\" import os import warnings from typing import Dict, Type import numpy as np import torch from torch import nn from transformers import AutoConfig, AutoModelForCausalLM, PretrainedConfig, MistralForSequenceClassification from verl.models.registry import ModelRegistry # pad input_ids_rmpad, cu_seqlens and max_seqlen_in_batch to be divisible by tp"},{"id":"verl/utils/model.py#LambdaLayer","name":"LambdaLayer","kind":"class","path":"verl/utils/model.py","span":"L28-L35","signature":"class LambdaLayer(nn.Module)","exported":true,"origin":"ast","body_hash":"39d993de56092c0ccacd9342a8b10bfa273d23432c6a779c3a299f0cc0d25477","body_text":"class LambdaLayer(nn.Module): def __init__(self, fn): super().__init__() self.fn = fn def forward(self, *args, **kwargs): return self.fn(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#LambdaLayer.__init__","name":"__init__","kind":"method","path":"verl/utils/model.py","span":"L30-L32","signature":"def __init__(self, fn)","exported":false,"origin":"ast","body_hash":"07827991a4186c8e286bd288750cd94aa990df95d07b5a9c2d7d6fb35d9fad97","body_text":"def __init__(self, fn): super().__init__() self.fn = fn","summary_state":"pending","summary":null,"crux":null,"owner":"LambdaLayer"},{"id":"verl/utils/model.py#LambdaLayer.forward","name":"forward","kind":"method","path":"verl/utils/model.py","span":"L34-L35","signature":"def forward(self, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"6c3e6b4bc2d407b378e6c140f81c2652ff90431b49582295fbeea93858f2abbd","body_text":"def forward(self, *args, **kwargs): return self.fn(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"LambdaLayer"},{"id":"verl/utils/model.py#squeeze","name":"squeeze","kind":"function","path":"verl/utils/model.py","span":"L38-L39","signature":"def squeeze(x)","exported":true,"origin":"ast","body_hash":"4ee6d7104f6eb052cd8fc0bec2c36eea7fc4073d2ffa3999a7905a6d15171b84","body_text":"def squeeze(x): return torch.squeeze(x, dim=-1)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#update_model_config","name":"update_model_config","kind":"function","path":"verl/utils/model.py","span":"L42-L44","signature":"def update_model_config(module_config, override_config_kwargs)","exported":true,"origin":"ast","body_hash":"b01f754dff392707460475a583ca53afecfcf1b7d9e25ee2f41d2d93d0969b76","body_text":"def update_model_config(module_config, override_config_kwargs): for key, val in override_config_kwargs.items(): setattr(module_config, key, val)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#get_huggingface_actor_config","name":"get_huggingface_actor_config","kind":"function","path":"verl/utils/model.py","span":"L47-L55","signature":"def get_huggingface_actor_config(model_name: str, override_config_kwargs=None, trust_remote_code=False) -> Dict","exported":true,"origin":"ast","body_hash":"bf1b971cc7835f33ae2238a04834717347cff51f9b4f77a1ac627de903155cdf","body_text":"def get_huggingface_actor_config(model_name: str, override_config_kwargs=None, trust_remote_code=False) -> Dict: if override_config_kwargs is None: override_config_kwargs = {} assert isinstance(override_config_kwargs, Dict), \\ f'override_config_kwargs must be a dict, got {type(override_config_kwargs)}' module_config = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code) update_model_config(module_config, override_config_kwargs) return module_config","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#create_huggingface_actor","name":"create_huggingface_actor","kind":"function","path":"verl/utils/model.py","span":"L58-L78","signature":"def create_huggingface_actor(model_name: str, override_config_kwargs=None, automodel_kwargs=None) -> nn.Module","exported":true,"origin":"ast","body_hash":"3c1e689d0e853ece27bb4038e620db790c542adf764a5bdd8fca3698b584ef16","body_text":"def create_huggingface_actor(model_name: str, override_config_kwargs=None, automodel_kwargs=None) -> nn.Module: \"\"\" Args: model_name: actor_override_config_kwargs: Returns: \"\"\" if override_config_kwargs is None: override_config_kwargs = {} if automodel_kwargs is None: automodel_kwargs = {} assert isinstance(override_config_kwargs, Dict), \\ f'override_config_kwargs must be a dict, got {type(override_config_kwargs)}' module_config = get_huggingface_actor_config(model_name, override_config_kwargs, trust_remote_code=automodel_kwargs.get('trust_remote_code', False)) module: nn.Module = AutoModelForCausalLM.from_config(module_config, **automodel_kwargs) return module","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#create_huggingface_critic","name":"create_huggingface_critic","kind":"function","path":"verl/utils/model.py","span":"L81-L99","signature":"def create_huggingface_critic(model_name: str, override_config_kwargs=None, automodel_kwargs=None) -> nn.Module","exported":true,"origin":"ast","body_hash":"3adbeab2ed3644f21fc75da32b8ada1c9beef787f1fe073b41eeb52f478324d5","body_text":"def create_huggingface_critic(model_name: str, override_config_kwargs=None, automodel_kwargs=None) -> nn.Module: \"\"\" Args: model_name: override_config_kwargs: Returns: \"\"\" critic_module: nn.Module = create_huggingface_actor(model_name, override_config_kwargs=override_config_kwargs, automodel_kwargs=automodel_kwargs) if automodel_kwargs is None: automodel_kwargs = {} torch_dtype = automodel_kwargs.get('torch_dtype', torch.float32) critic_module.lm_head = nn.Sequential(nn.Linear(critic_module.config.hidden_size, 1, dtype=torch_dtype), LambdaLayer(fn=squeeze)) return critic_module","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#get_model_size","name":"get_model_size","kind":"function","path":"verl/utils/model.py","span":"L102-L126","signature":"def get_model_size(model: nn.Module, scale='auto')","exported":true,"origin":"ast","body_hash":"80d4dbbe9977fe7ec9953daa437b151725baf137b741d682faa02d07c331fb91","body_text":"def get_model_size(model: nn.Module, scale='auto'): n_params = sum(p.numel() for p in model.parameters()) if scale == 'auto': if n_params > 1e9: scale = 'B' elif n_params > 1e6: scale = 'M' elif n_params > 1e3: scale = 'K' else: scale = '' if scale == 'B': n_params = n_params / 1e9 elif scale == 'M': n_params = n_params / 1e6 elif scale == 'K': n_params = n_params / 1e3 elif scale == '': pass else: raise NotImplemented(f'Unknown scale {scale}') return n_params, scale","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#print_model_size","name":"print_model_size","kind":"function","path":"verl/utils/model.py","span":"L129-L133","signature":"def print_model_size(model: nn.Module, name: str = None)","exported":true,"origin":"ast","body_hash":"1de6db13cda35553f4988948baa7ac70652186c238eee89ea581a3baeee51b75","body_text":"def print_model_size(model: nn.Module, name: str = None): n_params, scale = get_model_size(model, scale='auto') if name is None: name = model.__class__.__name__ print(f'{name} contains {n_params:.2f}{scale} parameters')","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#create_random_mask","name":"create_random_mask","kind":"function","path":"verl/utils/model.py","span":"L136-L174","signature":"def create_random_mask(input_ids: torch.Tensor, max_ratio_of_valid_token: float, max_ratio_of_left_padding: float, min_ratio_of_valid_token: float = 0)","exported":true,"origin":"ast","body_hash":"c87f9949d214e1555fa32b5ba83bca420e9b3e2275d8a6e901c851fc9352c53d","body_text":"def create_random_mask(input_ids: torch.Tensor, max_ratio_of_valid_token: float, max_ratio_of_left_padding: float, min_ratio_of_valid_token: float = 0): \"\"\"Create a random mask given input_ids. Support left padding and right padding. Process: - Sample valid token length - Sample left_padding length - Generate padding Args: input_ids: shape (batch_size, seq_len) Returns: \"\"\" assert max_ratio_of_valid_token > 0 and max_ratio_of_valid_token <= 1. assert max_ratio_of_left_padding >= 0 and max_ratio_of_left_padding < 1. assert min_ratio_of_valid_token <= max_ratio_of_valid_token batch_size, sequence_length = input_ids.shape max_num_valid_tokens = int(sequence_length * max_ratio_of_valid_token) min_num_valid_tokens = max(1, int(sequence_length * min_ratio_of_valid_token)) max_left_padding = int(sequence_length * max_ratio_of_left_padding) assert max_num_valid_tokens + max_left_padding <= sequence_length assert max_num_valid_tokens > 0 and max_ratio_of_valid_token <= sequence_length masks = torch.ones_like(input_ids, dtype=torch.int64) # TODO: we can make this faster for i in range(batch_size): num_left_padding = np.random.randint(low=0, high=max_left_padding + 1, dtype=np.int64) num_valid = np.random.randint(low=min_num_valid_tokens, high=max_num_valid_tokens + 1, dtype=np.int64) for index in range(num_left_padding): masks[i, index] = 0 for index in range(num_left_padding + num_valid, sequence_length): masks[i, index] = 0 return masks","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#compute_position_id_with_mask","name":"compute_position_id_with_mask","kind":"function","path":"verl/utils/model.py","span":"L177-L178","signature":"def compute_position_id_with_mask(mask)","exported":true,"origin":"ast","body_hash":"0d44d282f40de8f70052eaaf5decaffb8b95eea97c67d938652da409710e81ef","body_text":"def compute_position_id_with_mask(mask): return torch.clip(torch.cumsum(mask, dim=-1) - 1, min=0, max=None)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#normalize_pp_vpp_params","name":"normalize_pp_vpp_params","kind":"function","path":"verl/utils/model.py","span":"L181-L231","signature":"def normalize_pp_vpp_params(params, num_hidden_layers, layer_name='layers')","exported":true,"origin":"ast","body_hash":"c64e125b9f86df9c7b79848aacccc84b9bd0b2319f08fde3dba383f51ab99be3","body_text":"def normalize_pp_vpp_params(params, num_hidden_layers, layer_name='layers'): \"\"\" Normalize the pp vpp params into a complete named parameters. This is useful when gather parameters from pp ranks and passed to a model without pp params: List[List[Dict[str, param]]] params contains a list of pp, with a list of vpp named_parameters in each vpp chunk. output: Dict[str, param] \"\"\" def normalize_model_name(name, pp_rank, vpp_rank, pp_size, vpp_size, num_layers): \"\"\" Transform the model name in each model_chunk in each pp stage into the name in inference engine \"\"\" if vpp_size > 1: # print(f'try to bind vpp params to inference engine...') layers_per_pp = num_layers // pp_size layers_per_vpp = layers_per_pp // vpp_size pp_offset = layers_per_vpp * pp_rank vpp_offset = (layers_per_vpp * pp_size) * vpp_rank layer_offset = pp_offset + vpp_offset else: layers_per_pp = num_layers // pp_size layer_offset = layers_per_pp * pp_rank if layer_name in name: # belong to an intermediate layer split_name = name.split('.') # find the num next to split_name for i, name in enumerate(split_name): if name == layer_name: break layer_num_idx = i + 1 # check the name assert len(split_name) >= layer_num_idx + 1, f'split_name = {split_name}' assert split_name[layer_num_idx].isdigit(), f'split_name = {split_name}' # increment layer_num_idx by layer_offset split_name[layer_num_idx] = str(int(split_name[layer_num_idx]) + layer_offset) name = '.'.join(split_name) # weight name in inference_tp_model return name pp_size = len(params) normalized_name_to_param = {} for pp_rank in range(len(params)): vpp_size = len(params[pp_rank]) for vpp_rank in range(vpp_size): for name, param in params[pp_rank][vpp_rank].items(): normalized_name = normalize_model_name(name, pp_rank, vpp_rank, pp_size, vpp_size, num_hidden_layers) normalized_name_to_param[normalized_name] = param return normalized_name_to_param","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#normalize_pp_vpp_params.normalize_model_name","name":"normalize_model_name","kind":"function","path":"verl/utils/model.py","span":"L192-L220","signature":"def normalize_model_name(name, pp_rank, vpp_rank, pp_size, vpp_size, num_layers)","exported":true,"origin":"ast","body_hash":"b9a5fa5e36b5f8625dba4156b226b1dac769290e7c1e20f7d3bee05f9d4ef873","body_text":"def normalize_model_name(name, pp_rank, vpp_rank, pp_size, vpp_size, num_layers): \"\"\" Transform the model name in each model_chunk in each pp stage into the name in inference engine \"\"\" if vpp_size > 1: # print(f'try to bind vpp params to inference engine...') layers_per_pp = num_layers // pp_size layers_per_vpp = layers_per_pp // vpp_size pp_offset = layers_per_vpp * pp_rank vpp_offset = (layers_per_vpp * pp_size) * vpp_rank layer_offset = pp_offset + vpp_offset else: layers_per_pp = num_layers // pp_size layer_offset = layers_per_pp * pp_rank if layer_name in name: # belong to an intermediate layer split_name = name.split('.') # find the num next to split_name for i, name in enumerate(split_name): if name == layer_name: break layer_num_idx = i + 1 # check the name assert len(split_name) >= layer_num_idx + 1, f'split_name = {split_name}' assert split_name[layer_num_idx].isdigit(), f'split_name = {split_name}' # increment layer_num_idx by layer_offset split_name[layer_num_idx] = str(int(split_name[layer_num_idx]) + layer_offset) name = '.'.join(split_name) # weight name in inference_tp_model return name","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#get_parallel_model_from_config","name":"get_parallel_model_from_config","kind":"function","path":"verl/utils/model.py","span":"L234-L240","signature":"def get_parallel_model_from_config(config, megatron_config, pre_process=None, post_process=None, value=False)","exported":true,"origin":"ast","body_hash":"41cbe14f8f4b2e68f45f44bb853ae9aed89a99c5aaba12e84eb3e4c282cee25c","body_text":"def get_parallel_model_from_config(config, megatron_config, pre_process=None, post_process=None, value=False): from megatron.core import ModelParallelConfig assert isinstance(megatron_config, ModelParallelConfig) model_class = _get_parallel_model_architecture_from_config(config, value) model = model_class(config, megatron_config, pre_process=pre_process, post_process=post_process) return model","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#_get_parallel_model_architecture_from_config","name":"_get_parallel_model_architecture_from_config","kind":"function","path":"verl/utils/model.py","span":"L243-L250","signature":"def _get_parallel_model_architecture_from_config(config: PretrainedConfig, value=False) -> Type[nn.Module]","exported":false,"origin":"ast","body_hash":"f1dbadac697fa4a85d311060a0eebe3a7b7b1c472b028507022393ca6a649aa2","body_text":"def _get_parallel_model_architecture_from_config(config: PretrainedConfig, value=False) -> Type[nn.Module]: architectures = getattr(config, \"architectures\", []) for arch in architectures: model_cls = ModelRegistry.load_model_cls(arch, value) if model_cls is not None: return model_cls raise ValueError(f\"Model architectures {architectures} are not supported for now. \" f\"Supported architectures: {ModelRegistry.get_supported_archs()}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#load_megatron_model_weights","name":"load_megatron_model_weights","kind":"function","path":"verl/utils/model.py","span":"L253-L295","signature":"def load_megatron_model_weights(config, model_config, parallel_model, params_dtype, is_value_model=False, local_cache_path='~/.cache/verl/rlhf')","exported":true,"origin":"ast","body_hash":"358fc6fde500d9588723ea457281b19b59e3749dc3c1a07465cdccecffa582ed","body_text":"def load_megatron_model_weights(config, model_config, parallel_model, params_dtype, is_value_model=False, local_cache_path='~/.cache/verl/rlhf'): assert hasattr(model_config, \"architectures\"), \"architectures cannot be empty when load weight!\" architectures = getattr(model_config, \"architectures\", []) local_cache_path = os.path.expanduser(local_cache_path) if config.model.path.startswith(\"hdfs:\"): from verl.utils.fs import copy_local_path_from_hdfs print(f'start download from {config.model.path}') local_model_path = copy_local_path_from_hdfs(src=config.model.path, cache_dir=local_cache_path) print('finish download') else: print(f\"load from local dir {config.model.path}\") local_model_path = config.model.path # TODO: to find a better way to load mistral7b-rm lm_head if 'mistral7b-rm' in config.model.path: model = MistralForSequenceClassification.from_pretrained(local_model_path) # use score head instead of lm_head state_dict = model.state_dict() state_dict['lm_head.weight'] = state_dict['score.weight'] state_dict['model.embed_tokens.weight'] = state_dict[ 'model.embed_tokens.weight'][:32000] # workaround, 32001 -> 32000 is_value_model = True else: with warnings.catch_warnings(): warnings.simplefilter(\"ignore\") model = AutoModelForCausalLM.from_pretrained(local_model_path) state_dict = model.state_dict() from verl.models.weight_loader_registry import get_weight_loader print(f'before weight loader: architectures = {architectures}...') for arch in architectures: print(f'call weight loader arch = {arch}, model config = {model.config}') weight_loader = get_weight_loader(arch) weight_loader(state_dict=state_dict, wrapped_models=parallel_model, config=model.config, params_dtype=params_dtype, is_value_model=is_value_model)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/model.py#pad_packed_inputs","name":"pad_packed_inputs","kind":"function","path":"verl/utils/model.py","span":"L299-L332","signature":"def pad_packed_inputs(unpad_tokens: torch.Tensor, cu_seqlens, max_seqlen_in_batch, size)","exported":true,"origin":"ast","body_hash":"20e0a86dea346bce9aa0b497cc149e1e010483f3f8b33c691c6edf518cb2d9cb","body_text":"def pad_packed_inputs(unpad_tokens: torch.Tensor, cu_seqlens, max_seqlen_in_batch, size): \"\"\"pad the tokens such that the total length is a multiple of size. This function is useful when applying sequence parallel and context parallel Args: unpad_tokens: (total_nnz, ...). Tokens after removing padding cu_seqlens: (total_nnz + 1,) max_seqlen_in_batch: int Returns: \"\"\" F = nn.functional total_nnz = unpad_tokens.shape[0] if total_nnz % size == 0: pad_size = 0 else: pad_size = size - total_nnz % size # we assume adding a new data in the batch with seqlen pad_size if pad_size > 0: if unpad_tokens.ndim == 1: unpad_tokens = F.pad(unpad_tokens, (0, pad_size)) elif unpad_tokens.ndim == 2: unpad_tokens = F.pad(unpad_tokens, (0, 0, 0, pad_size)) else: raise NotImplementedError(f'Padding dim {unpad_tokens.ndim()} is not supported') cu_seqlens = F.pad(cu_seqlens, (0, 1), value=pad_size + cu_seqlens[-1]) max_seqlen_in_batch = max(max_seqlen_in_batch, pad_size) return unpad_tokens, cu_seqlens, max_seqlen_in_batch","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/model.py","relation":"imports","specifier":"os","file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"imports","specifier":"warnings","file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"imports","specifier":"typing","file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"imports","specifier":"torch","file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"imports","specifier":"torch","file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"imports","specifier":"transformers","file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"imports","specifier":"verl.models.registry","file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#LambdaLayer","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#LambdaLayer","relation":"contains","targetId":"verl/utils/model.py#LambdaLayer.__init__","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#LambdaLayer.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#LambdaLayer.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#LambdaLayer","relation":"contains","targetId":"verl/utils/model.py#LambdaLayer.forward","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#LambdaLayer.forward","relation":"calls","name":"fn","viaMember":true,"file":"verl/utils/model.py","recvType":"LambdaLayer"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#squeeze","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#squeeze","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#update_model_config","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#update_model_config","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#update_model_config","relation":"calls","name":"setattr","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#get_huggingface_actor_config","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_huggingface_actor_config","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_huggingface_actor_config","relation":"calls","name":"type","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_huggingface_actor_config","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_huggingface_actor_config","relation":"calls","name":"update_model_config","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#create_huggingface_actor","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_huggingface_actor","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_huggingface_actor","relation":"calls","name":"type","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_huggingface_actor","relation":"calls","name":"get_huggingface_actor_config","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_huggingface_actor","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_huggingface_actor","relation":"calls","name":"from_config","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#create_huggingface_critic","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_huggingface_critic","relation":"calls","name":"create_huggingface_actor","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_huggingface_critic","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_huggingface_critic","relation":"calls","name":"Sequential","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_huggingface_critic","relation":"calls","name":"Linear","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_huggingface_critic","relation":"calls","name":"LambdaLayer","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#get_model_size","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_model_size","relation":"calls","name":"sum","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_model_size","relation":"calls","name":"numel","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_model_size","relation":"calls","name":"parameters","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_model_size","relation":"calls","name":"NotImplemented","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#print_model_size","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#print_model_size","relation":"calls","name":"get_model_size","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#print_model_size","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#create_random_mask","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_random_mask","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_random_mask","relation":"calls","name":"max","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_random_mask","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_random_mask","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_random_mask","relation":"calls","name":"ones_like","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_random_mask","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_random_mask","relation":"calls","name":"randint","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_random_mask","relation":"calls","name":"randint","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_random_mask","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#create_random_mask","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#compute_position_id_with_mask","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#compute_position_id_with_mask","relation":"calls","name":"clip","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#compute_position_id_with_mask","relation":"calls","name":"cumsum","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#normalize_pp_vpp_params","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params","relation":"contains","targetId":"verl/utils/model.py#normalize_pp_vpp_params.normalize_model_name","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params.normalize_model_name","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params.normalize_model_name","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params.normalize_model_name","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params.normalize_model_name","relation":"calls","name":"isdigit","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params.normalize_model_name","relation":"calls","name":"str","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params.normalize_model_name","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params.normalize_model_name","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#normalize_pp_vpp_params","relation":"calls","name":"normalize_model_name","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#get_parallel_model_from_config","file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"imports","specifier":"megatron.core","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_parallel_model_from_config","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_parallel_model_from_config","relation":"calls","name":"_get_parallel_model_architecture_from_config","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#get_parallel_model_from_config","relation":"calls","name":"model_class","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#_get_parallel_model_architecture_from_config","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#_get_parallel_model_architecture_from_config","relation":"calls","name":"getattr","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#_get_parallel_model_architecture_from_config","relation":"calls","name":"load_model_cls","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#_get_parallel_model_architecture_from_config","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#_get_parallel_model_architecture_from_config","relation":"calls","name":"get_supported_archs","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#load_megatron_model_weights","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"getattr","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"expanduser","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"startswith","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"imports","specifier":"verl.utils.fs","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"state_dict","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"catch_warnings","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"simplefilter","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"state_dict","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"imports","specifier":"verl.models.weight_loader_registry","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"get_weight_loader","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#load_megatron_model_weights","relation":"calls","name":"weight_loader","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py","relation":"contains","targetId":"verl/utils/model.py#pad_packed_inputs","file":"verl/utils/model.py"},{"source":"verl/utils/model.py#pad_packed_inputs","relation":"calls","name":"pad","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#pad_packed_inputs","relation":"calls","name":"pad","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#pad_packed_inputs","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#pad_packed_inputs","relation":"calls","name":"ndim","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#pad_packed_inputs","relation":"calls","name":"pad","viaMember":true,"file":"verl/utils/model.py"},{"source":"verl/utils/model.py#pad_packed_inputs","relation":"calls","name":"max","viaMember":false,"file":"verl/utils/model.py"}]},"verl/utils/py_functional.py":{"size":1656,"mtimeMs":1789275830324.203,"hash":"2b7cb3329422c6d5efef8e922eabf4f0942dfc330c04728d6d99159e953a3e6f","nodes":[{"id":"verl/utils/py_functional.py","name":"py_functional.py","kind":"file","path":"verl/utils/py_functional.py","span":"L1-L57","signature":null,"exported":true,"origin":"ast","body_hash":"2b7cb3329422c6d5efef8e922eabf4f0942dfc330c04728d6d99159e953a3e6f","chars":1656,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 python utility functions \"\"\" from typing import Dict from types import SimpleNamespace"},{"id":"verl/utils/py_functional.py#union_two_dict","name":"union_two_dict","kind":"function","path":"verl/utils/py_functional.py","span":"L22-L38","signature":"def union_two_dict(dict1: Dict, dict2: Dict)","exported":true,"origin":"ast","body_hash":"5cb730c231686a179bbca1c1b64ba2ff0029c3fc25a95538b3fceb67b0a16dc5","body_text":"def union_two_dict(dict1: Dict, dict2: Dict): \"\"\"Union two dict. Will throw an error if there is an item not the same object with the same key. Args: dict1: dict2: Returns: \"\"\" for key, val in dict2.items(): if key in dict1: assert dict2[key] == dict1[key], \\ f'{key} in meta_dict1 and meta_dict2 are not the same object' dict1[key] = val return dict1","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/py_functional.py#append_to_dict","name":"append_to_dict","kind":"function","path":"verl/utils/py_functional.py","span":"L41-L45","signature":"def append_to_dict(data: Dict, new_data: Dict)","exported":true,"origin":"ast","body_hash":"693ea2f52d694afd7db966cb0a9cb1c2802f01f542991c4a2f8ab61f5916eb31","body_text":"def append_to_dict(data: Dict, new_data: Dict): for key, val in new_data.items(): if key not in data: data[key] = [] data[key].append(val)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/py_functional.py#NestedNamespace","name":"NestedNamespace","kind":"class","path":"verl/utils/py_functional.py","span":"L48-L56","signature":"class NestedNamespace(SimpleNamespace)","exported":true,"origin":"ast","body_hash":"d5baed6885fd0448b8f867769cfe63644c12d4a90b1347f919366d69d9223571","body_text":"class NestedNamespace(SimpleNamespace): def __init__(self, dictionary, **kwargs): super().__init__(**kwargs) for key, value in dictionary.items(): if isinstance(value, dict): self.__setattr__(key, NestedNamespace(value)) else: self.__setattr__(key, value)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/py_functional.py#NestedNamespace.__init__","name":"__init__","kind":"method","path":"verl/utils/py_functional.py","span":"L50-L56","signature":"def __init__(self, dictionary, **kwargs)","exported":false,"origin":"ast","body_hash":"927436a8d86387260b5a4f03d7d0195725d633d6fe0964ff3ba944e9ba37fa04","body_text":"def __init__(self, dictionary, **kwargs): super().__init__(**kwargs) for key, value in dictionary.items(): if isinstance(value, dict): self.__setattr__(key, NestedNamespace(value)) else: self.__setattr__(key, value)","summary_state":"pending","summary":null,"crux":null,"owner":"NestedNamespace"}],"rawEdges":[{"source":"verl/utils/py_functional.py","relation":"imports","specifier":"typing","file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py","relation":"imports","specifier":"types","file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py","relation":"contains","targetId":"verl/utils/py_functional.py#union_two_dict","file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py#union_two_dict","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/py_functional.py","recvType":"Dict"},{"source":"verl/utils/py_functional.py","relation":"contains","targetId":"verl/utils/py_functional.py#append_to_dict","file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py#append_to_dict","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/py_functional.py","recvType":"Dict"},{"source":"verl/utils/py_functional.py#append_to_dict","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py","relation":"contains","targetId":"verl/utils/py_functional.py#NestedNamespace","file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py#NestedNamespace","relation":"extends","name":"SimpleNamespace","file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py#NestedNamespace","relation":"contains","targetId":"verl/utils/py_functional.py#NestedNamespace.__init__","file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py#NestedNamespace.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py#NestedNamespace.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py#NestedNamespace.__init__","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py#NestedNamespace.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py#NestedNamespace.__init__","relation":"calls","name":"__setattr__","viaMember":true,"file":"verl/utils/py_functional.py","recvType":"NestedNamespace"},{"source":"verl/utils/py_functional.py#NestedNamespace.__init__","relation":"calls","name":"NestedNamespace","viaMember":false,"file":"verl/utils/py_functional.py"},{"source":"verl/utils/py_functional.py#NestedNamespace.__init__","relation":"calls","name":"__setattr__","viaMember":true,"file":"verl/utils/py_functional.py","recvType":"NestedNamespace"}]},"verl/utils/ray_utils.py":{"size":1401,"mtimeMs":1789275830324.203,"hash":"77ef654419c1c1de59a6ec2b32b457eca7dd8cc73fe74a0d98278a3255835b58","nodes":[{"id":"verl/utils/ray_utils.py","name":"ray_utils.py","kind":"file","path":"verl/utils/ray_utils.py","span":"L1-L44","signature":null,"exported":true,"origin":"ast","body_hash":"77ef654419c1c1de59a6ec2b32b457eca7dd8cc73fe74a0d98278a3255835b58","chars":1401,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Contains commonly used utilities for ray \"\"\" import ray import concurrent.futures"},{"id":"verl/utils/ray_utils.py#parallel_put","name":"parallel_put","kind":"function","path":"verl/utils/ray_utils.py","span":"L23-L43","signature":"def parallel_put(data_list, max_workers=None)","exported":true,"origin":"ast","body_hash":"b95156a2a9905dea647e39bd81a834f84618e480b9a6e60cddc8f62b588df9b5","body_text":"def parallel_put(data_list, max_workers=None): def put_data(index, data): return index, ray.put(data) if max_workers is None: max_workers = min(len(data_list), 16) with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor: data_list_f = [executor.submit(put_data, i, data) for i, data in enumerate(data_list)] res_lst = [] for future in concurrent.futures.as_completed(data_list_f): res_lst.append(future.result()) # reorder based on index output = [None for _ in range(len(data_list))] for res in res_lst: index, data_ref = res output[index] = data_ref return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ray_utils.py#parallel_put.put_data","name":"put_data","kind":"function","path":"verl/utils/ray_utils.py","span":"L25-L26","signature":"def put_data(index, data)","exported":true,"origin":"ast","body_hash":"3c950e305097ebc2dcb6e67bcb91cd36af775d0c9d7cad7b339097618eba41e1","body_text":"def put_data(index, data): return index, ray.put(data)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/ray_utils.py","relation":"imports","specifier":"ray","file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py","relation":"imports","specifier":"concurrent.futures","file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py","relation":"contains","targetId":"verl/utils/ray_utils.py#parallel_put","file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"contains","targetId":"verl/utils/ray_utils.py#parallel_put.put_data","file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put.put_data","relation":"calls","name":"put","viaMember":true,"file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"calls","name":"min","viaMember":false,"file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"calls","name":"ThreadPoolExecutor","viaMember":true,"file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"calls","name":"submit","viaMember":true,"file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"calls","name":"as_completed","viaMember":true,"file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"calls","name":"result","viaMember":true,"file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/ray_utils.py"},{"source":"verl/utils/ray_utils.py#parallel_put","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/ray_utils.py"}]},"verl/utils/rendezvous/__init__.py":{"size":600,"mtimeMs":1789275830325.2031,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/utils/rendezvous/__init__.py","name":"__init__.py","kind":"file","path":"verl/utils/rendezvous/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/utils/rendezvous/ray_backend.py":{"size":2720,"mtimeMs":1789275830325.2031,"hash":"902cda6fef0dd91061454e2f2461e3766eaa6b3227d4f7a9e829ab96b3e0b63d","nodes":[{"id":"verl/utils/rendezvous/ray_backend.py","name":"ray_backend.py","kind":"file","path":"verl/utils/rendezvous/ray_backend.py","span":"L1-L78","signature":null,"exported":true,"origin":"ast","body_hash":"902cda6fef0dd91061454e2f2461e3766eaa6b3227d4f7a9e829ab96b3e0b63d","chars":2720,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 time from cupy.cuda.nccl import NcclCommunicator, get_unique_id import ray from ray.util import list_named_actors @ray.remote"},{"id":"verl/utils/rendezvous/ray_backend.py#NCCLIDStore","name":"NCCLIDStore","kind":"class","path":"verl/utils/rendezvous/ray_backend.py","span":"L25-L31","signature":"class NCCLIDStore","exported":true,"origin":"ast","body_hash":"62be3c0f890883e49051767ccb9650d654cecd944fbe90fc2f9254fe1897e2aa","body_text":"class NCCLIDStore: def __init__(self, nccl_id): self._nccl_id = nccl_id def get(self): return self._nccl_id","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/rendezvous/ray_backend.py#NCCLIDStore.__init__","name":"__init__","kind":"method","path":"verl/utils/rendezvous/ray_backend.py","span":"L27-L28","signature":"def __init__(self, nccl_id)","exported":false,"origin":"ast","body_hash":"b7453119788201cb3c8afe136406461fefcfa3608affbf019196b2b94c39f175","body_text":"def __init__(self, nccl_id): self._nccl_id = nccl_id","summary_state":"pending","summary":null,"crux":null,"owner":"NCCLIDStore"},{"id":"verl/utils/rendezvous/ray_backend.py#NCCLIDStore.get","name":"get","kind":"method","path":"verl/utils/rendezvous/ray_backend.py","span":"L30-L31","signature":"def get(self)","exported":true,"origin":"ast","body_hash":"f1ea2a502abd14faea2ebf961c087fa4370e20719c5ed2b37a7782b40b5df27e","body_text":"def get(self): return self._nccl_id","summary_state":"pending","summary":null,"crux":null,"owner":"NCCLIDStore"},{"id":"verl/utils/rendezvous/ray_backend.py#get_nccl_id_store_by_name","name":"get_nccl_id_store_by_name","kind":"function","path":"verl/utils/rendezvous/ray_backend.py","span":"L34-L44","signature":"def get_nccl_id_store_by_name(name)","exported":true,"origin":"ast","body_hash":"60960e184e7198dad1ee6d71f7108303f4966349584afc8c7546ad0fb7e473f9","body_text":"def get_nccl_id_store_by_name(name): all_actors = list_named_actors(all_namespaces=True) matched_actors = [actor for actor in all_actors if actor.get(\"name\", None) == name] if len(matched_actors) == 1: actor = matched_actors[0] return ray.get_actor(**actor) elif len(matched_actors) > 1: logging.warning(f\"multiple actors with same name found: {matched_actors}\") elif len(matched_actors) == 0: logging.info(f\"failed to get any actor named {name}\") return None","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","name":"create_nccl_communicator_in_ray","kind":"function","path":"verl/utils/rendezvous/ray_backend.py","span":"L47-L77","signature":"def create_nccl_communicator_in_ray(rank: int, world_size: int, group_name: str, max_retries: int = 100, interval_s: int = 5)","exported":true,"origin":"ast","body_hash":"e369586bcc39ecb088add16595b78ddcc9e67fb5cb7d7cafe7a220a6761900e2","body_text":"def create_nccl_communicator_in_ray(rank: int, world_size: int, group_name: str, max_retries: int = 100, interval_s: int = 5): if rank == 0: nccl_id = get_unique_id() nccl_id_store = NCCLIDStore.options(name=group_name).remote(nccl_id) assert ray.get(nccl_id_store.get.remote()) == nccl_id communicator = NcclCommunicator( ndev=world_size, commId=nccl_id, rank=0, ) return communicator else: for i in range(max_retries): nccl_id_store = get_nccl_id_store_by_name(group_name) if nccl_id_store is not None: logging.info(f\"nccl_id_store {group_name} got\") nccl_id = ray.get(nccl_id_store.get.remote()) logging.info(f\"nccl id for {group_name} got: {nccl_id}\") communicator = NcclCommunicator( ndev=world_size, commId=nccl_id, rank=rank, ) return communicator logging.info(f\"failed to get nccl_id for {i+1} time, sleep for {interval_s} seconds\") time.sleep(interval_s)","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/rendezvous/ray_backend.py","relation":"imports","specifier":"logging","file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py","relation":"imports","specifier":"time","file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py","relation":"imports","specifier":"cupy.cuda.nccl","file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py","relation":"imports","specifier":"ray","file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py","relation":"imports","specifier":"ray.util","file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py","relation":"contains","targetId":"verl/utils/rendezvous/ray_backend.py#NCCLIDStore","file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#NCCLIDStore","relation":"contains","targetId":"verl/utils/rendezvous/ray_backend.py#NCCLIDStore.__init__","file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#NCCLIDStore","relation":"contains","targetId":"verl/utils/rendezvous/ray_backend.py#NCCLIDStore.get","file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py","relation":"contains","targetId":"verl/utils/rendezvous/ray_backend.py#get_nccl_id_store_by_name","file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#get_nccl_id_store_by_name","relation":"calls","name":"list_named_actors","viaMember":false,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#get_nccl_id_store_by_name","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#get_nccl_id_store_by_name","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#get_nccl_id_store_by_name","relation":"calls","name":"get_actor","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#get_nccl_id_store_by_name","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#get_nccl_id_store_by_name","relation":"calls","name":"warning","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#get_nccl_id_store_by_name","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#get_nccl_id_store_by_name","relation":"calls","name":"info","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py","relation":"contains","targetId":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"get_unique_id","viaMember":false,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"remote","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"options","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"remote","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"NcclCommunicator","viaMember":false,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"get_nccl_id_store_by_name","viaMember":false,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"info","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"remote","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"info","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"NcclCommunicator","viaMember":false,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"info","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"},{"source":"verl/utils/rendezvous/ray_backend.py#create_nccl_communicator_in_ray","relation":"calls","name":"sleep","viaMember":true,"file":"verl/utils/rendezvous/ray_backend.py"}]},"verl/utils/reward_score/__init__.py":{"size":600,"mtimeMs":1789275830325.2031,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/utils/reward_score/__init__.py","name":"__init__.py","kind":"file","path":"verl/utils/reward_score/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/utils/reward_score/countdown.py":{"size":3778,"mtimeMs":1789275830325.2031,"hash":"1e94bdc94c9cf9b013c3eb2a1f5046323e51f9fbf9481dac0d329dec80adff56","nodes":[{"id":"verl/utils/reward_score/countdown.py","name":"countdown.py","kind":"file","path":"verl/utils/reward_score/countdown.py","span":"L1-L111","signature":null,"exported":true,"origin":"ast","body_hash":"1e94bdc94c9cf9b013c3eb2a1f5046323e51f9fbf9481dac0d329dec80adff56","chars":3778,"summary_state":"pending","summary":null,"crux":null,"body_text":"import re import random import ast import operator"},{"id":"verl/utils/reward_score/countdown.py#extract_solution","name":"extract_solution","kind":"function","path":"verl/utils/reward_score/countdown.py","span":"L7-L25","signature":"def extract_solution(solution_str)","exported":true,"origin":"ast","body_hash":"867b2e90540dbb264646686c398b68473758bd8465b94a9fa0817c11a2f728f4","body_text":"def extract_solution(solution_str): \"\"\"Extract the equation from the solution string.\"\"\" # Remove everything before the first \"Assistant:\" if \"Assistant:\" in solution_str: solution_str = solution_str.split(\"Assistant:\", 1)[1] elif \"<|im_start|>assistant\" in solution_str: solution_str = solution_str.split(\"<|im_start|>assistant\", 1)[1] else: return None solution_str = solution_str.split('\\n')[-1] answer_pattern = r'<answer>(.*?)</answer>' match = re.finditer(answer_pattern, solution_str) matches = list(match) if matches: final_answer = matches[-1].group(1).strip() else: final_answer = None return final_answer","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/countdown.py#validate_equation","name":"validate_equation","kind":"function","path":"verl/utils/reward_score/countdown.py","span":"L28-L41","signature":"def validate_equation(equation_str, available_numbers)","exported":true,"origin":"ast","body_hash":"6183b9336aa8eae92645c106eebde87a5af4a86e6353e06417f16cc3a5e53ec8","body_text":"def validate_equation(equation_str, available_numbers): \"\"\"Validate that equation only uses available numbers and each number once.\"\"\" try: # Extract all numbers from the equation numbers_in_eq = [int(n) for n in re.findall(r'\\d+', equation_str)] # Check if all numbers in equation are available available_numbers = sorted(available_numbers) numbers_in_eq = sorted(numbers_in_eq) # Each number should be used exactly once return numbers_in_eq == available_numbers except: return False","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/countdown.py#evaluate_equation","name":"evaluate_equation","kind":"function","path":"verl/utils/reward_score/countdown.py","span":"L44-L56","signature":"def evaluate_equation(equation_str)","exported":true,"origin":"ast","body_hash":"189334995b7fd8a3e95234189811c3d0e1f2d3e1e82ba4d89e4c5f9f89e4fbc5","body_text":"def evaluate_equation(equation_str): \"\"\"Safely evaluate the arithmetic equation using eval() with precautions.\"\"\" try: # Define a regex pattern that only allows numbers, operators, parentheses, and whitespace allowed_pattern = r'^[\\d+\\-*/().\\s]+$' if not re.match(allowed_pattern, equation_str): raise ValueError(\"Invalid characters in equation.\") # Evaluate the equation with restricted globals and locals result = eval(equation_str, {\"__builtins__\": None}, {}) return result except Exception as e: return None","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/countdown.py#compute_score","name":"compute_score","kind":"function","path":"verl/utils/reward_score/countdown.py","span":"L59-L111","signature":"def compute_score(solution_str, ground_truth, method='strict', format_score=0.1, score=1.)","exported":true,"origin":"ast","body_hash":"9f97c2a8a00287e4b4b1e1d7f0010bf4cf0444bfd128a94b9a80f1fb70bd1429","body_text":"def compute_score(solution_str, ground_truth, method='strict', format_score=0.1, score=1.): \"\"\"The scoring function for countdown task. Args: solution_str: the solution text ground_truth: dictionary containing target number and available numbers method: the method to extract the solution format_score: the score for correct format but wrong answer score: the score for the correct answer \"\"\" target = ground_truth['target'] numbers = ground_truth['numbers'] equation = extract_solution(solution_str=solution_str) do_print = random.randint(1, 64) == 1 if do_print: print(f\"--------------------------------\") print(f\"Target: {target} | Numbers: {numbers}\") print(f\"Extracted equation: {equation}\") print(f\"Solution string: {solution_str}\") if equation is None: if do_print: print(f\"No equation found\") return 0 # Validate equation uses correct numbers if not validate_equation(equation, numbers): if do_print: print(f\"Invalid equation\") return format_score # Evaluate equation try: result = evaluate_equation(equation) if result is None: if do_print: print(f\"Could not evaluate equation\") return format_score if abs(result - target) < 1e-5: # Account for floating point precision if do_print: print(f\"Correct equation: {equation} = {result}\") return score else: if do_print: print(f\"Wrong result: equation = {result}, target = {target}\") return format_score except: if do_print: print(f\"Error evaluating equation\") return format_score","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/reward_score/countdown.py","relation":"imports","specifier":"re","file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py","relation":"imports","specifier":"random","file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py","relation":"imports","specifier":"ast","file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py","relation":"imports","specifier":"operator","file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py","relation":"contains","targetId":"verl/utils/reward_score/countdown.py#extract_solution","file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#extract_solution","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#extract_solution","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#extract_solution","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#extract_solution","relation":"calls","name":"finditer","viaMember":true,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#extract_solution","relation":"calls","name":"list","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#extract_solution","relation":"calls","name":"strip","viaMember":true,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#extract_solution","relation":"calls","name":"group","viaMember":true,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py","relation":"contains","targetId":"verl/utils/reward_score/countdown.py#validate_equation","file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#validate_equation","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#validate_equation","relation":"calls","name":"findall","viaMember":true,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#validate_equation","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#validate_equation","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py","relation":"contains","targetId":"verl/utils/reward_score/countdown.py#evaluate_equation","file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#evaluate_equation","relation":"calls","name":"match","viaMember":true,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#evaluate_equation","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#evaluate_equation","relation":"calls","name":"eval","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py","relation":"contains","targetId":"verl/utils/reward_score/countdown.py#compute_score","file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"extract_solution","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"randint","viaMember":true,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"validate_equation","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"evaluate_equation","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"abs","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/countdown.py"},{"source":"verl/utils/reward_score/countdown.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/countdown.py"}]},"verl/utils/reward_score/gsm8k.py":{"size":2385,"mtimeMs":1789275830325.2031,"hash":"ad8c22fa6becda86f429a9cac622aa9a2e9178758ecde00e3a02a55e09b93b6e","nodes":[{"id":"verl/utils/reward_score/gsm8k.py","name":"gsm8k.py","kind":"file","path":"verl/utils/reward_score/gsm8k.py","span":"L1-L63","signature":null,"exported":true,"origin":"ast","body_hash":"ad8c22fa6becda86f429a9cac622aa9a2e9178758ecde00e3a02a55e09b93b6e","chars":2385,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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"},{"id":"verl/utils/reward_score/gsm8k.py#extract_solution","name":"extract_solution","kind":"function","path":"verl/utils/reward_score/gsm8k.py","span":"L18-L41","signature":"def extract_solution(solution_str, method='strict')","exported":true,"origin":"ast","body_hash":"fb61336255c4a65ab8abc4a0cfd3215d545e678d9b8c31c71e1fbe3f422393d5","body_text":"def extract_solution(solution_str, method='strict'): assert method in ['strict', 'flexible'] if method == 'strict': # this also tests the formatting of the model solution = re.search(\"#### (\\\\-?[0-9\\\\.\\\\,]+)\", solution_str) if solution is None: final_answer = None else: final_answer = solution.group(0) final_answer = final_answer.split('#### ')[1].replace(',', '').replace('$', '') elif method == 'flexible': answer = re.findall(\"(\\\\-?[0-9\\\\.\\\\,]+)\", solution_str) final_answer = None if len(answer) == 0: # no reward is there is no answer pass else: invalid_str = ['', '.'] # find the last number that is not '.' for final_answer in reversed(answer): if final_answer not in invalid_str: break return final_answer","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/gsm8k.py#compute_score","name":"compute_score","kind":"function","path":"verl/utils/reward_score/gsm8k.py","span":"L44-L63","signature":"def compute_score(solution_str, ground_truth, method='strict', format_score=0., score=1.)","exported":true,"origin":"ast","body_hash":"21bf91fa302855047b76503cf0358a5f9292a6287aa5c7a268583d0195826993","body_text":"def compute_score(solution_str, ground_truth, method='strict', format_score=0., score=1.): \"\"\"The scoring function for GSM8k. Reference: Trung, Luong, et al. \"Reft: Reasoning with reinforced fine-tuning.\" Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024. 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, method=method) if answer is None: return 0 else: if answer == ground_truth: return score else: return format_score","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/reward_score/gsm8k.py","relation":"imports","specifier":"re","file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py","relation":"contains","targetId":"verl/utils/reward_score/gsm8k.py#extract_solution","file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py#extract_solution","relation":"calls","name":"search","viaMember":true,"file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py#extract_solution","relation":"calls","name":"group","viaMember":true,"file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py#extract_solution","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py#extract_solution","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py#extract_solution","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py#extract_solution","relation":"calls","name":"findall","viaMember":true,"file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py#extract_solution","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py#extract_solution","relation":"calls","name":"reversed","viaMember":false,"file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py","relation":"contains","targetId":"verl/utils/reward_score/gsm8k.py#compute_score","file":"verl/utils/reward_score/gsm8k.py"},{"source":"verl/utils/reward_score/gsm8k.py#compute_score","relation":"calls","name":"extract_solution","viaMember":false,"file":"verl/utils/reward_score/gsm8k.py"}]},"verl/utils/reward_score/math.py":{"size":6717,"mtimeMs":1789275830325.2031,"hash":"fd0eb6f4702e249e3386e91ac03964214bd3deafa5a0b3d9f9a05849930735c7","nodes":[{"id":"verl/utils/reward_score/math.py","name":"math.py","kind":"file","path":"verl/utils/reward_score/math.py","span":"L1-L228","signature":null,"exported":true,"origin":"ast","body_hash":"fd0eb6f4702e249e3386e91ac03964214bd3deafa5a0b3d9f9a05849930735c7","chars":6717,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates # Copyright 2022 EleutherAI and the HuggingFace Inc. team. 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. # Adapted from https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/hendrycks_math/utils.py # string normalization from https://github.com/EleutherAI/lm-evaluation-harness/blob/master/lm_eval/tasks/hendrycks_math.py"},{"id":"verl/utils/reward_score/math.py#compute_score","name":"compute_score","kind":"function","path":"verl/utils/reward_score/math.py","span":"L17-L28","signature":"def compute_score(solution_str, ground_truth) -> float","exported":true,"origin":"ast","body_hash":"4d1370e5b406fa47ea4df7d3c98232c9e9d1766470f4ec69d992d7029afdeb01","body_text":"def compute_score(solution_str, ground_truth) -> float: retval = 0. try: string_in_last_boxed = last_boxed_only_string(solution_str) if string_in_last_boxed is not None: answer = remove_boxed(string_in_last_boxed) if is_equiv(answer, ground_truth): retval = 1. except Exception as e: print(e) return retval","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/math.py#is_equiv","name":"is_equiv","kind":"function","path":"verl/utils/reward_score/math.py","span":"L32-L46","signature":"def is_equiv(str1, str2, verbose=False)","exported":true,"origin":"ast","body_hash":"a9e0962951c5f950a42219296c44a9dd564aaaa9b67ff17ef099c478744c8ddc","body_text":"def is_equiv(str1, str2, verbose=False): if str1 is None and str2 is None: print(\"WARNING: Both None\") return True if str1 is None or str2 is None: return False try: ss1 = strip_string(str1) ss2 = strip_string(str2) if verbose: print(ss1, ss2) return ss1 == ss2 except Exception: return str1 == str2","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/math.py#remove_boxed","name":"remove_boxed","kind":"function","path":"verl/utils/reward_score/math.py","span":"L49-L60","signature":"def remove_boxed(s)","exported":true,"origin":"ast","body_hash":"7d9b4545a960d2519ecfe549ac6181ffe65598576790703386f7d2f141a85b58","body_text":"def remove_boxed(s): if \"\\\\boxed \" in s: left = \"\\\\boxed \" assert s[:len(left)] == left return s[len(left):] left = \"\\\\boxed{\" assert s[:len(left)] == left assert s[-1] == \"}\" return s[len(left):-1]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/math.py#last_boxed_only_string","name":"last_boxed_only_string","kind":"function","path":"verl/utils/reward_score/math.py","span":"L63-L90","signature":"def last_boxed_only_string(string)","exported":true,"origin":"ast","body_hash":"d7e05000ad7d6546a8fc45790868d700b2b754ffc403c97747c7d5f6b42835c8","body_text":"def last_boxed_only_string(string): idx = string.rfind(\"\\\\boxed\") if \"\\\\boxed \" in string: return \"\\\\boxed \" + string.split(\"\\\\boxed \")[-1].split(\"$\")[0] if idx < 0: idx = string.rfind(\"\\\\fbox\") if idx < 0: return None i = idx right_brace_idx = None num_left_braces_open = 0 while i < len(string): if string[i] == \"{\": num_left_braces_open += 1 if string[i] == \"}\": num_left_braces_open -= 1 if num_left_braces_open == 0: right_brace_idx = i break i += 1 if right_brace_idx is None: retval = None else: retval = string[idx:right_brace_idx + 1] return retval","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/math.py#fix_fracs","name":"fix_fracs","kind":"function","path":"verl/utils/reward_score/math.py","span":"L93-L122","signature":"def fix_fracs(string)","exported":true,"origin":"ast","body_hash":"7ae7145c40356c71b4d7774c7ac5ad0dd13a4b70aec70171428c0f14f73f9bfe","body_text":"def fix_fracs(string): substrs = string.split(\"\\\\frac\") new_str = substrs[0] if len(substrs) > 1: substrs = substrs[1:] for substr in substrs: new_str += \"\\\\frac\" if substr[0] == \"{\": new_str += substr else: try: assert len(substr) >= 2 except AssertionError: return string a = substr[0] b = substr[1] if b != \"{\": if len(substr) > 2: post_substr = substr[2:] new_str += \"{\" + a + \"}{\" + b + \"}\" + post_substr else: new_str += \"{\" + a + \"}{\" + b + \"}\" else: if len(substr) > 2: post_substr = substr[2:] new_str += \"{\" + a + \"}\" + b + post_substr else: new_str += \"{\" + a + \"}\" + b string = new_str return string","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/math.py#fix_a_slash_b","name":"fix_a_slash_b","kind":"function","path":"verl/utils/reward_score/math.py","span":"L125-L137","signature":"def fix_a_slash_b(string)","exported":true,"origin":"ast","body_hash":"df6b791a3be15b9684d4fbea064cc634dbe9a225e412041a3278bafe4818a141","body_text":"def fix_a_slash_b(string): if len(string.split(\"/\")) != 2: return string a = string.split(\"/\")[0] b = string.split(\"/\")[1] try: a = int(a) b = int(b) assert string == \"{}/{}\".format(a, b) new_string = \"\\\\frac{\" + str(a) + \"}{\" + str(b) + \"}\" return new_string except AssertionError: return string","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/math.py#remove_right_units","name":"remove_right_units","kind":"function","path":"verl/utils/reward_score/math.py","span":"L140-L147","signature":"def remove_right_units(string): # \"\\\\text{ \" only ever occurs (at least in the val set) when describing units","exported":true,"origin":"ast","body_hash":"7ebdfd4935766a8cb6460b9c76cc4cf9165763b25ce5bbabafb4b8a23173bd8d","body_text":"def remove_right_units(string): # \"\\\\text{ \" only ever occurs (at least in the val set) when describing units if \"\\\\text{ \" in string: splits = string.split(\"\\\\text{ \") assert len(splits) == 2 return splits[0] else: return string","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/math.py#fix_sqrt","name":"fix_sqrt","kind":"function","path":"verl/utils/reward_score/math.py","span":"L150-L162","signature":"def fix_sqrt(string)","exported":true,"origin":"ast","body_hash":"1775e40d55618fa093aa35d9f49b144fde57ffb91e5da728855237cf7ea236b2","body_text":"def fix_sqrt(string): if \"\\\\sqrt\" not in string: return string splits = string.split(\"\\\\sqrt\") new_string = splits[0] for split in splits[1:]: if split[0] != \"{\": a = split[0] new_substr = \"\\\\sqrt{\" + a + \"}\" + split[1:] else: new_substr = \"\\\\sqrt\" + split new_string += new_substr return new_string","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/math.py#strip_string","name":"strip_string","kind":"function","path":"verl/utils/reward_score/math.py","span":"L165-L227","signature":"def strip_string(string): # linebreaks","exported":true,"origin":"ast","body_hash":"235f50dd2af031bd38287ee72faba3d27b926ede5b9d5aa7dd3a203588ca6d7f","body_text":"def strip_string(string): # linebreaks string = string.replace(\"\\n\", \"\") # remove inverse spaces string = string.replace(\"\\\\!\", \"\") # replace \\\\ with \\ string = string.replace(\"\\\\\\\\\", \"\\\\\") # replace tfrac and dfrac with frac string = string.replace(\"tfrac\", \"frac\") string = string.replace(\"dfrac\", \"frac\") # remove \\left and \\right string = string.replace(\"\\\\left\", \"\") string = string.replace(\"\\\\right\", \"\") # Remove circ (degrees) string = string.replace(\"^{\\\\circ}\", \"\") string = string.replace(\"^\\\\circ\", \"\") # remove dollar signs string = string.replace(\"\\\\$\", \"\") # remove units (on the right) string = remove_right_units(string) # remove percentage string = string.replace(\"\\\\%\", \"\") string = string.replace(\"\\%\", \"\") # noqa: W605 # \" 0.\" equivalent to \" .\" and \"{0.\" equivalent to \"{.\" Alternatively, add \"0\" if \".\" is the start of the string string = string.replace(\" .\", \" 0.\") string = string.replace(\"{.\", \"{0.\") # if empty, return empty string if len(string) == 0: return string if string[0] == \".\": string = \"0\" + string # to consider: get rid of e.g. \"k = \" or \"q = \" at beginning if len(string.split(\"=\")) == 2: if len(string.split(\"=\")[0]) <= 2: string = string.split(\"=\")[1] # fix sqrt3 --> sqrt{3} string = fix_sqrt(string) # remove spaces string = string.replace(\" \", \"\") # \\frac1b or \\frac12 --> \\frac{1}{b} and \\frac{1}{2}, etc. Even works with \\frac1{72} (but not \\frac{72}1). Also does a/b --> \\\\frac{a}{b} string = fix_fracs(string) # manually change 0.5 --> \\frac{1}{2} if string == \"0.5\": string = \"\\\\frac{1}{2}\" # NOTE: X/Y changed to \\frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y string = fix_a_slash_b(string) return string","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/reward_score/math.py","relation":"contains","targetId":"verl/utils/reward_score/math.py#compute_score","file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#compute_score","relation":"calls","name":"last_boxed_only_string","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#compute_score","relation":"calls","name":"remove_boxed","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#compute_score","relation":"calls","name":"is_equiv","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py","relation":"contains","targetId":"verl/utils/reward_score/math.py#is_equiv","file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#is_equiv","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#is_equiv","relation":"calls","name":"strip_string","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#is_equiv","relation":"calls","name":"strip_string","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#is_equiv","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py","relation":"contains","targetId":"verl/utils/reward_score/math.py#remove_boxed","file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#remove_boxed","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#remove_boxed","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#remove_boxed","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#remove_boxed","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py","relation":"contains","targetId":"verl/utils/reward_score/math.py#last_boxed_only_string","file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#last_boxed_only_string","relation":"calls","name":"rfind","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#last_boxed_only_string","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#last_boxed_only_string","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#last_boxed_only_string","relation":"calls","name":"rfind","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#last_boxed_only_string","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py","relation":"contains","targetId":"verl/utils/reward_score/math.py#fix_fracs","file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_fracs","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_fracs","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_fracs","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_fracs","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_fracs","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py","relation":"contains","targetId":"verl/utils/reward_score/math.py#fix_a_slash_b","file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_a_slash_b","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_a_slash_b","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_a_slash_b","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_a_slash_b","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_a_slash_b","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_a_slash_b","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_a_slash_b","relation":"calls","name":"format","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_a_slash_b","relation":"calls","name":"str","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_a_slash_b","relation":"calls","name":"str","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py","relation":"contains","targetId":"verl/utils/reward_score/math.py#remove_right_units","file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#remove_right_units","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#remove_right_units","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py","relation":"contains","targetId":"verl/utils/reward_score/math.py#fix_sqrt","file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#fix_sqrt","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py","relation":"contains","targetId":"verl/utils/reward_score/math.py#strip_string","file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"remove_right_units","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"fix_sqrt","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"replace","viaMember":true,"file":"verl/utils/reward_score/math.py","recvType":"fix_a_slash_b"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"fix_fracs","viaMember":false,"file":"verl/utils/reward_score/math.py"},{"source":"verl/utils/reward_score/math.py#strip_string","relation":"calls","name":"fix_a_slash_b","viaMember":false,"file":"verl/utils/reward_score/math.py"}]},"verl/utils/reward_score/multiply.py":{"size":1977,"mtimeMs":1789275830325.2031,"hash":"c42d763bcce3700fc0b7c17f94cce7eee973d550b61410d9d3f041a26787a180","nodes":[{"id":"verl/utils/reward_score/multiply.py","name":"multiply.py","kind":"file","path":"verl/utils/reward_score/multiply.py","span":"L1-L59","signature":null,"exported":true,"origin":"ast","body_hash":"c42d763bcce3700fc0b7c17f94cce7eee973d550b61410d9d3f041a26787a180","chars":1977,"summary_state":"pending","summary":null,"crux":null,"body_text":"import re import random"},{"id":"verl/utils/reward_score/multiply.py#extract_solution","name":"extract_solution","kind":"function","path":"verl/utils/reward_score/multiply.py","span":"L5-L24","signature":"def extract_solution(solution_str): # Remove everything before the first \"Assistant:\"","exported":true,"origin":"ast","body_hash":"ce310a53e84823359e805fc18cf4a6543a8dc0e7d194a39f777de594f6e4b2a7","body_text":"def extract_solution(solution_str): # Remove everything before the first \"Assistant:\" if \"Assistant:\" in solution_str: solution_str = solution_str.split(\"Assistant:\", 1)[1] else: return None answer_pattern = r'<answer>(.*?)</answer>' match = re.finditer(answer_pattern, solution_str) matches = list(match) if matches: final_answer = matches[-1].group(1).strip() else: final_answer = None if final_answer is not None: try: int_final_answer = int(final_answer) except ValueError: final_answer = None return final_answer","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/multiply.py#compute_score","name":"compute_score","kind":"function","path":"verl/utils/reward_score/multiply.py","span":"L27-L58","signature":"def compute_score(solution_str, ground_truth, method='strict', format_score=0.1, score=1.)","exported":true,"origin":"ast","body_hash":"6253e73abd50309b8865cabe2fd911bac402409005c15fff881795673024133c","body_text":"def compute_score(solution_str, ground_truth, method='strict', format_score=0.1, score=1.): \"\"\"The scoring function for GSM8k. Reference: Trung, Luong, et al. \"Reft: Reasoning with reinforced fine-tuning.\" Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024. 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(f\"--------------------------------\") print(f\"Ground truth: {ground_truth} | Extracted answer: {answer}\") print(f\"Solution string: {solution_str}\") if answer is None: if do_print: print(f\"No answer found\") return 0 else: if int(answer) == int(ground_truth): if do_print: print(f\"Correct answer: {answer}\") return score else: if do_print: print(f\"Incorrect answer {answer} | Ground truth: {ground_truth}\") return format_score","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/reward_score/multiply.py","relation":"imports","specifier":"re","file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py","relation":"imports","specifier":"random","file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py","relation":"contains","targetId":"verl/utils/reward_score/multiply.py#extract_solution","file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#extract_solution","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#extract_solution","relation":"calls","name":"finditer","viaMember":true,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#extract_solution","relation":"calls","name":"list","viaMember":false,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#extract_solution","relation":"calls","name":"strip","viaMember":true,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#extract_solution","relation":"calls","name":"group","viaMember":true,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#extract_solution","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py","relation":"contains","targetId":"verl/utils/reward_score/multiply.py#compute_score","file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#compute_score","relation":"calls","name":"extract_solution","viaMember":false,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#compute_score","relation":"calls","name":"randint","viaMember":true,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#compute_score","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#compute_score","relation":"calls","name":"int","viaMember":false,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/multiply.py"},{"source":"verl/utils/reward_score/multiply.py#compute_score","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/multiply.py"}]},"verl/utils/reward_score/qa_em.py":{"size":4727,"mtimeMs":1789275830325.2031,"hash":"2ec449e96a29cad66e220d9192ae0781fb066391bcca821db1fdf8ed802e7095","nodes":[{"id":"verl/utils/reward_score/qa_em.py","name":"qa_em.py","kind":"file","path":"verl/utils/reward_score/qa_em.py","span":"L1-L142","signature":null,"exported":true,"origin":"ast","body_hash":"2ec449e96a29cad66e220d9192ae0781fb066391bcca821db1fdf8ed802e7095","chars":4659,"summary_state":"pending","summary":null,"crux":null,"body_text":"# Copyright 2024 Bytedance Ltd. and/or its affiliates #修改后的answer提取机制 # 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 string import random"},{"id":"verl/utils/reward_score/qa_em.py#normalize_answer","name":"normalize_answer","kind":"function","path":"verl/utils/reward_score/qa_em.py","span":"L19-L33","signature":"def normalize_answer(s)","exported":true,"origin":"ast","body_hash":"ce803b99766628c4cb95d8b6a24e4f3e0866b62b683f5387f71c8ddda4baf8c9","body_text":"def normalize_answer(s): def remove_articles(text): return re.sub(r\"\\b(a|an|the)\\b\", \" \", text) def white_space_fix(text): return \" \".join(text.split()) def remove_punc(text): exclude = set(string.punctuation) return \"\".join(ch for ch in text if ch not in exclude) def lower(text): return text.lower() return white_space_fix(remove_articles(remove_punc(lower(s))))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_em.py#normalize_answer.remove_articles","name":"remove_articles","kind":"function","path":"verl/utils/reward_score/qa_em.py","span":"L20-L21","signature":"def remove_articles(text)","exported":true,"origin":"ast","body_hash":"379104c4e84a485bfe70e068ad3ab86a952e68ee257ff0948ef420779d7300fc","body_text":"def remove_articles(text): return re.sub(r\"\\b(a|an|the)\\b\", \" \", text)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_em.py#normalize_answer.white_space_fix","name":"white_space_fix","kind":"function","path":"verl/utils/reward_score/qa_em.py","span":"L23-L24","signature":"def white_space_fix(text)","exported":true,"origin":"ast","body_hash":"8ad0a68aa23e4563d3c147a4533421497e137461d860b3750ec4698e15ca18f5","body_text":"def white_space_fix(text): return \" \".join(text.split())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_em.py#normalize_answer.remove_punc","name":"remove_punc","kind":"function","path":"verl/utils/reward_score/qa_em.py","span":"L26-L28","signature":"def remove_punc(text)","exported":true,"origin":"ast","body_hash":"ec4bf1cdca5299f73e895b061ef28504bd52dfc218a3d6bd0609d914f3436f80","body_text":"def remove_punc(text): exclude = set(string.punctuation) return \"\".join(ch for ch in text if ch not in exclude)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_em.py#normalize_answer.lower","name":"lower","kind":"function","path":"verl/utils/reward_score/qa_em.py","span":"L30-L31","signature":"def lower(text)","exported":true,"origin":"ast","body_hash":"b70faaa318a21a90ebbdc5db764d69ab6aef83e2c575e07993e0e9e248ba0b9a","body_text":"def lower(text): return text.lower()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_em.py#em_check","name":"em_check","kind":"function","path":"verl/utils/reward_score/qa_em.py","span":"L36-L46","signature":"def em_check(prediction, golden_answers)","exported":true,"origin":"ast","body_hash":"f710d2e0b95fe11031cf46fa077af29e8877e16608ff71c2a7d9133c98b4f278","body_text":"def em_check(prediction, golden_answers): if isinstance(golden_answers, str): golden_answers = [golden_answers] normalized_prediction = normalize_answer(prediction) score = 0 for golden_answer in golden_answers: golden_answer = normalize_answer(golden_answer) if golden_answer == normalized_prediction: score = 1 break return score","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_em.py#subem_check","name":"subem_check","kind":"function","path":"verl/utils/reward_score/qa_em.py","span":"L49-L59","signature":"def subem_check(prediction, golden_answers)","exported":true,"origin":"ast","body_hash":"45977ea28fd310ea788ac6b47db8e1135d0bb3db6e6dd99a996a7885bf87f934","body_text":"def subem_check(prediction, golden_answers): if isinstance(golden_answers, str): golden_answers = [golden_answers] normalized_prediction = normalize_answer(prediction) score = 0 for golden_answer in golden_answers: golden_answer = normalize_answer(golden_answer) if golden_answer in normalized_prediction: score = 1 break return score","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_em.py#extract_solution","name":"extract_solution","kind":"function","path":"verl/utils/reward_score/qa_em.py","span":"L62-L85","signature":"def extract_solution(solution_str)","exported":true,"origin":"ast","body_hash":"520a1c71cb4040fddf6a8b8d831ab090fcd20c3c382bf6a8c990a82c6a60fa50","body_text":"def extract_solution(solution_str): \"\"\"Extract the equation from the solution string.\"\"\" # Remove everything before the first \"Assistant:\" # if \"Assistant:\" in solution_str: # solution_str = solution_str.split(\"Assistant:\", 1)[1] # elif \"<|im_start|>assistant\" in solution_str: # solution_str = solution_str.split(\"<|im_start|>assistant\", 1)[1] # else: # return None # solution_str = solution_str.split('\\n')[-1] answer_pattern = r'<answer>(.*?)</answer>' match = re.finditer(answer_pattern, solution_str, re.DOTALL) matches = list(match) # If there are 0 or exactly 1 matches, return None #这里要改 if len(matches) == 0: return None # 空答案也认为是none ans = matches[-1].group(1).strip() # 1个或多个都取最后一个 return ans if ans else None # 空答案也当 None #return matches[-1].group(1).strip()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_em.py#compute_score_em","name":"compute_score_em","kind":"function","path":"verl/utils/reward_score/qa_em.py","span":"L88-L113","signature":"def compute_score_em(solution_str, ground_truth, method='strict', format_score=0., score=1.)","exported":true,"origin":"ast","body_hash":"cf36b90cf362b59bc77d7e63e905ee42c7332677aa964cc908b8a248eb8d8bc6","body_text":"def compute_score_em(solution_str, ground_truth, method='strict', format_score=0., score=1.): \"\"\"The scoring function for 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(f\"--------------------------------\") 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 em_check(answer, ground_truth['target']): return score else: return format_score","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_em.py#compute_score_subem","name":"compute_score_subem","kind":"function","path":"verl/utils/reward_score/qa_em.py","span":"L116-L141","signature":"def compute_score_subem(solution_str, ground_truth, method='strict', format_score=0., score=1.)","exported":true,"origin":"ast","body_hash":"ef7203f143fa53a45db8c8f1293f0a77a12085bc70fb6d2ba7b037ccf7c6bbf3","body_text":"def compute_score_subem(solution_str, ground_truth, method='strict', format_score=0., score=1.): \"\"\"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(f\"--------------------------------\") 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","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/reward_score/qa_em.py","relation":"imports","specifier":"re","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py","relation":"imports","specifier":"string","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py","relation":"imports","specifier":"random","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py","relation":"contains","targetId":"verl/utils/reward_score/qa_em.py#normalize_answer","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer","relation":"contains","targetId":"verl/utils/reward_score/qa_em.py#normalize_answer.remove_articles","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer.remove_articles","relation":"calls","name":"sub","viaMember":true,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer","relation":"contains","targetId":"verl/utils/reward_score/qa_em.py#normalize_answer.white_space_fix","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer.white_space_fix","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer.white_space_fix","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer","relation":"contains","targetId":"verl/utils/reward_score/qa_em.py#normalize_answer.remove_punc","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer.remove_punc","relation":"calls","name":"set","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer.remove_punc","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer","relation":"contains","targetId":"verl/utils/reward_score/qa_em.py#normalize_answer.lower","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer.lower","relation":"calls","name":"lower","viaMember":true,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer","relation":"calls","name":"white_space_fix","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer","relation":"calls","name":"remove_articles","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer","relation":"calls","name":"remove_punc","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#normalize_answer","relation":"calls","name":"lower","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py","relation":"contains","targetId":"verl/utils/reward_score/qa_em.py#em_check","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#em_check","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#em_check","relation":"calls","name":"normalize_answer","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#em_check","relation":"calls","name":"normalize_answer","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py","relation":"contains","targetId":"verl/utils/reward_score/qa_em.py#subem_check","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#subem_check","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#subem_check","relation":"calls","name":"normalize_answer","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#subem_check","relation":"calls","name":"normalize_answer","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py","relation":"contains","targetId":"verl/utils/reward_score/qa_em.py#extract_solution","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#extract_solution","relation":"calls","name":"finditer","viaMember":true,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#extract_solution","relation":"calls","name":"list","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#extract_solution","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#extract_solution","relation":"calls","name":"strip","viaMember":true,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#extract_solution","relation":"calls","name":"group","viaMember":true,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py","relation":"contains","targetId":"verl/utils/reward_score/qa_em.py#compute_score_em","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_em","relation":"calls","name":"extract_solution","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_em","relation":"calls","name":"randint","viaMember":true,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_em","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_em","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_em","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_em","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_em","relation":"calls","name":"em_check","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py","relation":"contains","targetId":"verl/utils/reward_score/qa_em.py#compute_score_subem","file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_subem","relation":"calls","name":"extract_solution","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_subem","relation":"calls","name":"randint","viaMember":true,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_subem","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_subem","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_subem","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_subem","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"},{"source":"verl/utils/reward_score/qa_em.py#compute_score_subem","relation":"calls","name":"subem_check","viaMember":false,"file":"verl/utils/reward_score/qa_em.py"}]},"verl/utils/reward_score/qa_f1.py":{"size":3450,"mtimeMs":1789275830325.2031,"hash":"520830de56d01f3b07c5f5e68a14c59f5cff746bd7bd81db808e188c44e5bce1","nodes":[{"id":"verl/utils/reward_score/qa_f1.py","name":"qa_f1.py","kind":"file","path":"verl/utils/reward_score/qa_f1.py","span":"L1-L113","signature":null,"exported":true,"origin":"ast","body_hash":"520830de56d01f3b07c5f5e68a14c59f5cff746bd7bd81db808e188c44e5bce1","chars":3446,"summary_state":"pending","summary":null,"crux":null,"body_text":"# F1 score implementation for QA reward, following the same structure and style # as qa_em.py. import re import string import random # ----------------------------------------------------------------------------- # Normalization (identical to qa_em.py) # ----------------------------------------------------------------------------- # ----------------------------------------------------------------------------- # Token-level F1 score (SQuAD-style) # ----------------------------------------------------------------------------- # ----------------------------------------------------------------------------- # Extract solution (<answer>...</answer>) # Reuse the same function from qa_em # ----------------------------------------------------------------------------- from . import qa_em extract_solution = qa_em.extract_solution # ----------------------------------------------------------------------------- # Main scoring function: compute F1-only reward # -----------------------------------------------------------------------------"},{"id":"verl/utils/reward_score/qa_f1.py#normalize_answer","name":"normalize_answer","kind":"function","path":"verl/utils/reward_score/qa_f1.py","span":"L11-L25","signature":"def normalize_answer(s)","exported":true,"origin":"ast","body_hash":"ce803b99766628c4cb95d8b6a24e4f3e0866b62b683f5387f71c8ddda4baf8c9","body_text":"def normalize_answer(s): def remove_articles(text): return re.sub(r\"\\b(a|an|the)\\b\", \" \", text) def white_space_fix(text): return \" \".join(text.split()) def remove_punc(text): exclude = set(string.punctuation) return \"\".join(ch for ch in text if ch not in exclude) def lower(text): return text.lower() return white_space_fix(remove_articles(remove_punc(lower(s))))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_f1.py#normalize_answer.remove_articles","name":"remove_articles","kind":"function","path":"verl/utils/reward_score/qa_f1.py","span":"L12-L13","signature":"def remove_articles(text)","exported":true,"origin":"ast","body_hash":"379104c4e84a485bfe70e068ad3ab86a952e68ee257ff0948ef420779d7300fc","body_text":"def remove_articles(text): return re.sub(r\"\\b(a|an|the)\\b\", \" \", text)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_f1.py#normalize_answer.white_space_fix","name":"white_space_fix","kind":"function","path":"verl/utils/reward_score/qa_f1.py","span":"L15-L16","signature":"def white_space_fix(text)","exported":true,"origin":"ast","body_hash":"8ad0a68aa23e4563d3c147a4533421497e137461d860b3750ec4698e15ca18f5","body_text":"def white_space_fix(text): return \" \".join(text.split())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_f1.py#normalize_answer.remove_punc","name":"remove_punc","kind":"function","path":"verl/utils/reward_score/qa_f1.py","span":"L18-L20","signature":"def remove_punc(text)","exported":true,"origin":"ast","body_hash":"ec4bf1cdca5299f73e895b061ef28504bd52dfc218a3d6bd0609d914f3436f80","body_text":"def remove_punc(text): exclude = set(string.punctuation) return \"\".join(ch for ch in text if ch not in exclude)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_f1.py#normalize_answer.lower","name":"lower","kind":"function","path":"verl/utils/reward_score/qa_f1.py","span":"L22-L23","signature":"def lower(text)","exported":true,"origin":"ast","body_hash":"b70faaa318a21a90ebbdc5db764d69ab6aef83e2c575e07993e0e9e248ba0b9a","body_text":"def lower(text): return text.lower()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_f1.py#f1_score","name":"f1_score","kind":"function","path":"verl/utils/reward_score/qa_f1.py","span":"L31-L69","signature":"def f1_score(prediction, golden_answers)","exported":true,"origin":"ast","body_hash":"12b2e84f44b9558e0ea0dcb3ce96ce6f485f7adc5ea0321a23c4f21d66b796df","body_text":"def f1_score(prediction, golden_answers): if isinstance(golden_answers, str): golden_answers = [golden_answers] pred = normalize_answer(prediction) pred_tokens = pred.split() if len(pred_tokens) == 0: return 0.0 best_f1 = 0.0 for golden_answer in golden_answers: gold = normalize_answer(golden_answer) gold_tokens = gold.split() if len(gold_tokens) == 0: continue # Count overlaps (multiset intersection) gold_counts = {} for t in gold_tokens: gold_counts[t] = gold_counts.get(t, 0) + 1 common = 0 for t in pred_tokens: if gold_counts.get(t, 0) > 0: common += 1 gold_counts[t] -= 1 if common == 0: continue precision = common / len(pred_tokens) recall = common / len(gold_tokens) f1 = 2 * precision * recall / (precision + recall) best_f1 = max(best_f1, f1) return best_f1","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/reward_score/qa_f1.py#compute_score_f1","name":"compute_score_f1","kind":"function","path":"verl/utils/reward_score/qa_f1.py","span":"L83-L112","signature":"def compute_score_f1(solution_str, ground_truth, method='strict', format_score=0., score=1.)","exported":true,"origin":"ast","body_hash":"24f059b990325c44eed359e6de5d63aee9af952dc0f2028fdf1dfe45ae5f6547","body_text":"def compute_score_f1(solution_str, ground_truth, method='strict', format_score=0., score=1.): \"\"\" F1-only scoring function. Args: solution_str: decoded full sequence ground_truth: dict, contains `target` method: kept for compatibility format_score: not used in F1-only setting (kept for compatibility) score: max score for fully correct answer (usually 1.0) \"\"\" answer = extract_solution(solution_str=solution_str) do_print = random.randint(1, 64) == 1 if do_print: print(f\"--------------------------------\") print(f\"Golden answers: {ground_truth['target']}\") print(f\"Extracted answer: {answer}\") print(f\"Solution string: {solution_str}\") # No answer → format error → reward = 0 if answer is None: return 0.0 # Compute token-level F1 f1 = f1_score(answer, ground_truth['target']) # Reward = f1 * score return score * f1","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/reward_score/qa_f1.py","relation":"imports","specifier":"re","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py","relation":"imports","specifier":"string","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py","relation":"imports","specifier":"random","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py","relation":"contains","targetId":"verl/utils/reward_score/qa_f1.py#normalize_answer","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer","relation":"contains","targetId":"verl/utils/reward_score/qa_f1.py#normalize_answer.remove_articles","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer.remove_articles","relation":"calls","name":"sub","viaMember":true,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer","relation":"contains","targetId":"verl/utils/reward_score/qa_f1.py#normalize_answer.white_space_fix","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer.white_space_fix","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer.white_space_fix","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer","relation":"contains","targetId":"verl/utils/reward_score/qa_f1.py#normalize_answer.remove_punc","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer.remove_punc","relation":"calls","name":"set","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer.remove_punc","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer","relation":"contains","targetId":"verl/utils/reward_score/qa_f1.py#normalize_answer.lower","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer.lower","relation":"calls","name":"lower","viaMember":true,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer","relation":"calls","name":"white_space_fix","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer","relation":"calls","name":"remove_articles","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer","relation":"calls","name":"remove_punc","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#normalize_answer","relation":"calls","name":"lower","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py","relation":"contains","targetId":"verl/utils/reward_score/qa_f1.py#f1_score","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"normalize_answer","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/qa_f1.py","recvType":"normalize_answer"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"normalize_answer","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/reward_score/qa_f1.py","recvType":"normalize_answer"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#f1_score","relation":"calls","name":"max","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py","relation":"imports","specifier":".","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py","relation":"contains","targetId":"verl/utils/reward_score/qa_f1.py#compute_score_f1","file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#compute_score_f1","relation":"calls","name":"extract_solution","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#compute_score_f1","relation":"calls","name":"randint","viaMember":true,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#compute_score_f1","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#compute_score_f1","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#compute_score_f1","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#compute_score_f1","relation":"calls","name":"print","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"},{"source":"verl/utils/reward_score/qa_f1.py#compute_score_f1","relation":"calls","name":"f1_score","viaMember":false,"file":"verl/utils/reward_score/qa_f1.py"}]},"verl/utils/seqlen_balancing.py":{"size":10327,"mtimeMs":1789275830325.2031,"hash":"836a381b88a7c2185aea684b8e8c26e7eb0c25c6183e8f2f4d1686c85636fefb","nodes":[{"id":"verl/utils/seqlen_balancing.py","name":"seqlen_balancing.py","kind":"file","path":"verl/utils/seqlen_balancing.py","span":"L1-L266","signature":null,"exported":true,"origin":"ast","body_hash":"836a381b88a7c2185aea684b8e8c26e7eb0c25c6183e8f2f4d1686c85636fefb","chars":10327,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 typing import List, Tuple, Callable import heapq import torch from torch import distributed as dist from tensordict import TensorDict import copy"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp","name":"karmarkar_karp","kind":"function","path":"verl/utils/seqlen_balancing.py","span":"L25-L130","signature":"def karmarkar_karp(seqlen_list: List[int], k_partitions: int, equal_size: bool): # see: https://en.wikipedia.org/wiki/Largest_differencing_method","exported":true,"origin":"ast","body_hash":"ce8fa18daca205c037c2b4fbca576404b2250722a299d38b1c4dfcc98e7fb663","body_text":"def karmarkar_karp(seqlen_list: List[int], k_partitions: int, equal_size: bool): # 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 spread(self): return self.sets[0].sum - self.sets[-1].sum 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set","name":"Set","kind":"class","path":"verl/utils/seqlen_balancing.py","span":"L27-L47","signature":"class Set","exported":true,"origin":"ast","body_hash":"4052b5ec7dfd7fe086a7ce298db2c6e867ba6564865d3f2fa37d11d819b78da7","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.__init__","name":"__init__","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L29-L31","signature":"def __init__(self) -> None","exported":false,"origin":"ast","body_hash":"f42178848e7897312127642edf7f2ea4831a65f9ccc2c9da5176b8a02a21946f","body_text":"def __init__(self) -> None: self.sum = 0 self.items = []","summary_state":"pending","summary":null,"crux":null,"owner":"Set"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.add","name":"add","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L33-L35","signature":"def add(self, idx: int, val: int)","exported":true,"origin":"ast","body_hash":"f94c8736c1148f9543c1613e933918b42b98b6d58eb3aa01e4ad31c65c8a1939","body_text":"def add(self, idx: int, val: int): self.items.append((idx, val)) self.sum += val","summary_state":"pending","summary":null,"crux":null,"owner":"Set"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.merge","name":"merge","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L37-L40","signature":"def merge(self, other)","exported":true,"origin":"ast","body_hash":"8cf6a256099890e2ca627bd13142fad6124b84965f6f5beeb0ef744efbbe0e40","body_text":"def merge(self, other): for idx, val in other.items: self.items.append((idx, val)) self.sum += val","summary_state":"pending","summary":null,"crux":null,"owner":"Set"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.__lt__","name":"__lt__","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L42-L47","signature":"def __lt__(self, other)","exported":false,"origin":"ast","body_hash":"8f92662c5e10259eb0eadd5d855fff3347cad9afcd5a5417eb0a00665e562261","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"Set"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.State","name":"State","kind":"class","path":"verl/utils/seqlen_balancing.py","span":"L49-L101","signature":"class State","exported":true,"origin":"ast","body_hash":"560b2c9235750356e41bbe473606fcb5e60b9b349643ad48b1c62e110a42f48a","body_text":"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 spread(self): return self.sets[0].sum - self.sets[-1].sum 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__init__","name":"__init__","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L51-L58","signature":"def __init__(self, items: List[Tuple[int, int]], k: int) -> None","exported":false,"origin":"ast","body_hash":"21726cd524567568d2db731f571293b27bafa8b4a532f8b84a65583d3809405f","body_text":"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)","summary_state":"pending","summary":null,"crux":null,"owner":"State"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.spread","name":"spread","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L60-L61","signature":"def spread(self)","exported":true,"origin":"ast","body_hash":"39bcb7136c1a5ff41efaa738f2d52f4aae0eb6902809d807ca0d26922b8a66a5","body_text":"def spread(self): return self.sets[0].sum - self.sets[-1].sum","summary_state":"pending","summary":null,"crux":null,"owner":"State"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.get_partitions","name":"get_partitions","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L63-L70","signature":"def get_partitions(self)","exported":true,"origin":"ast","body_hash":"3b0a54e07a7fec36593202a32a48dd49e81934995318fc617e24b0fad979c073","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"State"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.merge","name":"merge","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L72-L75","signature":"def merge(self, other)","exported":true,"origin":"ast","body_hash":"4d287d1b10bcae960dcecc2b5e3fc1f1bbf06ca049161d5746c128e0bf282845","body_text":"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)","summary_state":"pending","summary":null,"crux":null,"owner":"State"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.spread~2","name":"spread","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L78-L79","signature":"def spread(self) -> int","exported":true,"origin":"ast","body_hash":"cac59f13ba1558582cfa461dee44d9a481efdf7026f06fcf50f21f80f44c1c59","body_text":"def spread(self) -> int: return self.sets[0].sum - self.sets[-1].sum","summary_state":"pending","summary":null,"crux":null,"owner":"State"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__lt__","name":"__lt__","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L81-L87","signature":"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.","exported":false,"origin":"ast","body_hash":"aba8018a052fda91b9c4f2879dbae04bf41d4d2bf6c7c47abc2e94a521f8b29a","body_text":"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]","summary_state":"pending","summary":null,"crux":null,"owner":"State"},{"id":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__repr__","name":"__repr__","kind":"method","path":"verl/utils/seqlen_balancing.py","span":"L89-L101","signature":"def __repr__(self) -> str","exported":false,"origin":"ast","body_hash":"72767039f26fb824a89bfeea8b0505c35ac9139f940c7d4b17c37091e35e4fa3","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"State"},{"id":"verl/utils/seqlen_balancing.py#greedy_partition","name":"greedy_partition","kind":"function","path":"verl/utils/seqlen_balancing.py","span":"L133-L149","signature":"def greedy_partition(seqlen_list: List[int], k_partitions: int, equal_size: bool)","exported":true,"origin":"ast","body_hash":"54850a28c0e5161b9e737a09243da4898c5ef5a0fec65161ab4c826be0300b2f","body_text":"def greedy_partition(seqlen_list: List[int], k_partitions: int, equal_size: bool): 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions","name":"get_seqlen_balanced_partitions","kind":"function","path":"verl/utils/seqlen_balancing.py","span":"L152-L183","signature":"def get_seqlen_balanced_partitions(seqlen_list: List[int], k_partitions: int, equal_size: bool)","exported":true,"origin":"ast","body_hash":"53d34d9ef31d0b531e7c71d8f91ed62bf8f7eb63776e5de97acc4343114a4734","body_text":"def get_seqlen_balanced_partitions(seqlen_list: List[int], k_partitions: int, equal_size: bool): \"\"\" get order of seq lengths to make partitions balanced, this is used in balacing sum of seqlength across dp ranks and microbatches Parameters: seqlen_list (List[int]): seq lengths of each items k_partitions (int): resulting number of partitions equal_size (bool): if True, number of items in each partitions must be equal. if False, only consider balancing the sum, each partition can have variable number of items Returns: partitions (List[List[int]]): return k_partitions list containing the index of items. \"\"\" 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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","name":"_check_and_sort_partitions","kind":"function","path":"verl/utils/seqlen_balancing.py","span":"L170-L180","signature":"def _check_and_sort_partitions(partitions)","exported":false,"origin":"ast","body_hash":"af848299c066ce6451bcf7e6972594c3a52bc7cfd381d9942326856d88e37191","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","name":"log_seqlen_unbalance","kind":"function","path":"verl/utils/seqlen_balancing.py","span":"L186-L217","signature":"def log_seqlen_unbalance(seqlen_list: List[int], partitions: List[List[int]], prefix): # add some metrics of seqlen sum on dp ranks","exported":true,"origin":"ast","body_hash":"a638c25f8740eb304391618dd2a86cd7b80657782cd449fda3c75508106891db","body_text":"def log_seqlen_unbalance(seqlen_list: List[int], partitions: List[List[int]], prefix): # add some metrics of seqlen sum on dp ranks 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 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) }","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/seqlen_balancing.py#ceildiv","name":"ceildiv","kind":"function","path":"verl/utils/seqlen_balancing.py","span":"L220-L221","signature":"def ceildiv(a, b)","exported":true,"origin":"ast","body_hash":"5b4c12ea00e96fc92d67b1334dd9e85278de3bb8a71afde949ca346e68bfbb82","body_text":"def ceildiv(a, b): return -(a // -b)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","name":"rearrange_micro_batches","kind":"function","path":"verl/utils/seqlen_balancing.py","span":"L224-L256","signature":"def rearrange_micro_batches(batch: TensorDict, max_token_len, dp_group=None)","exported":true,"origin":"ast","body_hash":"6b778b89332e27b7168bb21e030473bd890300ee7e1713c63a4c0e9123c98d6b","body_text":"def rearrange_micro_batches(batch: TensorDict, max_token_len, dp_group=None): \"\"\"Split the batch into a list of micro_batches, where the max_token_len is smaller than max_token_len and the number of valid tokens in each micro batch is well balanced. \"\"\" # this is per local micro_bsz max_seq_len = batch['attention_mask'].shape[-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=}' seq_len_effective: torch.Tensor = batch['attention_mask'].sum(dim=1) total_seqlen = seq_len_effective.sum().item() num_micro_batches = ceildiv(total_seqlen, max_token_len) if dist.is_initialized(): num_micro_batches = torch.tensor([num_micro_batches], device='cuda') dist.all_reduce(num_micro_batches, op=dist.ReduceOp.MAX, group=dp_group) num_micro_batches = num_micro_batches.cpu().item() seq_len_effective = seq_len_effective.tolist() assert num_micro_batches <= len(seq_len_effective) micro_bsz_idx = get_seqlen_balanced_partitions(seq_len_effective, num_micro_batches, equal_size=False) micro_batches = [] for partition in micro_bsz_idx: curr_micro_batch = [] for idx in partition: curr_micro_batch.append(batch[idx:idx + 1]) curr_micro_batch = torch.cat(curr_micro_batch) micro_batches.append(curr_micro_batch) return micro_batches, micro_bsz_idx","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/seqlen_balancing.py#get_reverse_idx","name":"get_reverse_idx","kind":"function","path":"verl/utils/seqlen_balancing.py","span":"L259-L265","signature":"def get_reverse_idx(idx_map)","exported":true,"origin":"ast","body_hash":"1300c8c5935163f6bb2c239aa7076595a280d1f3f70774d99abbe9afcaedc60a","body_text":"def get_reverse_idx(idx_map): reverse_idx_map = copy.deepcopy(idx_map) for i, idx in enumerate(idx_map): reverse_idx_map[idx] = i return reverse_idx_map","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/seqlen_balancing.py","relation":"imports","specifier":"typing","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"imports","specifier":"heapq","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"imports","specifier":"torch","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"imports","specifier":"torch","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"imports","specifier":"tensordict","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"imports","specifier":"copy","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.__init__","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.add","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.add","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.merge","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.merge","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.__lt__","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.__lt__","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.__lt__","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.__lt__","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.Set.__lt__","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.State","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__init__","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__init__","relation":"calls","name":"Set","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__init__","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__init__","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__init__","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__init__","relation":"calls","name":"add","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__init__","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.spread","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.get_partitions","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.get_partitions","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.get_partitions","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.get_partitions","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.get_partitions","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.merge","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.merge","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.merge","relation":"calls","name":"merge","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.merge","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.spread~2","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__lt__","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__repr__","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__repr__","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__repr__","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp.State.__repr__","relation":"calls","name":"str","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"heappush","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"State","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"heappush","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"State","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"heappop","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"heappop","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"merge","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"heappush","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"get_partitions","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#karmarkar_karp","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#greedy_partition","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"sum","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#greedy_partition","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","relation":"calls","name":"set","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","relation":"calls","name":"add","viaMember":true,"file":"verl/utils/seqlen_balancing.py","recvType":"set"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","relation":"calls","name":"set","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions._check_and_sort_partitions","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions","relation":"calls","name":"karmarkar_karp","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_seqlen_balanced_partitions","relation":"calls","name":"_check_and_sort_partitions","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","relation":"calls","name":"sum","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","relation":"calls","name":"sum","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","relation":"calls","name":"min","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","relation":"calls","name":"max","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#log_seqlen_unbalance","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#ceildiv","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"item","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"ceildiv","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"tensor","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"all_reduce","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"item","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"cpu","viaMember":true,"file":"verl/utils/seqlen_balancing.py","recvType":"ceildiv"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"tolist","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"get_seqlen_balanced_partitions","viaMember":false,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#rearrange_micro_batches","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py","relation":"contains","targetId":"verl/utils/seqlen_balancing.py#get_reverse_idx","file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_reverse_idx","relation":"calls","name":"deepcopy","viaMember":true,"file":"verl/utils/seqlen_balancing.py"},{"source":"verl/utils/seqlen_balancing.py#get_reverse_idx","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/seqlen_balancing.py"}]},"verl/utils/tokenizer.py":{"size":2384,"mtimeMs":1789275830326.2031,"hash":"839dd008294614789194d3d97bedbd2ae381038fc288d7ed1f600bca7dc3470b","nodes":[{"id":"verl/utils/tokenizer.py","name":"tokenizer.py","kind":"file","path":"verl/utils/tokenizer.py","span":"L1-L58","signature":null,"exported":true,"origin":"ast","body_hash":"839dd008294614789194d3d97bedbd2ae381038fc288d7ed1f600bca7dc3470b","chars":2384,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 warnings __all__ = ['hf_tokenizer']"},{"id":"verl/utils/tokenizer.py#set_pad_token_id","name":"set_pad_token_id","kind":"function","path":"verl/utils/tokenizer.py","span":"L20-L32","signature":"def set_pad_token_id(tokenizer)","exported":true,"origin":"ast","body_hash":"7dfa5e66c136a70395c18e98d38c062f241392ceec4e86ab75c89df7b046ffdd","body_text":"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}') 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}')","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/tokenizer.py#hf_tokenizer","name":"hf_tokenizer","kind":"function","path":"verl/utils/tokenizer.py","span":"L35-L58","signature":"def hf_tokenizer(name_or_path, correct_pad_token=True, correct_gemma2=True, **kwargs)","exported":true,"origin":"ast","body_hash":"96a12bd89c71481f4072b8397c8d90b68cd899cd7e0b2d693c19eaf1a0826e4f","body_text":"def hf_tokenizer(name_or_path, correct_pad_token=True, correct_gemma2=True, **kwargs): \"\"\"Create a huggingface pretrained tokenizer. 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. **kwargs: The keyword arguments for the 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 <end_of_turn> and 107.') kwargs['eos_token'] = '<end_of_turn>' kwargs['eos_token_id'] = 107 tokenizer = AutoTokenizer.from_pretrained(name_or_path, **kwargs) if correct_pad_token: set_pad_token_id(tokenizer) return tokenizer","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/tokenizer.py","relation":"imports","specifier":"warnings","file":"verl/utils/tokenizer.py"},{"source":"verl/utils/tokenizer.py","relation":"contains","targetId":"verl/utils/tokenizer.py#set_pad_token_id","file":"verl/utils/tokenizer.py"},{"source":"verl/utils/tokenizer.py#set_pad_token_id","relation":"calls","name":"warn","viaMember":true,"file":"verl/utils/tokenizer.py"},{"source":"verl/utils/tokenizer.py#set_pad_token_id","relation":"calls","name":"warn","viaMember":true,"file":"verl/utils/tokenizer.py"},{"source":"verl/utils/tokenizer.py","relation":"contains","targetId":"verl/utils/tokenizer.py#hf_tokenizer","file":"verl/utils/tokenizer.py"},{"source":"verl/utils/tokenizer.py","relation":"imports","specifier":"transformers","file":"verl/utils/tokenizer.py"},{"source":"verl/utils/tokenizer.py#hf_tokenizer","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/tokenizer.py"},{"source":"verl/utils/tokenizer.py#hf_tokenizer","relation":"calls","name":"warn","viaMember":true,"file":"verl/utils/tokenizer.py"},{"source":"verl/utils/tokenizer.py#hf_tokenizer","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/utils/tokenizer.py"},{"source":"verl/utils/tokenizer.py#hf_tokenizer","relation":"calls","name":"set_pad_token_id","viaMember":false,"file":"verl/utils/tokenizer.py"}]},"verl/utils/torch_dtypes.py":{"size":2217,"mtimeMs":1789275830326.2031,"hash":"db8a3796a51fed225f3781dcfe5341a68ba83811133f6ef0b41aba28958d92d0","nodes":[{"id":"verl/utils/torch_dtypes.py","name":"torch_dtypes.py","kind":"file","path":"verl/utils/torch_dtypes.py","span":"L1-L83","signature":null,"exported":true,"origin":"ast","body_hash":"db8a3796a51fed225f3781dcfe5341a68ba83811133f6ef0b41aba28958d92d0","chars":2217,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 from typing import Union HALF_LIST = [16, \"16\", \"fp16\", \"float16\"] FLOAT_LIST = [32, \"32\", \"fp32\", \"float32\"] BFLOAT_LIST = [\"bf16\", \"bfloat16\"]"},{"id":"verl/utils/torch_dtypes.py#PrecisionType","name":"PrecisionType","kind":"class","path":"verl/utils/torch_dtypes.py","span":"L27-L82","signature":"class PrecisionType(object)","exported":true,"origin":"ast","body_hash":"814e429c9fd5c1a5341f9d2ecb0d996282ad3b7d064196679f11c967d22efe11","body_text":"class PrecisionType(object): \"\"\"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: Union[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}\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_dtypes.py#PrecisionType.supported_type","name":"supported_type","kind":"method","path":"verl/utils/torch_dtypes.py","span":"L43-L44","signature":"def supported_type(precision: Union[str, int]) -> bool","exported":true,"origin":"ast","body_hash":"4363c68c2e24bc77f662b9336bf14ac64ffd4174d503ae66269b6fa66c7a1713","body_text":"def supported_type(precision: Union[str, int]) -> bool: return any(x == precision for x in PrecisionType)","summary_state":"pending","summary":null,"crux":null,"owner":"PrecisionType"},{"id":"verl/utils/torch_dtypes.py#PrecisionType.supported_types","name":"supported_types","kind":"method","path":"verl/utils/torch_dtypes.py","span":"L47-L48","signature":"def supported_types() -> list[str]","exported":true,"origin":"ast","body_hash":"b617e1a339f12dcb46d75dc2d6b59f767b675239b9d31e3104ea16ec18bcb04e","body_text":"def supported_types() -> list[str]: return [x.value for x in PrecisionType]","summary_state":"pending","summary":null,"crux":null,"owner":"PrecisionType"},{"id":"verl/utils/torch_dtypes.py#PrecisionType.is_fp16","name":"is_fp16","kind":"method","path":"verl/utils/torch_dtypes.py","span":"L51-L52","signature":"def is_fp16(precision)","exported":true,"origin":"ast","body_hash":"ae70fdb17d3298c417082721ae7b7b4f5184c5ed4d6c56e4080a85b6a874a161","body_text":"def is_fp16(precision): return precision in HALF_LIST","summary_state":"pending","summary":null,"crux":null,"owner":"PrecisionType"},{"id":"verl/utils/torch_dtypes.py#PrecisionType.is_fp32","name":"is_fp32","kind":"method","path":"verl/utils/torch_dtypes.py","span":"L55-L56","signature":"def is_fp32(precision)","exported":true,"origin":"ast","body_hash":"df9edf686059f124b70e47e934397dd34fe57584bfc5f01a763913e894924b2e","body_text":"def is_fp32(precision): return precision in FLOAT_LIST","summary_state":"pending","summary":null,"crux":null,"owner":"PrecisionType"},{"id":"verl/utils/torch_dtypes.py#PrecisionType.is_bf16","name":"is_bf16","kind":"method","path":"verl/utils/torch_dtypes.py","span":"L59-L60","signature":"def is_bf16(precision)","exported":true,"origin":"ast","body_hash":"bf9e6c6516a2fd9493c9e85ec65ddf946763f9ab3bd122ca2d5bcb83779d07bb","body_text":"def is_bf16(precision): return precision in BFLOAT_LIST","summary_state":"pending","summary":null,"crux":null,"owner":"PrecisionType"},{"id":"verl/utils/torch_dtypes.py#PrecisionType.to_dtype","name":"to_dtype","kind":"method","path":"verl/utils/torch_dtypes.py","span":"L63-L71","signature":"def to_dtype(precision)","exported":true,"origin":"ast","body_hash":"9689ae56c9b435d890032f977242dbce5c4da30fb7b1ea50b119b56f2df1e10c","body_text":"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}\")","summary_state":"pending","summary":null,"crux":null,"owner":"PrecisionType"},{"id":"verl/utils/torch_dtypes.py#PrecisionType.to_str","name":"to_str","kind":"method","path":"verl/utils/torch_dtypes.py","span":"L74-L82","signature":"def to_str(precision)","exported":true,"origin":"ast","body_hash":"feb63ec83fe66bfb14b785206c5fe164adf5224a4d53d38f00c0a6405a2fddc4","body_text":"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}\")","summary_state":"pending","summary":null,"crux":null,"owner":"PrecisionType"}],"rawEdges":[{"source":"verl/utils/torch_dtypes.py","relation":"imports","specifier":"torch","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py","relation":"imports","specifier":"typing","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py","relation":"contains","targetId":"verl/utils/torch_dtypes.py#PrecisionType","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType","relation":"extends","name":"object","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType","relation":"contains","targetId":"verl/utils/torch_dtypes.py#PrecisionType.supported_type","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType.supported_type","relation":"calls","name":"any","viaMember":false,"file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType","relation":"contains","targetId":"verl/utils/torch_dtypes.py#PrecisionType.supported_types","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType","relation":"contains","targetId":"verl/utils/torch_dtypes.py#PrecisionType.is_fp16","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType","relation":"contains","targetId":"verl/utils/torch_dtypes.py#PrecisionType.is_fp32","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType","relation":"contains","targetId":"verl/utils/torch_dtypes.py#PrecisionType.is_bf16","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType","relation":"contains","targetId":"verl/utils/torch_dtypes.py#PrecisionType.to_dtype","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType.to_dtype","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType","relation":"contains","targetId":"verl/utils/torch_dtypes.py#PrecisionType.to_str","file":"verl/utils/torch_dtypes.py"},{"source":"verl/utils/torch_dtypes.py#PrecisionType.to_str","relation":"calls","name":"RuntimeError","viaMember":false,"file":"verl/utils/torch_dtypes.py"}]},"verl/utils/torch_functional.py":{"size":18349,"mtimeMs":1789275830326.2031,"hash":"e6f0578b0f52b530d52e9fc69968108d67badeb8374f18f39258e3550ac35454","nodes":[{"id":"verl/utils/torch_functional.py","name":"torch_functional.py","kind":"file","path":"verl/utils/torch_functional.py","span":"L1-L493","signature":null,"exported":true,"origin":"ast","body_hash":"e6f0578b0f52b530d52e9fc69968108d67badeb8374f18f39258e3550ac35454","chars":18349,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 \"\"\" from typing import Dict, Union, List, Optional import os import torch import torch.distributed import torch.nn.functional as F from tensordict import TensorDict from torch import nn 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 from transformers import PreTrainedTokenizer from transformers.generation.logits_process import (TemperatureLogitsWarper, TopKLogitsWarper, TopPLogitsWarper) \"\"\" Optimizer related \"\"\" from torch.optim import Optimizer from torch.optim.lr_scheduler import LambdaLR import math # Copied from transformers.models.bart.modeling_bart._make_causal_mask # Copied from transformers.models.bart.modeling_bart._expand_mask"},{"id":"verl/utils/torch_functional.py#gather_from_labels","name":"gather_from_labels","kind":"function","path":"verl/utils/torch_functional.py","span":"L34-L46","signature":"def gather_from_labels(data, label)","exported":true,"origin":"ast","body_hash":"9f0bf473aa445f0bb63d9f5677716622e5af191e5632e2417d1b96a7f0197bb5","body_text":"def gather_from_labels(data, label): \"\"\"Gather the label from data. The value in label should be [0, vocab_size) Args: data: (..., vocab_size) label (torch.IntTensor) : (...,) Returns: \"\"\" output = torch.gather(data, -1, label.unsqueeze(-1)).squeeze(-1) return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#logprobs_from_logits","name":"logprobs_from_logits","kind":"function","path":"verl/utils/torch_functional.py","span":"L49-L62","signature":"def logprobs_from_logits(logits, labels)","exported":true,"origin":"ast","body_hash":"79975d357eaffd1eb8125526910aa15557f5b40da9fc5b1a6e98413ec0abab9a","body_text":"def logprobs_from_logits(logits, labels): \"\"\" See: https://github.com/pytorch/pytorch/issues/563#issuecomment-330103591 \"\"\" 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) output = output.view(*batch_dim) else: output = logprobs_from_logits_naive(logits, labels) return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#logprobs_from_logits_flash_attn","name":"logprobs_from_logits_flash_attn","kind":"function","path":"verl/utils/torch_functional.py","span":"L65-L67","signature":"def logprobs_from_logits_flash_attn(logits, labels)","exported":true,"origin":"ast","body_hash":"d4583c81f7ed002c6d6727455d6832142beb72e76ac071e25ff610ac8689f0e7","body_text":"def logprobs_from_logits_flash_attn(logits, labels): output = -cross_entropy_loss(logits, labels)[0] return output","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#logprobs_from_logits_naive","name":"logprobs_from_logits_naive","kind":"function","path":"verl/utils/torch_functional.py","span":"L70-L73","signature":"def logprobs_from_logits_naive(logits, labels)","exported":true,"origin":"ast","body_hash":"cdb224954d652aea2b29759ff8590c825cd316d2d6e03fddb1948551ed174980","body_text":"def logprobs_from_logits_naive(logits, labels): logp = F.log_softmax(logits, dim=-1) logpy = gather_from_labels(logp, labels) return logpy","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#logprobs_of_labels_v2","name":"logprobs_of_labels_v2","kind":"function","path":"verl/utils/torch_functional.py","span":"L76-L83","signature":"def logprobs_of_labels_v2(logits: torch.FloatTensor, labels)","exported":true,"origin":"ast","body_hash":"62e30c577569eb4327b3e8d9799f5065df3e805cd558f6754d2cd33c5229fb7d","body_text":"def logprobs_of_labels_v2(logits: torch.FloatTensor, labels): \"\"\" A memory efficient implementation of logprobs_from_logits \"\"\" assert logits.dtype == torch.float32, 'Using bf16 logits with logprobs_of_labels_v2 may lead to divergence' logprobs_labels = torch.gather(logits, dim=-1, index=labels.unsqueeze(-1)) logprobs_labels = logprobs_labels - torch.logsumexp(logits, dim=-1, keepdim=True) return logprobs_labels.squeeze(-1)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#clip_by_value","name":"clip_by_value","kind":"function","path":"verl/utils/torch_functional.py","span":"L86-L92","signature":"def clip_by_value(x, tensor_min, tensor_max)","exported":true,"origin":"ast","body_hash":"8690acf7ea492f2d50c807432754c3a0b714c93c1a00bf3823115fbbdf0fd429","body_text":"def clip_by_value(x, tensor_min, tensor_max): \"\"\" Tensor extenstion to torch.clamp https://github.com/pytorch/pytorch/issues/2793#issuecomment-428784713 \"\"\" clipped = torch.max(torch.min(x, tensor_max), tensor_min) return clipped","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#entropy_from_logits","name":"entropy_from_logits","kind":"function","path":"verl/utils/torch_functional.py","span":"L95-L99","signature":"def entropy_from_logits(logits: torch.Tensor)","exported":true,"origin":"ast","body_hash":"2e17ccb30bfc6930bdd83c71ace95830a5a71bc3057e2dea72b7d1c55cc28465","body_text":"def entropy_from_logits(logits: torch.Tensor): \"\"\"Calculate entropy from logits.\"\"\" pd = torch.nn.functional.softmax(logits, dim=-1) entropy = torch.logsumexp(logits, dim=-1) - torch.sum(pd * logits, dim=-1) return entropy","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#masked_sum","name":"masked_sum","kind":"function","path":"verl/utils/torch_functional.py","span":"L102-L104","signature":"def masked_sum(values, mask, axis=None)","exported":true,"origin":"ast","body_hash":"e5306b3c3d772c7de527574ab1ddf17b72a5f8f61c62f34fc3c2b96e46e25912","body_text":"def masked_sum(values, mask, axis=None): \"\"\"Compute mean of tensor with a masked values.\"\"\" return (values * mask).sum(axis=axis)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#masked_mean","name":"masked_mean","kind":"function","path":"verl/utils/torch_functional.py","span":"L107-L109","signature":"def masked_mean(values, mask, axis=None)","exported":true,"origin":"ast","body_hash":"5b9d3ba4196063f1ed5f38cc188f709cbe53bb19c21e374ff6a4d3b515b6f118","body_text":"def masked_mean(values, mask, axis=None): \"\"\"Compute mean of tensor with a masked values.\"\"\" return (values * mask).sum(axis=axis) / mask.sum(axis=axis)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#masked_var","name":"masked_var","kind":"function","path":"verl/utils/torch_functional.py","span":"L112-L127","signature":"def masked_var(values, mask, unbiased=True)","exported":true,"origin":"ast","body_hash":"9582c6ca30c8a862c5ca0c7e12b80a1a2ecffd7b0cde762edd55a4f2d282a0b2","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#masked_whiten","name":"masked_whiten","kind":"function","path":"verl/utils/torch_functional.py","span":"L130-L136","signature":"def masked_whiten(values, mask, shift_mean=True)","exported":true,"origin":"ast","body_hash":"c200dde1e644f490fc6e3e89cdd8ec55fec8550544629061b93fc25097c6efdb","body_text":"def masked_whiten(values, mask, shift_mean=True): \"\"\"Whiten values with masked 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#get_eos_mask","name":"get_eos_mask","kind":"function","path":"verl/utils/torch_functional.py","span":"L139-L148","signature":"def get_eos_mask(response_id: torch.Tensor, eos_token: int = 2, dtype=torch.int64)","exported":true,"origin":"ast","body_hash":"d65ec4d477606c72f9bf629f52c82b1dae824b5a240cd25b245cfa33b1288923","body_text":"def get_eos_mask(response_id: torch.Tensor, eos_token: int = 2, dtype=torch.int64): ''' e.g. end of sentence token=1 response_id: [0, 0, 2, 42, 3, 5, 1, 0, 0] eos_mask: [1, 1, 1, 1, 1, 1, 1, 0, 0] ''' eos_mask = response_id.eq(eos_token).long() eos_mask = (torch.cumsum(eos_mask, dim=1) - eos_mask).bool() eos_mask = torch.logical_not(eos_mask).to(dtype) return eos_mask","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#compute_grad_norm","name":"compute_grad_norm","kind":"function","path":"verl/utils/torch_functional.py","span":"L151-L157","signature":"def compute_grad_norm(model: nn.Module)","exported":true,"origin":"ast","body_hash":"73a67edc879852b141ea73632045198ba05412c6b1ecebcc2d1f4b78bcc2a70b","body_text":"def compute_grad_norm(model: nn.Module): total_grad_square = 0 total_params = 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#broadcast_dict_tensor","name":"broadcast_dict_tensor","kind":"function","path":"verl/utils/torch_functional.py","span":"L160-L166","signature":"def broadcast_dict_tensor(tensors: Union[Dict[str, torch.Tensor], TensorDict], src, group)","exported":true,"origin":"ast","body_hash":"ee62d62721d4481108015ea2926910eb7e7c6b0c2fff6b29cf62fea0e696011b","body_text":"def broadcast_dict_tensor(tensors: Union[Dict[str, torch.Tensor], TensorDict], src, group): \"\"\" TODO: optimize this. Technically, we only need one broadcast \"\"\" for key in tensors.sorted_keys: torch.distributed.broadcast(tensors[key], src=src, group=group, async_op=False)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#allgather_dict_tensors","name":"allgather_dict_tensors","kind":"function","path":"verl/utils/torch_functional.py","span":"L169-L200","signature":"def allgather_dict_tensors(tensors: Union[Dict[str, torch.Tensor], TensorDict], size, group, dim=0)","exported":true,"origin":"ast","body_hash":"9d475e3571c334933db84da68c6070263d2583500c6a79667a9b6d0487c21edb","body_text":"def allgather_dict_tensors(tensors: Union[Dict[str, torch.Tensor], TensorDict], size, group, dim=0): \"\"\" TODO: optimize this. - We can use async ops - We can use only one allgather Args: tensors: size: group: Returns: \"\"\" 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#split_dict_tensor_into_batches","name":"split_dict_tensor_into_batches","kind":"function","path":"verl/utils/torch_functional.py","span":"L203-L206","signature":"def split_dict_tensor_into_batches(tensors: TensorDict, batch_size) -> List[TensorDict]","exported":true,"origin":"ast","body_hash":"a8956f89b90127c58c4c38ec93bab3e4e44bf06c9671e90be567b9af83a40b1d","body_text":"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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#pad_sequence_to_length","name":"pad_sequence_to_length","kind":"function","path":"verl/utils/torch_functional.py","span":"L209-L219","signature":"def pad_sequence_to_length(tensors, max_seq_len, pad_token_id, left_pad=False)","exported":true,"origin":"ast","body_hash":"3fae4f0085407a2734f31d251f465b833613168b1ec8314f8192ca3e41207ace","body_text":"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] (0, max_seq_len - tensors.shape[-1]) means right pad to max_seq_length and no left pad \"\"\" if tensors.shape[-1] >= max_seq_len: return tensors 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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#tokenize_and_postprocess_data","name":"tokenize_and_postprocess_data","kind":"function","path":"verl/utils/torch_functional.py","span":"L225-L266","signature":"def tokenize_and_postprocess_data(prompt: str, tokenizer: PreTrainedTokenizer, max_length: int, pad_token_id: int, left_pad=True, truncation='error')","exported":true,"origin":"ast","body_hash":"85635a315247a52c3c98a8f2f92a7c029a0f08175e388d3ccf55c326eeb30062","body_text":"def tokenize_and_postprocess_data(prompt: str, tokenizer: PreTrainedTokenizer, max_length: int, pad_token_id: int, left_pad=True, truncation='error'): \"\"\" input_data is the output from tokenizer. \"\"\" assert truncation in ['left', 'right', 'error'] input_data = tokenizer(prompt, return_tensors='pt', add_special_tokens=False) input_ids = input_data['input_ids'] attention_mask = input_data['attention_mask'] 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 == 'error': raise NotImplementedError(f'{sequence_length=} is larger than {max_length=}') else: raise NotImplementedError(f'Unknown truncation method {truncation}') return input_ids, attention_mask","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#remove_pad_token","name":"remove_pad_token","kind":"function","path":"verl/utils/torch_functional.py","span":"L269-L281","signature":"def remove_pad_token(input_ids: torch.Tensor, attention_mask: torch.Tensor)","exported":true,"origin":"ast","body_hash":"a70133f3b55334687a24781f2293fa9874c12ee7937d44d25bb6a138b6fb0d66","body_text":"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): no_padding_batch.append((ids[len(ids) - mask.sum():]).cpu().numpy().tolist()) return no_padding_batch","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#log_probs_from_logits_response","name":"log_probs_from_logits_response","kind":"function","path":"verl/utils/torch_functional.py","span":"L284-L297","signature":"def log_probs_from_logits_response(input_ids, logits, response_length)","exported":true,"origin":"ast","body_hash":"311cc0cd3a88075650d01960a18abc8bf5504514a1c8d29a30a4de46fa40ffeb","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#log_probs_from_logits_response_rmpad","name":"log_probs_from_logits_response_rmpad","kind":"function","path":"verl/utils/torch_functional.py","span":"L300-L325","signature":"def log_probs_from_logits_response_rmpad(input_ids, attention_mask, logits_rmpad, response_length)","exported":true,"origin":"ast","body_hash":"98b6ed8890a779bc805568e45502c95950b965ef0a8ac51f2e14f59f5450c6d5","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#log_probs_from_logits_all_rmpad","name":"log_probs_from_logits_all_rmpad","kind":"function","path":"verl/utils/torch_functional.py","span":"L328-L353","signature":"def log_probs_from_logits_all_rmpad(input_ids_rmpad, logits_rmpad, indices, batch_size, seqlen, response_length)","exported":true,"origin":"ast","body_hash":"0d40a975b300b292b52740081a0df984b45e9829c925902fc78c1df4a0f9a131","body_text":"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 \"\"\" from flash_attn.bert_padding 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#post_process_logits","name":"post_process_logits","kind":"function","path":"verl/utils/torch_functional.py","span":"L359-L367","signature":"def post_process_logits(input_ids, logits, temperature, top_k, top_p)","exported":true,"origin":"ast","body_hash":"973a9d905cdb68893c53ac917f3d38945c337adcbca310b60f3e85129360efbd","body_text":"def post_process_logits(input_ids, logits, temperature, top_k, top_p): if temperature != 1.: 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup","name":"get_cosine_schedule_with_warmup","kind":"function","path":"verl/utils/torch_functional.py","span":"L379-L419","signature":"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, )","exported":true,"origin":"ast","body_hash":"1afccb07779812ad55c884e22b1f5b6ce05c59aed4d4e555cf2a8a9382979f5f","body_text":"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, ): \"\"\" 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. Return: :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. \"\"\" assert min_lr_ratio >= 0 and min_lr_ratio <= 1. coef = (1 - min_lr_ratio) * 0.5 intercept = (1 + min_lr_ratio) * 0.5 def lr_lambda(current_step): if current_step < num_warmup_steps: return 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(0.0, x * coef + intercept) return LambdaLR(optimizer, lr_lambda, last_epoch)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","name":"lr_lambda","kind":"function","path":"verl/utils/torch_functional.py","span":"L412-L417","signature":"def lr_lambda(current_step)","exported":true,"origin":"ast","body_hash":"d731dcccc0a4b083034210e78c82bfb45a6f423245b669ca7582a60e3a914230","body_text":"def lr_lambda(current_step): if current_step < num_warmup_steps: return 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(0.0, x * coef + intercept)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#get_constant_schedule_with_warmup","name":"get_constant_schedule_with_warmup","kind":"function","path":"verl/utils/torch_functional.py","span":"L422-L431","signature":"def get_constant_schedule_with_warmup( optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1, )","exported":true,"origin":"ast","body_hash":"e7a33a305cb70b765a96b222fd1df147f049ca619d797f02047469ec2dbc1a35","body_text":"def get_constant_schedule_with_warmup( optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1, ): def lr_lambda(current_step): return min(1, float(current_step) / float(max(1, num_warmup_steps))) return LambdaLR(optimizer, lr_lambda, last_epoch)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#get_constant_schedule_with_warmup.lr_lambda","name":"lr_lambda","kind":"function","path":"verl/utils/torch_functional.py","span":"L428-L429","signature":"def lr_lambda(current_step)","exported":true,"origin":"ast","body_hash":"72e1f00afd62ef088edeea92c2b233eaed43e1b26a032c69e22e19376b83dd9f","body_text":"def lr_lambda(current_step): return min(1, float(current_step) / float(max(1, num_warmup_steps)))","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#prepare_decoder_attention_mask","name":"prepare_decoder_attention_mask","kind":"function","path":"verl/utils/torch_functional.py","span":"L434-L452","signature":"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]","exported":true,"origin":"ast","body_hash":"d91771186a5d65f6bd75d815812d4923909f0986c4ab263fe58d030c1213f0b0","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#_make_causal_mask","name":"_make_causal_mask","kind":"function","path":"verl/utils/torch_functional.py","span":"L456-L465","signature":"def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device)","exported":false,"origin":"ast","body_hash":"4836799258639bfbfdc83054165da581e95fcb7e6c8522bf4651e42af051f1f7","body_text":"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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#_expand_mask","name":"_expand_mask","kind":"function","path":"verl/utils/torch_functional.py","span":"L469-L480","signature":"def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None)","exported":false,"origin":"ast","body_hash":"7eac0af7f47e8bd3c1587db698c5f5583656b24149ec5c2ea1a410e34ae82160","body_text":"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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/torch_functional.py#get_unpad_data","name":"get_unpad_data","kind":"function","path":"verl/utils/torch_functional.py","span":"L483-L492","signature":"def get_unpad_data(attention_mask)","exported":true,"origin":"ast","body_hash":"baad864a915277968edc11f260f7141aeb738c4938af98a79028dac270a38299","body_text":"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, )","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"typing","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"os","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"torch","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"torch.distributed","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"tensordict","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"torch","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"flash_attn.ops.triton.cross_entropy","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#gather_from_labels","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#gather_from_labels","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#gather_from_labels","relation":"calls","name":"gather","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#gather_from_labels","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#logprobs_from_logits","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_from_logits","relation":"calls","name":"reshape","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_from_logits","relation":"calls","name":"reshape","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_from_logits","relation":"calls","name":"logprobs_from_logits_flash_attn","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_from_logits","relation":"calls","name":"view","viaMember":true,"file":"verl/utils/torch_functional.py","recvType":"logprobs_from_logits_naive"},{"source":"verl/utils/torch_functional.py#logprobs_from_logits","relation":"calls","name":"logprobs_from_logits_naive","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#logprobs_from_logits_flash_attn","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_from_logits_flash_attn","relation":"calls","name":"cross_entropy_loss","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#logprobs_from_logits_naive","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_from_logits_naive","relation":"calls","name":"log_softmax","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_from_logits_naive","relation":"calls","name":"gather_from_labels","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#logprobs_of_labels_v2","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_of_labels_v2","relation":"calls","name":"gather","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_of_labels_v2","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_of_labels_v2","relation":"calls","name":"logsumexp","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#logprobs_of_labels_v2","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#clip_by_value","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#clip_by_value","relation":"calls","name":"max","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#clip_by_value","relation":"calls","name":"min","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#entropy_from_logits","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#entropy_from_logits","relation":"calls","name":"softmax","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#entropy_from_logits","relation":"calls","name":"logsumexp","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#entropy_from_logits","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#masked_sum","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_sum","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#masked_mean","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_mean","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_mean","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#masked_var","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_var","relation":"calls","name":"masked_mean","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_var","relation":"calls","name":"masked_mean","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_var","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_var","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_var","relation":"calls","name":"ValueError","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#masked_whiten","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_whiten","relation":"calls","name":"masked_mean","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_whiten","relation":"calls","name":"masked_var","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#masked_whiten","relation":"calls","name":"rsqrt","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#get_eos_mask","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_eos_mask","relation":"calls","name":"long","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_eos_mask","relation":"calls","name":"eq","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_eos_mask","relation":"calls","name":"bool","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_eos_mask","relation":"calls","name":"cumsum","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_eos_mask","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_eos_mask","relation":"calls","name":"logical_not","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#compute_grad_norm","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#compute_grad_norm","relation":"calls","name":"parameters","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#compute_grad_norm","relation":"calls","name":"item","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#compute_grad_norm","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#compute_grad_norm","relation":"calls","name":"square","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#compute_grad_norm","relation":"calls","name":"detach","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#broadcast_dict_tensor","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#broadcast_dict_tensor","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#allgather_dict_tensors","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#allgather_dict_tensors","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#allgather_dict_tensors","relation":"calls","name":"to_dict","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#allgather_dict_tensors","relation":"calls","name":"sorted","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#allgather_dict_tensors","relation":"calls","name":"keys","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#allgather_dict_tensors","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#allgather_dict_tensors","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#allgather_dict_tensors","relation":"calls","name":"all_gather","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#allgather_dict_tensors","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#allgather_dict_tensors","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#split_dict_tensor_into_batches","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#split_dict_tensor_into_batches","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/torch_functional.py","recvType":"TensorDict"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#pad_sequence_to_length","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#pad_sequence_to_length","relation":"calls","name":"pad","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"transformers","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#tokenize_and_postprocess_data","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#tokenize_and_postprocess_data","relation":"calls","name":"tokenizer","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#tokenize_and_postprocess_data","relation":"calls","name":"pad_sequence_to_length","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#tokenize_and_postprocess_data","relation":"calls","name":"pad_sequence_to_length","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#tokenize_and_postprocess_data","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#tokenize_and_postprocess_data","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#remove_pad_token","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#remove_pad_token","relation":"calls","name":"zip","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#remove_pad_token","relation":"calls","name":"append","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#remove_pad_token","relation":"calls","name":"tolist","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#remove_pad_token","relation":"calls","name":"numpy","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#remove_pad_token","relation":"calls","name":"cpu","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#remove_pad_token","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#remove_pad_token","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#log_probs_from_logits_response","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_response","relation":"calls","name":"logprobs_from_logits","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#log_probs_from_logits_response_rmpad","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"flash_attn.bert_padding","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_response_rmpad","relation":"calls","name":"unpad_input","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_response_rmpad","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_response_rmpad","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_response_rmpad","relation":"calls","name":"roll","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_response_rmpad","relation":"calls","name":"logprobs_from_logits","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_response_rmpad","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_response_rmpad","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/utils/torch_functional.py","recvType":"logprobs_from_logits"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_response_rmpad","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/utils/torch_functional.py","recvType":"pad_input"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#log_probs_from_logits_all_rmpad","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"flash_attn.bert_padding","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_all_rmpad","relation":"calls","name":"transpose","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_all_rmpad","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_all_rmpad","relation":"calls","name":"roll","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_all_rmpad","relation":"calls","name":"logprobs_from_logits","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_all_rmpad","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_all_rmpad","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/utils/torch_functional.py","recvType":"logprobs_from_logits"},{"source":"verl/utils/torch_functional.py#log_probs_from_logits_all_rmpad","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/utils/torch_functional.py","recvType":"pad_input"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"transformers.generation.logits_process","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#post_process_logits","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#post_process_logits","relation":"calls","name":"div_","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"torch.optim","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"torch.optim.lr_scheduler","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"imports","specifier":"math","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup","relation":"contains","targetId":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","relation":"calls","name":"float","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","relation":"calls","name":"float","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","relation":"calls","name":"max","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","relation":"calls","name":"float","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","relation":"calls","name":"float","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","relation":"calls","name":"max","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","relation":"calls","name":"cos","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","relation":"calls","name":"float","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup.lr_lambda","relation":"calls","name":"max","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_cosine_schedule_with_warmup","relation":"calls","name":"LambdaLR","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#get_constant_schedule_with_warmup","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_constant_schedule_with_warmup","relation":"contains","targetId":"verl/utils/torch_functional.py#get_constant_schedule_with_warmup.lr_lambda","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_constant_schedule_with_warmup.lr_lambda","relation":"calls","name":"min","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_constant_schedule_with_warmup.lr_lambda","relation":"calls","name":"float","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_constant_schedule_with_warmup.lr_lambda","relation":"calls","name":"float","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_constant_schedule_with_warmup.lr_lambda","relation":"calls","name":"max","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_constant_schedule_with_warmup","relation":"calls","name":"LambdaLR","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#prepare_decoder_attention_mask","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#prepare_decoder_attention_mask","relation":"calls","name":"_make_causal_mask","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#prepare_decoder_attention_mask","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#prepare_decoder_attention_mask","relation":"calls","name":"_expand_mask","viaMember":false,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#_make_causal_mask","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_make_causal_mask","relation":"calls","name":"full","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_make_causal_mask","relation":"calls","name":"finfo","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_make_causal_mask","relation":"calls","name":"arange","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_make_causal_mask","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_make_causal_mask","relation":"calls","name":"masked_fill_","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_make_causal_mask","relation":"calls","name":"view","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_make_causal_mask","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_make_causal_mask","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_make_causal_mask","relation":"calls","name":"expand","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#_expand_mask","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_expand_mask","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_expand_mask","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_expand_mask","relation":"calls","name":"expand","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_expand_mask","relation":"calls","name":"masked_fill","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_expand_mask","relation":"calls","name":"to","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#_expand_mask","relation":"calls","name":"finfo","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py","relation":"contains","targetId":"verl/utils/torch_functional.py#get_unpad_data","file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_unpad_data","relation":"calls","name":"sum","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_unpad_data","relation":"calls","name":"flatten","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_unpad_data","relation":"calls","name":"nonzero","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_unpad_data","relation":"calls","name":"flatten","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_unpad_data","relation":"calls","name":"item","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_unpad_data","relation":"calls","name":"max","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_unpad_data","relation":"calls","name":"pad","viaMember":true,"file":"verl/utils/torch_functional.py"},{"source":"verl/utils/torch_functional.py#get_unpad_data","relation":"calls","name":"cumsum","viaMember":true,"file":"verl/utils/torch_functional.py"}]},"verl/utils/tracking.py":{"size":4691,"mtimeMs":1789275830326.2031,"hash":"dd7537bdea3a69a11520e5f11ebe676cff2319b3a844d33ba58f7ce0b7175418","nodes":[{"id":"verl/utils/tracking.py","name":"tracking.py","kind":"file","path":"verl/utils/tracking.py","span":"L1-L126","signature":null,"exported":true,"origin":"ast","body_hash":"dd7537bdea3a69a11520e5f11ebe676cff2319b3a844d33ba58f7ce0b7175418","chars":4655,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 from enum import Enum from functools import partial from pathlib import Path from typing import List, Union, Dict, Any"},{"id":"verl/utils/tracking.py#Tracking","name":"Tracking","kind":"class","path":"verl/utils/tracking.py","span":"L24-L71","signature":"class Tracking(object)","exported":true,"origin":"ast","body_hash":"25464c2f944fd8508c34bb2ddd3bbfd345cc95b03585f3b3b1f3d79625d3d288","body_text":"class Tracking(object): supported_backend = ['wandb', 'mlflow', 'console', 'tensorboard'] def __init__(self, project_name, experiment_name, default_backend: Union[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) 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 wandb import os WANDB_API_KEY = os.environ.get(\"WANDB_API_KEY\", None) if WANDB_API_KEY: wandb.login(key=WANDB_API_KEY) wandb.init(project=project_name, name=experiment_name, config=config) self.logger['wandb'] = wandb if 'mlflow' in default_backend: import mlflow mlflow.start_run(run_name=experiment_name) mlflow.log_params(_compute_mlflow_params_from_objects(config)) self.logger['mlflow'] = _MlflowLoggingAdapter() if 'console' in default_backend: from verl.utils.logger.aggregate_logger import LocalLogger self.console_logger = LocalLogger(print_to_console=True) self.logger['console'] = self.console_logger if 'tensorboard' in default_backend: from torch.utils.tensorboard import SummaryWriter import os # 优先使用环境变量 TB_LOG_DIR，否则退回到相对路径 tb_base = os.environ.get('TB_LOG_DIR', 'tensorboard_logs') tb_log_dir = os.path.join(tb_base, project_name, experiment_name) os.makedirs(tb_log_dir, exist_ok=True) self.logger['tensorboard'] = _TensorboardLoggingAdapter(SummaryWriter(log_dir=tb_log_dir)) 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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/tracking.py#Tracking.__init__","name":"__init__","kind":"method","path":"verl/utils/tracking.py","span":"L27-L66","signature":"def __init__(self, project_name, experiment_name, default_backend: Union[str, List[str]] = 'console', config=None)","exported":false,"origin":"ast","body_hash":"0385b0dc68ec5a5754744a8cbd29acc165a043d20112323db62af12b5ffa132e","body_text":"def __init__(self, project_name, experiment_name, default_backend: Union[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) 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 wandb import os WANDB_API_KEY = os.environ.get(\"WANDB_API_KEY\", None) if WANDB_API_KEY: wandb.login(key=WANDB_API_KEY) wandb.init(project=project_name, name=experiment_name, config=config) self.logger['wandb'] = wandb if 'mlflow' in default_backend: import mlflow mlflow.start_run(run_name=experiment_name) mlflow.log_params(_compute_mlflow_params_from_objects(config)) self.logger['mlflow'] = _MlflowLoggingAdapter() if 'console' in default_backend: from verl.utils.logger.aggregate_logger import LocalLogger self.console_logger = LocalLogger(print_to_console=True) self.logger['console'] = self.console_logger if 'tensorboard' in default_backend: from torch.utils.tensorboard import SummaryWriter import os # 优先使用环境变量 TB_LOG_DIR，否则退回到相对路径 tb_base = os.environ.get('TB_LOG_DIR', 'tensorboard_logs') tb_log_dir = os.path.join(tb_base, project_name, experiment_name) os.makedirs(tb_log_dir, exist_ok=True) self.logger['tensorboard'] = _TensorboardLoggingAdapter(SummaryWriter(log_dir=tb_log_dir))","summary_state":"pending","summary":null,"crux":null,"owner":"Tracking"},{"id":"verl/utils/tracking.py#Tracking.log","name":"log","kind":"method","path":"verl/utils/tracking.py","span":"L68-L71","signature":"def log(self, data, step, backend=None)","exported":true,"origin":"ast","body_hash":"e2cfef826e3e293b4ff99c87184007e2838dd70a87bc948dd43f06f2726a5706","body_text":"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)","summary_state":"pending","summary":null,"crux":null,"owner":"Tracking"},{"id":"verl/utils/tracking.py#_MlflowLoggingAdapter","name":"_MlflowLoggingAdapter","kind":"class","path":"verl/utils/tracking.py","span":"L74-L78","signature":"class _MlflowLoggingAdapter","exported":false,"origin":"ast","body_hash":"015a594a418510cb192597d22ccd4e8e43534c78bd4a37ba6278199320456655","body_text":"class _MlflowLoggingAdapter: def log(self, data, step): import mlflow mlflow.log_metrics(metrics=data, step=step)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/tracking.py#_MlflowLoggingAdapter.log","name":"log","kind":"method","path":"verl/utils/tracking.py","span":"L76-L78","signature":"def log(self, data, step)","exported":true,"origin":"ast","body_hash":"246c7fd3a99c49d92a82b2475fdf5ea5c9a03f418bc5edbb0796bf171e8188b8","body_text":"def log(self, data, step): import mlflow mlflow.log_metrics(metrics=data, step=step)","summary_state":"pending","summary":null,"crux":null,"owner":"_MlflowLoggingAdapter"},{"id":"verl/utils/tracking.py#_TensorboardLoggingAdapter","name":"_TensorboardLoggingAdapter","kind":"class","path":"verl/utils/tracking.py","span":"L81-L91","signature":"class _TensorboardLoggingAdapter","exported":false,"origin":"ast","body_hash":"208c2a0bc0db7ce87db11275da1757433004b4ae7813bb1aabc2d00e75e59e64","body_text":"class _TensorboardLoggingAdapter: def __init__(self, writer): self.writer = writer def log(self, data, step): import numbers for k, v in data.items(): if isinstance(v, numbers.Number): self.writer.add_scalar(k, v, global_step=step) self.writer.flush()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/tracking.py#_TensorboardLoggingAdapter.__init__","name":"__init__","kind":"method","path":"verl/utils/tracking.py","span":"L83-L84","signature":"def __init__(self, writer)","exported":false,"origin":"ast","body_hash":"854fbe4a098084fe909bcb4a331ebd7932db0d4694e308e800e3e17fc86f41e7","body_text":"def __init__(self, writer): self.writer = writer","summary_state":"pending","summary":null,"crux":null,"owner":"_TensorboardLoggingAdapter"},{"id":"verl/utils/tracking.py#_TensorboardLoggingAdapter.log","name":"log","kind":"method","path":"verl/utils/tracking.py","span":"L86-L91","signature":"def log(self, data, step)","exported":true,"origin":"ast","body_hash":"44e98ae712741df144a718140a2f4dcf97c9274ea3261fb793ced3b04e500f4d","body_text":"def log(self, data, step): import numbers for k, v in data.items(): if isinstance(v, numbers.Number): self.writer.add_scalar(k, v, global_step=step) self.writer.flush()","summary_state":"pending","summary":null,"crux":null,"owner":"_TensorboardLoggingAdapter"},{"id":"verl/utils/tracking.py#_compute_mlflow_params_from_objects","name":"_compute_mlflow_params_from_objects","kind":"function","path":"verl/utils/tracking.py","span":"L94-L98","signature":"def _compute_mlflow_params_from_objects(params) -> Dict[str, Any]","exported":false,"origin":"ast","body_hash":"af44e4a9a7857d0a7a93baeed1f1d67269cd48b5ac78071c3c5dc9fc71367610","body_text":"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='/')","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/tracking.py#_transform_params_to_json_serializable","name":"_transform_params_to_json_serializable","kind":"function","path":"verl/utils/tracking.py","span":"L101-L118","signature":"def _transform_params_to_json_serializable(x, convert_list_to_dict: bool)","exported":false,"origin":"ast","body_hash":"884dfcc32f9313daa7ed0bb8a5a145e0733011dd9e462ffa6cfefe6dd50772ce","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/tracking.py#_flatten_dict","name":"_flatten_dict","kind":"function","path":"verl/utils/tracking.py","span":"L121-L125","signature":"def _flatten_dict(raw: Dict[str, Any], *, sep: str) -> Dict[str, Any]","exported":false,"origin":"ast","body_hash":"cbf73f1e064f42773b9071b095b718520f135f2ea78cef7906bc3601eb4c5459","body_text":"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","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/tracking.py","relation":"imports","specifier":"dataclasses","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"enum","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"functools","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"pathlib","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"typing","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"contains","targetId":"verl/utils/tracking.py#Tracking","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking","relation":"extends","name":"object","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking","relation":"contains","targetId":"verl/utils/tracking.py#Tracking.__init__","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"warnings","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"warn","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"wandb","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"os","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"login","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"init","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"mlflow","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"start_run","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"log_params","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"_compute_mlflow_params_from_objects","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"_MlflowLoggingAdapter","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"verl.utils.logger.aggregate_logger","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"LocalLogger","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"torch.utils.tensorboard","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"os","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"join","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"_TensorboardLoggingAdapter","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.__init__","relation":"calls","name":"SummaryWriter","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking","relation":"contains","targetId":"verl/utils/tracking.py#Tracking.log","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.log","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#Tracking.log","relation":"calls","name":"log","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"contains","targetId":"verl/utils/tracking.py#_MlflowLoggingAdapter","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_MlflowLoggingAdapter","relation":"contains","targetId":"verl/utils/tracking.py#_MlflowLoggingAdapter.log","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"mlflow","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_MlflowLoggingAdapter.log","relation":"calls","name":"log_metrics","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"contains","targetId":"verl/utils/tracking.py#_TensorboardLoggingAdapter","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_TensorboardLoggingAdapter","relation":"contains","targetId":"verl/utils/tracking.py#_TensorboardLoggingAdapter.__init__","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_TensorboardLoggingAdapter","relation":"contains","targetId":"verl/utils/tracking.py#_TensorboardLoggingAdapter.log","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"imports","specifier":"numbers","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_TensorboardLoggingAdapter.log","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_TensorboardLoggingAdapter.log","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_TensorboardLoggingAdapter.log","relation":"calls","name":"add_scalar","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_TensorboardLoggingAdapter.log","relation":"calls","name":"flush","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"contains","targetId":"verl/utils/tracking.py#_compute_mlflow_params_from_objects","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_compute_mlflow_params_from_objects","relation":"calls","name":"_flatten_dict","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_compute_mlflow_params_from_objects","relation":"calls","name":"_transform_params_to_json_serializable","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"contains","targetId":"verl/utils/tracking.py#_transform_params_to_json_serializable","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"partial","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"is_dataclass","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"_transform","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"asdict","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"_transform","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"items","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"_transform","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"_transform","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"str","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_transform_params_to_json_serializable","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py","relation":"contains","targetId":"verl/utils/tracking.py#_flatten_dict","file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_flatten_dict","relation":"calls","name":"to_dict","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_flatten_dict","relation":"calls","name":"json_normalize","viaMember":true,"file":"verl/utils/tracking.py"},{"source":"verl/utils/tracking.py#_flatten_dict","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/tracking.py"}]},"verl/utils/ulysses.py":{"size":10619,"mtimeMs":1789275830326.2031,"hash":"c47179787683e6ff73a50e09fd4acd6328bcd9d61d526afb827f234e538f238a","nodes":[{"id":"verl/utils/ulysses.py","name":"ulysses.py","kind":"file","path":"verl/utils/ulysses.py","span":"L1-L289","signature":null,"exported":true,"origin":"ast","body_hash":"c47179787683e6ff73a50e09fd4acd6328bcd9d61d526afb827f234e538f238a","chars":10619,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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/microsoft/DeepSpeed/blob/master/deepspeed/sequence/layer.py \"\"\" from typing import Any, Optional, List, Tuple import torch from torch import Tensor import torch.distributed as dist from torch.distributed import ProcessGroup _ULYSSES_SEQUENCE_PARALLEL_GROUP = None"},{"id":"verl/utils/ulysses.py#set_ulysses_sequence_parallel_group","name":"set_ulysses_sequence_parallel_group","kind":"function","path":"verl/utils/ulysses.py","span":"L29-L34","signature":"def set_ulysses_sequence_parallel_group(group: dist.ProcessGroup)","exported":true,"origin":"ast","body_hash":"2370249b6ffde2bc343356ebb2c712f8fd86f77b3a859efa620c2b9947a40bbc","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#get_ulysses_sequence_parallel_group","name":"get_ulysses_sequence_parallel_group","kind":"function","path":"verl/utils/ulysses.py","span":"L37-L42","signature":"def get_ulysses_sequence_parallel_group() -> Optional[dist.ProcessGroup]","exported":true,"origin":"ast","body_hash":"ede40bde9ffbc5ed692ed50d01996143ad440afaf28b8a9967ee7c19a502d195","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#get_ulysses_sequence_parallel_world_size","name":"get_ulysses_sequence_parallel_world_size","kind":"function","path":"verl/utils/ulysses.py","span":"L45-L50","signature":"def get_ulysses_sequence_parallel_world_size(group: ProcessGroup = None) -> int","exported":true,"origin":"ast","body_hash":"54c1a27535d7e257a7000b6539e0c816a62d7d5c1cdb71c4295c5407dd2e71e5","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#get_ulysses_sequence_parallel_rank","name":"get_ulysses_sequence_parallel_rank","kind":"function","path":"verl/utils/ulysses.py","span":"L53-L58","signature":"def get_ulysses_sequence_parallel_rank(group: ProcessGroup = None) -> int","exported":true,"origin":"ast","body_hash":"f7798f0f07bf603296dde656eda0149fdb0da164a0a60b22da339c43ec5d908d","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#gather_seq_scatter_heads","name":"gather_seq_scatter_heads","kind":"function","path":"verl/utils/ulysses.py","span":"L61-L82","signature":"def gather_seq_scatter_heads( x: Tensor, seq_dim: int, head_dim: int, unpadded_dim_size: int = 0, group: ProcessGroup = None, ) -> Tensor","exported":true,"origin":"ast","body_hash":"e3ecf360ce158b59fc83a9e851a709e5d2d6bf69f9df26b1e36cc757b9174c52","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#gather_heads_scatter_seq","name":"gather_heads_scatter_seq","kind":"function","path":"verl/utils/ulysses.py","span":"L85-L100","signature":"def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int, group: ProcessGroup = None) -> Tensor","exported":true,"origin":"ast","body_hash":"7f02a634795303f4d4e12cbe9797c5a8be128b8ae0e19dc6ae9990d37951e910","body_text":"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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#_pad_tensor","name":"_pad_tensor","kind":"function","path":"verl/utils/ulysses.py","span":"L103-L107","signature":"def _pad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor","exported":false,"origin":"ast","body_hash":"94eaf7a7ebf3dd1efa706f3e8df0c2fec3893db065f437f499142a7eda9fa2f5","body_text":"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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#_unpad_tensor","name":"_unpad_tensor","kind":"function","path":"verl/utils/ulysses.py","span":"L110-L113","signature":"def _unpad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor","exported":false,"origin":"ast","body_hash":"6969d4776bb36b54c94e25ff53f244b5b671c69d25b185d9d67d2947ea7d45e5","body_text":"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[slc]","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#slice_input_tensor","name":"slice_input_tensor","kind":"function","path":"verl/utils/ulysses.py","span":"L116-L129","signature":"def slice_input_tensor(x: Tensor, dim: int, padding: bool = True, group: ProcessGroup = None) -> Tensor","exported":true,"origin":"ast","body_hash":"3e03f49f5bb2d0030570d69b7690b7bef7383c3556a8502715a28e55710fbd85","body_text":"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[slc].contiguous()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#all_to_all_tensor","name":"all_to_all_tensor","kind":"function","path":"verl/utils/ulysses.py","span":"L132-L151","signature":"def all_to_all_tensor( local_input: Tensor, scatter_dim: int, gather_dim: int, group: Optional[dist.ProcessGroup] = None, async_op: bool = False, )","exported":true,"origin":"ast","body_hash":"580030347657a1cd210f8473dd218c586390f0c0417c79bdfeb8c7ab7b1ccd0f","body_text":"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()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#all_to_all_tensor.wait","name":"wait","kind":"function","path":"verl/utils/ulysses.py","span":"L146-L148","signature":"def wait()","exported":true,"origin":"ast","body_hash":"a57b2ac960ba1db9edd79ecae07b00e4be32f854c7e25484e18a22811fc9f648","body_text":"def wait(): comm.wait() return torch.cat(output_list, dim=gather_dim).contiguous()","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#all_gather_tensor","name":"all_gather_tensor","kind":"function","path":"verl/utils/ulysses.py","span":"L154-L161","signature":"def all_gather_tensor(local_tensor: Tensor, group: Optional[dist.ProcessGroup] = None, async_op: bool = False)","exported":true,"origin":"ast","body_hash":"0be18c1ae95269dab2a8c30855ae6037844d5eda822ba32454cd1751a1098831","body_text":"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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#SeqAllToAll","name":"SeqAllToAll","kind":"class","path":"verl/utils/ulysses.py","span":"L164-L194","signature":"class SeqAllToAll(torch.autograd.Function)","exported":true,"origin":"ast","body_hash":"ff028505d482d4b495ff63d1681bfc21ca3f3888244e7c6d6b0a00ce04858948","body_text":"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]: if ctx.async_op: input_t = torch.cat(grad_output[1:], dim=ctx.gather_dim).contiguous() else: input_t = grad_output[0] return ( None, all_to_all_tensor(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group, False), None, None, None, None, )","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#SeqAllToAll.forward","name":"forward","kind":"method","path":"verl/utils/ulysses.py","span":"L167-L179","signature":"def forward( ctx: Any, group: dist.ProcessGroup, local_input: Tensor, scatter_dim: int, gather_dim: int, async_op: bool = False, ) -> Tensor","exported":true,"origin":"ast","body_hash":"d9ee223f1bd158e7fb530920744a76f87c986a0eba3ef0108c72d310d122f681","body_text":"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)","summary_state":"pending","summary":null,"crux":null,"owner":"SeqAllToAll"},{"id":"verl/utils/ulysses.py#SeqAllToAll.backward","name":"backward","kind":"method","path":"verl/utils/ulysses.py","span":"L182-L194","signature":"def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]","exported":true,"origin":"ast","body_hash":"7aaea1fede11941738ba7a1b704d8419843ab49053a388f93bbc1a5fb3fe1852","body_text":"def backward(ctx: Any, *grad_output: Tensor) -> Tuple[None, Tensor, None, None]: if ctx.async_op: input_t = torch.cat(grad_output[1:], dim=ctx.gather_dim).contiguous() else: input_t = grad_output[0] return ( None, all_to_all_tensor(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group, False), None, None, None, None, )","summary_state":"pending","summary":null,"crux":null,"owner":"SeqAllToAll"},{"id":"verl/utils/ulysses.py#Gather","name":"Gather","kind":"class","path":"verl/utils/ulysses.py","span":"L197-L230","signature":"class Gather(torch.autograd.Function)","exported":true,"origin":"ast","body_hash":"f6149bf3e03f7e2ceeb73c1ec94f9f96579d307c6f8ae7e8a74bdd75a6cd7202","body_text":"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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#Gather.forward","name":"forward","kind":"method","path":"verl/utils/ulysses.py","span":"L200-L223","signature":"def forward(ctx: Any, group: dist.ProcessGroup, local_tensor: Tensor, gather_dim: int, grad_scaler: bool = True, async_op=False) -> Tensor","exported":true,"origin":"ast","body_hash":"b0a50f6fe7eb88144134a09e3bb50d2b3512f999b52eb3dd6d6febdb3e27d00e","body_text":"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)","summary_state":"pending","summary":null,"crux":null,"owner":"Gather"},{"id":"verl/utils/ulysses.py#Gather.backward","name":"backward","kind":"method","path":"verl/utils/ulysses.py","span":"L226-L230","signature":"def backward(ctx: Any, grad_output: Tensor) -> Any","exported":true,"origin":"ast","body_hash":"026f16e725f2c391fb23ff04133f062c883111b1018b400661906bf6ec8f635a","body_text":"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)","summary_state":"pending","summary":null,"crux":null,"owner":"Gather"},{"id":"verl/utils/ulysses.py#gather_outpus_and_unpad","name":"gather_outpus_and_unpad","kind":"function","path":"verl/utils/ulysses.py","span":"L233-L249","signature":"def gather_outpus_and_unpad(x: Tensor, gather_dim: int, unpad_dim: int = None, padding_size: int = 0, grad_scaler: bool = True, group: Optional[dist.ProcessGroup] = None)","exported":true,"origin":"ast","body_hash":"8a8fce63a5f2fc9d7dc475f0588d66bca577221a67a48ac10dde644474fd68c2","body_text":"def gather_outpus_and_unpad(x: Tensor, gather_dim: int, unpad_dim: int = None, padding_size: int = 0, grad_scaler: bool = True, group: Optional[dist.ProcessGroup] = None): group = get_ulysses_sequence_parallel_group() if group is None else group sp_size = get_ulysses_sequence_parallel_world_size() if group == 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/utils/ulysses.py#ulysses_pad_and_slice_inputs","name":"ulysses_pad_and_slice_inputs","kind":"function","path":"verl/utils/ulysses.py","span":"L252-L288","signature":"def ulysses_pad_and_slice_inputs(input_ids_rmpad: torch.Tensor, position_ids_rmpad: Optional[torch.Tensor] = None, sp_size: int = 1)","exported":true,"origin":"ast","body_hash":"eac54c8affff8df78b9047dd1414a916b6e03f79eee5826d02207fa34f364a18","body_text":"def ulysses_pad_and_slice_inputs(input_ids_rmpad: torch.Tensor, position_ids_rmpad: Optional[torch.Tensor] = None, sp_size: int = 1): \"\"\" 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, but only input_ids will be 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 Returns: torch.Tensor: padded and sliced input_ids torch.Tensor: padded and sliced position_ids int: pad size \"\"\" if position_ids_rmpad is not None: assert position_ids_rmpad.size(0) == 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=0) if position_ids_rmpad is not None: pad_pos_ids = torch.arange(pad_size, device=position_ids_rmpad.device).unsqueeze(0) position_ids_rmpad = torch.cat((position_ids_rmpad, pad_pos_ids), dim=-1) # we don't need to slice position ids input_ids_rmpad = slice_input_tensor(input_ids_rmpad, dim=1, padding=False) return input_ids_rmpad, position_ids_rmpad, pad_size","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/utils/ulysses.py","relation":"imports","specifier":"typing","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"imports","specifier":"torch","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"imports","specifier":"torch","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"imports","specifier":"torch.distributed","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#set_ulysses_sequence_parallel_group","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#get_ulysses_sequence_parallel_group","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#get_ulysses_sequence_parallel_world_size","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#get_ulysses_sequence_parallel_world_size","relation":"calls","name":"get_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#get_ulysses_sequence_parallel_world_size","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#get_ulysses_sequence_parallel_rank","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#get_ulysses_sequence_parallel_rank","relation":"calls","name":"get_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#get_ulysses_sequence_parallel_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#gather_seq_scatter_heads","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_seq_scatter_heads","relation":"calls","name":"get_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_seq_scatter_heads","relation":"calls","name":"get_ulysses_sequence_parallel_world_size","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_seq_scatter_heads","relation":"calls","name":"apply","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_seq_scatter_heads","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/ulysses.py","recvType":"_unpad_tensor"},{"source":"verl/utils/ulysses.py#gather_seq_scatter_heads","relation":"calls","name":"_unpad_tensor","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#gather_heads_scatter_seq","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_heads_scatter_seq","relation":"calls","name":"get_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_heads_scatter_seq","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/ulysses.py","recvType":"_pad_tensor"},{"source":"verl/utils/ulysses.py#gather_heads_scatter_seq","relation":"calls","name":"get_ulysses_sequence_parallel_world_size","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_heads_scatter_seq","relation":"calls","name":"_pad_tensor","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_heads_scatter_seq","relation":"calls","name":"apply","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#_pad_tensor","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#_pad_tensor","relation":"calls","name":"list","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#_pad_tensor","relation":"calls","name":"zeros","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#_pad_tensor","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#_unpad_tensor","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#_unpad_tensor","relation":"calls","name":"slice","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#_unpad_tensor","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#_unpad_tensor","relation":"calls","name":"slice","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#slice_input_tensor","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#slice_input_tensor","relation":"calls","name":"get_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#slice_input_tensor","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#slice_input_tensor","relation":"calls","name":"get_ulysses_sequence_parallel_rank","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#slice_input_tensor","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/ulysses.py","recvType":"_pad_tensor"},{"source":"verl/utils/ulysses.py#slice_input_tensor","relation":"calls","name":"_pad_tensor","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#slice_input_tensor","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/ulysses.py","recvType":"_pad_tensor"},{"source":"verl/utils/ulysses.py#slice_input_tensor","relation":"calls","name":"slice","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#slice_input_tensor","relation":"calls","name":"len","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#slice_input_tensor","relation":"calls","name":"slice","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#slice_input_tensor","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#all_to_all_tensor","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor","relation":"calls","name":"get_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor","relation":"calls","name":"tensor_split","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor","relation":"calls","name":"range","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor","relation":"calls","name":"all_to_all","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor","relation":"contains","targetId":"verl/utils/ulysses.py#all_to_all_tensor.wait","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor.wait","relation":"calls","name":"wait","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor.wait","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor.wait","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_to_all_tensor","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#all_gather_tensor","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_gather_tensor","relation":"calls","name":"get_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_gather_tensor","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_gather_tensor","relation":"calls","name":"list","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_gather_tensor","relation":"calls","name":"empty","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#all_gather_tensor","relation":"calls","name":"all_gather_into_tensor","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#SeqAllToAll","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#SeqAllToAll","relation":"contains","targetId":"verl/utils/ulysses.py#SeqAllToAll.forward","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#SeqAllToAll.forward","relation":"calls","name":"all_to_all_tensor","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#SeqAllToAll","relation":"contains","targetId":"verl/utils/ulysses.py#SeqAllToAll.backward","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#SeqAllToAll.backward","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#SeqAllToAll.backward","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#SeqAllToAll.backward","relation":"calls","name":"all_to_all_tensor","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#Gather","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#Gather","relation":"contains","targetId":"verl/utils/ulysses.py#Gather.forward","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#Gather.forward","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#Gather.forward","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#Gather.forward","relation":"calls","name":"list","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#Gather.forward","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/ulysses.py","recvType":"Tensor"},{"source":"verl/utils/ulysses.py#Gather.forward","relation":"calls","name":"all_gather_tensor","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#Gather.forward","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#Gather.forward","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/ulysses.py","recvType":"all_gather_tensor"},{"source":"verl/utils/ulysses.py#Gather","relation":"contains","targetId":"verl/utils/ulysses.py#Gather.backward","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#Gather.backward","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#Gather.backward","relation":"calls","name":"split","viaMember":true,"file":"verl/utils/ulysses.py","recvType":"Tensor"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#gather_outpus_and_unpad","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_outpus_and_unpad","relation":"calls","name":"get_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_outpus_and_unpad","relation":"calls","name":"get_ulysses_sequence_parallel_world_size","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_outpus_and_unpad","relation":"calls","name":"apply","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_outpus_and_unpad","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#gather_outpus_and_unpad","relation":"calls","name":"_unpad_tensor","viaMember":false,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py","relation":"contains","targetId":"verl/utils/ulysses.py#ulysses_pad_and_slice_inputs","file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#ulysses_pad_and_slice_inputs","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#ulysses_pad_and_slice_inputs","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/ulysses.py","recvType":"slice_input_tensor"},{"source":"verl/utils/ulysses.py#ulysses_pad_and_slice_inputs","relation":"calls","name":"size","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#ulysses_pad_and_slice_inputs","relation":"calls","name":"pad","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#ulysses_pad_and_slice_inputs","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#ulysses_pad_and_slice_inputs","relation":"calls","name":"arange","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#ulysses_pad_and_slice_inputs","relation":"calls","name":"cat","viaMember":true,"file":"verl/utils/ulysses.py"},{"source":"verl/utils/ulysses.py#ulysses_pad_and_slice_inputs","relation":"calls","name":"slice_input_tensor","viaMember":false,"file":"verl/utils/ulysses.py"}]},"verl/workers/__init__.py":{"size":600,"mtimeMs":1789275830326.2031,"hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","nodes":[{"id":"verl/workers/__init__.py","name":"__init__.py","kind":"file","path":"verl/workers/__init__.py","span":"L1-L14","signature":null,"exported":true,"origin":"ast","body_hash":"e65abd77db034b4f29eb1268c24867f84919dbce22e28d1997e563ed1d94179c","chars":600,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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."}],"rawEdges":[]},"verl/workers/actor/__init__.py":{"size":727,"mtimeMs":1789275830326.2031,"hash":"3053474474716dc57f7e1b5d80d23b28c7b146d9257f8b6bf3b7878de4eae219","nodes":[{"id":"verl/workers/actor/__init__.py","name":"__init__.py","kind":"file","path":"verl/workers/actor/__init__.py","span":"L1-L19","signature":null,"exported":true,"origin":"ast","body_hash":"3053474474716dc57f7e1b5d80d23b28c7b146d9257f8b6bf3b7878de4eae219","chars":727,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 BasePPOActor from .dp_actor import DataParallelPPOActor __all__ = [\"BasePPOActor\", \"DataParallelPPOActor\"]"}],"rawEdges":[{"source":"verl/workers/actor/__init__.py","relation":"imports","specifier":".base","file":"verl/workers/actor/__init__.py"},{"source":"verl/workers/actor/__init__.py","relation":"imports","specifier":".dp_actor","file":"verl/workers/actor/__init__.py"}]},"verl/workers/actor/base.py":{"size":2071,"mtimeMs":1789275830326.2031,"hash":"4afcfa27274dd1b58e40cc5484ca352890d5423ce2c162734f1ee6b4fb26890b","nodes":[{"id":"verl/workers/actor/base.py","name":"base.py","kind":"file","path":"verl/workers/actor/base.py","span":"L1-L67","signature":null,"exported":true,"origin":"ast","body_hash":"4afcfa27274dd1b58e40cc5484ca352890d5423ce2c162734f1ee6b4fb26890b","chars":2071,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 class for Actor \"\"\" from abc import ABC, abstractmethod from typing import Iterable, Dict from verl import DataProto import torch __all__ = ['BasePPOActor']"},{"id":"verl/workers/actor/base.py#BasePPOActor","name":"BasePPOActor","kind":"class","path":"verl/workers/actor/base.py","span":"L26-L66","signature":"class BasePPOActor(ABC)","exported":true,"origin":"ast","body_hash":"a7272c7b8df86a12e9c28f64f923d0765429e7df14e04525629d9e05eabe6ddc","body_text":"class BasePPOActor(ABC): def __init__(self, config): \"\"\"The base class for PPO actor Args: config (DictConfig): a config passed to the PPOActor. We expect the type to be DictConfig (https://omegaconf.readthedocs.io/), but it can be any namedtuple in general. \"\"\" super().__init__() self.config = config @abstractmethod def compute_log_prob(self, data: DataProto) -> torch.Tensor: \"\"\"Compute logits given a batch of data. Args: data (DataProto): a batch of data represented by DataProto. It must contain key ```input_ids```, ```attention_mask``` and ```position_ids```. Returns: DataProto: a DataProto containing the key ```log_probs``` \"\"\" pass @abstractmethod def update_policy(self, data: DataProto) -> Dict: \"\"\"Update the policy with an iterator of DataProto Args: data (DataProto): an iterator over the DataProto that returns by ```make_minibatch_iterator``` Returns: Dict: a dictionary contains anything. Typically, it contains the statistics during updating the model such as ```loss```, ```grad_norm```, etc,. \"\"\" pass","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/actor/base.py#BasePPOActor.__init__","name":"__init__","kind":"method","path":"verl/workers/actor/base.py","span":"L28-L36","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"d5a685efc62709fb34ff911424f23b918d9e639f17517d5df3d78afc6075abd1","body_text":"def __init__(self, config): \"\"\"The base class for PPO actor Args: config (DictConfig): a config passed to the PPOActor. We expect the type to be DictConfig (https://omegaconf.readthedocs.io/), but it can be any namedtuple in general. \"\"\" super().__init__() self.config = config","summary_state":"pending","summary":null,"crux":null,"owner":"BasePPOActor"},{"id":"verl/workers/actor/base.py#BasePPOActor.compute_log_prob","name":"compute_log_prob","kind":"method","path":"verl/workers/actor/base.py","span":"L39-L51","signature":"def compute_log_prob(self, data: DataProto) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"d25195bb4793dd1d6858b7104328a3d876f9147474c3243c13a446721ce109c7","body_text":"def compute_log_prob(self, data: DataProto) -> torch.Tensor: \"\"\"Compute logits given a batch of data. Args: data (DataProto): a batch of data represented by DataProto. It must contain key ```input_ids```, ```attention_mask``` and ```position_ids```. Returns: DataProto: a DataProto containing the key ```log_probs``` \"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"BasePPOActor"},{"id":"verl/workers/actor/base.py#BasePPOActor.update_policy","name":"update_policy","kind":"method","path":"verl/workers/actor/base.py","span":"L54-L66","signature":"def update_policy(self, data: DataProto) -> Dict","exported":true,"origin":"ast","body_hash":"a6308730ace71450cd82db991bdd3347cbd0f5040f296d0450efe986d334738b","body_text":"def update_policy(self, data: DataProto) -> Dict: \"\"\"Update the policy with an iterator of DataProto Args: data (DataProto): an iterator over the DataProto that returns by ```make_minibatch_iterator``` Returns: Dict: a dictionary contains anything. Typically, it contains the statistics during updating the model such as ```loss```, ```grad_norm```, etc,. \"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"BasePPOActor"}],"rawEdges":[{"source":"verl/workers/actor/base.py","relation":"imports","specifier":"abc","file":"verl/workers/actor/base.py"},{"source":"verl/workers/actor/base.py","relation":"imports","specifier":"typing","file":"verl/workers/actor/base.py"},{"source":"verl/workers/actor/base.py","relation":"imports","specifier":"verl","file":"verl/workers/actor/base.py"},{"source":"verl/workers/actor/base.py","relation":"imports","specifier":"torch","file":"verl/workers/actor/base.py"},{"source":"verl/workers/actor/base.py","relation":"contains","targetId":"verl/workers/actor/base.py#BasePPOActor","file":"verl/workers/actor/base.py"},{"source":"verl/workers/actor/base.py#BasePPOActor","relation":"extends","name":"ABC","file":"verl/workers/actor/base.py"},{"source":"verl/workers/actor/base.py#BasePPOActor","relation":"contains","targetId":"verl/workers/actor/base.py#BasePPOActor.__init__","file":"verl/workers/actor/base.py"},{"source":"verl/workers/actor/base.py#BasePPOActor.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/actor/base.py"},{"source":"verl/workers/actor/base.py#BasePPOActor.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/actor/base.py"},{"source":"verl/workers/actor/base.py#BasePPOActor","relation":"contains","targetId":"verl/workers/actor/base.py#BasePPOActor.compute_log_prob","file":"verl/workers/actor/base.py"},{"source":"verl/workers/actor/base.py#BasePPOActor","relation":"contains","targetId":"verl/workers/actor/base.py#BasePPOActor.update_policy","file":"verl/workers/actor/base.py"}]},"verl/workers/actor/dp_actor.py":{"size":14440,"mtimeMs":1789275830327.2031,"hash":"288e2494f787b515c11973cd30ac4ea1d6c73d03332f3fa650df70f3c9d035d8","nodes":[{"id":"verl/workers/actor/dp_actor.py","name":"dp_actor.py","kind":"file","path":"verl/workers/actor/dp_actor.py","span":"L1-L292","signature":null,"exported":true,"origin":"ast","body_hash":"288e2494f787b515c11973cd30ac4ea1d6c73d03332f3fa650df70f3c9d035d8","chars":14396,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Single Process Actor \"\"\" import itertools from typing import Iterable, Tuple import torch from torch import nn from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from verl import DataProto from verl.trainer.ppo import core_algos from verl.workers.actor import BasePPOActor from verl.utils.py_functional import append_to_dict from verl.utils.torch_functional import logprobs_from_logits, masked_mean from verl.utils.ulysses import ulysses_pad_and_slice_inputs, gather_outpus_and_unpad from verl.utils.seqlen_balancing import rearrange_micro_batches, get_reverse_idx import verl.utils.torch_functional as verl_F from flash_attn.bert_padding import pad_input, unpad_input, rearrange, index_first_axis __all__ = ['DataParallelPPOActor']"},{"id":"verl/workers/actor/dp_actor.py#DataParallelPPOActor","name":"DataParallelPPOActor","kind":"class","path":"verl/workers/actor/dp_actor.py","span":"L39-L291","signature":"class DataParallelPPOActor(BasePPOActor)","exported":true,"origin":"ast","body_hash":"044dbbfb3b271228ccdeca777d5c9cde3fef8ffb983d7d695c119f2dc4c6802d","body_text":"class DataParallelPPOActor(BasePPOActor): def __init__( self, config, actor_module: nn.Module, actor_optimizer: torch.optim.Optimizer = None, ): \"\"\"When optimizer is None, it is Reference Policy\"\"\" super().__init__(config) self.actor_module = actor_module self.actor_optimizer = actor_optimizer self.use_remove_padding = self.config.get('use_remove_padding', False) print(f'Actor use_remove_padding={self.use_remove_padding}') self.ulysses_sequence_parallel_size = self.config.ulysses_sequence_parallel_size self.use_ulysses_sp = self.ulysses_sequence_parallel_size > 1 # torch.compile 在部分集群环境（triton hash 不可用）会崩溃，改为直接调用 self.compute_entropy_from_logits = verl_F.entropy_from_logits def _forward_micro_batch(self, micro_batch, temperature) -> Tuple[torch.Tensor, torch.Tensor]: \"\"\" Returns: entropy: # (bs, response_len) log_probs: # (bs, response_len) \"\"\" response_length = micro_batch['responses'].size(-1) with torch.autocast(device_type='cuda', dtype=torch.bfloat16): input_ids = micro_batch['input_ids'] batch_size, seqlen = input_ids.shape attention_mask = micro_batch['attention_mask'] position_ids = micro_batch['position_ids'] if self.use_remove_padding: 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) # 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: input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs(input_ids_rmpad, \\ position_ids_rmpad, \\ sp_size=self.ulysses_sequence_parallel_size) input_ids_rmpad_rolled, _, _ = ulysses_pad_and_slice_inputs(input_ids_rmpad_rolled, None, self.ulysses_sequence_parallel_size) input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0) # ((total_nnz / sp) + pad) # only pass input_ids and position_ids to enable flash_attn_varlen output = self.actor_module(input_ids=input_ids_rmpad, attention_mask=None, position_ids=position_ids_rmpad, use_cache=False) # prevent model thinks we are generating logits_rmpad = output.logits.squeeze(0) # (total_nnz, vocab_size) logits_rmpad.div_(temperature) # compute entropy entropy_rmpad = self.compute_entropy_from_logits(logits_rmpad) # ((total_nnz / sp) + pad) # if use_sp: ((total_nnz / sp) + pad) ; if not use_sp: (batch, seqlen) log_probs = logprobs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled) # gather log_prob if sp > 1 if self.use_ulysses_sp: # gather and unpad for the ulysses sp log_probs = gather_outpus_and_unpad(log_probs, gather_dim=0, unpad_dim=0, padding_size=pad_size) entropy_rmpad = gather_outpus_and_unpad(entropy_rmpad, gather_dim=0, unpad_dim=0, padding_size=pad_size) # pad back to (bsz, seqlen) full_entropy = pad_input(hidden_states=entropy_rmpad.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen) full_log_probs = pad_input(hidden_states=log_probs.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen) # only return response part: entropy = full_entropy.squeeze(-1)[:, -response_length - 1:-1] # (bsz, response_length) log_probs = full_log_probs.squeeze(-1)[:, -response_length - 1:-1] # (bsz, response_length) else: # not using rmpad and no ulysses sp output = self.actor_module(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False) # prevent model thinks we are generating logits = output.logits logits.div_(temperature) logits = logits[:, -response_length - 1:-1] # (bsz, response_length) log_probs = logprobs_from_logits(logits, micro_batch['responses']) entropy = verl_F.entropy_from_logits(logits) # (bsz, response_length) return entropy, log_probs def _optimizer_step(self): assert self.config.grad_clip is not None if isinstance(self.actor_module, FSDP): grad_norm = self.actor_module.clip_grad_norm_(max_norm=self.config.grad_clip) else: grad_norm = torch.nn.utils.clip_grad_norm_(self.actor_module.parameters(), max_norm=self.config.grad_clip) self.actor_optimizer.step() return grad_norm def compute_log_prob(self, data: DataProto) -> torch.Tensor: \"\"\"Compute the log probability of the responses given input_ids, attention_mask and position_ids Args: data (DataProto): a DataProto containing keys ``input_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. Note that input_ids is the concatenation of prompt and response. Note that ``sequence_length = prompt_length + response_length``. ``attention_mask``: tensor of shape [batch_size, sequence_length]. torch.int64. ``position_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. ``responses``: tensor of shape [batch_size, response_length]. torch.int64. Returns: torch.Tensor: the log_prob tensor \"\"\" # set to eval self.actor_module.eval() micro","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.__init__","name":"__init__","kind":"method","path":"verl/workers/actor/dp_actor.py","span":"L41-L57","signature":"def __init__( self, config, actor_module: nn.Module, actor_optimizer: torch.optim.Optimizer = None, )","exported":false,"origin":"ast","body_hash":"398b651a365e70c3f56d99c7c2bfb2142a7e2d2720e8f7f83174ea30cbd07ec8","body_text":"def __init__( self, config, actor_module: nn.Module, actor_optimizer: torch.optim.Optimizer = None, ): \"\"\"When optimizer is None, it is Reference Policy\"\"\" super().__init__(config) self.actor_module = actor_module self.actor_optimizer = actor_optimizer self.use_remove_padding = self.config.get('use_remove_padding', False) print(f'Actor use_remove_padding={self.use_remove_padding}') self.ulysses_sequence_parallel_size = self.config.ulysses_sequence_parallel_size self.use_ulysses_sp = self.ulysses_sequence_parallel_size > 1 # torch.compile 在部分集群环境（triton hash 不可用）会崩溃，改为直接调用 self.compute_entropy_from_logits = verl_F.entropy_from_logits","summary_state":"pending","summary":null,"crux":null,"owner":"DataParallelPPOActor"},{"id":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","name":"_forward_micro_batch","kind":"method","path":"verl/workers/actor/dp_actor.py","span":"L59-L142","signature":"def _forward_micro_batch(self, micro_batch, temperature) -> Tuple[torch.Tensor, torch.Tensor]","exported":false,"origin":"ast","body_hash":"2089b91b722244366808914a443091331e71dbb4a2b40e7c085c9723c03e97d4","body_text":"def _forward_micro_batch(self, micro_batch, temperature) -> Tuple[torch.Tensor, torch.Tensor]: \"\"\" Returns: entropy: # (bs, response_len) log_probs: # (bs, response_len) \"\"\" response_length = micro_batch['responses'].size(-1) with torch.autocast(device_type='cuda', dtype=torch.bfloat16): input_ids = micro_batch['input_ids'] batch_size, seqlen = input_ids.shape attention_mask = micro_batch['attention_mask'] position_ids = micro_batch['position_ids'] if self.use_remove_padding: 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) # 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: input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs(input_ids_rmpad, \\ position_ids_rmpad, \\ sp_size=self.ulysses_sequence_parallel_size) input_ids_rmpad_rolled, _, _ = ulysses_pad_and_slice_inputs(input_ids_rmpad_rolled, None, self.ulysses_sequence_parallel_size) input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0) # ((total_nnz / sp) + pad) # only pass input_ids and position_ids to enable flash_attn_varlen output = self.actor_module(input_ids=input_ids_rmpad, attention_mask=None, position_ids=position_ids_rmpad, use_cache=False) # prevent model thinks we are generating logits_rmpad = output.logits.squeeze(0) # (total_nnz, vocab_size) logits_rmpad.div_(temperature) # compute entropy entropy_rmpad = self.compute_entropy_from_logits(logits_rmpad) # ((total_nnz / sp) + pad) # if use_sp: ((total_nnz / sp) + pad) ; if not use_sp: (batch, seqlen) log_probs = logprobs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled) # gather log_prob if sp > 1 if self.use_ulysses_sp: # gather and unpad for the ulysses sp log_probs = gather_outpus_and_unpad(log_probs, gather_dim=0, unpad_dim=0, padding_size=pad_size) entropy_rmpad = gather_outpus_and_unpad(entropy_rmpad, gather_dim=0, unpad_dim=0, padding_size=pad_size) # pad back to (bsz, seqlen) full_entropy = pad_input(hidden_states=entropy_rmpad.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen) full_log_probs = pad_input(hidden_states=log_probs.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen) # only return response part: entropy = full_entropy.squeeze(-1)[:, -response_length - 1:-1] # (bsz, response_length) log_probs = full_log_probs.squeeze(-1)[:, -response_length - 1:-1] # (bsz, response_length) else: # not using rmpad and no ulysses sp output = self.actor_module(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False) # prevent model thinks we are generating logits = output.logits logits.div_(temperature) logits = logits[:, -response_length - 1:-1] # (bsz, response_length) log_probs = logprobs_from_logits(logits, micro_batch['responses']) entropy = verl_F.entropy_from_logits(logits) # (bsz, response_length) return entropy, log_probs","summary_state":"pending","summary":null,"crux":null,"owner":"DataParallelPPOActor"},{"id":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._optimizer_step","name":"_optimizer_step","kind":"method","path":"verl/workers/actor/dp_actor.py","span":"L144-L152","signature":"def _optimizer_step(self)","exported":false,"origin":"ast","body_hash":"a2584f450b2bdc24abb6f2ac98ee4a471b8289485a87e0c65b42e3a7bb8aaa63","body_text":"def _optimizer_step(self): assert self.config.grad_clip is not None if isinstance(self.actor_module, FSDP): grad_norm = self.actor_module.clip_grad_norm_(max_norm=self.config.grad_clip) else: grad_norm = torch.nn.utils.clip_grad_norm_(self.actor_module.parameters(), max_norm=self.config.grad_clip) self.actor_optimizer.step() return grad_norm","summary_state":"pending","summary":null,"crux":null,"owner":"DataParallelPPOActor"},{"id":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","name":"compute_log_prob","kind":"method","path":"verl/workers/actor/dp_actor.py","span":"L154-L202","signature":"def compute_log_prob(self, data: DataProto) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"7c12aa1696035cb1cf9269760618067d43a50d8a6a39b41cd429b99fdcddecb2","body_text":"def compute_log_prob(self, data: DataProto) -> torch.Tensor: \"\"\"Compute the log probability of the responses given input_ids, attention_mask and position_ids Args: data (DataProto): a DataProto containing keys ``input_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. Note that input_ids is the concatenation of prompt and response. Note that ``sequence_length = prompt_length + response_length``. ``attention_mask``: tensor of shape [batch_size, sequence_length]. torch.int64. ``position_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. ``responses``: tensor of shape [batch_size, response_length]. torch.int64. Returns: torch.Tensor: the log_prob tensor \"\"\" # set to eval self.actor_module.eval() micro_batch_size = data.meta_info['micro_batch_size'] temperature = data.meta_info['temperature'] # temperature must be in the data.meta_info to avoid slient error use_dynamic_bsz = data.meta_info['use_dynamic_bsz'] select_keys = ['responses', 'input_ids', 'attention_mask', 'position_ids'] batch = data.select(batch_keys=select_keys).batch if use_dynamic_bsz: # split using dynamic bsz max_token_len = data.meta_info['max_token_len'] * self.ulysses_sequence_parallel_size micro_batches, indices = rearrange_micro_batches(batch=batch, max_token_len=max_token_len) else: micro_batches = batch.split(micro_batch_size) log_probs_lst = [] for micro_batch in micro_batches: with torch.no_grad(): _, log_probs = self._forward_micro_batch(micro_batch, temperature=temperature) log_probs_lst.append(log_probs) log_probs = torch.concat(log_probs_lst, dim=0) if use_dynamic_bsz: indices = list(itertools.chain.from_iterable(indices)) assert len(indices) == log_probs.size(0), f\"{len(indices)} vs. {log_probs.size()}\" revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long) log_probs = log_probs[revert_indices] return log_probs","summary_state":"pending","summary":null,"crux":null,"owner":"DataParallelPPOActor"},{"id":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","name":"update_policy","kind":"method","path":"verl/workers/actor/dp_actor.py","span":"L204-L291","signature":"def update_policy(self, data: DataProto): # make sure we are in training mode","exported":true,"origin":"ast","body_hash":"ab860d8b1afa146ad8ef64e1d0bc7e4648d60cd926aa51994ce3c13bf0ad1953","body_text":"def update_policy(self, data: DataProto): # make sure we are in training mode self.actor_module.train() assert self.config.ppo_mini_batch_size % self.config.ppo_micro_batch_size == 0 self.gradient_accumulation = self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size temperature = data.meta_info['temperature'] # temperature must be in the data.meta_info to avoid slient error select_keys = ['responses', 'input_ids', 'attention_mask', 'position_ids', 'old_log_probs', 'advantages'] if self.config.state_masking: select_keys.append('loss_mask') if self.config.use_kl_loss: select_keys.append('ref_log_prob') batch = data.select(batch_keys=select_keys).batch # Split to make minibatch iterator for updating the actor # See PPO paper for details. https://arxiv.org/abs/1707.06347 dataloader = batch.split(self.config.ppo_mini_batch_size) metrics = {} for batch_idx, data in enumerate(dataloader): # split batch into micro_batches mini_batch = data if self.config.use_dynamic_bsz: max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size micro_batches, _ = rearrange_micro_batches(batch=mini_batch, max_token_len=max_token_len) else: # split batch into micro_batches micro_batches = mini_batch.split(self.config.ppo_micro_batch_size) self.actor_optimizer.zero_grad() for data in micro_batches: data = data.cuda() # actor device is cpu when using offload responses = data['responses'] response_length = responses.size(1) attention_mask = data['attention_mask'] response_mask = attention_mask[:, -response_length:] if self.config.state_masking: response_mask = data['loss_mask'] old_log_prob = data['old_log_probs'] advantages = data['advantages'] clip_ratio = self.config.clip_ratio entropy_coeff = self.config.entropy_coeff # all return: (bsz, response_length) entropy, log_prob = self._forward_micro_batch(micro_batch=data, temperature=temperature) pg_loss, pg_clipfrac, ppo_kl = core_algos.compute_policy_loss(old_log_prob=old_log_prob, log_prob=log_prob, advantages=advantages, eos_mask=response_mask, cliprange=clip_ratio) # compute entropy loss from entropy entropy_loss = verl_F.masked_mean(entropy, response_mask) # compute policy loss policy_loss = pg_loss - entropy_loss * entropy_coeff if self.config.use_kl_loss: ref_log_prob = data['ref_log_prob'] # compute kl loss kld = core_algos.kl_penalty(logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type) kl_loss = masked_mean(kld, response_mask) policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef metrics['actor/kl_loss'] = kl_loss.detach().item() metrics['actor/kl_coef'] = self.config.kl_loss_coef loss = policy_loss / self.gradient_accumulation loss.backward() data = { 'actor/entropy_loss': entropy_loss.detach().item(), 'actor/pg_loss': pg_loss.detach().item(), 'actor/pg_clipfrac': pg_clipfrac.detach().item(), 'actor/ppo_kl': ppo_kl.detach().item(), } append_to_dict(metrics, data) grad_norm = self._optimizer_step() data = {'actor/grad_norm': grad_norm.detach().item()} append_to_dict(metrics, data) self.actor_optimizer.zero_grad() return metrics","summary_state":"pending","summary":null,"crux":null,"owner":"DataParallelPPOActor"}],"rawEdges":[{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"itertools","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"typing","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"torch","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"torch","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"torch.distributed.fsdp","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"verl","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"verl.trainer.ppo","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"verl.workers.actor","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"verl.utils.py_functional","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"verl.utils.ulysses","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"verl.utils.seqlen_balancing","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"imports","specifier":"flash_attn.bert_padding","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py","relation":"contains","targetId":"verl/workers/actor/dp_actor.py#DataParallelPPOActor","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor","relation":"extends","name":"BasePPOActor","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor","relation":"contains","targetId":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.__init__","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor","relation":"contains","targetId":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"autocast","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"unpad_input","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"transpose","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"transpose","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"index_first_axis","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"rearrange","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"roll","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"ulysses_pad_and_slice_inputs","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"ulysses_pad_and_slice_inputs","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"actor_module","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"DataParallelPPOActor"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"div_","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"compute_entropy_from_logits","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"DataParallelPPOActor"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"logprobs_from_logits","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"gather_outpus_and_unpad","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"gather_outpus_and_unpad","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"gather_outpus_and_unpad"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"logprobs_from_logits"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"pad_input"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"pad_input"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"actor_module","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"DataParallelPPOActor"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"div_","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"logprobs_from_logits","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._forward_micro_batch","relation":"calls","name":"entropy_from_logits","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor","relation":"contains","targetId":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._optimizer_step","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._optimizer_step","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._optimizer_step","relation":"calls","name":"clip_grad_norm_","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._optimizer_step","relation":"calls","name":"clip_grad_norm_","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._optimizer_step","relation":"calls","name":"parameters","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor._optimizer_step","relation":"calls","name":"step","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor","relation":"contains","targetId":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"eval","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"select","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"DataProto"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"rearrange_micro_batches","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"split","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"_forward_micro_batch","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"DataParallelPPOActor"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"concat","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"list","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"from_iterable","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"tensor","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.compute_log_prob","relation":"calls","name":"get_reverse_idx","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor","relation":"contains","targetId":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"train","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"select","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"DataProto"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"split","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"rearrange_micro_batches","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"split","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"zero_grad","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"DataProto"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"_forward_micro_batch","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"DataParallelPPOActor"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"compute_policy_loss","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"masked_mean","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"kl_penalty","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"masked_mean","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"masked_mean"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"backward","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"append_to_dict","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"_optimizer_step","viaMember":true,"file":"verl/workers/actor/dp_actor.py","recvType":"DataParallelPPOActor"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"append_to_dict","viaMember":false,"file":"verl/workers/actor/dp_actor.py"},{"source":"verl/workers/actor/dp_actor.py#DataParallelPPOActor.update_policy","relation":"calls","name":"zero_grad","viaMember":true,"file":"verl/workers/actor/dp_actor.py"}]},"verl/workers/actor/megatron_actor.py":{"size":18814,"mtimeMs":1789275830327.2031,"hash":"4e06adb8ae9802223b7d6c2ccd7658ec71c2286aaeb4139eeeebbf4319dc3df4","nodes":[{"id":"verl/workers/actor/megatron_actor.py","name":"megatron_actor.py","kind":"file","path":"verl/workers/actor/megatron_actor.py","span":"L1-L369","signature":null,"exported":true,"origin":"ast","body_hash":"4e06adb8ae9802223b7d6c2ccd7658ec71c2286aaeb4139eeeebbf4319dc3df4","chars":18814,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Megatron Actor. In megatron actor, the differences are: 1. We only make minibatch Note that our model doesn't have to be `MegatronModule` because we don't share embedding in the last layer \"\"\" from functools import partial from typing import Iterable, Dict import torch from torch import nn import torch.distributed # from megatron import get_args from megatron.optimizer import DistributedOptimizer from verl.utils.megatron.optimizer_config import OptimizerConfig from megatron.core import parallel_state as mpu from megatron.core import ModelParallelConfig from megatron.core.pipeline_parallel import get_forward_backward_func # from megatron.core.optimizer import DistributedOptimizer from omegaconf import OmegaConf from verl.utils.megatron.tensor_parallel import vocab_parallel_compute_entropy_loss, vocab_parallel_log_probs_from_logits from verl.utils.megatron.pipeline_parallel import (compute_transformers_input_shapes, make_batch_generator) from verl import DataProto from verl.trainer.ppo import core_algos from verl.workers.actor import BasePPOActor from verl.utils.py_functional import append_to_dict from verl.utils.torch_functional import logprobs_from_logits, broadcast_dict_tensor, split_dict_tensor_into_batches __all__ = ['MegatronPPOActor']"},{"id":"verl/workers/actor/megatron_actor.py#MegatronPPOActor","name":"MegatronPPOActor","kind":"class","path":"verl/workers/actor/megatron_actor.py","span":"L48-L368","signature":"class MegatronPPOActor(BasePPOActor)","exported":true,"origin":"ast","body_hash":"fdd5e7a73da8cd9dc379930daa3ac6ca521984010430ba057b8d8776a5b26d65","body_text":"class MegatronPPOActor(BasePPOActor): def __init__(self, config, model_config, megatron_config: ModelParallelConfig, actor_module: nn.ModuleList, actor_optimizer: DistributedOptimizer, actor_optimizer_config: OptimizerConfig): \"\"\"MeagtronPPOActor class. This class implements the simple PPO logics when the model is built with Megatron. Args: config (OmegaConf): the basic config that contains the hyper-parameters of PPO Actor. It must contain ``ppo_micro_batch_size``: minibatch size when updating ppo. ``ppo_mini_batch_size``: minibatch size when updating ppo using the batch data. ``ppo_epochs``: number of epochs to update the actor using the batch data. ``shuffle``: whether to shuffle the data after each ppo epoch. ``clip_ratio``: clip ratio of the ppo algorithm. See https://arxiv.org/abs/1707.06347. ``entropy_coeff``: entropy coefficient of the PPO loss. See https://arxiv.org/abs/1707.06347. model_config (OmegaConf): model configuration. It must contains ``model_config.vocab_size`` and ``model_config.hidden_size`` megatron_config (OmegaConf): megatron configuration. It must contains ``sequence_parallel_enabled``: whether the sequence parallel is enabled. ``param_dtype``: the dtype of the parameters. ``virtual_pipeline_model_parallel_size``: virtual pipeline model parallel size. a.k.a number of chunks in each pp stage. actor_module (nn.ModuleList): actor module is a ModuleList that contains a list of nn.Module in this pp stage. each nn.Module in this rank holds a vpp module chunk. See https://arxiv.org/pdf/2104.04473.pdf for more details. The actor module has some constraints to follow in order to use the updating logics implemented here 1. It must implement unpad_input before any computation and pad_input after all the computation. Remove padding is an optimization that removes the padding tokens. See unpad_input and pad_input function in flash-attn (https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py). 2. Each pp stage must return the hidden state with the same shape [total_nnz, 1, hidden_size], where total_nnz is the number of valid tokens in this batch. If sequence parallel is enabled, the size of the hidden state is [total_nnz // tp, 1, hidden_size]. actor_optimizer (DistributedOptimizer): currently, we only support DistributedOptimizer in Megatron. It implements zero1 optimizer that shards the optimizer state across dp ranks. >>> def megatron_actor_model_provider(pre_process, post_process): >>> vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() >>> parallel_model = ParallelMistralForCausalLMRmPadPP(config=actor_model_config, >>> megatron_config=megatron_config, >>> pre_process=pre_process, >>> post_process=post_process).cuda() >>> return parallel_model >>> from megatron.training import get_model >>> from megatron.optimizer import get_megatron_optimizer >>> actor_module = get_model(megatron_actor_model_provider, wrap_with_ddp=True) >>> actor_module = nn.ModuleList(actor_module) >>> actor_optimizer = get_megatron_optimizer(actor_module) >>> actor = MegatronPPOActor(config=config, >>> model_config=actor_model_config, >>> megatron_config=megatron_config, >>> actor_module=actor_module, >>> actor_optimizer=actor_optimizer) \"\"\" super().__init__(config) self.model_config = model_config self.megatron_config = megatron_config # self.megatron_args = get_args() self.actor_module = actor_module self.actor_optimizer: DistributedOptimizer = actor_optimizer self.actor_optimizer_config = actor_optimizer_config self.optimizer_step_args = OmegaConf.create({ 'skip_grad': None, 'overlap_dp_param_comm': False, 'overlap_dp_grad_comm': False, 'gradient_accumulation_steps': 1, 'sequence_parallel': self.megatron_config.sequence_parallel, 'DDP_impl': 'local', 'layernorm_allreduce_bucket_threshold': 0, 'pipeline_model_parallel_split_rank': None, 'reduce_grads_use_alltoall': False }) def compute_log_prob(self, data: DataProto) -> torch.Tensor: \"\"\"Compute the log probability of the responses given input_ids, attention_mask and position_ids Args: data (DataProto): a DataProto containing keys ``input_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. Note that input_ids is the concatenation of prompt and response. Note that ``sequence_length = prompt_length + response_length``. ``attention_mask``: tensor of shape [batch_size, sequence_length]. torch.int64. ``position_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. ``responses``: tensor of shape [batch_size, response_length]. torch.int64. Returns: DataProto: torch.Tensor: the log_prob tensor \"\"\" data.batch = data.batch.contiguous() def compute_logprobs_fn(output, data): response = data['responses'] response_length = response.size(1) logits = output['logits'] logits = logits[:, -response_length - 1:-1] log_probs = vocab_parallel_log_probs_from_logits(logits, response) return {'log_probs': log_probs} # We make recompute_old_log_prob by default here. # TODO (zhangchi.usc1992): actually, this function should only retur","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.__init__","name":"__init__","kind":"method","path":"verl/workers/actor/megatron_actor.py","span":"L50-L127","signature":"def __init__(self, config, model_config, megatron_config: ModelParallelConfig, actor_module: nn.ModuleList, actor_optimizer: DistributedOptimizer, actor_optimizer_config: OptimizerConfig)","exported":false,"origin":"ast","body_hash":"2fc98f2337f6707c846dc4d4545d5692a9864cbd615c9e0ba86162f7ae2621dd","body_text":"def __init__(self, config, model_config, megatron_config: ModelParallelConfig, actor_module: nn.ModuleList, actor_optimizer: DistributedOptimizer, actor_optimizer_config: OptimizerConfig): \"\"\"MeagtronPPOActor class. This class implements the simple PPO logics when the model is built with Megatron. Args: config (OmegaConf): the basic config that contains the hyper-parameters of PPO Actor. It must contain ``ppo_micro_batch_size``: minibatch size when updating ppo. ``ppo_mini_batch_size``: minibatch size when updating ppo using the batch data. ``ppo_epochs``: number of epochs to update the actor using the batch data. ``shuffle``: whether to shuffle the data after each ppo epoch. ``clip_ratio``: clip ratio of the ppo algorithm. See https://arxiv.org/abs/1707.06347. ``entropy_coeff``: entropy coefficient of the PPO loss. See https://arxiv.org/abs/1707.06347. model_config (OmegaConf): model configuration. It must contains ``model_config.vocab_size`` and ``model_config.hidden_size`` megatron_config (OmegaConf): megatron configuration. It must contains ``sequence_parallel_enabled``: whether the sequence parallel is enabled. ``param_dtype``: the dtype of the parameters. ``virtual_pipeline_model_parallel_size``: virtual pipeline model parallel size. a.k.a number of chunks in each pp stage. actor_module (nn.ModuleList): actor module is a ModuleList that contains a list of nn.Module in this pp stage. each nn.Module in this rank holds a vpp module chunk. See https://arxiv.org/pdf/2104.04473.pdf for more details. The actor module has some constraints to follow in order to use the updating logics implemented here 1. It must implement unpad_input before any computation and pad_input after all the computation. Remove padding is an optimization that removes the padding tokens. See unpad_input and pad_input function in flash-attn (https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py). 2. Each pp stage must return the hidden state with the same shape [total_nnz, 1, hidden_size], where total_nnz is the number of valid tokens in this batch. If sequence parallel is enabled, the size of the hidden state is [total_nnz // tp, 1, hidden_size]. actor_optimizer (DistributedOptimizer): currently, we only support DistributedOptimizer in Megatron. It implements zero1 optimizer that shards the optimizer state across dp ranks. >>> def megatron_actor_model_provider(pre_process, post_process): >>> vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() >>> parallel_model = ParallelMistralForCausalLMRmPadPP(config=actor_model_config, >>> megatron_config=megatron_config, >>> pre_process=pre_process, >>> post_process=post_process).cuda() >>> return parallel_model >>> from megatron.training import get_model >>> from megatron.optimizer import get_megatron_optimizer >>> actor_module = get_model(megatron_actor_model_provider, wrap_with_ddp=True) >>> actor_module = nn.ModuleList(actor_module) >>> actor_optimizer = get_megatron_optimizer(actor_module) >>> actor = MegatronPPOActor(config=config, >>> model_config=actor_model_config, >>> megatron_config=megatron_config, >>> actor_module=actor_module, >>> actor_optimizer=actor_optimizer) \"\"\" super().__init__(config) self.model_config = model_config self.megatron_config = megatron_config # self.megatron_args = get_args() self.actor_module = actor_module self.actor_optimizer: DistributedOptimizer = actor_optimizer self.actor_optimizer_config = actor_optimizer_config self.optimizer_step_args = OmegaConf.create({ 'skip_grad': None, 'overlap_dp_param_comm': False, 'overlap_dp_grad_comm': False, 'gradient_accumulation_steps': 1, 'sequence_parallel': self.megatron_config.sequence_parallel, 'DDP_impl': 'local', 'layernorm_allreduce_bucket_threshold': 0, 'pipeline_model_parallel_split_rank': None, 'reduce_grads_use_alltoall': False })","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronPPOActor"},{"id":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","name":"compute_log_prob","kind":"method","path":"verl/workers/actor/megatron_actor.py","span":"L129-L188","signature":"def compute_log_prob(self, data: DataProto) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"ea05a36e579b85baa4984ca064931af2324ef32f55017df441587e988af00425","body_text":"def compute_log_prob(self, data: DataProto) -> torch.Tensor: \"\"\"Compute the log probability of the responses given input_ids, attention_mask and position_ids Args: data (DataProto): a DataProto containing keys ``input_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. Note that input_ids is the concatenation of prompt and response. Note that ``sequence_length = prompt_length + response_length``. ``attention_mask``: tensor of shape [batch_size, sequence_length]. torch.int64. ``position_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. ``responses``: tensor of shape [batch_size, response_length]. torch.int64. Returns: DataProto: torch.Tensor: the log_prob tensor \"\"\" data.batch = data.batch.contiguous() def compute_logprobs_fn(output, data): response = data['responses'] response_length = response.size(1) logits = output['logits'] logits = logits[:, -response_length - 1:-1] log_probs = vocab_parallel_log_probs_from_logits(logits, response) return {'log_probs': log_probs} # We make recompute_old_log_prob by default here. # TODO (zhangchi.usc1992): actually, this function should only return log_prob and this logic should be handled by user outside recompute_old_log_prob = self.config.get('recompute_old_log_prob', True) if recompute_old_log_prob or 'old_log_probs' not in data.batch.keys(): select_keys = ['responses', 'input_ids', 'attention_mask', 'position_ids'] batch = data.select(batch_keys=select_keys).batch input_ids = batch['input_ids'] batch_size = input_ids.size(0) response = batch['responses'] response_length = response.size(1) with torch.no_grad(): output = self.forward_backward_batch(data, forward_only=True, post_process_fn=compute_logprobs_fn) if mpu.is_pipeline_last_stage(ignore_virtual=True): # only on last rank. It should be on every tp rank log_probs = torch.cat([o['log_probs'] for o in output], dim=0) # (bs, seq_size) log_probs = log_probs.to(torch.float32) else: log_probs = torch.empty(size=(batch_size, response_length), dtype=torch.float32, device=input_ids.device) # broadcast across pp ranks torch.distributed.broadcast(tensor=log_probs, src=mpu.get_pipeline_model_parallel_last_rank(), group=mpu.get_pipeline_model_parallel_group(), async_op=False) # add empty cache after each compute torch.cuda.empty_cache() return log_probs","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronPPOActor"},{"id":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob.compute_logprobs_fn","name":"compute_logprobs_fn","kind":"function","path":"verl/workers/actor/megatron_actor.py","span":"L149-L155","signature":"def compute_logprobs_fn(output, data)","exported":true,"origin":"ast","body_hash":"6261cf883520453a536b313bde62265dd3387b4bc23d63dd063ce98575c0d3c9","body_text":"def compute_logprobs_fn(output, data): response = data['responses'] response_length = response.size(1) logits = output['logits'] logits = logits[:, -response_length - 1:-1] log_probs = vocab_parallel_log_probs_from_logits(logits, response) return {'log_probs': log_probs}","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.make_minibatch_iterator","name":"make_minibatch_iterator","kind":"method","path":"verl/workers/actor/megatron_actor.py","span":"L190-L216","signature":"def make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]","exported":true,"origin":"ast","body_hash":"719224233743aa8c60b8a3df510b16f9da77a4865ff50bde05c67b2b1d95951f","body_text":"def make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]: \"\"\"Make minibatch iterator for updating the actor Args: data (DataProto): a DataProto containing keys ``input_ids``: tensor of shape [batch_size, sequence_length]. torch.int64, where ``sequence_length = prompt_length + response_length`` ``attention_mask``: tensor of shape [batch_size, sequence_length]. torch.int64 ``position_ids``: tensor of shape [batch_size, sequence_length]. torch.int64 ``responses``: tensor of shape [batch_size, response_length]. torch.int64. Note that responses = input_ids[:, -response_length:] ``old_log_probs``: tensor of shape [batch_size, response_length]. torch.float32. The log probability of responses. ``advantages``: tensor of shape [batch_size, response_length]. torch.float32. The advantages of responses. See PPO paper for details. https://arxiv.org/abs/1707.06347 Returns: \"\"\" select_keys = ['responses', 'input_ids', 'attention_mask', 'position_ids', 'old_log_probs', 'advantages'] data = data.select(batch_keys=select_keys) return data.make_iterator(mini_batch_size=self.config.ppo_mini_batch_size, epochs=self.config.ppo_epochs, dataloader_kwargs={'shuffle': self.config.shuffle})","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronPPOActor"},{"id":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","name":"forward_backward_batch","kind":"method","path":"verl/workers/actor/megatron_actor.py","span":"L218-L327","signature":"def forward_backward_batch(self, data: DataProto, forward_only=False, post_process_fn=None)","exported":true,"origin":"ast","body_hash":"63767fea86eb865ec0444a13293d65b61c123cbd7b86791c248f09879014e043","body_text":"def forward_backward_batch(self, data: DataProto, forward_only=False, post_process_fn=None): \"\"\" We assume: - The model takes input: (input_ids, attention_mask, position_ids). No rmpad for the input - The communication shape is (total_nnz_pad_to_sp // tp_size, 1, hidden_size) if sequence parallel is enabled \"\"\" # broadcast from last pp rank to all other pp ranks # TODO: actually, we just need to control the sampling order. broadcast_dict_tensor(data.batch, src=mpu.get_pipeline_model_parallel_last_rank(), group=mpu.get_pipeline_model_parallel_group()) # split into micro-batches data.batch['attention_mask'] = data.batch['attention_mask'].to(bool) if data.meta_info.get('micro_batch_size', None) is not None: batch_size = data.meta_info['micro_batch_size'] else: batch_size = self.config.ppo_micro_batch_size batches = split_dict_tensor_into_batches(data.batch, batch_size=batch_size) # compute input shapes for pp stages input_shapes = compute_transformers_input_shapes( batches, meta_info={ 'sequence_parallel': self.megatron_config.sequence_parallel, 'hidden_size': self.model_config.hidden_size }) n_micro_batch = len(batches) seq_len = batches[0]['input_ids'].shape[1] forward_backward_func = get_forward_backward_func() def loss_func(output, data, meta_info): if forward_only: if post_process_fn is None: return 1.0, {'logits': output.logits} else: return 1.0, post_process_fn(output, data) responses = data['responses'] response_length = responses.size(1) attention_mask = data['attention_mask'] response_mask = attention_mask[:, -response_length:] old_log_prob = data['old_log_probs'] advantages = data['advantages'] clip_ratio = meta_info['clip_ratio'] entropy_coeff = meta_info['entropy_coeff'] # compute policy loss logits = output.logits logits = logits[:, -response_length - 1:-1] log_prob = vocab_parallel_log_probs_from_logits(logits, responses) pg_loss, pg_clipfrac, ppo_kl = core_algos.compute_policy_loss(old_log_prob=old_log_prob, log_prob=log_prob, advantages=advantages, eos_mask=response_mask, cliprange=clip_ratio) entropy_loss = vocab_parallel_compute_entropy_loss(logits, eos_mask=response_mask) policy_loss = pg_loss - entropy_loss * entropy_coeff # return loss and stats stats = { 'actor/entropy_loss': entropy_loss.detach().item(), 'actor/pg_loss': pg_loss.detach().item(), 'actor/pg_clipfrac': pg_clipfrac.detach().item(), 'actor/ppo_kl': ppo_kl.detach().item() } return policy_loss, stats def forward_step(batch_iter, model): batch = next(batch_iter) input_ids = batch['input_ids'] attention_mask = batch['attention_mask'] position_ids = batch['position_ids'] output = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) if forward_only: meta_info = None else: meta_info = {'clip_ratio': self.config.clip_ratio, 'entropy_coeff': self.config.entropy_coeff} return output, partial(loss_func, data=batch, meta_info=meta_info) # batch should be a list of batches inside micro-batches batch_generator = make_batch_generator(batches, vpp_size=len(self.actor_module)) # TODO: we may use the new schedule instead # for flash-attn: (seq_len, batch_size, hidden_size) = (mbs*seq_len, 1, hidden_size) if mpu.get_pipeline_model_parallel_world_size() > 1: losses_reduced = forward_backward_func( forward_step_func=forward_step, data_iterator=batch_generator, model=self.actor_module, num_microbatches=n_micro_batch, input_shapes=input_shapes, # must set for flash-attn sequence packing seq_length=batch_size * seq_len, # no use when input_shapes was set hidden_size=self.model_config.hidden_size, # no use when input_shapes was set micro_batch_size=1, # no use when input_shapes was set forward_only=forward_only, ) else: losses_reduced = forward_backward_func( forward_step_func=forward_step, data_iterator=batch_generator, model=self.actor_module, num_microbatches=n_micro_batch, seq_length=batch_size * seq_len, # in use for pp = 1 hidden_size=self.model_config.hidden_size, # in use for pp = 1 micro_batch_size=1, # in use for pp = 1 forward_only=forward_only, ) # loss_reduces contains the stats returned from loss_func return losses_reduced","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronPPOActor"},{"id":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","name":"loss_func","kind":"function","path":"verl/workers/actor/megatron_actor.py","span":"L249-L284","signature":"def loss_func(output, data, meta_info)","exported":true,"origin":"ast","body_hash":"6cc03128046ec6e525184d22b4bc3e9d3bf1f7608a5fe2e9a40f8e5180be5f90","body_text":"def loss_func(output, data, meta_info): if forward_only: if post_process_fn is None: return 1.0, {'logits': output.logits} else: return 1.0, post_process_fn(output, data) responses = data['responses'] response_length = responses.size(1) attention_mask = data['attention_mask'] response_mask = attention_mask[:, -response_length:] old_log_prob = data['old_log_probs'] advantages = data['advantages'] clip_ratio = meta_info['clip_ratio'] entropy_coeff = meta_info['entropy_coeff'] # compute policy loss logits = output.logits logits = logits[:, -response_length - 1:-1] log_prob = vocab_parallel_log_probs_from_logits(logits, responses) pg_loss, pg_clipfrac, ppo_kl = core_algos.compute_policy_loss(old_log_prob=old_log_prob, log_prob=log_prob, advantages=advantages, eos_mask=response_mask, cliprange=clip_ratio) entropy_loss = vocab_parallel_compute_entropy_loss(logits, eos_mask=response_mask) policy_loss = pg_loss - entropy_loss * entropy_coeff # return loss and stats stats = { 'actor/entropy_loss': entropy_loss.detach().item(), 'actor/pg_loss': pg_loss.detach().item(), 'actor/pg_clipfrac': pg_clipfrac.detach().item(), 'actor/ppo_kl': ppo_kl.detach().item() } return policy_loss, stats","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.forward_step","name":"forward_step","kind":"function","path":"verl/workers/actor/megatron_actor.py","span":"L286-L296","signature":"def forward_step(batch_iter, model)","exported":true,"origin":"ast","body_hash":"0d1e7f5fec7eaed167d4e8706e16c348e4ab0323ccac317a0580dc032640db23","body_text":"def forward_step(batch_iter, model): batch = next(batch_iter) input_ids = batch['input_ids'] attention_mask = batch['attention_mask'] position_ids = batch['position_ids'] output = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) if forward_only: meta_info = None else: meta_info = {'clip_ratio': self.config.clip_ratio, 'entropy_coeff': self.config.entropy_coeff} return output, partial(loss_func, data=batch, meta_info=meta_info)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.update_policy","name":"update_policy","kind":"method","path":"verl/workers/actor/megatron_actor.py","span":"L329-L368","signature":"def update_policy(self, dataloader: Iterable[DataProto]) -> Dict","exported":true,"origin":"ast","body_hash":"f2e8dbbe35cb03a13fa8d8479a74752c7edb006d202b5cc280078bc57071cb65","body_text":"def update_policy(self, dataloader: Iterable[DataProto]) -> Dict: \"\"\"Update the policy with an iterator of DataProto Args: dataloader (Iterable[DataProto]): an iterator over the DataProto that returns by ``make_minibatch_iterator`` The keys of each data batch is described in the make_minibatch_iterator. Returns: Dict: a dictionary containing the statistics. Note that the statistics are only valid in the last pp stage and users have to combine the output in each dp rank manually. \"\"\" metrics = {} for data in dataloader: # data = data.batch.to(self.actor_module.device) self.actor_optimizer.zero_grad() # use use_contiguous_buffers_in_local_ddp and no overlap_dp_param_comm for chunk in self.actor_module: # if use distributed optimizer, zero grad buffer will be handled by optimizer chunk.zero_grad_buffer(zero_buffer=(not self.actor_optimizer_config.use_distributed_optimizer)) metric_micro_batch = self.forward_backward_batch(data) for metric in metric_micro_batch: append_to_dict(metrics, metric) # append the metric from this micro-batch to global metrics. update_successful, grad_norm, num_zeros_in_grad = self.actor_optimizer.step( self.megatron_config, self.megatron_config.timers) if update_successful: # allgather already execute in optimizer.step in new megatron pass else: raise NotImplementedError for metric in metric_micro_batch: append_to_dict(metrics, metric) # append the metric from this micro-batch to global metrics. # add empty cache after each compute torch.cuda.empty_cache() return metrics","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronPPOActor"}],"rawEdges":[{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"functools","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"typing","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"torch","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"torch","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"megatron.optimizer","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"verl.utils.megatron.optimizer_config","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"megatron.core","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"megatron.core","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"megatron.core.pipeline_parallel","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"omegaconf","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"verl.utils.megatron.tensor_parallel","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"verl.utils.megatron.pipeline_parallel","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"verl","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"verl.trainer.ppo","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"verl.workers.actor","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"verl.utils.py_functional","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py","relation":"contains","targetId":"verl/workers/actor/megatron_actor.py#MegatronPPOActor","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor","relation":"extends","name":"BasePPOActor","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor","relation":"contains","targetId":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.__init__","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.__init__","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor","relation":"contains","targetId":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"contains","targetId":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob.compute_logprobs_fn","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob.compute_logprobs_fn","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob.compute_logprobs_fn","relation":"calls","name":"vocab_parallel_log_probs_from_logits","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"keys","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"select","viaMember":true,"file":"verl/workers/actor/megatron_actor.py","recvType":"DataProto"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"forward_backward_batch","viaMember":true,"file":"verl/workers/actor/megatron_actor.py","recvType":"MegatronPPOActor"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"is_pipeline_last_stage","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"empty","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"get_pipeline_model_parallel_last_rank","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"get_pipeline_model_parallel_group","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.compute_log_prob","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor","relation":"contains","targetId":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.make_minibatch_iterator","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.make_minibatch_iterator","relation":"calls","name":"select","viaMember":true,"file":"verl/workers/actor/megatron_actor.py","recvType":"DataProto"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.make_minibatch_iterator","relation":"calls","name":"make_iterator","viaMember":true,"file":"verl/workers/actor/megatron_actor.py","recvType":"DataProto"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor","relation":"contains","targetId":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"broadcast_dict_tensor","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"get_pipeline_model_parallel_last_rank","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"get_pipeline_model_parallel_group","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"split_dict_tensor_into_batches","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"compute_transformers_input_shapes","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"get_forward_backward_func","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"contains","targetId":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"post_process_fn","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"vocab_parallel_log_probs_from_logits","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"compute_policy_loss","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"vocab_parallel_compute_entropy_loss","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/actor/megatron_actor.py","recvType":"vocab_parallel_compute_entropy_loss"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.loss_func","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"contains","targetId":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.forward_step","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.forward_step","relation":"calls","name":"next","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.forward_step","relation":"calls","name":"model","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch.forward_step","relation":"calls","name":"partial","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"make_batch_generator","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"forward_backward_func","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.forward_backward_batch","relation":"calls","name":"forward_backward_func","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor","relation":"contains","targetId":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.update_policy","file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.update_policy","relation":"calls","name":"zero_grad","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.update_policy","relation":"calls","name":"zero_grad_buffer","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.update_policy","relation":"calls","name":"forward_backward_batch","viaMember":true,"file":"verl/workers/actor/megatron_actor.py","recvType":"MegatronPPOActor"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.update_policy","relation":"calls","name":"append_to_dict","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.update_policy","relation":"calls","name":"step","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.update_policy","relation":"calls","name":"append_to_dict","viaMember":false,"file":"verl/workers/actor/megatron_actor.py"},{"source":"verl/workers/actor/megatron_actor.py#MegatronPPOActor.update_policy","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/actor/megatron_actor.py"}]},"verl/workers/critic/__init__.py":{"size":732,"mtimeMs":1789275830327.2031,"hash":"a61c624ef2dd6bec59cc59302d88d1ccdc151ea5beb776b9733ac19bcfbe8fb7","nodes":[{"id":"verl/workers/critic/__init__.py","name":"__init__.py","kind":"file","path":"verl/workers/critic/__init__.py","span":"L1-L19","signature":null,"exported":true,"origin":"ast","body_hash":"a61c624ef2dd6bec59cc59302d88d1ccdc151ea5beb776b9733ac19bcfbe8fb7","chars":732,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 BasePPOCritic from .dp_critic import DataParallelPPOCritic __all__ = [\"BasePPOCritic\", \"DataParallelPPOCritic\"]"}],"rawEdges":[{"source":"verl/workers/critic/__init__.py","relation":"imports","specifier":".base","file":"verl/workers/critic/__init__.py"},{"source":"verl/workers/critic/__init__.py","relation":"imports","specifier":".dp_critic","file":"verl/workers/critic/__init__.py"}]},"verl/workers/critic/base.py":{"size":1094,"mtimeMs":1789275830327.2031,"hash":"9609662cb80dc02afb53a424abb6a9ae93e0448d190a3dd5bc4cd5e0d92ea6a2","nodes":[{"id":"verl/workers/critic/base.py","name":"base.py","kind":"file","path":"verl/workers/critic/base.py","span":"L1-L41","signature":null,"exported":true,"origin":"ast","body_hash":"9609662cb80dc02afb53a424abb6a9ae93e0448d190a3dd5bc4cd5e0d92ea6a2","chars":1094,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Base class for a critic \"\"\" from abc import ABC, abstractmethod import torch from verl import DataProto __all__ = ['BasePPOCritic']"},{"id":"verl/workers/critic/base.py#BasePPOCritic","name":"BasePPOCritic","kind":"class","path":"verl/workers/critic/base.py","span":"L26-L40","signature":"class BasePPOCritic(ABC)","exported":true,"origin":"ast","body_hash":"32f57186bf844b37a9fe896effe157869eca8cadae245d3a3229bc472f933c49","body_text":"class BasePPOCritic(ABC): def __init__(self, config): super().__init__() self.config = config @abstractmethod def compute_values(self, data: DataProto) -> torch.Tensor: \"\"\"Compute values\"\"\" pass @abstractmethod def update_critic(self, data: DataProto): \"\"\"Update the critic\"\"\" pass","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/critic/base.py#BasePPOCritic.__init__","name":"__init__","kind":"method","path":"verl/workers/critic/base.py","span":"L28-L30","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"cf080aee1d9f1d0678095acc910739ef8971d44f656ba0d7db1249407408b911","body_text":"def __init__(self, config): super().__init__() self.config = config","summary_state":"pending","summary":null,"crux":null,"owner":"BasePPOCritic"},{"id":"verl/workers/critic/base.py#BasePPOCritic.compute_values","name":"compute_values","kind":"method","path":"verl/workers/critic/base.py","span":"L33-L35","signature":"def compute_values(self, data: DataProto) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"9c09a6cafc8f3c08f35cba23e61fc972cdd1191b28502e172640887d0ba94ff9","body_text":"def compute_values(self, data: DataProto) -> torch.Tensor: \"\"\"Compute values\"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"BasePPOCritic"},{"id":"verl/workers/critic/base.py#BasePPOCritic.update_critic","name":"update_critic","kind":"method","path":"verl/workers/critic/base.py","span":"L38-L40","signature":"def update_critic(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"ff4296921cdb634b825cf6c11fe2a71571b57aed72720c291f07854090cf7531","body_text":"def update_critic(self, data: DataProto): \"\"\"Update the critic\"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"BasePPOCritic"}],"rawEdges":[{"source":"verl/workers/critic/base.py","relation":"imports","specifier":"abc","file":"verl/workers/critic/base.py"},{"source":"verl/workers/critic/base.py","relation":"imports","specifier":"torch","file":"verl/workers/critic/base.py"},{"source":"verl/workers/critic/base.py","relation":"imports","specifier":"verl","file":"verl/workers/critic/base.py"},{"source":"verl/workers/critic/base.py","relation":"contains","targetId":"verl/workers/critic/base.py#BasePPOCritic","file":"verl/workers/critic/base.py"},{"source":"verl/workers/critic/base.py#BasePPOCritic","relation":"extends","name":"ABC","file":"verl/workers/critic/base.py"},{"source":"verl/workers/critic/base.py#BasePPOCritic","relation":"contains","targetId":"verl/workers/critic/base.py#BasePPOCritic.__init__","file":"verl/workers/critic/base.py"},{"source":"verl/workers/critic/base.py#BasePPOCritic.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/critic/base.py"},{"source":"verl/workers/critic/base.py#BasePPOCritic.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/critic/base.py"},{"source":"verl/workers/critic/base.py#BasePPOCritic","relation":"contains","targetId":"verl/workers/critic/base.py#BasePPOCritic.compute_values","file":"verl/workers/critic/base.py"},{"source":"verl/workers/critic/base.py#BasePPOCritic","relation":"contains","targetId":"verl/workers/critic/base.py#BasePPOCritic.update_critic","file":"verl/workers/critic/base.py"}]},"verl/workers/critic/dp_critic.py":{"size":9785,"mtimeMs":1789275830327.2031,"hash":"53bd4fcebb5aee477d24bb2e9d40d0f1fc94f57677cb918873cc2984234229a3","nodes":[{"id":"verl/workers/critic/dp_critic.py","name":"dp_critic.py","kind":"file","path":"verl/workers/critic/dp_critic.py","span":"L1-L205","signature":null,"exported":true,"origin":"ast","body_hash":"53bd4fcebb5aee477d24bb2e9d40d0f1fc94f57677cb918873cc2984234229a3","chars":9785,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Implement a multiprocess PPOCritic \"\"\" import itertools from typing import Iterable import torch import torch.distributed from torch import nn, optim from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from verl import DataProto from verl.trainer.ppo import core_algos from verl.workers.critic import BasePPOCritic from verl.utils.py_functional import append_to_dict from verl.utils.torch_functional import masked_mean from verl.utils.ulysses import ulysses_pad_and_slice_inputs, gather_outpus_and_unpad from verl.utils.seqlen_balancing import rearrange_micro_batches, get_reverse_idx from flash_attn.bert_padding import pad_input, unpad_input, rearrange, index_first_axis __all__ = ['DataParallelPPOCritic']"},{"id":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic","name":"DataParallelPPOCritic","kind":"class","path":"verl/workers/critic/dp_critic.py","span":"L39-L204","signature":"class DataParallelPPOCritic(BasePPOCritic)","exported":true,"origin":"ast","body_hash":"428279817a4e284e3a22737870f85212fc005b847d06274306b62e743b6c7c51","body_text":"class DataParallelPPOCritic(BasePPOCritic): def __init__(self, config, critic_module: nn.Module, critic_optimizer: optim.Optimizer): super().__init__(config=config) self.critic_module = critic_module self.critic_optimizer = critic_optimizer self.use_remove_padding = self.config.model.get('use_remove_padding', False) print(f'Critic use_remove_padding={self.use_remove_padding}') assert self.config.ppo_mini_batch_size % self.config.ppo_micro_batch_size == 0 self.gradient_accumulation = self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size self.ulysses_sequence_parallel_size = self.config.get('ulysses_sequence_parallel_size', 1) def _forward_micro_batch(self, micro_batch): response_length = micro_batch['responses'].size(-1) with torch.autocast(device_type='cuda', dtype=torch.bfloat16): input_ids = micro_batch['input_ids'] batch, seqlen = input_ids.shape attention_mask = micro_batch['attention_mask'] position_ids = micro_batch['position_ids'] if self.use_remove_padding: 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) # pad and slice the inputs if sp > 1 if self.ulysses_sequence_parallel_size > 1: input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs(input_ids_rmpad, \\ position_ids_rmpad, \\ sp_size=self.ulysses_sequence_parallel_size) # only pass input_ids and position_ids to enable flash_attn_varlen output = self.critic_module(input_ids=input_ids_rmpad, attention_mask=None, position_ids=position_ids_rmpad, use_cache=False) # prevent model thinks we are generating values_rmpad = output.logits values_rmpad = values_rmpad.squeeze(0) # (total_nnz) # gather output if sp > 1 if self.ulysses_sequence_parallel_size > 1: values_rmpad = gather_outpus_and_unpad(values_rmpad, gather_dim=0, unpad_dim=0, padding_size=pad_size) # pad it back values = pad_input(values_rmpad, indices=indices, batch=batch, seqlen=seqlen).squeeze(-1) values = values[:, -response_length - 1:-1] else: output = self.critic_module(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False) # prevent model thinks we are generating values = output.logits values = values[:, -response_length - 1:-1].squeeze(-1) return values def _optimizer_step(self): assert self.config.grad_clip is not None if isinstance(self.critic_module, FSDP): grad_norm = self.critic_module.clip_grad_norm_(self.config.grad_clip) else: grad_norm = torch.nn.utils.clip_grad_norm_(self.critic_module.parameters(), max_norm=self.config.grad_clip) self.critic_optimizer.step() return grad_norm def compute_values(self, data: DataProto) -> torch.Tensor: self.critic_module.eval() micro_batch_size = data.meta_info['micro_batch_size'] select_keys = ['responses', 'input_ids', 'attention_mask', 'position_ids'] batch = data.select(batch_keys=select_keys).batch use_dynamic_bsz = data.meta_info['use_dynamic_bsz'] if use_dynamic_bsz: # split using dynamic bsz max_token_len = data.meta_info['max_token_len'] * self.ulysses_sequence_parallel_size micro_batches, indices = rearrange_micro_batches(batch=batch, max_token_len=max_token_len) else: micro_batches = batch.split(micro_batch_size) values_lst = [] for micro_batch in micro_batches: with torch.no_grad(): values = self._forward_micro_batch(micro_batch) values_lst.append(values) values = torch.concat(values_lst, dim=0) responses = data.batch['responses'] attention_mask = data.batch['attention_mask'] response_length = responses.size(1) values = values * attention_mask[:, -response_length - 1:-1] if use_dynamic_bsz: indices = list(itertools.chain.from_iterable(indices)) assert len(indices) == values.size(0), f\"{len(indices)} vs. {values.size()}\" revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long) values = values[revert_indices] return values def update_critic(self, data: DataProto): # make sure we are in training mode self.critic_module.train() metrics = {} select_keys = ['input_ids', 'responses', 'attention_mask', 'position_ids', 'values', 'returns'] batch = data.select(batch_keys=select_keys).batch # Split to make minibatch iterator for updating the actor # See PPO paper for details. https://arxiv.org/abs/1707.06347 dataloader = batch.split(self.config.ppo_mini_batch_size) for batch_idx, data in enumerate(dataloader): # split batch into micro_batches mini_batch = data if self.config.use_dynamic_bsz: max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size micro_batches, _ = rearrange_micro_batches(batch=mini_batch, max_token_len=max_token_len) else: micro_batches = mini_batch.split(self.config.ppo_micro_batch_size) self.critic_optimizer.zero_grad() for data in micro_batches: data = data.cuda() # critic device is cpu when us","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.__init__","name":"__init__","kind":"method","path":"verl/workers/critic/dp_critic.py","span":"L41-L51","signature":"def __init__(self, config, critic_module: nn.Module, critic_optimizer: optim.Optimizer)","exported":false,"origin":"ast","body_hash":"7b274fad5ae21901f4f235ddb30a918933a9524c921d17c2843b79132e48e481","body_text":"def __init__(self, config, critic_module: nn.Module, critic_optimizer: optim.Optimizer): super().__init__(config=config) self.critic_module = critic_module self.critic_optimizer = critic_optimizer self.use_remove_padding = self.config.model.get('use_remove_padding', False) print(f'Critic use_remove_padding={self.use_remove_padding}') assert self.config.ppo_mini_batch_size % self.config.ppo_micro_batch_size == 0 self.gradient_accumulation = self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size self.ulysses_sequence_parallel_size = self.config.get('ulysses_sequence_parallel_size', 1)","summary_state":"pending","summary":null,"crux":null,"owner":"DataParallelPPOCritic"},{"id":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","name":"_forward_micro_batch","kind":"method","path":"verl/workers/critic/dp_critic.py","span":"L53-L101","signature":"def _forward_micro_batch(self, micro_batch)","exported":false,"origin":"ast","body_hash":"db76a1e4bde7640528abbc96e74e21f904df7d8be17ed48b05832d4d252785d6","body_text":"def _forward_micro_batch(self, micro_batch): response_length = micro_batch['responses'].size(-1) with torch.autocast(device_type='cuda', dtype=torch.bfloat16): input_ids = micro_batch['input_ids'] batch, seqlen = input_ids.shape attention_mask = micro_batch['attention_mask'] position_ids = micro_batch['position_ids'] if self.use_remove_padding: 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) # pad and slice the inputs if sp > 1 if self.ulysses_sequence_parallel_size > 1: input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs(input_ids_rmpad, \\ position_ids_rmpad, \\ sp_size=self.ulysses_sequence_parallel_size) # only pass input_ids and position_ids to enable flash_attn_varlen output = self.critic_module(input_ids=input_ids_rmpad, attention_mask=None, position_ids=position_ids_rmpad, use_cache=False) # prevent model thinks we are generating values_rmpad = output.logits values_rmpad = values_rmpad.squeeze(0) # (total_nnz) # gather output if sp > 1 if self.ulysses_sequence_parallel_size > 1: values_rmpad = gather_outpus_and_unpad(values_rmpad, gather_dim=0, unpad_dim=0, padding_size=pad_size) # pad it back values = pad_input(values_rmpad, indices=indices, batch=batch, seqlen=seqlen).squeeze(-1) values = values[:, -response_length - 1:-1] else: output = self.critic_module(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False) # prevent model thinks we are generating values = output.logits values = values[:, -response_length - 1:-1].squeeze(-1) return values","summary_state":"pending","summary":null,"crux":null,"owner":"DataParallelPPOCritic"},{"id":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._optimizer_step","name":"_optimizer_step","kind":"method","path":"verl/workers/critic/dp_critic.py","span":"L103-L111","signature":"def _optimizer_step(self)","exported":false,"origin":"ast","body_hash":"c50ec803eec6182c80e21ec98dcc3d9843fb8ec77023910fee039170c524f903","body_text":"def _optimizer_step(self): assert self.config.grad_clip is not None if isinstance(self.critic_module, FSDP): grad_norm = self.critic_module.clip_grad_norm_(self.config.grad_clip) else: grad_norm = torch.nn.utils.clip_grad_norm_(self.critic_module.parameters(), max_norm=self.config.grad_clip) self.critic_optimizer.step() return grad_norm","summary_state":"pending","summary":null,"crux":null,"owner":"DataParallelPPOCritic"},{"id":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","name":"compute_values","kind":"method","path":"verl/workers/critic/dp_critic.py","span":"L113-L144","signature":"def compute_values(self, data: DataProto) -> torch.Tensor","exported":true,"origin":"ast","body_hash":"b496079b39fcc1d192495463e2da91d65e1233b1181305f30a3fb9875b9a7071","body_text":"def compute_values(self, data: DataProto) -> torch.Tensor: self.critic_module.eval() micro_batch_size = data.meta_info['micro_batch_size'] select_keys = ['responses', 'input_ids', 'attention_mask', 'position_ids'] batch = data.select(batch_keys=select_keys).batch use_dynamic_bsz = data.meta_info['use_dynamic_bsz'] if use_dynamic_bsz: # split using dynamic bsz max_token_len = data.meta_info['max_token_len'] * self.ulysses_sequence_parallel_size micro_batches, indices = rearrange_micro_batches(batch=batch, max_token_len=max_token_len) else: micro_batches = batch.split(micro_batch_size) values_lst = [] for micro_batch in micro_batches: with torch.no_grad(): values = self._forward_micro_batch(micro_batch) values_lst.append(values) values = torch.concat(values_lst, dim=0) responses = data.batch['responses'] attention_mask = data.batch['attention_mask'] response_length = responses.size(1) values = values * attention_mask[:, -response_length - 1:-1] if use_dynamic_bsz: indices = list(itertools.chain.from_iterable(indices)) assert len(indices) == values.size(0), f\"{len(indices)} vs. {values.size()}\" revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long) values = values[revert_indices] return values","summary_state":"pending","summary":null,"crux":null,"owner":"DataParallelPPOCritic"},{"id":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","name":"update_critic","kind":"method","path":"verl/workers/critic/dp_critic.py","span":"L146-L204","signature":"def update_critic(self, data: DataProto): # make sure we are in training mode","exported":true,"origin":"ast","body_hash":"06067b711e5f95962cb947ac66c603451a1acea1370183405e1cb545199d42bd","body_text":"def update_critic(self, data: DataProto): # make sure we are in training mode self.critic_module.train() metrics = {} select_keys = ['input_ids', 'responses', 'attention_mask', 'position_ids', 'values', 'returns'] batch = data.select(batch_keys=select_keys).batch # Split to make minibatch iterator for updating the actor # See PPO paper for details. https://arxiv.org/abs/1707.06347 dataloader = batch.split(self.config.ppo_mini_batch_size) for batch_idx, data in enumerate(dataloader): # split batch into micro_batches mini_batch = data if self.config.use_dynamic_bsz: max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size micro_batches, _ = rearrange_micro_batches(batch=mini_batch, max_token_len=max_token_len) else: micro_batches = mini_batch.split(self.config.ppo_micro_batch_size) self.critic_optimizer.zero_grad() for data in micro_batches: data = data.cuda() # critic device is cpu when using offload input_ids = data['input_ids'] responses = data['responses'] attention_mask = data['attention_mask'] position_ids = data['position_ids'] values = data['values'] returns = data['returns'] response_length = responses.size(1) eos_mask = attention_mask[:, -response_length - 1:-1] vpreds = self._forward_micro_batch(data) # assert not torch.any(torch.isnan(vpreds)).item() vf_loss, vf_clipfrac = core_algos.compute_value_loss(vpreds=vpreds, values=values, returns=returns, eos_mask=eos_mask, cliprange_value=self.config.cliprange_value) loss = vf_loss / self.gradient_accumulation loss.backward() data = { 'critic/vf_loss': vf_loss.detach().item(), 'critic/vf_clipfrac': vf_clipfrac.detach().item(), 'critic/vpred_mean': masked_mean(vpreds, eos_mask).detach().item(), } append_to_dict(metrics, data) grad_norm = self._optimizer_step() data = {'critic/grad_norm': grad_norm.detach().item()} append_to_dict(metrics, data) self.critic_optimizer.zero_grad() return metrics","summary_state":"pending","summary":null,"crux":null,"owner":"DataParallelPPOCritic"}],"rawEdges":[{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"itertools","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"typing","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"torch","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"torch","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"torch.distributed.fsdp","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"verl","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"verl.trainer.ppo","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"verl.workers.critic","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"verl.utils.py_functional","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"verl.utils.ulysses","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"verl.utils.seqlen_balancing","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"imports","specifier":"flash_attn.bert_padding","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py","relation":"contains","targetId":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic","relation":"extends","name":"BasePPOCritic","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic","relation":"contains","targetId":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.__init__","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic","relation":"contains","targetId":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"autocast","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"unpad_input","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"transpose","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"transpose","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"index_first_axis","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"rearrange","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"ulysses_pad_and_slice_inputs","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"critic_module","viaMember":true,"file":"verl/workers/critic/dp_critic.py","recvType":"DataParallelPPOCritic"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/workers/critic/dp_critic.py","recvType":"gather_outpus_and_unpad"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"gather_outpus_and_unpad","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"critic_module","viaMember":true,"file":"verl/workers/critic/dp_critic.py","recvType":"DataParallelPPOCritic"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._forward_micro_batch","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic","relation":"contains","targetId":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._optimizer_step","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._optimizer_step","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._optimizer_step","relation":"calls","name":"clip_grad_norm_","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._optimizer_step","relation":"calls","name":"clip_grad_norm_","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._optimizer_step","relation":"calls","name":"parameters","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic._optimizer_step","relation":"calls","name":"step","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic","relation":"contains","targetId":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"eval","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"select","viaMember":true,"file":"verl/workers/critic/dp_critic.py","recvType":"DataProto"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"rearrange_micro_batches","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"split","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"_forward_micro_batch","viaMember":true,"file":"verl/workers/critic/dp_critic.py","recvType":"DataParallelPPOCritic"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"concat","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"list","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"from_iterable","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"tensor","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.compute_values","relation":"calls","name":"get_reverse_idx","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic","relation":"contains","targetId":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"train","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"select","viaMember":true,"file":"verl/workers/critic/dp_critic.py","recvType":"DataProto"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"split","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"enumerate","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"rearrange_micro_batches","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"split","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"zero_grad","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/critic/dp_critic.py","recvType":"DataProto"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"_forward_micro_batch","viaMember":true,"file":"verl/workers/critic/dp_critic.py","recvType":"DataParallelPPOCritic"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"compute_value_loss","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"backward","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"masked_mean","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"append_to_dict","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"_optimizer_step","viaMember":true,"file":"verl/workers/critic/dp_critic.py","recvType":"DataParallelPPOCritic"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"append_to_dict","viaMember":false,"file":"verl/workers/critic/dp_critic.py"},{"source":"verl/workers/critic/dp_critic.py#DataParallelPPOCritic.update_critic","relation":"calls","name":"zero_grad","viaMember":true,"file":"verl/workers/critic/dp_critic.py"}]},"verl/workers/critic/megatron_critic.py":{"size":10641,"mtimeMs":1789275830327.2031,"hash":"19fb92b0d5cca0024ba7ff77e90ef62772d629bf99210804182945393f1deebb","nodes":[{"id":"verl/workers/critic/megatron_critic.py","name":"megatron_critic.py","kind":"file","path":"verl/workers/critic/megatron_critic.py","span":"L1-L230","signature":null,"exported":true,"origin":"ast","body_hash":"19fb92b0d5cca0024ba7ff77e90ef62772d629bf99210804182945393f1deebb","chars":10641,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Implement a multiprocess PPOCritic \"\"\" from functools import partial from typing import Iterable import torch import torch.distributed from omegaconf import OmegaConf from torch import nn from verl import DataProto from verl.trainer.ppo import core_algos from verl.workers.critic import BasePPOCritic from verl.utils.megatron.pipeline_parallel import (compute_transformers_input_shapes, make_batch_generator) from verl.utils.py_functional import append_to_dict from verl.utils.torch_dtypes import PrecisionType from verl.utils.torch_functional import masked_mean, broadcast_dict_tensor, split_dict_tensor_into_batches from verl.utils.megatron import sequence_parallel as sp_utils from verl.utils.megatron.optimizer_config import OptimizerConfig from megatron.optimizer import DistributedOptimizer from megatron.core import parallel_state as mpu from megatron.core.pipeline_parallel import get_forward_backward_func"},{"id":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic","name":"MegatronPPOCritic","kind":"class","path":"verl/workers/critic/megatron_critic.py","span":"L41-L229","signature":"class MegatronPPOCritic(BasePPOCritic)","exported":true,"origin":"ast","body_hash":"d76717011abd45a6269165a1e37f00f0fdd48ec8755322e9c9a43a5c58253a8f","body_text":"class MegatronPPOCritic(BasePPOCritic): def __init__(self, config, model_config, megatron_config, critic_module: nn.ModuleList, critic_optimizer: DistributedOptimizer, critic_optimizer_config: OptimizerConfig): super().__init__(config=config) self.model_config = model_config self.megatron_config = megatron_config self.critic_module = critic_module self.critic_optimizer = critic_optimizer self.critic_optimizer_config = critic_optimizer_config # we create a separate nametuple for optimizer step so that global args won't affect it. self.optimizer_step_args = OmegaConf.create({ 'skip_grad': None, 'overlap_dp_param_comm': False, 'overlap_dp_grad_comm': False, 'gradient_accumulation_steps': 1, 'sequence_parallel': self.megatron_config.sequence_parallel, 'DDP_impl': 'local', 'layernorm_allreduce_bucket_threshold': 0, 'pipeline_model_parallel_split_rank': None, 'reduce_grads_use_alltoall': False }) if self.config.kl_ctrl.type == 'fixed': self.kl_ctrl = core_algos.FixedKLController(kl_coef=self.config.kl_ctrl.kl_coef) elif self.config.kl_ctrl.type == 'adaptive': assert self.config.kl_ctrl.horizon > 0, f'horizon must be larger than 0. Got {self.config.kl_ctrl.horizon}' self.kl_ctrl = core_algos.AdaptiveKLController(init_kl_coef=self.config.kl_ctrl.kl_coef, target_kl=self.config.kl_ctrl.target_kl, horizon=self.config.kl_ctrl.horizon) else: raise NotImplementedError def compute_values(self, data: DataProto) -> DataProto: # data.batch = data.batch.to(self.critic_module.module.device) responses = data.batch['responses'] attention_mask = data.batch['attention_mask'] response_length = responses.size(1) with torch.no_grad(): output = self.forward_backward_batch(data=data, forward_only=True) if mpu.is_pipeline_last_stage(ignore_virtual=True): # only on last rank. It should be on every tp rank values = torch.cat([o['vpreds'] for o in output], dim=0) # (bs, seq_size, vocal_size) values = values.to(torch.float32) else: values = torch.empty_like(attention_mask, dtype=torch.float32) # each tp ranks should contain the same value values = values * attention_mask values = values[:, -response_length - 1:-1] values = values.contiguous() # sync among pp ranks torch.distributed.broadcast(tensor=values, src=mpu.get_pipeline_model_parallel_last_rank(), group=mpu.get_pipeline_model_parallel_group()) # add empty cache after each compute torch.cuda.empty_cache() return values def make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]: select_keys = ['input_ids', 'responses', 'attention_mask', 'position_ids', 'values', 'returns'] data = data.select(batch_keys=select_keys) return data.make_iterator(mini_batch_size=self.config.ppo_mini_batch_size, epochs=self.config.ppo_epochs, dataloader_kwargs={'shuffle': self.config.shuffle}) def forward_backward_batch(self, data: DataProto, forward_only=False): # broadcast from last pp rank to all other pp ranks data.batch = data.batch.contiguous() broadcast_dict_tensor(data.batch, src=mpu.get_pipeline_model_parallel_last_rank(), group=mpu.get_pipeline_model_parallel_group()) # split into micro-batches data.batch['attention_mask'] = data.batch['attention_mask'].to(bool) batches = split_dict_tensor_into_batches(data.batch, batch_size=self.config.ppo_micro_batch_size) n_micro_batch = len(batches) seq_len = batches[0]['input_ids'].shape[1] # compute input shapes for pp stages input_shapes = compute_transformers_input_shapes( batches, meta_info={ 'sequence_parallel': self.megatron_config.sequence_parallel, 'hidden_size': self.model_config.hidden_size }) forward_backward_func = get_forward_backward_func() def loss_func(output, data, meta_info): if forward_only: return 1.0, {'vpreds': output.logits} responses = data['responses'] attention_mask = data['attention_mask'] values = data['values'] returns = data['returns'] response_length = responses.size(1) eos_mask = attention_mask[:, -response_length:] cliprange_value = self.config.cliprange_value vpreds = output.logits # (bs, sequence_length) vpreds = vpreds[:, -response_length - 1:-1] vf_loss, vf_clipfrac = core_algos.compute_value_loss(vpreds=vpreds, values=values, returns=returns, eos_mask=eos_mask, cliprange_value=cliprange_value) stats = { 'critic/vf_loss': vf_loss.detach().item(), 'critic/vf_clipfrac': vf_clipfrac.detach().item(), 'critic/vpred_mean': masked_mean(vpreds, eos_mask).detach().item(), } return vf_loss, stats def forward_step(batch_iter, model): batch = next(batch_iter) input_ids = batch['input_ids'] attention_mask = batch['attention_mask'] position_ids = batch['position_ids'] output = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) return output, partial(loss_func, data=batch, meta_info={}) # batch should be a list of batches inside micro-batches batch_generator = make_batch_generator(batches, vpp_size=len(self.critic_module)) # TODO: we may use the new schedule instead # for flash-attn: (seq_len, batch_size, hidden_size) = (mbs*seq_len, 1, hidden_size) if mpu.get_pipeline_model_parallel_world_s","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.__init__","name":"__init__","kind":"method","path":"verl/workers/critic/megatron_critic.py","span":"L43-L75","signature":"def __init__(self, config, model_config, megatron_config, critic_module: nn.ModuleList, critic_optimizer: DistributedOptimizer, critic_optimizer_config: OptimizerConfig)","exported":false,"origin":"ast","body_hash":"e33f2cb0146626e1e892b3380cc50df3cf1206386c05591325679a654a30ef45","body_text":"def __init__(self, config, model_config, megatron_config, critic_module: nn.ModuleList, critic_optimizer: DistributedOptimizer, critic_optimizer_config: OptimizerConfig): super().__init__(config=config) self.model_config = model_config self.megatron_config = megatron_config self.critic_module = critic_module self.critic_optimizer = critic_optimizer self.critic_optimizer_config = critic_optimizer_config # we create a separate nametuple for optimizer step so that global args won't affect it. self.optimizer_step_args = OmegaConf.create({ 'skip_grad': None, 'overlap_dp_param_comm': False, 'overlap_dp_grad_comm': False, 'gradient_accumulation_steps': 1, 'sequence_parallel': self.megatron_config.sequence_parallel, 'DDP_impl': 'local', 'layernorm_allreduce_bucket_threshold': 0, 'pipeline_model_parallel_split_rank': None, 'reduce_grads_use_alltoall': False }) if self.config.kl_ctrl.type == 'fixed': self.kl_ctrl = core_algos.FixedKLController(kl_coef=self.config.kl_ctrl.kl_coef) elif self.config.kl_ctrl.type == 'adaptive': assert self.config.kl_ctrl.horizon > 0, f'horizon must be larger than 0. Got {self.config.kl_ctrl.horizon}' self.kl_ctrl = core_algos.AdaptiveKLController(init_kl_coef=self.config.kl_ctrl.kl_coef, target_kl=self.config.kl_ctrl.target_kl, horizon=self.config.kl_ctrl.horizon) else: raise NotImplementedError","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronPPOCritic"},{"id":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","name":"compute_values","kind":"method","path":"verl/workers/critic/megatron_critic.py","span":"L77-L104","signature":"def compute_values(self, data: DataProto) -> DataProto: # data.batch = data.batch.to(self.critic_module.module.device)","exported":true,"origin":"ast","body_hash":"cc7c5ab91060465e730c63c6bd4ed449f1b51918fd599592cf19161ff71328d5","body_text":"def compute_values(self, data: DataProto) -> DataProto: # data.batch = data.batch.to(self.critic_module.module.device) responses = data.batch['responses'] attention_mask = data.batch['attention_mask'] response_length = responses.size(1) with torch.no_grad(): output = self.forward_backward_batch(data=data, forward_only=True) if mpu.is_pipeline_last_stage(ignore_virtual=True): # only on last rank. It should be on every tp rank values = torch.cat([o['vpreds'] for o in output], dim=0) # (bs, seq_size, vocal_size) values = values.to(torch.float32) else: values = torch.empty_like(attention_mask, dtype=torch.float32) # each tp ranks should contain the same value values = values * attention_mask values = values[:, -response_length - 1:-1] values = values.contiguous() # sync among pp ranks torch.distributed.broadcast(tensor=values, src=mpu.get_pipeline_model_parallel_last_rank(), group=mpu.get_pipeline_model_parallel_group()) # add empty cache after each compute torch.cuda.empty_cache() return values","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronPPOCritic"},{"id":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.make_minibatch_iterator","name":"make_minibatch_iterator","kind":"method","path":"verl/workers/critic/megatron_critic.py","span":"L106-L111","signature":"def make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]","exported":true,"origin":"ast","body_hash":"d62b5e5408740e95de13a21cff5940bba599256cd69fa58e7b2e3ea10d4e276e","body_text":"def make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]: select_keys = ['input_ids', 'responses', 'attention_mask', 'position_ids', 'values', 'returns'] data = data.select(batch_keys=select_keys) return data.make_iterator(mini_batch_size=self.config.ppo_mini_batch_size, epochs=self.config.ppo_epochs, dataloader_kwargs={'shuffle': self.config.shuffle})","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronPPOCritic"},{"id":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","name":"forward_backward_batch","kind":"method","path":"verl/workers/critic/megatron_critic.py","span":"L113-L202","signature":"def forward_backward_batch(self, data: DataProto, forward_only=False): # broadcast from last pp rank to all other pp ranks","exported":true,"origin":"ast","body_hash":"3f8089226c38083c2b62839d4d7e4c662bacd541d5b1172747af0a25cd93c928","body_text":"def forward_backward_batch(self, data: DataProto, forward_only=False): # broadcast from last pp rank to all other pp ranks data.batch = data.batch.contiguous() broadcast_dict_tensor(data.batch, src=mpu.get_pipeline_model_parallel_last_rank(), group=mpu.get_pipeline_model_parallel_group()) # split into micro-batches data.batch['attention_mask'] = data.batch['attention_mask'].to(bool) batches = split_dict_tensor_into_batches(data.batch, batch_size=self.config.ppo_micro_batch_size) n_micro_batch = len(batches) seq_len = batches[0]['input_ids'].shape[1] # compute input shapes for pp stages input_shapes = compute_transformers_input_shapes( batches, meta_info={ 'sequence_parallel': self.megatron_config.sequence_parallel, 'hidden_size': self.model_config.hidden_size }) forward_backward_func = get_forward_backward_func() def loss_func(output, data, meta_info): if forward_only: return 1.0, {'vpreds': output.logits} responses = data['responses'] attention_mask = data['attention_mask'] values = data['values'] returns = data['returns'] response_length = responses.size(1) eos_mask = attention_mask[:, -response_length:] cliprange_value = self.config.cliprange_value vpreds = output.logits # (bs, sequence_length) vpreds = vpreds[:, -response_length - 1:-1] vf_loss, vf_clipfrac = core_algos.compute_value_loss(vpreds=vpreds, values=values, returns=returns, eos_mask=eos_mask, cliprange_value=cliprange_value) stats = { 'critic/vf_loss': vf_loss.detach().item(), 'critic/vf_clipfrac': vf_clipfrac.detach().item(), 'critic/vpred_mean': masked_mean(vpreds, eos_mask).detach().item(), } return vf_loss, stats def forward_step(batch_iter, model): batch = next(batch_iter) input_ids = batch['input_ids'] attention_mask = batch['attention_mask'] position_ids = batch['position_ids'] output = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) return output, partial(loss_func, data=batch, meta_info={}) # batch should be a list of batches inside micro-batches batch_generator = make_batch_generator(batches, vpp_size=len(self.critic_module)) # TODO: we may use the new schedule instead # for flash-attn: (seq_len, batch_size, hidden_size) = (mbs*seq_len, 1, hidden_size) if mpu.get_pipeline_model_parallel_world_size() > 1: losses_reduced = forward_backward_func( forward_step_func=forward_step, data_iterator=batch_generator, model=self.critic_module, num_microbatches=n_micro_batch, input_shapes=input_shapes, # must set for flash-attn sequence packing seq_length=self.config.ppo_micro_batch_size * seq_len, # no use when input_shapes was set hidden_size=self.model_config.hidden_size, # no use when input_shapes was set micro_batch_size=1, # no use when input_shapes was set forward_only=forward_only, ) else: losses_reduced = forward_backward_func( forward_step_func=forward_step, data_iterator=batch_generator, model=self.critic_module, num_microbatches=n_micro_batch, seq_length=self.config.ppo_micro_batch_size * seq_len, # in use for pp = 1 hidden_size=self.model_config.hidden_size, # in use for pp = 1 micro_batch_size=1, # in use for pp = 1 forward_only=forward_only, ) # loss_reduces contains the stats returned from loss_func return losses_reduced","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronPPOCritic"},{"id":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","name":"loss_func","kind":"function","path":"verl/workers/critic/megatron_critic.py","span":"L135-L163","signature":"def loss_func(output, data, meta_info)","exported":true,"origin":"ast","body_hash":"fb3cc9a1d633411c1a394a4a4e87060104e726349f8aad8fa7fc758e52476caa","body_text":"def loss_func(output, data, meta_info): if forward_only: return 1.0, {'vpreds': output.logits} responses = data['responses'] attention_mask = data['attention_mask'] values = data['values'] returns = data['returns'] response_length = responses.size(1) eos_mask = attention_mask[:, -response_length:] cliprange_value = self.config.cliprange_value vpreds = output.logits # (bs, sequence_length) vpreds = vpreds[:, -response_length - 1:-1] vf_loss, vf_clipfrac = core_algos.compute_value_loss(vpreds=vpreds, values=values, returns=returns, eos_mask=eos_mask, cliprange_value=cliprange_value) stats = { 'critic/vf_loss': vf_loss.detach().item(), 'critic/vf_clipfrac': vf_clipfrac.detach().item(), 'critic/vpred_mean': masked_mean(vpreds, eos_mask).detach().item(), } return vf_loss, stats","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.forward_step","name":"forward_step","kind":"function","path":"verl/workers/critic/megatron_critic.py","span":"L165-L171","signature":"def forward_step(batch_iter, model)","exported":true,"origin":"ast","body_hash":"5aaae9c2450759884f575fd73746957dc2dfd0e365042322cbf606c610b92a79","body_text":"def forward_step(batch_iter, model): batch = next(batch_iter) input_ids = batch['input_ids'] attention_mask = batch['attention_mask'] position_ids = batch['position_ids'] output = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) return output, partial(loss_func, data=batch, meta_info={})","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.update_critic","name":"update_critic","kind":"method","path":"verl/workers/critic/megatron_critic.py","span":"L204-L229","signature":"def update_critic(self, dataloader: Iterable[DataProto])","exported":true,"origin":"ast","body_hash":"e0e0149e0e7af9959ab384f03f80928220e54bf1fbac8307edde0822c5e06eed","body_text":"def update_critic(self, dataloader: Iterable[DataProto]): metrics = {} for data in dataloader: # data = data.batch.to(self.critic_module.device) self.critic_optimizer.zero_grad() # use use_contiguous_buffers_in_local_ddp and no overlap_dp_param_comm for chunk in self.critic_module: chunk.zero_grad_buffer(zero_buffer=(not self.critic_optimizer_config.use_distributed_optimizer)) metric_micro_batch = self.forward_backward_batch(data) update_successful, grad_norm, num_zeros_in_grad = self.critic_optimizer.step( self.megatron_config, self.megatron_config.timers) if update_successful: # allgather already execute in optimizer.step in new megatron pass else: raise NotImplementedError for metric in metric_micro_batch: append_to_dict(metrics, metric) # append the metric from this micro-batch to global metrics. # add empty cache after each compute torch.cuda.empty_cache() return metrics","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronPPOCritic"}],"rawEdges":[{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"functools","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"typing","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"torch","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"omegaconf","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"torch","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"verl","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"verl.trainer.ppo","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"verl.workers.critic","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"verl.utils.megatron.pipeline_parallel","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"verl.utils.py_functional","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"verl.utils.torch_dtypes","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"verl.utils.megatron","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"verl.utils.megatron.optimizer_config","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"megatron.optimizer","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"megatron.core","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"imports","specifier":"megatron.core.pipeline_parallel","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py","relation":"contains","targetId":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic","relation":"extends","name":"BasePPOCritic","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic","relation":"contains","targetId":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.__init__","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.__init__","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.__init__","relation":"calls","name":"FixedKLController","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.__init__","relation":"calls","name":"AdaptiveKLController","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic","relation":"contains","targetId":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"forward_backward_batch","viaMember":true,"file":"verl/workers/critic/megatron_critic.py","recvType":"MegatronPPOCritic"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"is_pipeline_last_stage","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"get_pipeline_model_parallel_last_rank","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"get_pipeline_model_parallel_group","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.compute_values","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic","relation":"contains","targetId":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.make_minibatch_iterator","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.make_minibatch_iterator","relation":"calls","name":"select","viaMember":true,"file":"verl/workers/critic/megatron_critic.py","recvType":"DataProto"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.make_minibatch_iterator","relation":"calls","name":"make_iterator","viaMember":true,"file":"verl/workers/critic/megatron_critic.py","recvType":"DataProto"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic","relation":"contains","targetId":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"broadcast_dict_tensor","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"get_pipeline_model_parallel_last_rank","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"get_pipeline_model_parallel_group","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"split_dict_tensor_into_batches","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"compute_transformers_input_shapes","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"get_forward_backward_func","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"contains","targetId":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","relation":"calls","name":"compute_value_loss","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.loss_func","relation":"calls","name":"masked_mean","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"contains","targetId":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.forward_step","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.forward_step","relation":"calls","name":"next","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.forward_step","relation":"calls","name":"model","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch.forward_step","relation":"calls","name":"partial","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"make_batch_generator","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"forward_backward_func","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.forward_backward_batch","relation":"calls","name":"forward_backward_func","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic","relation":"contains","targetId":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.update_critic","file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.update_critic","relation":"calls","name":"zero_grad","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.update_critic","relation":"calls","name":"zero_grad_buffer","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.update_critic","relation":"calls","name":"forward_backward_batch","viaMember":true,"file":"verl/workers/critic/megatron_critic.py","recvType":"MegatronPPOCritic"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.update_critic","relation":"calls","name":"step","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.update_critic","relation":"calls","name":"append_to_dict","viaMember":false,"file":"verl/workers/critic/megatron_critic.py"},{"source":"verl/workers/critic/megatron_critic.py#MegatronPPOCritic.update_critic","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/critic/megatron_critic.py"}]},"verl/workers/fsdp_workers.py":{"size":54106,"mtimeMs":1789275830328.2031,"hash":"7240fc2a9334a6a521f7e51836a031d09d2b4422796b40d5bcb4e3e897867e26","nodes":[{"id":"verl/workers/fsdp_workers.py","name":"fsdp_workers.py","kind":"file","path":"verl/workers/fsdp_workers.py","span":"L1-L1055","signature":null,"exported":true,"origin":"ast","body_hash":"7240fc2a9334a6a521f7e51836a031d09d2b4422796b40d5bcb4e3e897867e26","chars":54106,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 main entry point to run the PPO algorithm \"\"\" import logging import os import warnings import torch import torch.distributed import verl.utils.hdfs_io as hdfs_io import verl.utils.torch_functional as verl_F from omegaconf import DictConfig, open_dict from verl import DataProto from verl.single_controller.base import Worker from verl.single_controller.base.decorator import register, Dispatch from verl.utils import hf_tokenizer from verl.utils.debug import log_gpu_memory_usage from verl.utils.fs import copy_local_path_from_hdfs from verl.utils.fsdp_utils import get_fsdp_wrap_policy, offload_fsdp_grad, init_fn, get_init_weight_context_manager from verl.utils.fsdp_utils import offload_fsdp_optimizer, offload_fsdp_param_and_grad, load_fsdp_optimizer, \\ load_fsdp_param_and_grad from verl.utils.import_utils import import_external_libs from verl.utils.model import compute_position_id_with_mask from verl.utils.flops_counter import FlopsCounter from verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager from codetiming import Timer logger = logging.getLogger(__file__) logger.setLevel(os.getenv('VERL_PPO_LOGGING_LEVEL', 'WARN')) # TODO(sgm): we may need to extract it to dp_reward_model.py"},{"id":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","name":"ActorRolloutRefWorker","kind":"class","path":"verl/workers/fsdp_workers.py","span":"L47-L535","signature":"class ActorRolloutRefWorker(Worker)","exported":true,"origin":"ast","body_hash":"b0e3c2603cb62ed637799f992569f96833e1381befaa4d12855cd1311a9800e4","body_text":"class ActorRolloutRefWorker(Worker): \"\"\" This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy or a hybrid engine based on the config.rollout \"\"\" def __init__(self, config: DictConfig, role: str): super().__init__() self.config = config import torch.distributed if not torch.distributed.is_initialized(): torch.distributed.init_process_group(backend=\"nccl\") # build device mesh for FSDP world_size = torch.distributed.get_world_size() from torch.distributed.device_mesh import init_device_mesh # TODO(sgm): support FSDP hybrid shard for larger model self.device_mesh = init_device_mesh('cuda', mesh_shape=(world_size,), mesh_dim_names=['fsdp']) # build device mesh for Ulysses Sequence Parallel self.ulysses_device_mesh = None self.ulysses_sequence_parallel_size = self.config.actor.get('ulysses_sequence_parallel_size', 1) dp = world_size // self.ulysses_sequence_parallel_size if self.ulysses_sequence_parallel_size > 1: self.ulysses_device_mesh = init_device_mesh('cuda', mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=['dp', 'sp']) self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) self.role = role 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'] self._is_offload_param = False self._is_offload_grad = False self._is_offload_optimizer = False if self._is_actor: self._is_offload_param = self.config.actor.fsdp_config.get('param_offload', False) self._is_offload_grad = self.config.actor.fsdp_config.get('grad_offload', False) self._is_offload_optimizer = self.config.actor.fsdp_config.get('optimizer_offload', False) elif self._is_ref: # TODO: it seems that manual offload is slowly than FSDP offload self._is_offload_param = self.config.ref.fsdp_config.get('param_offload', False) # normalize config if self._is_actor: self.config.actor.ppo_mini_batch_size //= (self.device_mesh.shape[0] // self.ulysses_sequence_parallel_size) self.config.actor.ppo_micro_batch_size //= (self.device_mesh.shape[0] // self.ulysses_sequence_parallel_size) self.config.actor.ppo_mini_batch_size *= self.config.rollout.n self.config.actor.ppo_micro_batch_size *= self.config.rollout.n if self._is_rollout: self.config.rollout.log_prob_micro_batch_size //= (self.device_mesh.shape[0] // self.ulysses_sequence_parallel_size) self.config.rollout.log_prob_micro_batch_size *= self.config.rollout.n if self._is_ref: self.config.ref.log_prob_micro_batch_size //= (self.device_mesh.shape[0] // self.ulysses_sequence_parallel_size) self.config.ref.log_prob_micro_batch_size *= self.config.rollout.n def _build_model_optimizer(self, model_path, fsdp_config, optim_config, override_model_config, use_remove_padding=False, enable_gradient_checkpointing=False, trust_remote_code=False): from verl.utils.model import print_model_size, update_model_config from verl.utils.torch_dtypes import PrecisionType from transformers import AutoModelForCausalLM, AutoConfig from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy, MixedPrecision from torch import optim log_gpu_memory_usage('Before init from HF AutoModel', logger=logger) local_path = copy_local_path_from_hdfs(model_path) # note that we have to create model in fp32. Otherwise, the optimizer is in bf16, which is incorrect # TODO(zhangchi.usc1992): 1. support create from random initialized model. 2. Support init with FSDP directly self.tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code) torch_dtype = fsdp_config.get('model_dtype', None) if torch_dtype is None: torch_dtype = torch.float32 if self._is_actor else torch.bfloat16 else: torch_dtype = PrecisionType.to_dtype(torch_dtype) # override model kwargs actor_model_config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code) if use_remove_padding: from verl.models.registry import check_model_support_rmpad check_model_support_rmpad(actor_model_config.model_type) if use_remove_padding and self.ulysses_sequence_parallel_size > 1: from verl.models.transformers.monkey_patch import apply_monkey_patch apply_monkey_patch(actor_model_config, verbose=True) override_config_kwargs = { 'bos_token_id': self.tokenizer.bos_token_id, 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, } override_config_kwargs.update(override_model_config) update_model_config(actor_model_config, override_config_kwargs=override_config_kwargs) if self.rank == 0: print(f'Model config after override: {actor_model_config}') # NOTE(fix me): tie_word_embedding causes meta_tensor init to hang init_context = get_init_weight_context_manager(use_meta_tensor=not actor_model_config.tie_word_embeddings) with init_context(), warnings.catch_warnings(): warnings.","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","name":"__init__","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L53-L109","signature":"def __init__(self, config: DictConfig, role: str)","exported":false,"origin":"ast","body_hash":"3b9abf02a861b2fd01935e8a666539d42d962769a783ced38ff1fb2b90893e34","body_text":"def __init__(self, config: DictConfig, role: str): super().__init__() self.config = config import torch.distributed if not torch.distributed.is_initialized(): torch.distributed.init_process_group(backend=\"nccl\") # build device mesh for FSDP world_size = torch.distributed.get_world_size() from torch.distributed.device_mesh import init_device_mesh # TODO(sgm): support FSDP hybrid shard for larger model self.device_mesh = init_device_mesh('cuda', mesh_shape=(world_size,), mesh_dim_names=['fsdp']) # build device mesh for Ulysses Sequence Parallel self.ulysses_device_mesh = None self.ulysses_sequence_parallel_size = self.config.actor.get('ulysses_sequence_parallel_size', 1) dp = world_size // self.ulysses_sequence_parallel_size if self.ulysses_sequence_parallel_size > 1: self.ulysses_device_mesh = init_device_mesh('cuda', mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=['dp', 'sp']) self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) self.role = role 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'] self._is_offload_param = False self._is_offload_grad = False self._is_offload_optimizer = False if self._is_actor: self._is_offload_param = self.config.actor.fsdp_config.get('param_offload', False) self._is_offload_grad = self.config.actor.fsdp_config.get('grad_offload', False) self._is_offload_optimizer = self.config.actor.fsdp_config.get('optimizer_offload', False) elif self._is_ref: # TODO: it seems that manual offload is slowly than FSDP offload self._is_offload_param = self.config.ref.fsdp_config.get('param_offload', False) # normalize config if self._is_actor: self.config.actor.ppo_mini_batch_size //= (self.device_mesh.shape[0] // self.ulysses_sequence_parallel_size) self.config.actor.ppo_micro_batch_size //= (self.device_mesh.shape[0] // self.ulysses_sequence_parallel_size) self.config.actor.ppo_mini_batch_size *= self.config.rollout.n self.config.actor.ppo_micro_batch_size *= self.config.rollout.n if self._is_rollout: self.config.rollout.log_prob_micro_batch_size //= (self.device_mesh.shape[0] // self.ulysses_sequence_parallel_size) self.config.rollout.log_prob_micro_batch_size *= self.config.rollout.n if self._is_ref: self.config.ref.log_prob_micro_batch_size //= (self.device_mesh.shape[0] // self.ulysses_sequence_parallel_size) self.config.ref.log_prob_micro_batch_size *= self.config.rollout.n","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","name":"_build_model_optimizer","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L111-L248","signature":"def _build_model_optimizer(self, model_path, fsdp_config, optim_config, override_model_config, use_remove_padding=False, enable_gradient_checkpointing=False, trust_remote_code=False)","exported":false,"origin":"ast","body_hash":"350c7a10097e8b0cbba36048e4a90dd6470366626fadb69943f1bef87869dd8f","body_text":"def _build_model_optimizer(self, model_path, fsdp_config, optim_config, override_model_config, use_remove_padding=False, enable_gradient_checkpointing=False, trust_remote_code=False): from verl.utils.model import print_model_size, update_model_config from verl.utils.torch_dtypes import PrecisionType from transformers import AutoModelForCausalLM, AutoConfig from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy, MixedPrecision from torch import optim log_gpu_memory_usage('Before init from HF AutoModel', logger=logger) local_path = copy_local_path_from_hdfs(model_path) # note that we have to create model in fp32. Otherwise, the optimizer is in bf16, which is incorrect # TODO(zhangchi.usc1992): 1. support create from random initialized model. 2. Support init with FSDP directly self.tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code) torch_dtype = fsdp_config.get('model_dtype', None) if torch_dtype is None: torch_dtype = torch.float32 if self._is_actor else torch.bfloat16 else: torch_dtype = PrecisionType.to_dtype(torch_dtype) # override model kwargs actor_model_config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code) if use_remove_padding: from verl.models.registry import check_model_support_rmpad check_model_support_rmpad(actor_model_config.model_type) if use_remove_padding and self.ulysses_sequence_parallel_size > 1: from verl.models.transformers.monkey_patch import apply_monkey_patch apply_monkey_patch(actor_model_config, verbose=True) override_config_kwargs = { 'bos_token_id': self.tokenizer.bos_token_id, 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, } override_config_kwargs.update(override_model_config) update_model_config(actor_model_config, override_config_kwargs=override_config_kwargs) if self.rank == 0: print(f'Model config after override: {actor_model_config}') # NOTE(fix me): tie_word_embedding causes meta_tensor init to hang init_context = get_init_weight_context_manager(use_meta_tensor=not actor_model_config.tie_word_embeddings) with init_context(), warnings.catch_warnings(): warnings.simplefilter(\"ignore\") actor_module = AutoModelForCausalLM.from_pretrained(pretrained_model_name_or_path=local_path, torch_dtype=torch_dtype, config=actor_model_config, attn_implementation='flash_attention_2', trust_remote_code=trust_remote_code) # some parameters may not in torch_dtype. TODO(zhangchi.usc1992) remove this after we switch to fsdp2 actor_module.to(torch_dtype) if enable_gradient_checkpointing: actor_module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) torch.distributed.barrier() if self.rank == 0: print_model_size(actor_module) log_gpu_memory_usage('After init from HF AutoModel', logger=logger) # We wrap FSDP for rollout as well mixed_precision_config = fsdp_config.get('mixed_precision', None) 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) if self._is_ref: mixed_precision = None auto_wrap_policy = get_fsdp_wrap_policy(module=actor_module, config=fsdp_config.get('wrap_policy', None)) if self._is_rollout and self.config.rollout.name == 'hf': # TODO(zhangchi.usc1992, shengguangming) fix me. Current, auto_wrap_policy causes HFRollout to hang in Gemma auto_wrap_policy = None print(f'wrap_policy: {auto_wrap_policy}') # TODO(sgm): support hybrid if auto_wrap_policy is None: sharding_strategy = ShardingStrategy.SHARD_GRAD_OP else: sharding_strategy = ShardingStrategy.FULL_SHARD # TODO: add transformer policy actor_module_fsdp = FSDP( actor_module, param_init_fn=init_fn, use_orig_params=False, auto_wrap_policy=auto_wrap_policy, device_id=torch.cuda.current_device(), sharding_strategy=sharding_strategy, # zero3 mixed_precision=mixed_precision, sync_module_states=True, device_mesh=self.device_mesh, forward_prefetch=False) log_gpu_memory_usage('After Actor FSDP init', logger=logger) # TODO: add more optimizer args into config if self._is_actor: from verl.utils.torch_functional import get_constant_schedule_with_warmup actor_optimizer = optim.AdamW(actor_module_fsdp.parameters(), lr=optim_config.lr, betas=optim_config.get('betas', (0.9, 0.999)), weight_decay=optim_config.get('weight_decay', 1e-2)) total_steps = optim_config.get('total_training_steps', 0) num_warmup_steps_ratio = optim_config.get('lr_warmup_steps_ratio', 0.) num_warmup_steps = int(num_warmup_steps_ratio * total_steps) print(f'Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}') actor_lr_scheduler = get_constant_schedul","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","name":"_build_rollout","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L250-L282","signature":"def _build_rollout(self)","exported":false,"origin":"ast","body_hash":"1aead75975d6d450e56bf221aa533aadac67a0ea33918d637d416dac54629589","body_text":"def _build_rollout(self): from torch.distributed.device_mesh import init_device_mesh # TODO(sgm): support FSDP hybrid shard for larger model infer_tp = self.config.rollout.tensor_model_parallel_size dp = self.world_size // infer_tp assert self.world_size % infer_tp == 0, f'rollout world_size: {self.world_size} is not divisible by infer_tp: {infer_tp}' rollout_device_mesh = init_device_mesh('cuda', mesh_shape=(dp, infer_tp), mesh_dim_names=['dp', 'infer_tp']) if self.config.rollout.name == 'hf': from verl.workers.rollout import HFRollout from verl.workers.sharding_manager import BaseShardingManager rollout = HFRollout(module=self.actor_module_fsdp, config=self.config.rollout) rollout_sharding_manager = BaseShardingManager() # TODO: a sharding manager that do nothing? elif self.config.rollout.name == 'vllm': from verl.workers.rollout.vllm_rollout import vLLMRollout from verl.workers.sharding_manager import FSDPVLLMShardingManager log_gpu_memory_usage('Before building vllm rollout', logger=None) rollout = vLLMRollout(actor_module=self.actor_module_fsdp, config=self.config.rollout, tokenizer=self.tokenizer, model_hf_config=self.actor_model_config) log_gpu_memory_usage('After building vllm rollout', logger=None) if torch.distributed.get_world_size() == 1: self.config.rollout.load_format = 'dummy_hf' rollout_sharding_manager = FSDPVLLMShardingManager(module=self.actor_module_fsdp, inference_engine=rollout.inference_engine, model_config=self.actor_model_config, full_params='hf' in self.config.rollout.load_format, device_mesh=rollout_device_mesh) log_gpu_memory_usage('After building sharding manager', logger=None) return rollout, rollout_sharding_manager","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","name":"init_model","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L285-L353","signature":"def init_model(self)","exported":true,"origin":"ast","body_hash":"1acaa544197b1d9e1d5c6154a955a62cdc97d416376b8dbd2e2666876e981358","body_text":"def init_model(self): from verl.workers.actor import DataParallelPPOActor # This is used to import external_lib into the huggingface systems import_external_libs(self.config.model.get('external_lib', None)) from omegaconf import OmegaConf override_model_config = OmegaConf.to_container(self.config.model.get('override_config', OmegaConf.create())) use_remove_padding = self.config.model.get('use_remove_padding', False) if self._is_actor or self._is_rollout: # we need the model for actor and rollout if self._is_actor: optim_config = self.config.actor.optim fsdp_config = self.config.actor.fsdp_config else: optim_config = None fsdp_config = OmegaConf.create() self.actor_module_fsdp, self.actor_optimizer, self.actor_lr_scheduler, self.actor_model_config = self._build_model_optimizer( model_path=self.config.model.path, fsdp_config=fsdp_config, optim_config=optim_config, override_model_config=override_model_config, use_remove_padding=use_remove_padding, enable_gradient_checkpointing=self.config.model.get('enable_gradient_checkpointing', False), trust_remote_code=self.config.model.get('trust_remote_code', False)) # get the original unwrapped module self.actor_module = self.actor_module_fsdp._fsdp_wrapped_module if self._is_offload_param: # param is require during state_dict in sharding manager offload_fsdp_grad(module=self.actor_module_fsdp) log_gpu_memory_usage('After offload actor grad during init', logger=logger) if self._is_offload_optimizer: offload_fsdp_optimizer(optimizer=self.actor_optimizer) log_gpu_memory_usage('After offload actor optimizer during init', logger=logger) # load from checkpoint if self._is_actor: OmegaConf.set_struct(self.config.actor, True) with open_dict(self.config.actor): self.config.actor.use_remove_padding = use_remove_padding self.actor = DataParallelPPOActor(config=self.config.actor, actor_module=self.actor_module_fsdp, actor_optimizer=self.actor_optimizer) if self._is_rollout: self.rollout, self.rollout_sharding_manager = self._build_rollout() if self._is_ref: self.ref_module_fsdp = self._build_model_optimizer(model_path=self.config.model.path, fsdp_config=self.config.ref.fsdp_config, optim_config=None, override_model_config=override_model_config, use_remove_padding=use_remove_padding, trust_remote_code=self.config.model.get( 'trust_remote_code', False))[0] if self._is_offload_param: offload_fsdp_param_and_grad(module=self.ref_module_fsdp, offload_grad=self._is_offload_grad) OmegaConf.set_struct(self.config.ref, True) with open_dict(self.config.ref): self.config.ref.use_remove_padding = use_remove_padding self.ref_policy = DataParallelPPOActor(config=self.config.ref, actor_module=self.ref_module_fsdp) if self._is_actor: self.flops_counter = FlopsCounter(self.actor_model_config) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","name":"update_actor","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L356-L398","signature":"def update_actor(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"3a2c73d2177ba1849fbf0ce1a0d56686e33f519e9a7916a03f533f1a4ca03aea","body_text":"def update_actor(self, data: DataProto): data = data.to('cuda') assert self._is_actor if self._is_offload_param: load_fsdp_param_and_grad(module=self.actor_module_fsdp, device_id=torch.cuda.current_device(), load_grad=self._is_offload_grad) if self._is_offload_optimizer: load_fsdp_optimizer(optimizer=self.actor_optimizer, device_id=torch.cuda.current_device()) data.batch = data.batch.cuda() log_gpu_memory_usage('Before update policy', logger=logger) with self.ulysses_sharding_manager: data = self.ulysses_sharding_manager.preprocess_data(data=data) # perform training with Timer(name='update_policy', logger=None) as timer: metrics = self.actor.update_policy(data=data) delta_time = timer.last global_num_tokens = data.meta_info['global_token_num'] estimated_flops, promised_flops = self.flops_counter.estimate_flops(global_num_tokens, delta_time) metrics['mfu/actor'] = estimated_flops * self.config.actor.ppo_epochs / promised_flops / self.world_size self.actor_lr_scheduler.step() lr = self.actor_lr_scheduler.get_last_lr()[0] metrics['actor/lr'] = lr log_gpu_memory_usage('After update policy', logger=logger) # TODO: here, we should return all metrics output = DataProto(meta_info={'metrics': metrics}) output = self.ulysses_sharding_manager.postprocess_data(data=output) output = output.to('cpu') if self._is_offload_param: offload_fsdp_param_and_grad(module=self.actor_module_fsdp, offload_grad=self._is_offload_grad) if self._is_offload_optimizer: offload_fsdp_optimizer(optimizer=self.actor_optimizer) torch.cuda.empty_cache() return output","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","name":"compute_log_prob","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L401-L429","signature":"def compute_log_prob(self, data: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"9e11c99528c202e92f2b323ff3b3fd62f38b5ce886c508f4ec241d1363d7229b","body_text":"def compute_log_prob(self, data: DataProto) -> DataProto: \"\"\"mostly copying from generate_sequences\"\"\" data = data.to('cuda') assert self._is_rollout if self._is_offload_param: load_fsdp_param_and_grad(module=self.actor_module_fsdp, device_id=torch.cuda.current_device(), load_grad=self._is_offload_grad) data.batch = data.batch.cuda() meta_info = {'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id} data.meta_info.update(meta_info) with self.ulysses_sharding_manager: data = self.ulysses_sharding_manager.preprocess_data(data) old_log_probs = self.actor.compute_log_prob(data=data) output = DataProto.from_dict(tensors={'old_log_probs': old_log_probs}) output = self.ulysses_sharding_manager.postprocess_data(output) output = output.to('cpu') if self._is_offload_param: # NOTE(sgm): the grad is already in CPU, only offload param here offload_fsdp_param_and_grad(module=self.actor_module_fsdp, offload_grad=self._is_offload_grad) # clear kv cache torch.cuda.empty_cache() log_gpu_memory_usage('After recompute log prob', logger=logger) return output","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","name":"generate_sequences","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L432-L477","signature":"def generate_sequences(self, prompts: DataProto)","exported":true,"origin":"ast","body_hash":"733191fa5ce4c0d66af801c78e340e7dc0b9b42149cae008a40662aa7461f413","body_text":"def generate_sequences(self, prompts: DataProto): prompts = prompts.to('cuda') # set to False if it is validation recompute_log_prob = prompts.meta_info.get('recompute_log_prob', True) assert self._is_rollout if self._is_offload_param: load_fsdp_param_and_grad(module=self.actor_module_fsdp, device_id=torch.cuda.current_device(), load_grad=self._is_offload_grad) prompts.batch = prompts.batch.cuda() meta_info = {'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id} prompts.meta_info.update(meta_info) with self.rollout_sharding_manager: log_gpu_memory_usage('After entering rollout sharding manager', logger=logger) prompts = self.rollout_sharding_manager.preprocess_data(prompts) output = self.rollout.generate_sequences(prompts=prompts) log_gpu_memory_usage('After rollout generation', logger=logger) output = self.rollout_sharding_manager.postprocess_data(output) if self._is_actor and recompute_log_prob: # we should always recompute old_log_probs when it is HybridEngine output.meta_info['micro_batch_size'] = self.config.rollout.log_prob_micro_batch_size output.meta_info['max_token_len'] = self.config.rollout.log_prob_max_token_len_per_gpu output.meta_info['use_dynamic_bsz'] = self.config.rollout.log_prob_use_dynamic_bsz output.meta_info['temperature'] = self.config.rollout.temperature # perform recompute log_prob with self.ulysses_sharding_manager: output = self.ulysses_sharding_manager.preprocess_data(output) old_log_probs = self.actor.compute_log_prob(data=output) output.batch['old_log_probs'] = old_log_probs output = self.ulysses_sharding_manager.postprocess_data(output) output = output.to('cpu') if self._is_offload_param: # NOTE(sgm): the grad is already in CPU, only offload param here offload_fsdp_param_and_grad(module=self.actor_module_fsdp, offload_grad=self._is_offload_grad) # clear kv cache torch.cuda.empty_cache() log_gpu_memory_usage('After recompute log prob', logger=logger) return output","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","name":"compute_ref_log_prob","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L480-L506","signature":"def compute_ref_log_prob(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"ba537dbba907f6d2d2267398d8d15e627db2e78ec704798e89502aaff5ecb295","body_text":"def compute_ref_log_prob(self, data: DataProto): assert self._is_ref data = data.to('cuda') if self._is_offload_param: load_fsdp_param_and_grad(module=self.ref_module_fsdp, device_id=torch.cuda.current_device(), load_grad=self._is_offload_grad) micro_batch_size = self.config.ref.log_prob_micro_batch_size data.meta_info['micro_batch_size'] = micro_batch_size data.meta_info['temperature'] = self.config.rollout.temperature data.meta_info['max_token_len'] = self.config.ref.log_prob_max_token_len_per_gpu data.meta_info['use_dynamic_bsz'] = self.config.ref.log_prob_use_dynamic_bsz with self.ulysses_sharding_manager: data = self.ulysses_sharding_manager.preprocess_data(data) output = self.ref_policy.compute_log_prob(data=data) output = DataProto.from_dict(tensors={'ref_log_prob': output}) output = self.ulysses_sharding_manager.postprocess_data(output) output = output.to('cpu') if self._is_offload_param: offload_fsdp_param_and_grad(module=self.ref_module_fsdp, offload_grad=self._is_offload_grad) torch.cuda.empty_cache() return output","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","name":"save_checkpoint","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L509-L535","signature":"def save_checkpoint(self, local_path, hdfs_path=None)","exported":true,"origin":"ast","body_hash":"3568777ef7917e213f1b8639a1137a84b4d1ed096d6c16f0a9db038a4d0724ae","body_text":"def save_checkpoint(self, local_path, hdfs_path=None): assert self._is_actor import torch if self._is_offload_param: load_fsdp_param_and_grad(module=self.actor_module_fsdp, device_id=torch.cuda.current_device(), load_grad=self._is_offload_grad) # TODO: support DCP and save sharded checkpoints import torch.distributed from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, StateDictType, FullStateDictConfig cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) with FSDP.state_dict_type(self.actor.actor_module, StateDictType.FULL_STATE_DICT, cfg): state_dict = self.actor.actor_module.state_dict() if self.rank == 0: print(f'Saving actor checkpoint to {local_path}') os.makedirs(local_path, exist_ok=True) self.actor_module.save_pretrained(local_path, state_dict=state_dict) self.tokenizer.save_pretrained(local_path) if hdfs_path is not None: print(f'Uploading actor checkpoint to {hdfs_path}') hdfs_io.makedirs(hdfs_path, exist_ok=True) hdfs_io.copy(src=local_path, dst=hdfs_path) torch.distributed.barrier() if self._is_offload_param: offload_fsdp_param_and_grad(module=self.actor_module_fsdp, offload_grad=self._is_offload_grad)","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/fsdp_workers.py#CriticWorker","name":"CriticWorker","kind":"class","path":"verl/workers/fsdp_workers.py","span":"L538-L792","signature":"class CriticWorker(Worker)","exported":true,"origin":"ast","body_hash":"a6d3dba0c4a5431bae6c1b2d42342322c99d4d7cb3ed670df223f98bdfcd9b03","body_text":"class CriticWorker(Worker): def __init__(self, config): super().__init__() import torch.distributed if not torch.distributed.is_initialized(): torch.distributed.init_process_group(backend=\"nccl\") self.config = config # build device mesh for Ulysses Sequence Parallel world_size = torch.distributed.get_world_size() from torch.distributed.device_mesh import init_device_mesh self.ulysses_device_mesh = None self.ulysses_sequence_parallel_size = self.config.get('ulysses_sequence_parallel_size', 1) dp = world_size // self.ulysses_sequence_parallel_size if self.ulysses_sequence_parallel_size > 1: self.ulysses_device_mesh = init_device_mesh('cuda', mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=['dp', 'sp']) self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) # set FSDP offload params self._is_offload_param = self.config.model.fsdp_config.param_offload self._is_offload_grad = self.config.model.fsdp_config.grad_offload self._is_offload_optimizer = self.config.model.fsdp_config.optimizer_offload # normalize config self.config.ppo_mini_batch_size //= (torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size) self.config.ppo_micro_batch_size //= (torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size) self.config.forward_micro_batch_size //= (torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size) def _build_critic_model_optimizer(self, config): # the following line is necessary from verl.utils.model import LambdaLayer, print_model_size, squeeze from verl.utils.torch_dtypes import PrecisionType from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy, MixedPrecision from torch import optim local_path = copy_local_path_from_hdfs(config.model.path) # note that the tokenizer between actor and critic may be different. So override tokenizer info with actor info # using random initialized model from any architecture. May not be the same as Actor. tokenizer_path = copy_local_path_from_hdfs(config.model.tokenizer_path) self.tokenizer = hf_tokenizer(tokenizer_path, trust_remote_code=config.model.get('trust_remote_code', False)) from omegaconf import OmegaConf override_config = OmegaConf.to_container(self.config.model.get('override_config', OmegaConf.create())) override_config_kwargs = { 'bos_token_id': self.tokenizer.bos_token_id, 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, } override_config_kwargs.update(override_config) if self.rank == 0: print(f'Critic overriding config {override_config_kwargs}') torch_dtype = self.config.model.fsdp_config.get('model_dtype', 'fp32') torch_dtype = PrecisionType.to_dtype(torch_dtype) from transformers import AutoConfig, AutoModelForTokenClassification from torch import nn trust_remote_code = False critic_model_config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code) critic_model_config.num_labels = 1 use_remove_padding = config.model.get('use_remove_padding', False) if use_remove_padding: from verl.models.registry import check_model_support_rmpad check_model_support_rmpad(critic_model_config.model_type) if use_remove_padding and self.ulysses_sequence_parallel_size > 1: from verl.models.transformers.monkey_patch import apply_monkey_patch apply_monkey_patch(critic_model_config, verbose=True) init_context = get_init_weight_context_manager() with init_context(), warnings.catch_warnings(): warnings.simplefilter(\"ignore\") setattr(critic_model_config, 'classifier_dropout', 0.) setattr(critic_model_config, 'hidden_dropout', '0') critic_module = AutoModelForTokenClassification.from_pretrained(pretrained_model_name_or_path=local_path, torch_dtype=torch_dtype, config=critic_model_config, attn_implementation='flash_attention_2', trust_remote_code=trust_remote_code) # some parameters may not in torch_dtype critic_module.to(torch_dtype) if config.model.get('enable_gradient_checkpointing', False): critic_module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) if self.rank == 0: print_model_size(critic_module) self.critic_model_config = critic_model_config fsdp_config = self.config.model.fsdp_config mixed_precision_config = fsdp_config.get('mixed_precision', None) 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) auto_wrap_policy = get_fsdp_wrap_policy(module=critic_module, config=self.config.model.fsdp_config.wrap_policy) log_gpu_memory_usage('Before critic FSDP', logger=None) critic_module = FSDP(critic_module","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/fsdp_workers.py#CriticWorker.__init__","name":"__init__","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L540-L569","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"dbcf8454d56f80ee886ebc2f2d7f3aff2959ed4fe2912a5c39549eac8a4a243c","body_text":"def __init__(self, config): super().__init__() import torch.distributed if not torch.distributed.is_initialized(): torch.distributed.init_process_group(backend=\"nccl\") self.config = config # build device mesh for Ulysses Sequence Parallel world_size = torch.distributed.get_world_size() from torch.distributed.device_mesh import init_device_mesh self.ulysses_device_mesh = None self.ulysses_sequence_parallel_size = self.config.get('ulysses_sequence_parallel_size', 1) dp = world_size // self.ulysses_sequence_parallel_size if self.ulysses_sequence_parallel_size > 1: self.ulysses_device_mesh = init_device_mesh('cuda', mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=['dp', 'sp']) self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) # set FSDP offload params self._is_offload_param = self.config.model.fsdp_config.param_offload self._is_offload_grad = self.config.model.fsdp_config.grad_offload self._is_offload_optimizer = self.config.model.fsdp_config.optimizer_offload # normalize config self.config.ppo_mini_batch_size //= (torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size) self.config.ppo_micro_batch_size //= (torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size) self.config.forward_micro_batch_size //= (torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size)","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","name":"_build_critic_model_optimizer","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L571-L680","signature":"def _build_critic_model_optimizer(self, config): # the following line is necessary","exported":false,"origin":"ast","body_hash":"df11f7be2270b736109deb66c7c379739c462eb04dac91ea8065a2a517500978","body_text":"def _build_critic_model_optimizer(self, config): # the following line is necessary from verl.utils.model import LambdaLayer, print_model_size, squeeze from verl.utils.torch_dtypes import PrecisionType from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy, MixedPrecision from torch import optim local_path = copy_local_path_from_hdfs(config.model.path) # note that the tokenizer between actor and critic may be different. So override tokenizer info with actor info # using random initialized model from any architecture. May not be the same as Actor. tokenizer_path = copy_local_path_from_hdfs(config.model.tokenizer_path) self.tokenizer = hf_tokenizer(tokenizer_path, trust_remote_code=config.model.get('trust_remote_code', False)) from omegaconf import OmegaConf override_config = OmegaConf.to_container(self.config.model.get('override_config', OmegaConf.create())) override_config_kwargs = { 'bos_token_id': self.tokenizer.bos_token_id, 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, } override_config_kwargs.update(override_config) if self.rank == 0: print(f'Critic overriding config {override_config_kwargs}') torch_dtype = self.config.model.fsdp_config.get('model_dtype', 'fp32') torch_dtype = PrecisionType.to_dtype(torch_dtype) from transformers import AutoConfig, AutoModelForTokenClassification from torch import nn trust_remote_code = False critic_model_config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code) critic_model_config.num_labels = 1 use_remove_padding = config.model.get('use_remove_padding', False) if use_remove_padding: from verl.models.registry import check_model_support_rmpad check_model_support_rmpad(critic_model_config.model_type) if use_remove_padding and self.ulysses_sequence_parallel_size > 1: from verl.models.transformers.monkey_patch import apply_monkey_patch apply_monkey_patch(critic_model_config, verbose=True) init_context = get_init_weight_context_manager() with init_context(), warnings.catch_warnings(): warnings.simplefilter(\"ignore\") setattr(critic_model_config, 'classifier_dropout', 0.) setattr(critic_model_config, 'hidden_dropout', '0') critic_module = AutoModelForTokenClassification.from_pretrained(pretrained_model_name_or_path=local_path, torch_dtype=torch_dtype, config=critic_model_config, attn_implementation='flash_attention_2', trust_remote_code=trust_remote_code) # some parameters may not in torch_dtype critic_module.to(torch_dtype) if config.model.get('enable_gradient_checkpointing', False): critic_module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={'use_reentrant': False}) if self.rank == 0: print_model_size(critic_module) self.critic_model_config = critic_model_config fsdp_config = self.config.model.fsdp_config mixed_precision_config = fsdp_config.get('mixed_precision', None) 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) auto_wrap_policy = get_fsdp_wrap_policy(module=critic_module, config=self.config.model.fsdp_config.wrap_policy) log_gpu_memory_usage('Before critic FSDP', logger=None) critic_module = FSDP(critic_module, param_init_fn=init_fn, use_orig_params=False, auto_wrap_policy=auto_wrap_policy, device_id=torch.cuda.current_device(), sharding_strategy=ShardingStrategy.FULL_SHARD, mixed_precision=mixed_precision, sync_module_states=True, forward_prefetch=False) log_gpu_memory_usage('After critic FSDP', logger=None) critic_optimizer = optim.AdamW(critic_module.parameters(), lr=config.optim.lr, betas=config.optim.get('betas', (0.9, 0.999)), weight_decay=config.optim.get('weight_decay', 1e-2)) total_steps = config.optim.get('total_training_steps', 0) num_warmup_steps_ratio = config.optim.get('lr_warmup_steps_ratio', 0.) num_warmup_steps = int(num_warmup_steps_ratio * total_steps) print(f'Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}') from verl.utils.torch_functional import get_constant_schedule_with_warmup critic_lr_scheduler = get_constant_schedule_with_warmup(optimizer=critic_optimizer, num_warmup_steps=num_warmup_steps) return critic_module, critic_optimizer, critic_lr_scheduler","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/fsdp_workers.py#CriticWorker.init_model","name":"init_model","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L683-L702","signature":"def init_model(self): # This is used to import external_lib into the huggingface systems","exported":true,"origin":"ast","body_hash":"451dadbafd9c10ec4caa08bcfdb959995844197b02befc605031a65d24207336","body_text":"def init_model(self): # This is used to import external_lib into the huggingface systems import_external_libs(self.config.model.get('external_lib', None)) from verl.workers.critic import DataParallelPPOCritic self.critic_module, self.critic_optimizer, self.critic_lr_scheduler = self._build_critic_model_optimizer( self.config) if self._is_offload_param: offload_fsdp_param_and_grad(module=self.critic_module, offload_grad=self._is_offload_grad) if self._is_offload_optimizer: offload_fsdp_optimizer(optimizer=self.critic_optimizer) self.critic = DataParallelPPOCritic(config=self.config, critic_module=self.critic_module, critic_optimizer=self.critic_optimizer) self.flops_counter = FlopsCounter(self.critic_model_config) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","name":"compute_values","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L705-L727","signature":"def compute_values(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"4b15b743ee7e4370212cac3089e38c0c3c4bb8d22844671625b576e045acca55","body_text":"def compute_values(self, data: DataProto): data = data.to('cuda') if self._is_offload_param: load_fsdp_param_and_grad(module=self.critic_module, device_id=torch.cuda.current_device(), load_grad=self._is_offload_grad) micro_batch_size = self.config.forward_micro_batch_size data.meta_info['micro_batch_size'] = micro_batch_size data.meta_info['max_token_len'] = self.config.forward_max_token_len_per_gpu data.meta_info['use_dynamic_bsz'] = self.config.use_dynamic_bsz # perform forward computation with self.ulysses_sharding_manager: data = self.ulysses_sharding_manager.preprocess_data(data=data) values = self.critic.compute_values(data=data) output = DataProto.from_dict(tensors={'values': values}) output = self.ulysses_sharding_manager.postprocess_data(data=output) output = output.to('cpu') if self._is_offload_param: offload_fsdp_param_and_grad(module=self.critic_module, offload_grad=self._is_offload_grad) torch.cuda.empty_cache() return output","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","name":"update_critic","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L730-L764","signature":"def update_critic(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"5f7d4a73e02343452e348dd3ebb6abfc89e5a9b4596fec1b2417106f229a6124","body_text":"def update_critic(self, data: DataProto): data = data.to('cuda') if self._is_offload_param: load_fsdp_param_and_grad(module=self.critic_module, device_id=torch.cuda.current_device(), load_grad=self._is_offload_grad) if self._is_offload_optimizer: load_fsdp_optimizer(optimizer=self.critic_optimizer, device_id=torch.cuda.current_device()) # perform forward computation with self.ulysses_sharding_manager: data = self.ulysses_sharding_manager.preprocess_data(data=data) with Timer(name='update_critic', logger=None) as timer: metrics = self.critic.update_critic(data=data) delta_time = timer.last global_num_tokens = data.meta_info['global_token_num'] estimated_flops, promised_flops = self.flops_counter.estimate_flops(global_num_tokens, delta_time) metrics['mfu/critic'] = estimated_flops * self.config.ppo_epochs / promised_flops / self.world_size self.critic_lr_scheduler.step() lr = self.critic_lr_scheduler.get_last_lr()[0] metrics['critic/lr'] = lr output = DataProto(batch=None, meta_info={'metrics': metrics}) output = self.ulysses_sharding_manager.postprocess_data(data=output) if self._is_offload_param: offload_fsdp_param_and_grad(module=self.critic_module, offload_grad=self._is_offload_grad) if self._is_offload_optimizer: offload_fsdp_optimizer(optimizer=self.critic_optimizer) torch.cuda.empty_cache() output = output.to('cpu') return output","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","name":"save_checkpoint","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L767-L792","signature":"def save_checkpoint(self, local_path, hdfs_path=None)","exported":true,"origin":"ast","body_hash":"4ae333b7eb09d2d740beddd6a6342d03c1c2f25a90861fc71843035d28ac2db8","body_text":"def save_checkpoint(self, local_path, hdfs_path=None): import torch if self._is_offload_param: load_fsdp_param_and_grad(module=self.critic_module, device_id=torch.cuda.current_device(), load_grad=self._is_offload_grad) # TODO: support DCP and save sharded checkpoints import torch.distributed from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, StateDictType, FullStateDictConfig cfg = FullStateDictConfig(offload_to_cpu=True, rank0_only=True) with FSDP.state_dict_type(self.critic_module, StateDictType.FULL_STATE_DICT, cfg): state_dict = self.critic_module.state_dict() if self.rank == 0: print(f'Saving critic checkpoint to {local_path}') os.makedirs(local_path, exist_ok=True) self.critic_module._fsdp_wrapped_module.save_pretrained(local_path, state_dict=state_dict) self.tokenizer.save_pretrained(local_path) if hdfs_path is not None: print(f'Uploading critic checkpoint to {hdfs_path}') hdfs_io.makedirs(hdfs_path, exist_ok=True) hdfs_io.copy(src=local_path, dst=hdfs_path) torch.distributed.barrier() if self._is_offload_param: offload_fsdp_param_and_grad(module=self.critic_module, offload_grad=self._is_offload_grad)","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/fsdp_workers.py#RewardModelWorker","name":"RewardModelWorker","kind":"class","path":"verl/workers/fsdp_workers.py","span":"L796-L1054","signature":"class RewardModelWorker(Worker)","exported":true,"origin":"ast","body_hash":"b20ef126039468290f0d0f61601a5525493765f5aded59682e95693da7b8751e","body_text":"class RewardModelWorker(Worker): \"\"\" Note that we only implement the reward model that is subclass of AutoModelForTokenClassification. \"\"\" def __init__(self, config): super().__init__() import torch.distributed if not torch.distributed.is_initialized(): torch.distributed.init_process_group(backend=\"nccl\") self.config = config # build device mesh for Ulysses Sequence Parallel world_size = torch.distributed.get_world_size() from torch.distributed.device_mesh import init_device_mesh self.ulysses_device_mesh = None self.ulysses_sequence_parallel_size = self.config.get('ulysses_sequence_parallel_size', 1) dp = world_size // self.ulysses_sequence_parallel_size if self.ulysses_sequence_parallel_size > 1: self.ulysses_device_mesh = init_device_mesh('cuda', mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=['dp', 'sp']) self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) self.use_remove_padding = self.config.model.get('use_remove_padding', False) self.config.micro_batch_size //= torch.distributed.get_world_size() def _build_model(self, config): # the following line is necessary from transformers import AutoModelForTokenClassification, AutoConfig from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy, CPUOffload # download the checkpoint from hdfs local_path = copy_local_path_from_hdfs(config.model.path) if self.config.model.input_tokenizer is None: self._do_switch_chat_template = False else: self._do_switch_chat_template = True input_tokenizer_local_path = copy_local_path_from_hdfs(config.model.input_tokenizer) self.input_tokenizer = hf_tokenizer(input_tokenizer_local_path, trust_remote_code=config.model.get('trust_remote_code', False)) self.tokenizer = hf_tokenizer(local_path, trust_remote_code=config.model.get('trust_remote_code', False)) trust_remote_code = config.model.get('trust_remote_code', False) model_config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code) model_config.num_labels = 1 use_remove_padding = config.model.get('use_remove_padding', False) if use_remove_padding: from verl.models.registry import check_model_support_rmpad check_model_support_rmpad(model_config.model_type) if use_remove_padding and self.ulysses_sequence_parallel_size > 1: from verl.models.transformers.monkey_patch import apply_monkey_patch apply_monkey_patch(model_config, verbose=True) # note that we have to create model in fp32. Otherwise, the optimizer is in bf16, which is incorrect init_context = get_init_weight_context_manager(use_meta_tensor=not model_config.tie_word_embeddings) with init_context(), warnings.catch_warnings(): warnings.simplefilter(\"ignore\") setattr(model_config, 'classifier_dropout', 0.) reward_module = AutoModelForTokenClassification.from_pretrained(pretrained_model_name_or_path=local_path, config=model_config, torch_dtype=torch.bfloat16, attn_implementation='flash_attention_2', trust_remote_code=trust_remote_code) reward_module.to(torch.bfloat16) auto_wrap_policy = get_fsdp_wrap_policy(module=reward_module, config=self.config.model.fsdp_config) reward_module = FSDP( reward_module, param_init_fn=init_fn, use_orig_params=False, auto_wrap_policy=auto_wrap_policy, device_id=torch.cuda.current_device(), sharding_strategy=ShardingStrategy.FULL_SHARD, # zero3 sync_module_states=True, cpu_offload=CPUOffload(offload_params=self.config.model.fsdp_config.param_offload), forward_prefetch=False) return reward_module @register(dispatch_mode=Dispatch.ONE_TO_ALL) def init_model(self): # This is used to import external_lib into the huggingface systems import_external_libs(self.config.model.get('external_lib', None)) self.reward_module = self._build_model(config=self.config) torch.cuda.empty_cache() def _forward_micro_batch(self, micro_batch): from flash_attn.bert_padding import pad_input, unpad_input, index_first_axis, rearrange from verl.utils.ulysses import ulysses_pad_and_slice_inputs, gather_outpus_and_unpad with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16): input_ids = micro_batch['input_ids'] batch_size, seqlen = input_ids.shape attention_mask = micro_batch['attention_mask'] position_ids = micro_batch['position_ids'] if self.use_remove_padding: 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) # pad and slice the inputs if sp > 1 if self.ulysses_sequence_parallel_size > 1: input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs(input_ids_rmpad, \\ position_ids_rmpad, \\ sp_size=self.ulysses_sequence_parallel_size) # only pass input_ids and position_ids to enable flash_attn_varlen output = self.reward_module(input_ids=input_ids_rmpad, attention_mask=None, position","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","name":"__init__","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L801-L822","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"1707509d2f11cec2a6da8726f248594aeec2ffe9acbe015b4356e87b8c6d7d49","body_text":"def __init__(self, config): super().__init__() import torch.distributed if not torch.distributed.is_initialized(): torch.distributed.init_process_group(backend=\"nccl\") self.config = config # build device mesh for Ulysses Sequence Parallel world_size = torch.distributed.get_world_size() from torch.distributed.device_mesh import init_device_mesh self.ulysses_device_mesh = None self.ulysses_sequence_parallel_size = self.config.get('ulysses_sequence_parallel_size', 1) dp = world_size // self.ulysses_sequence_parallel_size if self.ulysses_sequence_parallel_size > 1: self.ulysses_device_mesh = init_device_mesh('cuda', mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=['dp', 'sp']) self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) self.use_remove_padding = self.config.model.get('use_remove_padding', False) self.config.micro_batch_size //= torch.distributed.get_world_size()","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"},{"id":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","name":"_build_model","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L824-L879","signature":"def _build_model(self, config): # the following line is necessary","exported":false,"origin":"ast","body_hash":"8640a92a32ada7cec433206fa43d0679da42b0c9fa14049daad4f4b4641a23c3","body_text":"def _build_model(self, config): # the following line is necessary from transformers import AutoModelForTokenClassification, AutoConfig from torch.distributed.fsdp import FullyShardedDataParallel as FSDP, ShardingStrategy, CPUOffload # download the checkpoint from hdfs local_path = copy_local_path_from_hdfs(config.model.path) if self.config.model.input_tokenizer is None: self._do_switch_chat_template = False else: self._do_switch_chat_template = True input_tokenizer_local_path = copy_local_path_from_hdfs(config.model.input_tokenizer) self.input_tokenizer = hf_tokenizer(input_tokenizer_local_path, trust_remote_code=config.model.get('trust_remote_code', False)) self.tokenizer = hf_tokenizer(local_path, trust_remote_code=config.model.get('trust_remote_code', False)) trust_remote_code = config.model.get('trust_remote_code', False) model_config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code) model_config.num_labels = 1 use_remove_padding = config.model.get('use_remove_padding', False) if use_remove_padding: from verl.models.registry import check_model_support_rmpad check_model_support_rmpad(model_config.model_type) if use_remove_padding and self.ulysses_sequence_parallel_size > 1: from verl.models.transformers.monkey_patch import apply_monkey_patch apply_monkey_patch(model_config, verbose=True) # note that we have to create model in fp32. Otherwise, the optimizer is in bf16, which is incorrect init_context = get_init_weight_context_manager(use_meta_tensor=not model_config.tie_word_embeddings) with init_context(), warnings.catch_warnings(): warnings.simplefilter(\"ignore\") setattr(model_config, 'classifier_dropout', 0.) reward_module = AutoModelForTokenClassification.from_pretrained(pretrained_model_name_or_path=local_path, config=model_config, torch_dtype=torch.bfloat16, attn_implementation='flash_attention_2', trust_remote_code=trust_remote_code) reward_module.to(torch.bfloat16) auto_wrap_policy = get_fsdp_wrap_policy(module=reward_module, config=self.config.model.fsdp_config) reward_module = FSDP( reward_module, param_init_fn=init_fn, use_orig_params=False, auto_wrap_policy=auto_wrap_policy, device_id=torch.cuda.current_device(), sharding_strategy=ShardingStrategy.FULL_SHARD, # zero3 sync_module_states=True, cpu_offload=CPUOffload(offload_params=self.config.model.fsdp_config.param_offload), forward_prefetch=False) return reward_module","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"},{"id":"verl/workers/fsdp_workers.py#RewardModelWorker.init_model","name":"init_model","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L882-L886","signature":"def init_model(self): # This is used to import external_lib into the huggingface systems","exported":true,"origin":"ast","body_hash":"c235e9d31de0de2b3f81d1c4e89fdb5a443b818ad6ae615c7a074dd056626254","body_text":"def init_model(self): # This is used to import external_lib into the huggingface systems import_external_libs(self.config.model.get('external_lib', None)) self.reward_module = self._build_model(config=self.config) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"},{"id":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","name":"_forward_micro_batch","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L888-L940","signature":"def _forward_micro_batch(self, micro_batch)","exported":false,"origin":"ast","body_hash":"093881710ccea4d56ebfedc00da6d964aadbeef0dd966e3c7e80ccc783855b70","body_text":"def _forward_micro_batch(self, micro_batch): from flash_attn.bert_padding import pad_input, unpad_input, index_first_axis, rearrange from verl.utils.ulysses import ulysses_pad_and_slice_inputs, gather_outpus_and_unpad with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.bfloat16): input_ids = micro_batch['input_ids'] batch_size, seqlen = input_ids.shape attention_mask = micro_batch['attention_mask'] position_ids = micro_batch['position_ids'] if self.use_remove_padding: 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) # pad and slice the inputs if sp > 1 if self.ulysses_sequence_parallel_size > 1: input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs(input_ids_rmpad, \\ position_ids_rmpad, \\ sp_size=self.ulysses_sequence_parallel_size) # only pass input_ids and position_ids to enable flash_attn_varlen output = self.reward_module(input_ids=input_ids_rmpad, attention_mask=None, position_ids=position_ids_rmpad, use_cache=False) # prevent model thinks we are generating reward_rmpad = output.logits reward_rmpad = reward_rmpad.squeeze(0) # (total_nnz) # gather output if sp > 1 if self.ulysses_sequence_parallel_size > 1: reward_rmpad = gather_outpus_and_unpad(reward_rmpad, gather_dim=0, unpad_dim=0, padding_size=pad_size) # pad it back rm_score = pad_input(reward_rmpad, indices=indices, batch=batch_size, seqlen=seqlen).squeeze(-1) else: output = self.reward_module(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) rm_score = output.logits # (batch_size, seq_len, 1) rm_score = rm_score.squeeze(-1) # extract the result of the last valid token eos_mask_idx = torch.argmax(position_ids * attention_mask, dim=-1) # (bsz,) rm_score = rm_score[torch.arange(batch_size), eos_mask_idx] return rm_score","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"},{"id":"verl/workers/fsdp_workers.py#RewardModelWorker._expand_to_token_level","name":"_expand_to_token_level","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L942-L955","signature":"def _expand_to_token_level(self, data: DataProto, scores: torch.Tensor)","exported":false,"origin":"ast","body_hash":"d275cf2afd0e71d24cdb396800177afe9062ebcffcf0cefa16c32041a8f84694","body_text":"def _expand_to_token_level(self, data: DataProto, scores: torch.Tensor): batch_size = data.batch.batch_size[0] # expand as token_level_reward attention_mask = data.batch['attention_mask'] position_ids = data.batch['position_ids'] response_length = data.batch['responses'].shape[-1] eos_mask_idx = torch.argmax(position_ids * attention_mask, dim=-1) # (bsz,) token_level_scores = torch.zeros_like(attention_mask, dtype=scores.dtype) # (bsz, seqlen) token_level_scores[torch.arange(batch_size), eos_mask_idx] = scores # select the response part token_level_scores = token_level_scores[:, -response_length:] return token_level_scores","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"},{"id":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","name":"_switch_chat_template","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L957-L1012","signature":"def _switch_chat_template(self, data: DataProto)","exported":false,"origin":"ast","body_hash":"c697e926b0846c698f8ec90018d547959deec8e5e07dbeec228f4bc5ad0f4b29","body_text":"def _switch_chat_template(self, data: DataProto): src_max_length = data.batch['attention_mask'].shape[-1] src_tokenizer = self.input_tokenizer target_tokenizer = self.tokenizer rm_input_ids = [] rm_attention_mask = [] for i in range(data.batch.batch_size[0]): # extract raw prompt chat: list = data.non_tensor_batch['raw_prompt'][i].tolist() # extract response response_ids = data.batch['responses'][i] response_length = response_ids.shape[-1] valid_response_length = data.batch['attention_mask'][i][-response_length:].sum() valid_response_ids = response_ids[:valid_response_length] # decode response = src_tokenizer.decode(valid_response_ids) # remove bos and eos response = response.replace(src_tokenizer.eos_token, '') chat.append({'role': 'assistant', 'content': response}) prompt_with_chat_template = target_tokenizer.apply_chat_template(chat, add_generation_prompt=False, tokenize=False) if self.rank == 0 and i == 0: # for debugging purpose print(f'Switch template. chat: {prompt_with_chat_template}') # the maximum length is actually determined by the reward model itself max_length = self.config.get('max_length', src_max_length) if max_length is None: max_length = src_max_length input_ids, attention_mask = verl_F.tokenize_and_postprocess_data( prompt=prompt_with_chat_template, tokenizer=target_tokenizer, max_length=max_length, pad_token_id=target_tokenizer.pad_token_id, left_pad=False, # right padding truncation=self.config.get('truncation', 'right')) # truncate from the right rm_input_ids.append(input_ids) rm_attention_mask.append(attention_mask) rm_input_ids = torch.cat(rm_input_ids, dim=0) rm_attention_mask = torch.cat(rm_attention_mask, dim=0) rm_position_ids = compute_position_id_with_mask(rm_attention_mask) rm_inputs = {'input_ids': rm_input_ids, 'attention_mask': rm_attention_mask, 'position_ids': rm_position_ids} return DataProto.from_dict(rm_inputs)","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"},{"id":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","name":"compute_rm_score","kind":"method","path":"verl/workers/fsdp_workers.py","span":"L1015-L1054","signature":"def compute_rm_score(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"da94e90aef38a332e2c5c389e3cffce3c09a153e60124fcdc6fe61f18d7056ab","body_text":"def compute_rm_score(self, data: DataProto): import itertools from verl.utils.seqlen_balancing import rearrange_micro_batches, get_reverse_idx data = data.to('cuda') if self._do_switch_chat_template: rm_data = self._switch_chat_template(data) rm_data.batch = rm_data.batch.cuda() # perform forward computation with self.ulysses_sharding_manager: rm_data = self.ulysses_sharding_manager.preprocess_data(data=rm_data) data = self.ulysses_sharding_manager.preprocess_data(data=data) use_dynamic_bsz = self.config.use_dynamic_bsz if use_dynamic_bsz: max_token_len = self.config.forward_max_token_len_per_gpu * self.ulysses_sequence_parallel_size micro_batches, indices = rearrange_micro_batches(batch=rm_data.batch, max_token_len=max_token_len) else: micro_batches = rm_data.batch.split(self.config.micro_batch_size) output = [] for micro_batch in micro_batches: rm_score = self._forward_micro_batch(micro_batch) output.append(rm_score) scores = torch.cat(output, dim=0) # (batch_size) if use_dynamic_bsz: indices = list(itertools.chain.from_iterable(indices)) assert len(indices) == scores.size(0), f\"{len(indices)} vs. {scores.size()}\" revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long) scores = scores[revert_indices] token_level_scores = self._expand_to_token_level(data, scores) # Note that this is only the scores, may not be the final rewards used to train RL output = DataProto.from_dict(tensors={'rm_scores': token_level_scores}) output = self.ulysses_sharding_manager.postprocess_data(data=output) output = output.to('cpu') torch.cuda.empty_cache() return output","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"}],"rawEdges":[{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"logging","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"os","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"warnings","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"omegaconf","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.single_controller.base","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.single_controller.base.decorator","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.debug","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.fs","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.fsdp_utils","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.fsdp_utils","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.import_utils","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.flops_counter","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.workers.sharding_manager.fsdp_ulysses","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"codetiming","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"calls","name":"getLogger","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"calls","name":"setLevel","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"calls","name":"getenv","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"contains","targetId":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"extends","name":"Worker","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"init_process_group","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed.device_mesh","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"init_device_mesh","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"init_device_mesh","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"FSDPUlyssesShardingManager","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.torch_dtypes","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"transformers","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed.fsdp","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.models.registry","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"check_model_support_rmpad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.models.transformers.monkey_patch","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"apply_monkey_patch","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"update","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"update_model_config","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get_init_weight_context_manager","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"init_context","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"catch_warnings","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"simplefilter","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"gradient_checkpointing_enable","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"barrier","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"print_model_size","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"MixedPrecision","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get_fsdp_wrap_policy","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"FSDP","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"AdamW","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"parameters","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FullyShardedDataParallel"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"int","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get_constant_schedule_with_warmup","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed.device_mesh","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"init_device_mesh","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.workers.rollout","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.workers.sharding_manager","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"HFRollout","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"BaseShardingManager","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.workers.rollout.vllm_rollout","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.workers.sharding_manager","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"vLLMRollout","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"FSDPVLLMShardingManager","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.workers.actor","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"import_external_libs","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"omegaconf","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"to_container","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"_build_model_optimizer","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"ActorRolloutRefWorker"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"offload_fsdp_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"offload_fsdp_optimizer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"set_struct","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"open_dict","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"DataParallelPPOActor","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"_build_rollout","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"ActorRolloutRefWorker"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"_build_model_optimizer","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"ActorRolloutRefWorker"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"offload_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"set_struct","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"open_dict","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"DataParallelPPOActor","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"FlopsCounter","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataProto"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"load_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"load_fsdp_optimizer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"preprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"Timer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"update_policy","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataParallelPPOActor"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"estimate_flops","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FlopsCounter"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"step","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"get_last_lr","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"postprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataProto"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"offload_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"offload_fsdp_optimizer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataProto"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"load_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"update","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"preprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"compute_log_prob","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataParallelPPOActor"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"from_dict","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"postprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"offload_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_log_prob","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataProto"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"load_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"update","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"preprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"generate_sequences","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"postprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"preprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"compute_log_prob","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataParallelPPOActor"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"postprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"offload_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataProto"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"load_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"preprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"compute_log_prob","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataParallelPPOActor"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"from_dict","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"postprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"offload_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"load_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed.fsdp","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"FullStateDictConfig","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"state_dict_type","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"state_dict","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"save_pretrained","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"save_pretrained","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"hf_tokenizer"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"copy","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"barrier","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#ActorRolloutRefWorker.save_checkpoint","relation":"calls","name":"offload_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"contains","targetId":"verl/workers/fsdp_workers.py#CriticWorker","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"extends","name":"Worker","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#CriticWorker.__init__","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"init_process_group","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed.device_mesh","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"init_device_mesh","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"FSDPUlyssesShardingManager","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.__init__","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.torch_dtypes","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed.fsdp","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"omegaconf","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"to_container","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"update","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"transformers","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.models.registry","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"check_model_support_rmpad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.models.transformers.monkey_patch","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"apply_monkey_patch","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get_init_weight_context_manager","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"init_context","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"catch_warnings","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"simplefilter","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"setattr","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"setattr","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FullyShardedDataParallel"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"gradient_checkpointing_enable","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FullyShardedDataParallel"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"print_model_size","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"to_dtype","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"MixedPrecision","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get_fsdp_wrap_policy","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"FSDP","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"AdamW","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"parameters","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FullyShardedDataParallel"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"int","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get_constant_schedule_with_warmup","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#CriticWorker.init_model","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.init_model","relation":"calls","name":"import_external_libs","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.workers.critic","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.init_model","relation":"calls","name":"_build_critic_model_optimizer","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"CriticWorker"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.init_model","relation":"calls","name":"offload_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.init_model","relation":"calls","name":"offload_fsdp_optimizer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.init_model","relation":"calls","name":"DataParallelPPOCritic","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.init_model","relation":"calls","name":"FlopsCounter","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.init_model","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataProto"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","relation":"calls","name":"load_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","relation":"calls","name":"preprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","relation":"calls","name":"compute_values","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataParallelPPOCritic"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","relation":"calls","name":"from_dict","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","relation":"calls","name":"postprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","relation":"calls","name":"offload_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.compute_values","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataProto"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"load_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"load_fsdp_optimizer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"preprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"Timer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"update_critic","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataParallelPPOCritic"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"estimate_flops","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FlopsCounter"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"step","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"get_last_lr","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"postprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"offload_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"offload_fsdp_optimizer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.update_critic","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataProto"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"load_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed.fsdp","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"FullStateDictConfig","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"state_dict_type","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"state_dict","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"save_pretrained","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"save_pretrained","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"hf_tokenizer"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"makedirs","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"copy","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"barrier","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#CriticWorker.save_checkpoint","relation":"calls","name":"offload_fsdp_param_and_grad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"contains","targetId":"verl/workers/fsdp_workers.py#RewardModelWorker","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker","relation":"extends","name":"Worker","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","relation":"calls","name":"init_process_group","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed.device_mesh","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","relation":"calls","name":"init_device_mesh","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","relation":"calls","name":"FSDPUlyssesShardingManager","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.__init__","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"transformers","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"torch.distributed.fsdp","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.models.registry","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"check_model_support_rmpad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.models.transformers.monkey_patch","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"apply_monkey_patch","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"get_init_weight_context_manager","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"init_context","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"catch_warnings","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"simplefilter","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"setattr","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FullyShardedDataParallel"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"get_fsdp_wrap_policy","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"FSDP","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._build_model","relation":"calls","name":"CPUOffload","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#RewardModelWorker.init_model","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.init_model","relation":"calls","name":"import_external_libs","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.init_model","relation":"calls","name":"_build_model","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"RewardModelWorker"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.init_model","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"flash_attn.bert_padding","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.ulysses","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"autocast","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"unpad_input","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"transpose","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"transpose","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"index_first_axis","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"rearrange","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"ulysses_pad_and_slice_inputs","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"reward_module","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"RewardModelWorker"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"gather_outpus_and_unpad"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"gather_outpus_and_unpad","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"pad_input","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"reward_module","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"RewardModelWorker"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"squeeze","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"argmax","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._forward_micro_batch","relation":"calls","name":"arange","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#RewardModelWorker._expand_to_token_level","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._expand_to_token_level","relation":"calls","name":"argmax","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._expand_to_token_level","relation":"calls","name":"zeros_like","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._expand_to_token_level","relation":"calls","name":"arange","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"tolist","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"sum","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"decode","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"replace","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"list"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"apply_chat_template","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"tokenize_and_postprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"compute_position_id_with_mask","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker._switch_chat_template","relation":"calls","name":"from_dict","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"itertools","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py","relation":"imports","specifier":"verl.utils.seqlen_balancing","file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"DataProto"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"_switch_chat_template","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"RewardModelWorker"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"preprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"preprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"rearrange_micro_batches","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"split","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"_forward_micro_batch","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"RewardModelWorker"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"list","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"from_iterable","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"tensor","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"get_reverse_idx","viaMember":false,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"_expand_to_token_level","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"RewardModelWorker"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"from_dict","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"postprocess_data","viaMember":true,"file":"verl/workers/fsdp_workers.py","recvType":"FSDPUlyssesShardingManager"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/fsdp_workers.py"},{"source":"verl/workers/fsdp_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/fsdp_workers.py"}]},"verl/workers/megatron_workers.py":{"size":36430,"mtimeMs":1789275830328.2031,"hash":"bf136fe2d2595dd78a1df92234ecd6bd6b9a0ed0252e8915ebdd7caa88449fc0","nodes":[{"id":"verl/workers/megatron_workers.py","name":"megatron_workers.py","kind":"file","path":"verl/workers/megatron_workers.py","span":"L1-L736","signature":null,"exported":true,"origin":"ast","body_hash":"bf136fe2d2595dd78a1df92234ecd6bd6b9a0ed0252e8915ebdd7caa88449fc0","chars":36430,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 main entry point to run the PPO algorithm \"\"\" import os import logging import ray import torch import torch.distributed import torch.nn as nn from omegaconf import DictConfig from verl.single_controller.base.megatron.worker import MegatronWorker from verl.workers.actor.megatron_actor import MegatronPPOActor from verl.workers.critic.megatron_critic import MegatronPPOCritic from verl.workers.sharding_manager import AllGatherPPModel from verl.workers.reward_model.megatron.reward_model import MegatronRewardModel from verl.single_controller.base.decorator import register, Dispatch from verl import DataProto from verl.utils.fs import copy_local_path_from_hdfs from verl.utils.debug import log_gpu_memory_usage from verl.utils.model import load_megatron_model_weights from verl.utils.megatron_utils import init_model_parallel_config from verl.utils.megatron_utils import offload_megatron_param_and_grad, load_megatron_param_and_grad from verl.utils import hf_tokenizer from megatron.core import parallel_state as mpu from megatron.core import ModelParallelConfig logger = logging.getLogger(__file__) logger.setLevel(os.getenv('VERL_PPO_LOGGING_LEVEL', 'WARN'))"},{"id":"verl/workers/megatron_workers.py#set_random_seed","name":"set_random_seed","kind":"function","path":"verl/workers/megatron_workers.py","span":"L47-L60","signature":"def set_random_seed(seed)","exported":true,"origin":"ast","body_hash":"75e7be754ab3aabd20dd53e3b74f68fe37f499c17ad565f9d86e63d1423e9bf0","body_text":"def set_random_seed(seed): import torch import numpy as np import random torch.manual_seed(seed) np.random.seed(seed) random.seed(seed) if torch.cuda.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'","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","name":"ActorRolloutRefWorker","kind":"class","path":"verl/workers/megatron_workers.py","span":"L63-L413","signature":"class ActorRolloutRefWorker(MegatronWorker)","exported":true,"origin":"ast","body_hash":"3ef7b644b3bb6456c99f95deafdd5c1084f6d2d6485a2c36c5a62fd3475dc24f","body_text":"class ActorRolloutRefWorker(MegatronWorker): \"\"\" This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy or a hybrid engine based on the config.rollout \"\"\" def __init__(self, config: DictConfig, role: str): super().__init__() self.config = config # NOTE(sgm): We utilize colocate WorkerGroup by default. # As a result, Workers for different model share the same process. # Therefore, we only require one distribute initialization. # To utilize different parallel startegy in different models: # 1, users should disable WorkerDict; 2.assign different ResourcePool to different models, # 3. and apply the following patch in ray==2.10, https://github.com/ray-project/ray/pull/44385 if not torch.distributed.is_initialized(): rank = int(os.environ['LOCAL_RANK']) torch.distributed.init_process_group(backend=\"nccl\") torch.cuda.set_device(rank) if self.config.actor.megatron.sequence_parallel: os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' mpu.initialize_model_parallel( tensor_model_parallel_size=self.config.actor.megatron.tensor_model_parallel_size, pipeline_model_parallel_size=self.config.actor.megatron.pipeline_model_parallel_size, virtual_pipeline_model_parallel_size=None, pipeline_model_parallel_split_rank=None, use_sharp=False, context_parallel_size=1, expert_model_parallel_size=1, nccl_communicator_config_path=None, ) set_random_seed(seed=self.config.actor.megatron.seed) self.role = role 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'] # TODO(sgm): Currently, we only support reference model param offload # will support other offload later self._is_offload_param = False self._is_offload_grad = False self._is_offload_optimizer = False # normalize config if self._is_actor and self._is_rollout: self.config.actor.ppo_mini_batch_size //= mpu.get_data_parallel_world_size() self.config.actor.ppo_micro_batch_size //= mpu.get_data_parallel_world_size() self.config.rollout.log_prob_micro_batch_size //= mpu.get_data_parallel_world_size() self._is_offload_param = self.config.actor.get('param_offload', False) self._is_offload_grad = self.config.actor.get('grad_offload', False) self._is_offload_optimizer = self.config.actor.get('optimizer_offload', False) elif self._is_ref: self.config.ref.log_prob_micro_batch_size //= mpu.get_data_parallel_world_size() self._is_offload_param = self.config.ref.get('param_offload', False) def _build_model_optimizer(self, model_path, megatron_config: ModelParallelConfig, optim_config, override_model_config, enable_gradient_checkpointing=False): from verl.utils.megatron.optimizer import get_megatron_optimizer from megatron.core.models.gpt.gpt_model import ModelType from verl.utils.model import print_model_size, update_model_config from verl.utils.megatron_utils import get_model, init_megatron_optim_config from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig # Step 1: initialize the tokenizer local_path = copy_local_path_from_hdfs(model_path) self.tokenizer = hf_tokenizer(local_path) # Step 2: get the actor_model_config actor_model_config = AutoConfig.from_pretrained(local_path) override_config_kwargs = { 'bos_token_id': self.tokenizer.bos_token_id, 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, } override_config_kwargs.update(override_model_config) update_model_config(actor_model_config, override_config_kwargs=override_config_kwargs) if self.rank == 0: print(f'Model config after override: {actor_model_config}') def megatron_actor_model_provider(pre_process, post_process): from verl.utils.model import get_parallel_model_from_config # vpp is not supported yet because it will hang for some reason. Need debugging vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() # this will be set inside get_model # this_megatron_config = copy.deepcopy(megatron_config) # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank parallel_model = get_parallel_model_from_config(config=actor_model_config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process, value=False) parallel_model.cuda() return parallel_model # Step 3: initialize the megatron model if self._is_actor and self._is_rollout: # Initialize the 3D HybridEngine hybrid_engine = AllGatherPPModel(model_provider=megatron_actor_model_provider) # Fetch the model at current rank actor_module = hybrid_engine.this_rank_models if isinstance(actor_module, nn.ModuleList): actor_module = [actor_module[0]] if self.config.actor.load_weight: load_megatron_model_weights(self.config, actor_model_config, actor_module, params_dtype=megatron_config.params_dtype, is_value_model=False) if self.rank == 0: print_model_size(actor_module[0]) log_gpu_memory_usage('After AllGa","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","name":"__init__","kind":"method","path":"verl/workers/megatron_workers.py","span":"L69-L122","signature":"def __init__(self, config: DictConfig, role: str)","exported":false,"origin":"ast","body_hash":"cb66ddb97622b356f79d365d1741bd489c417467bbd3b78e07bd4bb8607377ba","body_text":"def __init__(self, config: DictConfig, role: str): super().__init__() self.config = config # NOTE(sgm): We utilize colocate WorkerGroup by default. # As a result, Workers for different model share the same process. # Therefore, we only require one distribute initialization. # To utilize different parallel startegy in different models: # 1, users should disable WorkerDict; 2.assign different ResourcePool to different models, # 3. and apply the following patch in ray==2.10, https://github.com/ray-project/ray/pull/44385 if not torch.distributed.is_initialized(): rank = int(os.environ['LOCAL_RANK']) torch.distributed.init_process_group(backend=\"nccl\") torch.cuda.set_device(rank) if self.config.actor.megatron.sequence_parallel: os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' mpu.initialize_model_parallel( tensor_model_parallel_size=self.config.actor.megatron.tensor_model_parallel_size, pipeline_model_parallel_size=self.config.actor.megatron.pipeline_model_parallel_size, virtual_pipeline_model_parallel_size=None, pipeline_model_parallel_split_rank=None, use_sharp=False, context_parallel_size=1, expert_model_parallel_size=1, nccl_communicator_config_path=None, ) set_random_seed(seed=self.config.actor.megatron.seed) self.role = role 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'] # TODO(sgm): Currently, we only support reference model param offload # will support other offload later self._is_offload_param = False self._is_offload_grad = False self._is_offload_optimizer = False # normalize config if self._is_actor and self._is_rollout: self.config.actor.ppo_mini_batch_size //= mpu.get_data_parallel_world_size() self.config.actor.ppo_micro_batch_size //= mpu.get_data_parallel_world_size() self.config.rollout.log_prob_micro_batch_size //= mpu.get_data_parallel_world_size() self._is_offload_param = self.config.actor.get('param_offload', False) self._is_offload_grad = self.config.actor.get('grad_offload', False) self._is_offload_optimizer = self.config.actor.get('optimizer_offload', False) elif self._is_ref: self.config.ref.log_prob_micro_batch_size //= mpu.get_data_parallel_world_size() self._is_offload_param = self.config.ref.get('param_offload', False)","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","name":"_build_model_optimizer","kind":"method","path":"verl/workers/megatron_workers.py","span":"L124-L214","signature":"def _build_model_optimizer(self, model_path, megatron_config: ModelParallelConfig, optim_config, override_model_config, enable_gradient_checkpointing=False)","exported":false,"origin":"ast","body_hash":"36ebeb2e232e0c8835e650228b53037fdb107d48fe3a98fb395144b67f20836b","body_text":"def _build_model_optimizer(self, model_path, megatron_config: ModelParallelConfig, optim_config, override_model_config, enable_gradient_checkpointing=False): from verl.utils.megatron.optimizer import get_megatron_optimizer from megatron.core.models.gpt.gpt_model import ModelType from verl.utils.model import print_model_size, update_model_config from verl.utils.megatron_utils import get_model, init_megatron_optim_config from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig # Step 1: initialize the tokenizer local_path = copy_local_path_from_hdfs(model_path) self.tokenizer = hf_tokenizer(local_path) # Step 2: get the actor_model_config actor_model_config = AutoConfig.from_pretrained(local_path) override_config_kwargs = { 'bos_token_id': self.tokenizer.bos_token_id, 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, } override_config_kwargs.update(override_model_config) update_model_config(actor_model_config, override_config_kwargs=override_config_kwargs) if self.rank == 0: print(f'Model config after override: {actor_model_config}') def megatron_actor_model_provider(pre_process, post_process): from verl.utils.model import get_parallel_model_from_config # vpp is not supported yet because it will hang for some reason. Need debugging vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() # this will be set inside get_model # this_megatron_config = copy.deepcopy(megatron_config) # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank parallel_model = get_parallel_model_from_config(config=actor_model_config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process, value=False) parallel_model.cuda() return parallel_model # Step 3: initialize the megatron model if self._is_actor and self._is_rollout: # Initialize the 3D HybridEngine hybrid_engine = AllGatherPPModel(model_provider=megatron_actor_model_provider) # Fetch the model at current rank actor_module = hybrid_engine.this_rank_models if isinstance(actor_module, nn.ModuleList): actor_module = [actor_module[0]] if self.config.actor.load_weight: load_megatron_model_weights(self.config, actor_model_config, actor_module, params_dtype=megatron_config.params_dtype, is_value_model=False) if self.rank == 0: print_model_size(actor_module[0]) log_gpu_memory_usage('After AllGatherPPModel init', logger=logger) elif self._is_ref: print(f'self.config.ref.load_weight: {self.config.ref.load_weight}') ref_module = get_model(model_provider_func=megatron_actor_model_provider, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=False) # ref_module = nn.ModuleList(ref_module) if self.config.ref.load_weight: # should align with the actor: assert self.config.actor.load_weight == self.config.ref.load_weight print(f'load ref weight start') load_megatron_model_weights(self.config, actor_model_config, ref_module, params_dtype=megatron_config.params_dtype, is_value_model=False) log_gpu_memory_usage('After ref module init', logger=logger) return ref_module, actor_model_config # TODO: add more optimizer args into config if self._is_actor: optim_config = init_megatron_optim_config(optim_config) actor_optimizer = get_megatron_optimizer(model=actor_module, config=optim_config) else: optim_config = None actor_optimizer = None log_gpu_memory_usage('After actor optimizer init', logger=logger) return actor_module, hybrid_engine, actor_optimizer, actor_model_config, optim_config","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer.megatron_actor_model_provider","name":"megatron_actor_model_provider","kind":"function","path":"verl/workers/megatron_workers.py","span":"L154-L166","signature":"def megatron_actor_model_provider(pre_process, post_process)","exported":true,"origin":"ast","body_hash":"2ac83f88ba10eebf33f7c0f8c3172a73a95024409e4bc42d7c23f73d2baa3bf7","body_text":"def megatron_actor_model_provider(pre_process, post_process): from verl.utils.model import get_parallel_model_from_config # vpp is not supported yet because it will hang for some reason. Need debugging vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() # this will be set inside get_model # this_megatron_config = copy.deepcopy(megatron_config) # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank parallel_model = get_parallel_model_from_config(config=actor_model_config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process, value=False) parallel_model.cuda() return parallel_model","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","name":"_build_rollout","kind":"method","path":"verl/workers/megatron_workers.py","span":"L216-L258","signature":"def _build_rollout(self)","exported":false,"origin":"ast","body_hash":"5e7fd53808199f2d874442ec81a5b9b061e9cbba5787275d3d844efbf67fd3d7","body_text":"def _build_rollout(self): if self.config.rollout.name == 'vllm': from verl.workers.rollout.vllm_rollout import vLLMRollout from verl.workers.sharding_manager import MegatronVLLMShardingManager from verl.utils.model import normalize_pp_vpp_params # NOTE(sgm): If the QKV and gate_up projection layer are concate together in actor, # we will reorganize their weight format when resharding from actor to rollout. layer_name_mapping = { \"qkv_layer_name\": self.config.rollout.layer_name_map.get(\"qkv_layer_name\", \"qkv\"), \"gate_proj_layer_name\": self.config.rollout.layer_name_map.get(\"gate_proj_layer_name\", \"linear_fc1.weight\"), } # reshard the weight partition from actor to rollout to initialize the rollout class # create a new cuda space for parameters not in this pp rank self.hybrid_engine.load_params_to_cuda() # broadcast the parameters from pp rank to other ranks self.hybrid_engine.allgather_params() # obtain name to parameters in pp/vpp params = self.hybrid_engine.get_all_params() # update the param name for the params = normalize_pp_vpp_params(params=params, num_hidden_layers=self.actor_model_config.num_hidden_layers, layer_name='layers') rollout = vLLMRollout(actor_module=params, config=self.config.rollout, tokenizer=self.tokenizer, model_hf_config=self.actor_model_config, train_tp=mpu.get_tensor_model_parallel_world_size()) log_gpu_memory_usage('After building vllm rollout', logger=logger) # perform weight resharding between actor and rollout sharding_manager = MegatronVLLMShardingManager(module=self.hybrid_engine, inference_engine=rollout.inference_engine, model_config=self.actor_model_config, layer_name_mapping=layer_name_mapping) log_gpu_memory_usage('After building sharding manager', logger=logger) else: NotImplementedError('Only vllmRollout is supported with Megatron now') return rollout, sharding_manager","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","name":"init_model","kind":"method","path":"verl/workers/megatron_workers.py","span":"L261-L323","signature":"def init_model(self)","exported":true,"origin":"ast","body_hash":"a5404b80fa39504c2d353308be4729d3659f334c0017da13d1515284a609b689","body_text":"def init_model(self): if self.config.model.get('external_lib', None) is not None: # This is used to import external_lib into the huggingface systems import importlib importlib.import_module(self.config.model.external_lib) from omegaconf import OmegaConf from verl.utils.torch_dtypes import PrecisionType override_model_config = OmegaConf.to_container(self.config.model.get('override_config', OmegaConf.create())) torch_dtype = torch.bfloat16 megatron_config = OmegaConf.create({ 'sequence_parallel': self.config.actor.megatron.get('sequence_parallel', True), 'param_dtype': PrecisionType.to_str(torch_dtype), 'tensor_model_parallel_size': mpu.get_tensor_model_parallel_world_size(), 'pipeline_model_parallel_rank': mpu.get_pipeline_model_parallel_rank(), 'pipeline_model_parallel_size': mpu.get_pipeline_model_parallel_world_size(), 'virtual_pipeline_model_parallel_rank': mpu.get_virtual_pipeline_model_parallel_rank(), 'virtual_pipeline_model_parallel_size': mpu.get_virtual_pipeline_model_parallel_world_size() }) megatron_config = init_model_parallel_config(megatron_config) if self._is_actor or self._is_rollout: # we need the model for actor and rollout if self._is_actor: optim_config = self.config.actor.optim else: optim_config = None self.actor_module, self.hybrid_engine, self.actor_optimizer, \\ self.actor_model_config, self.actor_optim_config = self._build_model_optimizer( model_path=self.config.model.path, megatron_config=megatron_config, optim_config=optim_config, override_model_config=override_model_config, ) if self._is_actor: self.actor = MegatronPPOActor(config=self.config.actor, model_config=self.actor_model_config, megatron_config=megatron_config, actor_module=self.actor_module, actor_optimizer=self.actor_optimizer, actor_optimizer_config=self.actor_optim_config) if self._is_rollout: self.rollout, self.sharding_manager = self._build_rollout() if self._is_ref: self.ref_module, self.ref_model_config = self._build_model_optimizer( model_path=self.config.model.path, megatron_config=megatron_config, optim_config=None, override_model_config=override_model_config, ) self.ref_policy = MegatronPPOActor(config=self.config.ref, model_config=self.ref_model_config, megatron_config=megatron_config, actor_module=self.ref_module, actor_optimizer=None, actor_optimizer_config=None) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.update_actor","name":"update_actor","kind":"method","path":"verl/workers/megatron_workers.py","span":"L326-L342","signature":"def update_actor(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"a811999552eb4fba3ab7a616b2567654f63909c6790b85ab9237b6337eea8d78","body_text":"def update_actor(self, data: DataProto): assert self._is_actor data.batch = data.batch.cuda() log_gpu_memory_usage('Before update policy', logger=logger) dataloader = self.actor.make_minibatch_iterator(data=data) metrics = self.actor.update_policy(dataloader=dataloader) log_gpu_memory_usage('After update policy', logger=logger) # TODO: here, we should return all metrics output = DataProto(meta_info={'metrics': metrics}) output = output.to('cpu') torch.cuda.empty_cache() return output","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","name":"generate_sequences","kind":"method","path":"verl/workers/megatron_workers.py","span":"L353-L381","signature":"def generate_sequences(self, prompts: DataProto)","exported":true,"origin":"ast","body_hash":"1f18dbb87d8eb69a29323cdb05e59463e31e388f21499144d7fa0b14d6e16a1c","body_text":"def generate_sequences(self, prompts: DataProto): assert self._is_rollout prompts.batch = prompts.batch.cuda() meta_info = {'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id} prompts.meta_info.update(meta_info) with self.sharding_manager: log_gpu_memory_usage('After entering sharding manager', logger=logger) prompts = self.sharding_manager.preprocess_data(prompts) output = self.rollout.generate_sequences(prompts=prompts) log_gpu_memory_usage('After rollout generation', logger=logger) output = self.sharding_manager.postprocess_data(output) validate = prompts.meta_info.get('validate', False) if self._is_actor and not validate: # we should always recompute old_log_probs when it is HybridEngine output.meta_info['micro_batch_size'] = self.config.rollout.log_prob_micro_batch_size output.meta_info['temperature'] = self.config.rollout.temperature old_log_probs = self.actor.compute_log_prob(data=output) output.batch['old_log_probs'] = old_log_probs output = output.to('cpu') # clear kv cache torch.cuda.empty_cache() log_gpu_memory_usage('After recompute log prob', logger=logger) return output","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","name":"compute_ref_log_prob","kind":"method","path":"verl/workers/megatron_workers.py","span":"L384-L400","signature":"def compute_ref_log_prob(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"445018072fb27c04167ca5cabe87709b49dbc987522db9fee6a009a9033e2a2c","body_text":"def compute_ref_log_prob(self, data: DataProto): data = data.to('cuda') assert self._is_ref if self._is_offload_param: load_megatron_param_and_grad(self.ref_module, torch.cuda.current_device(), self._is_offload_grad) micro_batch_size = self.config.rollout.log_prob_micro_batch_size data.meta_info['micro_batch_size'] = micro_batch_size data.meta_info['temperature'] = self.config.rollout.temperature output = self.ref_policy.compute_log_prob(data=data) output = DataProto.from_dict(tensors={'ref_log_prob': output}) output = output.to('cpu') if self._is_offload_param: offload_megatron_param_and_grad(self.ref_module, self._is_offload_grad) torch.cuda.empty_cache() return output","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.load_checkpoint","name":"load_checkpoint","kind":"method","path":"verl/workers/megatron_workers.py","span":"L403-L404","signature":"def load_checkpoint(self, checkpoint_path)","exported":true,"origin":"ast","body_hash":"4ed36be5305905dea521c3cc29d885054d41a3fb64ed765fc69e56b045131353","body_text":"def load_checkpoint(self, checkpoint_path): pass","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.load_pretrained_model","name":"load_pretrained_model","kind":"method","path":"verl/workers/megatron_workers.py","span":"L407-L408","signature":"def load_pretrained_model(self, checkpoint_path)","exported":true,"origin":"ast","body_hash":"ea566dd116d312eec00aacac919eccaf5c6574b8e39857bd4510f90aa9d197a3","body_text":"def load_pretrained_model(self, checkpoint_path): pass","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.save_checkpoint","name":"save_checkpoint","kind":"method","path":"verl/workers/megatron_workers.py","span":"L411-L413","signature":"def save_checkpoint(self, checkpoint_path)","exported":true,"origin":"ast","body_hash":"3faf8083a98c8106dab262b498469a45b0ba9b2a020eb2c56a1c5fd8e5e00e16","body_text":"def save_checkpoint(self, checkpoint_path): assert self._is_actor pass","summary_state":"pending","summary":null,"crux":null,"owner":"ActorRolloutRefWorker"},{"id":"verl/workers/megatron_workers.py#CriticWorker","name":"CriticWorker","kind":"class","path":"verl/workers/megatron_workers.py","span":"L416-L581","signature":"class CriticWorker(MegatronWorker)","exported":true,"origin":"ast","body_hash":"10b940934641fa635d3eff444108a0de5e22b27f4603babd1705e8a6227a82ec","body_text":"class CriticWorker(MegatronWorker): def __init__(self, config): super().__init__() self.config = config # NOTE(sgm): We utilize colocate WorkerGroup by default. # As a result, Workers for different model share the same process. # Therefore, we only require one distribute initialization. # To utilize different parallel startegy in different models: # 1, users should disable WorkerDict; 2.assign different ResourcePool to different models, # 3. and apply the following patch in ray==2.10, https://github.com/ray-project/ray/pull/44385 if not torch.distributed.is_initialized(): rank = int(os.environ['LOCAL_RANK']) torch.distributed.init_process_group(backend=\"nccl\") torch.cuda.set_device(rank) if self.config.megatron.sequence_parallel: os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' mpu.initialize_model_parallel( tensor_model_parallel_size=self.config.megatron.tensor_model_parallel_size, pipeline_model_parallel_size=self.config.megatron.pipeline_model_parallel_size, virtual_pipeline_model_parallel_size=None, pipeline_model_parallel_split_rank=None, use_sharp=False, context_parallel_size=1, expert_model_parallel_size=1, nccl_communicator_config_path=None, ) set_random_seed(seed=self.config.megatron.seed) # normalize config self.config.ppo_mini_batch_size //= mpu.get_data_parallel_world_size() self.config.ppo_micro_batch_size //= mpu.get_data_parallel_world_size() # TODO(sgm): support critic model offload def _build_critic_model_optimizer(self, model_path, megatron_config: ModelParallelConfig, optim_config, override_model_config, enable_gradient_checkpointing=False): from megatron.core.models.gpt.gpt_model import ModelType from verl.utils.model import print_model_size, update_model_config from verl.utils.megatron.optimizer import get_megatron_optimizer from verl.utils.megatron_utils import get_model, init_megatron_optim_config, init_model_parallel_config from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig # Step 1: initialize the tokenizer local_path = copy_local_path_from_hdfs(model_path) self.tokenizer = hf_tokenizer(local_path) # Step 2: get the actor_model_config critic_model_config = AutoConfig.from_pretrained(local_path) override_config_kwargs = { 'bos_token_id': self.tokenizer.bos_token_id, 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, } override_config_kwargs.update(override_model_config) update_model_config(critic_model_config, override_config_kwargs=override_config_kwargs) if self.rank == 0: print(f'Model config after override: {critic_model_config}') def megatron_critic_model_provider(pre_process, post_process): from verl.utils.model import get_parallel_model_from_config # TODO: support vpp here # vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() # this will be set inside get_model # this_megatron_config = copy.deepcopy(megatron_config) # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank parallel_model = get_parallel_model_from_config(config=critic_model_config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process, value=True) parallel_model.cuda() return parallel_model # Step 3: initialize the megatron model critic_module = get_model(model_provider_func=megatron_critic_model_provider, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True) # note that here critic_module will be a list to be compatible with the construction of interleaved pp (vpp). # but here, we do not use pp (vpp) yet. For simplicity, we remove the list # critic_module = nn.ModuleList(critic_module) if self.config.load_weight: load_megatron_model_weights(self.config, critic_model_config, critic_module, params_dtype=megatron_config.params_dtype, is_value_model=True) if self.rank == 0: print_model_size(critic_module[0]) # TODO: add more optimizer args into config optim_config = init_megatron_optim_config(optim_config) critic_optimizer = get_megatron_optimizer(model=critic_module, config=optim_config) torch.cuda.empty_cache() return critic_module, critic_optimizer, critic_model_config, optim_config @register(dispatch_mode=Dispatch.ONE_TO_ALL) def init_model(self): # create critic from omegaconf import OmegaConf from verl.utils.torch_dtypes import PrecisionType if self.config.model.get('external_lib', None) is not None: # This is used to import external_lib into the huggingface systems import importlib importlib.import_module(self.config.model.external_lib) override_model_config = OmegaConf.to_container(self.config.model.get('override_config', OmegaConf.create())) torch_dtype = torch.bfloat16 megatron_config = OmegaConf.create({ 'sequence_parallel': self.config.megatron.get('sequence_parallel', True), 'param_dtype': PrecisionType.to_str(torch_dtype), 'tensor_model_parallel_size': mpu.get_tensor_model_parallel_world_size(), 'pipeline_model_parallel_rank': mpu.get_pipeline_model_parallel_rank(), 'pipeline_model_parallel_size': mpu.get_pipeline_model_parallel_world_size(), 'virtual_pipeline_model_parallel_rank': mpu.","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/megatron_workers.py#CriticWorker.__init__","name":"__init__","kind":"method","path":"verl/workers/megatron_workers.py","span":"L418-L452","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"be9f6a1c69b6f04bd5b512631bc0087245bf26fe264971adf40af2357ddde172","body_text":"def __init__(self, config): super().__init__() self.config = config # NOTE(sgm): We utilize colocate WorkerGroup by default. # As a result, Workers for different model share the same process. # Therefore, we only require one distribute initialization. # To utilize different parallel startegy in different models: # 1, users should disable WorkerDict; 2.assign different ResourcePool to different models, # 3. and apply the following patch in ray==2.10, https://github.com/ray-project/ray/pull/44385 if not torch.distributed.is_initialized(): rank = int(os.environ['LOCAL_RANK']) torch.distributed.init_process_group(backend=\"nccl\") torch.cuda.set_device(rank) if self.config.megatron.sequence_parallel: os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' mpu.initialize_model_parallel( tensor_model_parallel_size=self.config.megatron.tensor_model_parallel_size, pipeline_model_parallel_size=self.config.megatron.pipeline_model_parallel_size, virtual_pipeline_model_parallel_size=None, pipeline_model_parallel_split_rank=None, use_sharp=False, context_parallel_size=1, expert_model_parallel_size=1, nccl_communicator_config_path=None, ) set_random_seed(seed=self.config.megatron.seed) # normalize config self.config.ppo_mini_batch_size //= mpu.get_data_parallel_world_size() self.config.ppo_micro_batch_size //= mpu.get_data_parallel_world_size() # TODO(sgm): support critic model offload","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","name":"_build_critic_model_optimizer","kind":"method","path":"verl/workers/megatron_workers.py","span":"L454-L519","signature":"def _build_critic_model_optimizer(self, model_path, megatron_config: ModelParallelConfig, optim_config, override_model_config, enable_gradient_checkpointing=False)","exported":false,"origin":"ast","body_hash":"ecbaa1ae59aed94d310f147f95b47ead695609c37a1f95819d2f83d9d0dea6ab","body_text":"def _build_critic_model_optimizer(self, model_path, megatron_config: ModelParallelConfig, optim_config, override_model_config, enable_gradient_checkpointing=False): from megatron.core.models.gpt.gpt_model import ModelType from verl.utils.model import print_model_size, update_model_config from verl.utils.megatron.optimizer import get_megatron_optimizer from verl.utils.megatron_utils import get_model, init_megatron_optim_config, init_model_parallel_config from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig # Step 1: initialize the tokenizer local_path = copy_local_path_from_hdfs(model_path) self.tokenizer = hf_tokenizer(local_path) # Step 2: get the actor_model_config critic_model_config = AutoConfig.from_pretrained(local_path) override_config_kwargs = { 'bos_token_id': self.tokenizer.bos_token_id, 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, } override_config_kwargs.update(override_model_config) update_model_config(critic_model_config, override_config_kwargs=override_config_kwargs) if self.rank == 0: print(f'Model config after override: {critic_model_config}') def megatron_critic_model_provider(pre_process, post_process): from verl.utils.model import get_parallel_model_from_config # TODO: support vpp here # vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() # this will be set inside get_model # this_megatron_config = copy.deepcopy(megatron_config) # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank parallel_model = get_parallel_model_from_config(config=critic_model_config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process, value=True) parallel_model.cuda() return parallel_model # Step 3: initialize the megatron model critic_module = get_model(model_provider_func=megatron_critic_model_provider, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=True) # note that here critic_module will be a list to be compatible with the construction of interleaved pp (vpp). # but here, we do not use pp (vpp) yet. For simplicity, we remove the list # critic_module = nn.ModuleList(critic_module) if self.config.load_weight: load_megatron_model_weights(self.config, critic_model_config, critic_module, params_dtype=megatron_config.params_dtype, is_value_model=True) if self.rank == 0: print_model_size(critic_module[0]) # TODO: add more optimizer args into config optim_config = init_megatron_optim_config(optim_config) critic_optimizer = get_megatron_optimizer(model=critic_module, config=optim_config) torch.cuda.empty_cache() return critic_module, critic_optimizer, critic_model_config, optim_config","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer.megatron_critic_model_provider","name":"megatron_critic_model_provider","kind":"function","path":"verl/workers/megatron_workers.py","span":"L484-L496","signature":"def megatron_critic_model_provider(pre_process, post_process)","exported":true,"origin":"ast","body_hash":"9dde905bcec61792c0481e23a1bfc85034805c130cccc72f1245dd215c53d6e0","body_text":"def megatron_critic_model_provider(pre_process, post_process): from verl.utils.model import get_parallel_model_from_config # TODO: support vpp here # vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() # this will be set inside get_model # this_megatron_config = copy.deepcopy(megatron_config) # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank parallel_model = get_parallel_model_from_config(config=critic_model_config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process, value=True) parallel_model.cuda() return parallel_model","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/megatron_workers.py#CriticWorker.init_model","name":"init_model","kind":"method","path":"verl/workers/megatron_workers.py","span":"L522-L556","signature":"def init_model(self): # create critic","exported":true,"origin":"ast","body_hash":"ce9575f6b7d85c7af77cc240fed7b61bd1b422ef7c2382d85b5d7e98fffb72f7","body_text":"def init_model(self): # create critic from omegaconf import OmegaConf from verl.utils.torch_dtypes import PrecisionType if self.config.model.get('external_lib', None) is not None: # This is used to import external_lib into the huggingface systems import importlib importlib.import_module(self.config.model.external_lib) override_model_config = OmegaConf.to_container(self.config.model.get('override_config', OmegaConf.create())) torch_dtype = torch.bfloat16 megatron_config = OmegaConf.create({ 'sequence_parallel': self.config.megatron.get('sequence_parallel', True), 'param_dtype': PrecisionType.to_str(torch_dtype), 'tensor_model_parallel_size': mpu.get_tensor_model_parallel_world_size(), 'pipeline_model_parallel_rank': mpu.get_pipeline_model_parallel_rank(), 'pipeline_model_parallel_size': mpu.get_pipeline_model_parallel_world_size(), 'virtual_pipeline_model_parallel_rank': mpu.get_virtual_pipeline_model_parallel_rank(), 'virtual_pipeline_model_parallel_size': mpu.get_virtual_pipeline_model_parallel_world_size() }) megatron_config = init_model_parallel_config(megatron_config) critic_module, critic_optimizer, critic_model_config, critic_optimizer_config = self._build_critic_model_optimizer( model_path=self.config.model.path, megatron_config=megatron_config, optim_config=self.config.optim, override_model_config=override_model_config) self.critic = MegatronPPOCritic(config=self.config, model_config=critic_model_config, megatron_config=megatron_config, critic_module=critic_module, critic_optimizer=critic_optimizer, critic_optimizer_config=critic_optimizer_config)","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/megatron_workers.py#CriticWorker.compute_values","name":"compute_values","kind":"method","path":"verl/workers/megatron_workers.py","span":"L559-L564","signature":"def compute_values(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"d82d9f3c8dfc6b59788d0627563a7bdc63d810177f5d0c792514f5deb80acf52","body_text":"def compute_values(self, data: DataProto): data = data.to('cuda') values = self.critic.compute_values(data=data) output = DataProto.from_dict(tensors={'values': values}) output = output.to('cpu') return output","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/megatron_workers.py#CriticWorker.update_critic","name":"update_critic","kind":"method","path":"verl/workers/megatron_workers.py","span":"L567-L573","signature":"def update_critic(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"f277c319ab61a7d647f5a1217d05803e481a1fe8c2cb79ec59b6961038666d0f","body_text":"def update_critic(self, data: DataProto): data = data.to('cuda') dataloader = self.critic.make_minibatch_iterator(data) metrics = self.critic.update_critic(dataloader=dataloader) output = DataProto(batch=None, meta_info={'metrics': metrics}) output = output.to('cpu') return output","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/megatron_workers.py#CriticWorker.load_checkpoint","name":"load_checkpoint","kind":"method","path":"verl/workers/megatron_workers.py","span":"L576-L577","signature":"def load_checkpoint(self, checkpoint_path)","exported":true,"origin":"ast","body_hash":"4ed36be5305905dea521c3cc29d885054d41a3fb64ed765fc69e56b045131353","body_text":"def load_checkpoint(self, checkpoint_path): pass","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/megatron_workers.py#CriticWorker.save_checkpoint","name":"save_checkpoint","kind":"method","path":"verl/workers/megatron_workers.py","span":"L580-L581","signature":"def save_checkpoint(self, checkpoint_path)","exported":true,"origin":"ast","body_hash":"0f86f9b40d600f9ce4be917bc1ec07554636032355ff2ad3f20a44d3aead5d57","body_text":"def save_checkpoint(self, checkpoint_path): pass","summary_state":"pending","summary":null,"crux":null,"owner":"CriticWorker"},{"id":"verl/workers/megatron_workers.py#RewardModelWorker","name":"RewardModelWorker","kind":"class","path":"verl/workers/megatron_workers.py","span":"L584-L735","signature":"class RewardModelWorker(MegatronWorker)","exported":true,"origin":"ast","body_hash":"b3bf587f6dde24887d89831973f0d6c1eb5f1755ef452732487f58d9411617f6","body_text":"class RewardModelWorker(MegatronWorker): \"\"\" Note that we only implement the reward model that is subclass of AutoModelForSequenceClassification. \"\"\" def __init__(self, config): super().__init__() self.config = config # NOTE(sgm): We utilize colocate WorkerGroup by default. # As a result, Workers for different model share the same process. # Therefore, we only require one distribute initialization. # To utilize different parallel startegy in different models: # 1, users should disable WorkerDict; 2.assign different ResourcePool to different models, # 3. and apply the following patch in ray==2.10, https://github.com/ray-project/ray/pull/44385 if not torch.distributed.is_initialized(): rank = int(os.environ['LOCAL_RANK']) torch.distributed.init_process_group(backend=\"nccl\") torch.cuda.set_device(rank) if self.config.megatron.sequence_parallel: os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' mpu.initialize_model_parallel( tensor_model_parallel_size=self.config.megatron.tensor_model_parallel_size, pipeline_model_parallel_size=self.config.megatron.pipeline_model_parallel_size, virtual_pipeline_model_parallel_size=None, pipeline_model_parallel_split_rank=None, use_sharp=False, context_parallel_size=1, expert_model_parallel_size=1, nccl_communicator_config_path=None, ) set_random_seed(seed=self.config.megatron.seed) # normalize config self.config.micro_batch_size //= mpu.get_data_parallel_world_size() def _build_rm_model(self, model_path, megatron_config: ModelParallelConfig, override_model_config): from megatron.core.models.gpt.gpt_model import ModelType from verl.utils.model import print_model_size, update_model_config from verl.utils.megatron_utils import get_model from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig # Step 1: initialize the tokenizer local_path = copy_local_path_from_hdfs(model_path) self.tokenizer = hf_tokenizer(local_path) # Step 2: get the actor_model_config rm_model_config = AutoConfig.from_pretrained(local_path) override_config_kwargs = { 'bos_token_id': self.tokenizer.bos_token_id, 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, } override_config_kwargs.update(override_model_config) update_model_config(rm_model_config, override_config_kwargs=override_config_kwargs) if self.rank == 0: print(f'Model config after override: {rm_model_config}') def megatron_rm_model_provider(pre_process, post_process): from verl.utils.model import get_parallel_model_from_config # vpp is not supported yet because it will hang for some reason. Need debugging vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() # this will be set inside get_model # this_megatron_config = copy.deepcopy(megatron_config) # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank parallel_model = get_parallel_model_from_config(config=rm_model_config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process, value=True) parallel_model.cuda() return parallel_model # Step 3: initialize the megatron model reward_model = get_model(model_provider_func=megatron_rm_model_provider, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=False) # note that here critic_module will be a list to be compatible with the construction of interleaved pp (vpp). # but here, we do not use pp (vpp) yet. For simplicity, we remove the list # reward_model = nn.ModuleList(reward_model) if self.config.load_weight: load_megatron_model_weights(self.config, rm_model_config, reward_model, params_dtype=megatron_config.params_dtype, is_value_model=True) # TODO: add more optimizer args into config torch.cuda.empty_cache() return reward_model, rm_model_config @register(dispatch_mode=Dispatch.ONE_TO_ALL) def init_model(self): # create critic from omegaconf import OmegaConf from verl.utils.torch_dtypes import PrecisionType from transformers import AutoTokenizer if self.config.model.get('external_lib', None) is not None: # This is used to import external_lib into the huggingface systems import importlib importlib.import_module(self.config.model.external_lib) override_model_config = OmegaConf.to_container(self.config.model.get('override_config', OmegaConf.create())) sft_tokenizer_local_path = copy_local_path_from_hdfs(self.config.model.input_tokenizer) sft_tokenizer = hf_tokenizer(sft_tokenizer_local_path) rm_tokenizer_path = self.config.model.get('rm_tokenizer', None) rm_tokenizer = None if rm_tokenizer_path is not None: rm_tokenizer_local_path = copy_local_path_from_hdfs(rm_tokenizer_path) rm_tokenizer = hf_tokenizer(rm_tokenizer_local_path) torch_dtype = torch.bfloat16 megatron_config = OmegaConf.create({ 'sequence_parallel': self.config.megatron.get('sequence_parallel', True), 'param_dtype': PrecisionType.to_str(torch_dtype), 'tensor_model_parallel_size': mpu.get_tensor_model_parallel_world_size(), 'pipeline_model_parallel_rank': mpu.get_pipeline_model_parallel_rank(), 'pipeline_model_parallel_size': mpu.get_pipeline_model_parallel_world_size(), 'virtual_pipeline_","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","name":"__init__","kind":"method","path":"verl/workers/megatron_workers.py","span":"L589-L620","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"33eb28386e1e1f0ccdcd9b5542eaf5fdc5f0fd8b5b6b25c0fddc5736274794e2","body_text":"def __init__(self, config): super().__init__() self.config = config # NOTE(sgm): We utilize colocate WorkerGroup by default. # As a result, Workers for different model share the same process. # Therefore, we only require one distribute initialization. # To utilize different parallel startegy in different models: # 1, users should disable WorkerDict; 2.assign different ResourcePool to different models, # 3. and apply the following patch in ray==2.10, https://github.com/ray-project/ray/pull/44385 if not torch.distributed.is_initialized(): rank = int(os.environ['LOCAL_RANK']) torch.distributed.init_process_group(backend=\"nccl\") torch.cuda.set_device(rank) if self.config.megatron.sequence_parallel: os.environ['CUDA_DEVICE_MAX_CONNECTIONS'] = '1' mpu.initialize_model_parallel( tensor_model_parallel_size=self.config.megatron.tensor_model_parallel_size, pipeline_model_parallel_size=self.config.megatron.pipeline_model_parallel_size, virtual_pipeline_model_parallel_size=None, pipeline_model_parallel_split_rank=None, use_sharp=False, context_parallel_size=1, expert_model_parallel_size=1, nccl_communicator_config_path=None, ) set_random_seed(seed=self.config.megatron.seed) # normalize config self.config.micro_batch_size //= mpu.get_data_parallel_world_size()","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"},{"id":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","name":"_build_rm_model","kind":"method","path":"verl/workers/megatron_workers.py","span":"L622-L677","signature":"def _build_rm_model(self, model_path, megatron_config: ModelParallelConfig, override_model_config)","exported":false,"origin":"ast","body_hash":"9ae3bf8a4817c0f6aee709f0cca8f6753744ae90c80c68cf046edd0657f13f1a","body_text":"def _build_rm_model(self, model_path, megatron_config: ModelParallelConfig, override_model_config): from megatron.core.models.gpt.gpt_model import ModelType from verl.utils.model import print_model_size, update_model_config from verl.utils.megatron_utils import get_model from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig # Step 1: initialize the tokenizer local_path = copy_local_path_from_hdfs(model_path) self.tokenizer = hf_tokenizer(local_path) # Step 2: get the actor_model_config rm_model_config = AutoConfig.from_pretrained(local_path) override_config_kwargs = { 'bos_token_id': self.tokenizer.bos_token_id, 'eos_token_id': self.tokenizer.eos_token_id, 'pad_token_id': self.tokenizer.pad_token_id, } override_config_kwargs.update(override_model_config) update_model_config(rm_model_config, override_config_kwargs=override_config_kwargs) if self.rank == 0: print(f'Model config after override: {rm_model_config}') def megatron_rm_model_provider(pre_process, post_process): from verl.utils.model import get_parallel_model_from_config # vpp is not supported yet because it will hang for some reason. Need debugging vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() # this will be set inside get_model # this_megatron_config = copy.deepcopy(megatron_config) # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank parallel_model = get_parallel_model_from_config(config=rm_model_config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process, value=True) parallel_model.cuda() return parallel_model # Step 3: initialize the megatron model reward_model = get_model(model_provider_func=megatron_rm_model_provider, model_type=ModelType.encoder_or_decoder, wrap_with_ddp=False) # note that here critic_module will be a list to be compatible with the construction of interleaved pp (vpp). # but here, we do not use pp (vpp) yet. For simplicity, we remove the list # reward_model = nn.ModuleList(reward_model) if self.config.load_weight: load_megatron_model_weights(self.config, rm_model_config, reward_model, params_dtype=megatron_config.params_dtype, is_value_model=True) # TODO: add more optimizer args into config torch.cuda.empty_cache() return reward_model, rm_model_config","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"},{"id":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model.megatron_rm_model_provider","name":"megatron_rm_model_provider","kind":"function","path":"verl/workers/megatron_workers.py","span":"L646-L658","signature":"def megatron_rm_model_provider(pre_process, post_process)","exported":true,"origin":"ast","body_hash":"08b8bf8094f36d117b7510c8f111ddff8e0d4792631d5bfc3f9476130c0087be","body_text":"def megatron_rm_model_provider(pre_process, post_process): from verl.utils.model import get_parallel_model_from_config # vpp is not supported yet because it will hang for some reason. Need debugging vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank() # this will be set inside get_model # this_megatron_config = copy.deepcopy(megatron_config) # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank parallel_model = get_parallel_model_from_config(config=rm_model_config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process, value=True) parallel_model.cuda() return parallel_model","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","name":"init_model","kind":"method","path":"verl/workers/megatron_workers.py","span":"L680-L726","signature":"def init_model(self): # create critic","exported":true,"origin":"ast","body_hash":"3770ed997ee0b3e5889587a93dbf30b9c0e0361af5ab3f504aaf02d2b364cbda","body_text":"def init_model(self): # create critic from omegaconf import OmegaConf from verl.utils.torch_dtypes import PrecisionType from transformers import AutoTokenizer if self.config.model.get('external_lib', None) is not None: # This is used to import external_lib into the huggingface systems import importlib importlib.import_module(self.config.model.external_lib) override_model_config = OmegaConf.to_container(self.config.model.get('override_config', OmegaConf.create())) sft_tokenizer_local_path = copy_local_path_from_hdfs(self.config.model.input_tokenizer) sft_tokenizer = hf_tokenizer(sft_tokenizer_local_path) rm_tokenizer_path = self.config.model.get('rm_tokenizer', None) rm_tokenizer = None if rm_tokenizer_path is not None: rm_tokenizer_local_path = copy_local_path_from_hdfs(rm_tokenizer_path) rm_tokenizer = hf_tokenizer(rm_tokenizer_local_path) torch_dtype = torch.bfloat16 megatron_config = OmegaConf.create({ 'sequence_parallel': self.config.megatron.get('sequence_parallel', True), 'param_dtype': PrecisionType.to_str(torch_dtype), 'tensor_model_parallel_size': mpu.get_tensor_model_parallel_world_size(), 'pipeline_model_parallel_rank': mpu.get_pipeline_model_parallel_rank(), 'pipeline_model_parallel_size': mpu.get_pipeline_model_parallel_world_size(), 'virtual_pipeline_model_parallel_rank': mpu.get_virtual_pipeline_model_parallel_rank(), 'virtual_pipeline_model_parallel_size': mpu.get_virtual_pipeline_model_parallel_world_size() }) megatron_config = init_model_parallel_config(megatron_config) reward_model_module, reward_model_config = self._build_rm_model( model_path=self.config.model.path, megatron_config=megatron_config, override_model_config=override_model_config, ) # FIXME(sgm): reward model param offload is implemented in MegatronRewardModel # should be implemented in workers self.rm = MegatronRewardModel(config=self.config, reward_model_module=reward_model_module, model_config=reward_model_config, megatron_config=megatron_config, sft_tokenizer=sft_tokenizer, rm_tokenizer=rm_tokenizer)","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"},{"id":"verl/workers/megatron_workers.py#RewardModelWorker.compute_rm_score","name":"compute_rm_score","kind":"method","path":"verl/workers/megatron_workers.py","span":"L731-L735","signature":"def compute_rm_score(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"832ac5cb6f5b9faa5d402ce9d8b26a45d2dc8a41a33ee856e76fe047656f7996","body_text":"def compute_rm_score(self, data: DataProto): data.batch = data.batch.cuda() output = self.rm.compute_reward(data) output = output.to('cpu') return output","summary_state":"pending","summary":null,"crux":null,"owner":"RewardModelWorker"}],"rawEdges":[{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"os","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"logging","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"ray","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"omegaconf","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.single_controller.base.megatron.worker","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.workers.actor.megatron_actor","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.workers.critic.megatron_critic","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.workers.sharding_manager","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.workers.reward_model.megatron.reward_model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.single_controller.base.decorator","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.fs","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.debug","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.megatron_utils","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.megatron_utils","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"megatron.core","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"megatron.core","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"calls","name":"getLogger","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"calls","name":"setLevel","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"calls","name":"getenv","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"contains","targetId":"verl/workers/megatron_workers.py#set_random_seed","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"random","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#set_random_seed","relation":"calls","name":"manual_seed","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#set_random_seed","relation":"calls","name":"seed","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#set_random_seed","relation":"calls","name":"seed","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#set_random_seed","relation":"calls","name":"device_count","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"megatron.core","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#set_random_seed","relation":"calls","name":"model_parallel_cuda_manual_seed","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"extends","name":"MegatronWorker","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"int","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"init_process_group","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"set_device","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"initialize_model_parallel","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"set_random_seed","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get_data_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get_data_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get_data_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get_data_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.megatron.optimizer","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"megatron.core.models.gpt.gpt_model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.megatron_utils","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"transformers","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"update","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"update_model_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer.megatron_actor_model_provider","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer.megatron_actor_model_provider","relation":"calls","name":"get_virtual_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer.megatron_actor_model_provider","relation":"calls","name":"get_parallel_model_from_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer.megatron_actor_model_provider","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"get_parallel_model_from_config"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"AllGatherPPModel","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"load_megatron_model_weights","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"print_model_size","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get_model","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"load_megatron_model_weights","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"init_megatron_optim_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"get_megatron_optimizer","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_model_optimizer","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.workers.rollout.vllm_rollout","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.workers.sharding_manager","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"load_params_to_cuda","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"allgather_params","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"get_all_params","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"normalize_pp_vpp_params","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"vLLMRollout","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"MegatronVLLMShardingManager","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker._build_rollout","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"importlib","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"import_module","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"omegaconf","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.torch_dtypes","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"to_container","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"to_str","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get_virtual_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"init_model_parallel_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"_build_model_optimizer","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"ActorRolloutRefWorker"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"MegatronPPOActor","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"_build_rollout","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"ActorRolloutRefWorker"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"_build_model_optimizer","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"ActorRolloutRefWorker"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"MegatronPPOActor","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.init_model","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.update_actor","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"make_minibatch_iterator","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"MegatronPPOActor"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"update_policy","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"MegatronPPOActor"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"DataProto"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.update_actor","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"update","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"preprocess_data","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"generate_sequences","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"postprocess_data","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"compute_log_prob","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"MegatronPPOActor"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.generate_sequences","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"DataProto"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"load_megatron_param_and_grad","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"compute_log_prob","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"MegatronPPOActor"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"from_dict","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"offload_megatron_param_and_grad","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.compute_ref_log_prob","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.load_checkpoint","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.load_pretrained_model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#ActorRolloutRefWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#ActorRolloutRefWorker.save_checkpoint","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"contains","targetId":"verl/workers/megatron_workers.py#CriticWorker","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"extends","name":"MegatronWorker","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#CriticWorker.__init__","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.__init__","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.__init__","relation":"calls","name":"int","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.__init__","relation":"calls","name":"init_process_group","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.__init__","relation":"calls","name":"set_device","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.__init__","relation":"calls","name":"initialize_model_parallel","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.__init__","relation":"calls","name":"set_random_seed","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.__init__","relation":"calls","name":"get_data_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.__init__","relation":"calls","name":"get_data_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"megatron.core.models.gpt.gpt_model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.megatron.optimizer","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.megatron_utils","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"transformers","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"update","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"update_model_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"contains","targetId":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer.megatron_critic_model_provider","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer.megatron_critic_model_provider","relation":"calls","name":"get_parallel_model_from_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer.megatron_critic_model_provider","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"get_parallel_model_from_config"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get_model","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"load_megatron_model_weights","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"print_model_size","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"init_megatron_optim_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"get_megatron_optimizer","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker._build_critic_model_optimizer","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#CriticWorker.init_model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"omegaconf","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.torch_dtypes","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"importlib","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"import_module","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"to_container","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"to_str","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"get_virtual_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"init_model_parallel_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"_build_critic_model_optimizer","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"CriticWorker"},{"source":"verl/workers/megatron_workers.py#CriticWorker.init_model","relation":"calls","name":"MegatronPPOCritic","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#CriticWorker.compute_values","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.compute_values","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"DataProto"},{"source":"verl/workers/megatron_workers.py#CriticWorker.compute_values","relation":"calls","name":"compute_values","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"MegatronPPOCritic"},{"source":"verl/workers/megatron_workers.py#CriticWorker.compute_values","relation":"calls","name":"from_dict","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.compute_values","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#CriticWorker.update_critic","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.update_critic","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"DataProto"},{"source":"verl/workers/megatron_workers.py#CriticWorker.update_critic","relation":"calls","name":"make_minibatch_iterator","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"MegatronPPOCritic"},{"source":"verl/workers/megatron_workers.py#CriticWorker.update_critic","relation":"calls","name":"update_critic","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"MegatronPPOCritic"},{"source":"verl/workers/megatron_workers.py#CriticWorker.update_critic","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker.update_critic","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"DataProto"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#CriticWorker.load_checkpoint","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#CriticWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#CriticWorker.save_checkpoint","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"contains","targetId":"verl/workers/megatron_workers.py#RewardModelWorker","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker","relation":"extends","name":"MegatronWorker","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","relation":"calls","name":"is_initialized","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","relation":"calls","name":"int","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","relation":"calls","name":"init_process_group","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","relation":"calls","name":"set_device","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","relation":"calls","name":"initialize_model_parallel","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","relation":"calls","name":"set_random_seed","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.__init__","relation":"calls","name":"get_data_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"megatron.core.models.gpt.gpt_model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.megatron_utils","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"transformers","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","relation":"calls","name":"update","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","relation":"calls","name":"update_model_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","relation":"contains","targetId":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model.megatron_rm_model_provider","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model.megatron_rm_model_provider","relation":"calls","name":"get_virtual_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model.megatron_rm_model_provider","relation":"calls","name":"get_parallel_model_from_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model.megatron_rm_model_provider","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"get_parallel_model_from_config"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","relation":"calls","name":"get_model","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","relation":"calls","name":"load_megatron_model_weights","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker._build_rm_model","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"omegaconf","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"verl.utils.torch_dtypes","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"transformers","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py","relation":"imports","specifier":"importlib","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"import_module","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"to_container","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"copy_local_path_from_hdfs","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"hf_tokenizer","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"create","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"to_str","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"get_virtual_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"init_model_parallel_config","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"_build_rm_model","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"RewardModelWorker"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.init_model","relation":"calls","name":"MegatronRewardModel","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker","relation":"contains","targetId":"verl/workers/megatron_workers.py#RewardModelWorker.compute_rm_score","file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/megatron_workers.py"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"compute_reward","viaMember":true,"file":"verl/workers/megatron_workers.py","recvType":"MegatronRewardModel"},{"source":"verl/workers/megatron_workers.py#RewardModelWorker.compute_rm_score","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/megatron_workers.py"}]},"verl/workers/retriever_workers.py":{"size":13198,"mtimeMs":1789275830328.2031,"hash":"548ec6d688c66a6a5823dfe758c981a9d8046431c6e2dc5af243d3a0c5eba5df","nodes":[{"id":"verl/workers/retriever_workers.py","name":"retriever_workers.py","kind":"file","path":"verl/workers/retriever_workers.py","span":"L1-L384","signature":null,"exported":true,"origin":"ast","body_hash":"548ec6d688c66a6a5823dfe758c981a9d8046431c6e2dc5af243d3a0c5eba5df","chars":13198,"summary_state":"pending","summary":null,"crux":null,"body_text":"import os import torch from verl.single_controller.base import Worker from verl.single_controller.base.decorator import register, Dispatch # from search_r1.env.search.retrieval import get_retriever import json import os import warnings from typing import List, Dict import functools from tqdm import tqdm from multiprocessing import Pool import faiss import torch import numpy as np from transformers import AutoConfig, AutoTokenizer, AutoModel import argparse import datasets import ray import faiss import torch @ray.remote(num_gpus=8) # Allocate all GPUs"},{"id":"verl/workers/retriever_workers.py#load_corpus","name":"load_corpus","kind":"function","path":"verl/workers/retriever_workers.py","span":"L22-L28","signature":"def load_corpus(corpus_path: str)","exported":true,"origin":"ast","body_hash":"63642a3531f6506e9fae80592a2170b344a1b6179a3d27a3002081529a9e3d23","body_text":"def load_corpus(corpus_path: str): corpus = datasets.load_dataset( 'json', data_files=corpus_path, split=\"train\", num_proc=4) return corpus","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#read_jsonl","name":"read_jsonl","kind":"function","path":"verl/workers/retriever_workers.py","span":"L31-L38","signature":"def read_jsonl(file_path)","exported":true,"origin":"ast","body_hash":"4c47880b2c2ac231addf857ba43f05bdfc1be3f19de4f48067418ca65a407f1d","body_text":"def read_jsonl(file_path): data = [] with open(file_path, \"r\") as f: readin = f.readlines() for line in readin: data.append(json.loads(line)) return data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#load_docs","name":"load_docs","kind":"function","path":"verl/workers/retriever_workers.py","span":"L41-L44","signature":"def load_docs(corpus, doc_idxs)","exported":true,"origin":"ast","body_hash":"6391624ce3919bef69c4bd7c04e58574246a106871a10aa208bcde98337a2922","body_text":"def load_docs(corpus, doc_idxs): results = [corpus[int(idx)] for idx in doc_idxs] return results","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#load_model","name":"load_model","kind":"function","path":"verl/workers/retriever_workers.py","span":"L47-L59","signature":"def load_model( model_path: str, use_fp16: bool = False )","exported":true,"origin":"ast","body_hash":"2d4e1f08f900b69638da6c1525eb3006b46ed0822e5e3f12fb86d8faa60ab34b","body_text":"def load_model( model_path: str, use_fp16: bool = False ): model_config = AutoConfig.from_pretrained(model_path, trust_remote_code=True) model = AutoModel.from_pretrained(model_path, trust_remote_code=True) model.eval() model.cuda() if use_fp16: model = model.half() tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, trust_remote_code=True) return model, tokenizer","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#pooling","name":"pooling","kind":"function","path":"verl/workers/retriever_workers.py","span":"L62-L76","signature":"def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = \"mean\" )","exported":true,"origin":"ast","body_hash":"6e356a5fbac440ef6226d2119161f4cf7a669ceb3505baddff6331a4a50030ec","body_text":"def pooling( pooler_output, last_hidden_state, attention_mask = None, pooling_method = \"mean\" ): if pooling_method == \"mean\": last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0) return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None] elif pooling_method == \"cls\": return last_hidden_state[:, 0] elif pooling_method == \"pooler\": return pooler_output else: raise NotImplementedError(\"Pooling method not implemented!\")","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#Encoder","name":"Encoder","kind":"class","path":"verl/workers/retriever_workers.py","span":"L79-L135","signature":"class Encoder","exported":true,"origin":"ast","body_hash":"a513ae91bd7e48309b1f39780b3beac631d16168574d922913ff8662dd199b11","body_text":"class Encoder: def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16): self.model_name = model_name self.model_path = model_path self.pooling_method = pooling_method self.max_length = max_length self.use_fp16 = use_fp16 self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16) @torch.no_grad() def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders if isinstance(query_list, str): query_list = [query_list] if \"e5\" in self.model_name.lower(): if is_query: query_list = [f\"query: {query}\" for query in query_list] else: query_list = [f\"passage: {query}\" for query in query_list] if \"bge\" in self.model_name.lower(): if is_query: query_list = [f\"Represent this sentence for searching relevant passages: {query}\" for query in query_list] inputs = self.tokenizer(query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors=\"pt\" ) inputs = {k: v.cuda() for k, v in inputs.items()} if \"T5\" in type(self.model).__name__: # T5-based retrieval model decoder_input_ids = torch.zeros( (inputs['input_ids'].shape[0], 1), dtype=torch.long ).to(inputs['input_ids'].device) output = self.model( **inputs, decoder_input_ids=decoder_input_ids, return_dict=True ) query_emb = output.last_hidden_state[:, 0, :] else: output = self.model(**inputs, return_dict=True) query_emb = pooling(output.pooler_output, output.last_hidden_state, inputs['attention_mask'], self.pooling_method) if \"dpr\" not in self.model_name.lower(): query_emb = torch.nn.functional.normalize(query_emb, dim=-1) query_emb = query_emb.detach().cpu().numpy() query_emb = query_emb.astype(np.float32, order=\"C\") return query_emb","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#Encoder.__init__","name":"__init__","kind":"method","path":"verl/workers/retriever_workers.py","span":"L80-L88","signature":"def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16)","exported":false,"origin":"ast","body_hash":"91ad5eec7d331441c6f1b920061f211d1e755e85e3509cf6d2b1b47a3215c1dd","body_text":"def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16): self.model_name = model_name self.model_path = model_path self.pooling_method = pooling_method self.max_length = max_length self.use_fp16 = use_fp16 self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16)","summary_state":"pending","summary":null,"crux":null,"owner":"Encoder"},{"id":"verl/workers/retriever_workers.py#Encoder.encode","name":"encode","kind":"method","path":"verl/workers/retriever_workers.py","span":"L91-L135","signature":"def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders","exported":true,"origin":"ast","body_hash":"6a6511e1300b1aa62cedb981fcdf443894d0ee025705cf692573931f36fd8dc0","body_text":"def encode(self, query_list: List[str], is_query=True) -> np.ndarray: # processing query for different encoders if isinstance(query_list, str): query_list = [query_list] if \"e5\" in self.model_name.lower(): if is_query: query_list = [f\"query: {query}\" for query in query_list] else: query_list = [f\"passage: {query}\" for query in query_list] if \"bge\" in self.model_name.lower(): if is_query: query_list = [f\"Represent this sentence for searching relevant passages: {query}\" for query in query_list] inputs = self.tokenizer(query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors=\"pt\" ) inputs = {k: v.cuda() for k, v in inputs.items()} if \"T5\" in type(self.model).__name__: # T5-based retrieval model decoder_input_ids = torch.zeros( (inputs['input_ids'].shape[0], 1), dtype=torch.long ).to(inputs['input_ids'].device) output = self.model( **inputs, decoder_input_ids=decoder_input_ids, return_dict=True ) query_emb = output.last_hidden_state[:, 0, :] else: output = self.model(**inputs, return_dict=True) query_emb = pooling(output.pooler_output, output.last_hidden_state, inputs['attention_mask'], self.pooling_method) if \"dpr\" not in self.model_name.lower(): query_emb = torch.nn.functional.normalize(query_emb, dim=-1) query_emb = query_emb.detach().cpu().numpy() query_emb = query_emb.astype(np.float32, order=\"C\") return query_emb","summary_state":"pending","summary":null,"crux":null,"owner":"Encoder"},{"id":"verl/workers/retriever_workers.py#BaseRetriever","name":"BaseRetriever","kind":"class","path":"verl/workers/retriever_workers.py","span":"L138-L168","signature":"class BaseRetriever","exported":true,"origin":"ast","body_hash":"1ce69eef8b2362b2435db716c275651dd2e03bf182aa743b91f8230edd7fabef","body_text":"class BaseRetriever: \"\"\"Base object for all retrievers.\"\"\" def __init__(self, config): self.config = config self.retrieval_method = config.retrieval_method self.topk = config.retrieval_topk self.index_path = config.index_path self.corpus_path = config.corpus_path # self.cache_save_path = os.path.join(config.save_dir, 'retrieval_cache.json') def _search(self, query: str, num: int, return_score:bool) -> List[Dict[str, str]]: r\"\"\"Retrieve topk relevant documents in corpus. Return: list: contains information related to the document, including: contents: used for building index title: (if provided) text: (if provided) \"\"\" pass def _batch_search(self, query_list, num, return_score): pass def search(self, *args, **kwargs): return self._search(*args, **kwargs) def batch_search(self, *args, **kwargs): return self._batch_search(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#BaseRetriever.__init__","name":"__init__","kind":"method","path":"verl/workers/retriever_workers.py","span":"L141-L149","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"1592cd9ea0215fbf4afe28ec72ebc8027028107edbd1f2bfb86b7b06a1204680","body_text":"def __init__(self, config): self.config = config self.retrieval_method = config.retrieval_method self.topk = config.retrieval_topk self.index_path = config.index_path self.corpus_path = config.corpus_path # self.cache_save_path = os.path.join(config.save_dir, 'retrieval_cache.json')","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"verl/workers/retriever_workers.py#BaseRetriever._search","name":"_search","kind":"method","path":"verl/workers/retriever_workers.py","span":"L151-L159","signature":"def _search(self, query: str, num: int, return_score:bool) -> List[Dict[str, str]]","exported":false,"origin":"ast","body_hash":"a1863af52b58a6982768a8c05eb0ceae5946ccc8ab516251836104adf4c6d0a3","body_text":"def _search(self, query: str, num: int, return_score:bool) -> List[Dict[str, str]]: r\"\"\"Retrieve topk relevant documents in corpus. Return: list: contains information related to the document, including: contents: used for building index title: (if provided) text: (if provided) \"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"verl/workers/retriever_workers.py#BaseRetriever._batch_search","name":"_batch_search","kind":"method","path":"verl/workers/retriever_workers.py","span":"L161-L162","signature":"def _batch_search(self, query_list, num, return_score)","exported":false,"origin":"ast","body_hash":"3a4873d6a4beb109c53552d834cfc803b0c161cbc908217f0aea5c67d4667fb1","body_text":"def _batch_search(self, query_list, num, return_score): pass","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"verl/workers/retriever_workers.py#BaseRetriever.search","name":"search","kind":"method","path":"verl/workers/retriever_workers.py","span":"L164-L165","signature":"def search(self, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"3d2bd614b882a4bd4196108a68168e90abd4d275d61e4c7e493e3800af3888a1","body_text":"def search(self, *args, **kwargs): return self._search(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"verl/workers/retriever_workers.py#BaseRetriever.batch_search","name":"batch_search","kind":"method","path":"verl/workers/retriever_workers.py","span":"L167-L168","signature":"def batch_search(self, *args, **kwargs)","exported":true,"origin":"ast","body_hash":"d37a6c77b70ed1c7532c9a1a534aa2bb6e6e942eeacc42bc02faa2e163da8629","body_text":"def batch_search(self, *args, **kwargs): return self._batch_search(*args, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRetriever"},{"id":"verl/workers/retriever_workers.py#BM25Retriever","name":"BM25Retriever","kind":"class","path":"verl/workers/retriever_workers.py","span":"L171-L231","signature":"class BM25Retriever(BaseRetriever)","exported":true,"origin":"ast","body_hash":"8dbd4c324303ae0e233bc2a9f5a38663f171aecc4896d8e0bb931fd5872c2f0b","body_text":"class BM25Retriever(BaseRetriever): r\"\"\"BM25 retriever based on pre-built pyserini index.\"\"\" def __init__(self, config): super().__init__(config) raise NotImplementedError from pyserini.search.lucene import LuceneSearcher self.searcher = LuceneSearcher(self.index_path) self.contain_doc = self._check_contain_doc() if not self.contain_doc: self.corpus = load_corpus(self.corpus_path) self.max_process_num = 8 def _check_contain_doc(self): r\"\"\"Check if the index contains document content \"\"\" return self.searcher.doc(0).raw() is not None def _search(self, query: str, num: int = None, return_score = False) -> List[Dict[str, str]]: if num is None: num = self.topk hits = self.searcher.search(query, num) if len(hits) < 1: if return_score: return [],[] else: return [] scores = [hit.score for hit in hits] if len(hits) < num: warnings.warn('Not enough documents retrieved!') else: hits = hits[:num] if self.contain_doc: all_contents = [json.loads(self.searcher.doc(hit.docid).raw())['contents'] for hit in hits] results = [{'title': content.split(\"\\n\")[0].strip(\"\\\"\"), 'text': \"\\n\".join(content.split(\"\\n\")[1:]), 'contents': content} for content in all_contents] else: results = load_docs(self.corpus, [hit.docid for hit in hits]) if return_score: return results, scores else: return results def _batch_search(self, query_list, num: int = None, return_score = False): # TODO: modify batch method results = [] scores = [] for query in query_list: item_result, item_score = self._search(query, num,True) results.append(item_result) scores.append(item_score) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#BM25Retriever.__init__","name":"__init__","kind":"method","path":"verl/workers/retriever_workers.py","span":"L174-L182","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"6f0a024980f061fd4a1296e405731137ff33a00581bb5f82040e27f735759ce7","body_text":"def __init__(self, config): super().__init__(config) raise NotImplementedError from pyserini.search.lucene import LuceneSearcher self.searcher = LuceneSearcher(self.index_path) self.contain_doc = self._check_contain_doc() if not self.contain_doc: self.corpus = load_corpus(self.corpus_path) self.max_process_num = 8","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"verl/workers/retriever_workers.py#BM25Retriever._check_contain_doc","name":"_check_contain_doc","kind":"method","path":"verl/workers/retriever_workers.py","span":"L184-L187","signature":"def _check_contain_doc(self)","exported":false,"origin":"ast","body_hash":"f7a574ad883e23e81357c53f3a2947ddd8491adbc804954c0d23985aa507d3c4","body_text":"def _check_contain_doc(self): r\"\"\"Check if the index contains document content \"\"\" return self.searcher.doc(0).raw() is not None","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"verl/workers/retriever_workers.py#BM25Retriever._search","name":"_search","kind":"method","path":"verl/workers/retriever_workers.py","span":"L189-L217","signature":"def _search(self, query: str, num: int = None, return_score = False) -> List[Dict[str, str]]","exported":false,"origin":"ast","body_hash":"a07f442a9561ebd63dc4c1e3f79003636e2f38b206f554402e2ee3148fd9a28f","body_text":"def _search(self, query: str, num: int = None, return_score = False) -> List[Dict[str, str]]: if num is None: num = self.topk hits = self.searcher.search(query, num) if len(hits) < 1: if return_score: return [],[] else: return [] scores = [hit.score for hit in hits] if len(hits) < num: warnings.warn('Not enough documents retrieved!') else: hits = hits[:num] if self.contain_doc: all_contents = [json.loads(self.searcher.doc(hit.docid).raw())['contents'] for hit in hits] results = [{'title': content.split(\"\\n\")[0].strip(\"\\\"\"), 'text': \"\\n\".join(content.split(\"\\n\")[1:]), 'contents': content} for content in all_contents] else: results = load_docs(self.corpus, [hit.docid for hit in hits]) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"verl/workers/retriever_workers.py#BM25Retriever._batch_search","name":"_batch_search","kind":"method","path":"verl/workers/retriever_workers.py","span":"L219-L231","signature":"def _batch_search(self, query_list, num: int = None, return_score = False): # TODO: modify batch method","exported":false,"origin":"ast","body_hash":"b569a216808f8f2a47a87bac80d0ffb727e6e0c4e1cf265dcd4d35ed86f12b0a","body_text":"def _batch_search(self, query_list, num: int = None, return_score = False): # TODO: modify batch method results = [] scores = [] for query in query_list: item_result, item_score = self._search(query, num,True) results.append(item_result) scores.append(item_score) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"BM25Retriever"},{"id":"verl/workers/retriever_workers.py#DenseRetriever","name":"DenseRetriever","kind":"class","path":"verl/workers/retriever_workers.py","span":"L234-L309","signature":"class DenseRetriever(BaseRetriever)","exported":true,"origin":"ast","body_hash":"312d2476de0066ae09ea8b20263fb9366be9b13413f30456513421b3e397cc72","body_text":"class DenseRetriever(BaseRetriever): r\"\"\"Dense retriever based on pre-built faiss index.\"\"\" def __init__(self, config: dict, index): super().__init__(config) self.index = index # self.index = faiss.read_index(self.index_path) # if config.faiss_gpu: # co = faiss.GpuMultipleClonerOptions() # co.useFloat16 = True # co.shard = True # self.index = faiss.index_cpu_to_all_gpus(self.index, co=co) # # self.index = faiss.index_cpu_to_all_gpus(self.index) self.corpus = load_corpus(self.corpus_path) self.encoder = Encoder( model_name = self.retrieval_method, model_path = config.retrieval_model_path, pooling_method = config.retrieval_pooling_method, max_length = config.retrieval_query_max_length, use_fp16 = config.retrieval_use_fp16 ) self.topk = config.retrieval_topk self.batch_size = self.config.retrieval_batch_size def _search(self, query: str, num: int = None, return_score = False): raise NotImplementedError if num is None: num = self.topk query_emb = self.encoder.encode(query) scores, idxs = self.index.search(query_emb, k=num) idxs = idxs[0] scores = scores[0] results = load_docs(self.corpus, idxs) if return_score: return results, scores else: return results def _batch_search(self, query_list: List[str], num: int = None, return_score = False): if isinstance(query_list, str): query_list = [query_list] if num is None: num = self.topk batch_size = self.batch_size results = [] scores = [] for start_idx in tqdm(range(0, len(query_list), batch_size), desc='Retrieval process: '): query_batch = query_list[start_idx:start_idx + batch_size] # from time import time # a = time() batch_emb = self.encoder.encode(query_batch) # b = time() # print(f'################### encode time {b-a} #####################') batch_scores, batch_idxs = ray.get(self.index.batch_search.remote(batch_emb, k=num)) batch_scores = batch_scores.tolist() batch_idxs = batch_idxs.tolist() # print(f'################### search time {time()-b} #####################') # exit() flat_idxs = sum(batch_idxs, []) batch_results = load_docs(self.corpus, flat_idxs) batch_results = [batch_results[i*num : (i+1)*num] for i in range(len(batch_idxs))] scores.extend(batch_scores) results.extend(batch_results) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#DenseRetriever.__init__","name":"__init__","kind":"method","path":"verl/workers/retriever_workers.py","span":"L237-L257","signature":"def __init__(self, config: dict, index)","exported":false,"origin":"ast","body_hash":"e36182e79edee473ddc23b0d9cd71a7f3637161982c5e0c6f013dfed386505df","body_text":"def __init__(self, config: dict, index): super().__init__(config) self.index = index # self.index = faiss.read_index(self.index_path) # if config.faiss_gpu: # co = faiss.GpuMultipleClonerOptions() # co.useFloat16 = True # co.shard = True # self.index = faiss.index_cpu_to_all_gpus(self.index, co=co) # # self.index = faiss.index_cpu_to_all_gpus(self.index) self.corpus = load_corpus(self.corpus_path) self.encoder = Encoder( model_name = self.retrieval_method, model_path = config.retrieval_model_path, pooling_method = config.retrieval_pooling_method, max_length = config.retrieval_query_max_length, use_fp16 = config.retrieval_use_fp16 ) self.topk = config.retrieval_topk self.batch_size = self.config.retrieval_batch_size","summary_state":"pending","summary":null,"crux":null,"owner":"DenseRetriever"},{"id":"verl/workers/retriever_workers.py#DenseRetriever._search","name":"_search","kind":"method","path":"verl/workers/retriever_workers.py","span":"L259-L272","signature":"def _search(self, query: str, num: int = None, return_score = False)","exported":false,"origin":"ast","body_hash":"0d57ba5ba406327befde2733fcd86103326cab6341d6a32df14c0b21814a104c","body_text":"def _search(self, query: str, num: int = None, return_score = False): raise NotImplementedError if num is None: num = self.topk query_emb = self.encoder.encode(query) scores, idxs = self.index.search(query_emb, k=num) idxs = idxs[0] scores = scores[0] results = load_docs(self.corpus, idxs) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"DenseRetriever"},{"id":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","name":"_batch_search","kind":"method","path":"verl/workers/retriever_workers.py","span":"L274-L309","signature":"def _batch_search(self, query_list: List[str], num: int = None, return_score = False)","exported":false,"origin":"ast","body_hash":"f9066ec2e79c87df3f047b886af4e539c108a0cf867cf696dd83bd7b6761745f","body_text":"def _batch_search(self, query_list: List[str], num: int = None, return_score = False): if isinstance(query_list, str): query_list = [query_list] if num is None: num = self.topk batch_size = self.batch_size results = [] scores = [] for start_idx in tqdm(range(0, len(query_list), batch_size), desc='Retrieval process: '): query_batch = query_list[start_idx:start_idx + batch_size] # from time import time # a = time() batch_emb = self.encoder.encode(query_batch) # b = time() # print(f'################### encode time {b-a} #####################') batch_scores, batch_idxs = ray.get(self.index.batch_search.remote(batch_emb, k=num)) batch_scores = batch_scores.tolist() batch_idxs = batch_idxs.tolist() # print(f'################### search time {time()-b} #####################') # exit() flat_idxs = sum(batch_idxs, []) batch_results = load_docs(self.corpus, flat_idxs) batch_results = [batch_results[i*num : (i+1)*num] for i in range(len(batch_idxs))] scores.extend(batch_scores) results.extend(batch_results) if return_score: return results, scores else: return results","summary_state":"pending","summary":null,"crux":null,"owner":"DenseRetriever"},{"id":"verl/workers/retriever_workers.py#get_retriever","name":"get_retriever","kind":"function","path":"verl/workers/retriever_workers.py","span":"L311-L324","signature":"def get_retriever(config, index)","exported":true,"origin":"ast","body_hash":"7d599aa8707797b7cc945e80cb58cab1bbbded8bc7fd4d0a8a91f6dbf133dd4d","body_text":"def get_retriever(config, index): r\"\"\"Automatically select retriever class based on config's retrieval method Args: config (dict): configuration with 'retrieval_method' key Returns: Retriever: retriever instance \"\"\" if config.retrieval_method == \"bm25\": raise NotImplementedError return BM25Retriever(config) else: return DenseRetriever(config, index)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#RetrieveWorker","name":"RetrieveWorker","kind":"class","path":"verl/workers/retriever_workers.py","span":"L328-L345","signature":"class RetrieveWorker(Worker)","exported":true,"origin":"ast","body_hash":"134537ce7d547e3f8939115b099dce466fee06337e08d0eeb0bcb56d4e14aa6a","body_text":"class RetrieveWorker(Worker): \"\"\"Environment worker that handles GPU-based environment operations.\"\"\" def __init__(self, config, faiss_server): super().__init__() config.index_path = os.path.join(config.index_path, f'{config.retrieval_method}_Flat.index') if config.retrieval_method != 'bm25' else os.path.join(config.index_path, 'bm25') self.config = config # Initialize environment later self.faiss_server = faiss_server @register(dispatch_mode=Dispatch.ONE_TO_ALL) def init_model(self): self.retriever = get_retriever(self.config, self.faiss_server) torch.cuda.empty_cache() @register(dispatch_mode=Dispatch.ALL_TO_ALL) def batch_search(self, queries): return self.retriever.batch_search(queries)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#RetrieveWorker.__init__","name":"__init__","kind":"method","path":"verl/workers/retriever_workers.py","span":"L331-L336","signature":"def __init__(self, config, faiss_server)","exported":false,"origin":"ast","body_hash":"eac5f538b84cfd4a364f66aff39b7451ffa7de36ff321b0c4be27617b36e1f95","body_text":"def __init__(self, config, faiss_server): super().__init__() config.index_path = os.path.join(config.index_path, f'{config.retrieval_method}_Flat.index') if config.retrieval_method != 'bm25' else os.path.join(config.index_path, 'bm25') self.config = config # Initialize environment later self.faiss_server = faiss_server","summary_state":"pending","summary":null,"crux":null,"owner":"RetrieveWorker"},{"id":"verl/workers/retriever_workers.py#RetrieveWorker.init_model","name":"init_model","kind":"method","path":"verl/workers/retriever_workers.py","span":"L339-L341","signature":"def init_model(self)","exported":true,"origin":"ast","body_hash":"012a737bfea81310f2f960a1d36698388dd0f51921d0b22b029a3f6f177596fc","body_text":"def init_model(self): self.retriever = get_retriever(self.config, self.faiss_server) torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null,"owner":"RetrieveWorker"},{"id":"verl/workers/retriever_workers.py#RetrieveWorker.batch_search","name":"batch_search","kind":"method","path":"verl/workers/retriever_workers.py","span":"L344-L345","signature":"def batch_search(self, queries)","exported":true,"origin":"ast","body_hash":"9c9d4693ffe79dbdb040376e17128b5454c3009b60b79b0b2aeba9451125950f","body_text":"def batch_search(self, queries): return self.retriever.batch_search(queries)","summary_state":"pending","summary":null,"crux":null,"owner":"RetrieveWorker"},{"id":"verl/workers/retriever_workers.py#FAISSIndexServer","name":"FAISSIndexServer","kind":"class","path":"verl/workers/retriever_workers.py","span":"L353-L383","signature":"class FAISSIndexServer","exported":true,"origin":"ast","body_hash":"6eba3d0fefb3630dfa70263093335e487a945ccb3add2db3513cbe3bb2dce1f7","body_text":"class FAISSIndexServer: \"\"\"Ray Actor that loads and serves a shared FAISS index with FAISS GPU optimization.\"\"\" def __init__(self, config): \"\"\"Initialize the FAISS index only once.\"\"\" print(\"[FAISSIndexServer] Loading FAISS index...\") self.config = config self.index = self.load_index(config) def load_index(self, config): \"\"\"Loads the FAISS index into GPU memory with sharding.\"\"\" index_path = os.path.join(config.index_path, f'{config.retrieval_method}_Flat.index') index = faiss.read_index(index_path) if self.config.faiss_gpu: # Apply FAISS GPU settings co = faiss.GpuMultipleClonerOptions() co.useFloat16 = True # Reduce memory footprint co.shard = True # Distribute index across all GPUs print(\"[FAISSIndexServer] Moving FAISS index to all GPUs with sharding enabled...\") index = faiss.index_cpu_to_all_gpus(index, co=co) print(\"[FAISSIndexServer] FAISS index successfully moved to GPUs.\") return index def batch_search(self, batch_emb, k): \"\"\"Perform batch search on the FAISS index.\"\"\" print(f\"[FAISSIndexServer] Received {len(batch_emb)} queries.\") return self.index.search(batch_emb, k) # Adjust 'k' as needed","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/retriever_workers.py#FAISSIndexServer.__init__","name":"__init__","kind":"method","path":"verl/workers/retriever_workers.py","span":"L356-L360","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"2ddc71189b77903ddba4bd01620a676c4762d3c520bb272436e18ec58615e0ae","body_text":"def __init__(self, config): \"\"\"Initialize the FAISS index only once.\"\"\" print(\"[FAISSIndexServer] Loading FAISS index...\") self.config = config self.index = self.load_index(config)","summary_state":"pending","summary":null,"crux":null,"owner":"FAISSIndexServer"},{"id":"verl/workers/retriever_workers.py#FAISSIndexServer.load_index","name":"load_index","kind":"method","path":"verl/workers/retriever_workers.py","span":"L362-L378","signature":"def load_index(self, config)","exported":true,"origin":"ast","body_hash":"b0b0172bdfc64da20ec852665bbfd846d30780265780096ae442e61e5270d0ee","body_text":"def load_index(self, config): \"\"\"Loads the FAISS index into GPU memory with sharding.\"\"\" index_path = os.path.join(config.index_path, f'{config.retrieval_method}_Flat.index') index = faiss.read_index(index_path) if self.config.faiss_gpu: # Apply FAISS GPU settings co = faiss.GpuMultipleClonerOptions() co.useFloat16 = True # Reduce memory footprint co.shard = True # Distribute index across all GPUs print(\"[FAISSIndexServer] Moving FAISS index to all GPUs with sharding enabled...\") index = faiss.index_cpu_to_all_gpus(index, co=co) print(\"[FAISSIndexServer] FAISS index successfully moved to GPUs.\") return index","summary_state":"pending","summary":null,"crux":null,"owner":"FAISSIndexServer"},{"id":"verl/workers/retriever_workers.py#FAISSIndexServer.batch_search","name":"batch_search","kind":"method","path":"verl/workers/retriever_workers.py","span":"L380-L383","signature":"def batch_search(self, batch_emb, k)","exported":true,"origin":"ast","body_hash":"7338814b55736c7ba4b87a76bd356722a4bec4139594b3b917815f25841516a1","body_text":"def batch_search(self, batch_emb, k): \"\"\"Perform batch search on the FAISS index.\"\"\" print(f\"[FAISSIndexServer] Received {len(batch_emb)} queries.\") return self.index.search(batch_emb, k) # Adjust 'k' as needed","summary_state":"pending","summary":null,"crux":null,"owner":"FAISSIndexServer"}],"rawEdges":[{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"os","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"verl.single_controller.base","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"verl.single_controller.base.decorator","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"json","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"os","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"warnings","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"typing","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"functools","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"tqdm","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"multiprocessing","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"faiss","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"transformers","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"argparse","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"datasets","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#load_corpus","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#load_corpus","relation":"calls","name":"load_dataset","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#read_jsonl","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#read_jsonl","relation":"calls","name":"open","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#read_jsonl","relation":"calls","name":"readlines","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#read_jsonl","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#read_jsonl","relation":"calls","name":"loads","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#load_docs","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#load_docs","relation":"calls","name":"int","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#load_model","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#load_model","relation":"calls","name":"eval","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#load_model","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#load_model","relation":"calls","name":"half","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#load_model","relation":"calls","name":"from_pretrained","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#pooling","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#pooling","relation":"calls","name":"masked_fill","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#pooling","relation":"calls","name":"bool","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#pooling","relation":"calls","name":"sum","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#pooling","relation":"calls","name":"sum","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#pooling","relation":"calls","name":"NotImplementedError","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#Encoder","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder","relation":"contains","targetId":"verl/workers/retriever_workers.py#Encoder.__init__","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.__init__","relation":"calls","name":"load_model","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder","relation":"contains","targetId":"verl/workers/retriever_workers.py#Encoder.encode","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"lower","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"lower","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"tokenizer","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"Encoder"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"items","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"type","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"zeros","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"model","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"Encoder"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"model","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"Encoder"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"pooling","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"lower","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"normalize","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"numpy","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"cpu","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"detach","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"pooling"},{"source":"verl/workers/retriever_workers.py#Encoder.encode","relation":"calls","name":"astype","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"pooling"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#BaseRetriever","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BaseRetriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#BaseRetriever.__init__","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BaseRetriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#BaseRetriever._search","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BaseRetriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#BaseRetriever._batch_search","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BaseRetriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#BaseRetriever.search","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BaseRetriever.search","relation":"calls","name":"_search","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"BaseRetriever"},{"source":"verl/workers/retriever_workers.py#BaseRetriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#BaseRetriever.batch_search","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BaseRetriever.batch_search","relation":"calls","name":"_batch_search","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"BaseRetriever"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#BM25Retriever","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever","relation":"extends","name":"BaseRetriever","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#BM25Retriever.__init__","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"pyserini.search.lucene","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever.__init__","relation":"calls","name":"LuceneSearcher","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever.__init__","relation":"calls","name":"_check_contain_doc","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"BM25Retriever"},{"source":"verl/workers/retriever_workers.py#BM25Retriever.__init__","relation":"calls","name":"load_corpus","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#BM25Retriever._check_contain_doc","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._check_contain_doc","relation":"calls","name":"raw","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._check_contain_doc","relation":"calls","name":"doc","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"LuceneSearcher"},{"source":"verl/workers/retriever_workers.py#BM25Retriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#BM25Retriever._search","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"search","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"LuceneSearcher"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"warn","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"loads","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"raw","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"doc","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"LuceneSearcher"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"strip","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"split","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"join","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"split","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._search","relation":"calls","name":"load_docs","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#BM25Retriever._batch_search","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._batch_search","relation":"calls","name":"_search","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"BM25Retriever"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._batch_search","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#BM25Retriever._batch_search","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#DenseRetriever","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever","relation":"extends","name":"BaseRetriever","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#DenseRetriever.__init__","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever.__init__","relation":"calls","name":"load_corpus","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever.__init__","relation":"calls","name":"Encoder","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#DenseRetriever._search","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._search","relation":"calls","name":"encode","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"Encoder"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._search","relation":"calls","name":"search","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._search","relation":"calls","name":"load_docs","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever","relation":"contains","targetId":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"tqdm","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"encode","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"Encoder"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"remote","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"tolist","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"tolist","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"sum","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"load_docs","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"extend","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#DenseRetriever._batch_search","relation":"calls","name":"extend","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#get_retriever","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#get_retriever","relation":"calls","name":"BM25Retriever","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#get_retriever","relation":"calls","name":"DenseRetriever","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#RetrieveWorker","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker","relation":"extends","name":"Worker","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker","relation":"contains","targetId":"verl/workers/retriever_workers.py#RetrieveWorker.__init__","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker.__init__","relation":"calls","name":"join","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker.__init__","relation":"calls","name":"join","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker","relation":"contains","targetId":"verl/workers/retriever_workers.py#RetrieveWorker.init_model","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker.init_model","relation":"calls","name":"get_retriever","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker.init_model","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker","relation":"calls","name":"register","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker","relation":"contains","targetId":"verl/workers/retriever_workers.py#RetrieveWorker.batch_search","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#RetrieveWorker.batch_search","relation":"calls","name":"batch_search","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"get_retriever"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"ray","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"faiss","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"imports","specifier":"torch","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"calls","name":"remote","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py","relation":"contains","targetId":"verl/workers/retriever_workers.py#FAISSIndexServer","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer","relation":"contains","targetId":"verl/workers/retriever_workers.py#FAISSIndexServer.__init__","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.__init__","relation":"calls","name":"load_index","viaMember":true,"file":"verl/workers/retriever_workers.py","recvType":"FAISSIndexServer"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer","relation":"contains","targetId":"verl/workers/retriever_workers.py#FAISSIndexServer.load_index","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.load_index","relation":"calls","name":"join","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.load_index","relation":"calls","name":"read_index","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.load_index","relation":"calls","name":"GpuMultipleClonerOptions","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.load_index","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.load_index","relation":"calls","name":"index_cpu_to_all_gpus","viaMember":true,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.load_index","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer","relation":"contains","targetId":"verl/workers/retriever_workers.py#FAISSIndexServer.batch_search","file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.batch_search","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.batch_search","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/retriever_workers.py"},{"source":"verl/workers/retriever_workers.py#FAISSIndexServer.batch_search","relation":"calls","name":"search","viaMember":true,"file":"verl/workers/retriever_workers.py"}]},"verl/workers/reward_model/__init__.py":{"size":638,"mtimeMs":1789275830328.2031,"hash":"969b640ba9c36d51ddf12380041963eb2255b094318027e76053c191339bd825","nodes":[{"id":"verl/workers/reward_model/__init__.py","name":"__init__.py","kind":"file","path":"verl/workers/reward_model/__init__.py","span":"L1-L16","signature":null,"exported":true,"origin":"ast","body_hash":"969b640ba9c36d51ddf12380041963eb2255b094318027e76053c191339bd825","chars":638,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 BasePPORewardModel"}],"rawEdges":[{"source":"verl/workers/reward_model/__init__.py","relation":"imports","specifier":".base","file":"verl/workers/reward_model/__init__.py"}]},"verl/workers/reward_model/base.py":{"size":1710,"mtimeMs":1789275830328.2031,"hash":"7f6e0ccd8edaca5a5edc02e8f545297bef4acaea84d0858b870340a157ed0983","nodes":[{"id":"verl/workers/reward_model/base.py","name":"base.py","kind":"file","path":"verl/workers/reward_model/base.py","span":"L1-L46","signature":null,"exported":true,"origin":"ast","body_hash":"7f6e0ccd8edaca5a5edc02e8f545297bef4acaea84d0858b870340a157ed0983","chars":1710,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 class for reward model \"\"\" from abc import ABC, abstractmethod from verl import DataProto"},{"id":"verl/workers/reward_model/base.py#BasePPORewardModel","name":"BasePPORewardModel","kind":"class","path":"verl/workers/reward_model/base.py","span":"L23-L45","signature":"class BasePPORewardModel(ABC)","exported":true,"origin":"ast","body_hash":"5899f0b13c0a80cf1b4c7a1d3753e2fb0309a500a657cc877cdc77f1e03eecba","body_text":"class BasePPORewardModel(ABC): def __init__(self, config): self.config = config @abstractmethod def compute_reward(self, data: DataProto) -> DataProto: \"\"\"Computing reward given input_ids. The transformers should output a tensor with shape [batch_size, sequence_length], and the value at [EOS] mask should be gathered. Args: data: must contain keys \"input_ids\", \"attention_mask\" and \"position_ids\". - input_ids: [batch_size, sequence_length] - attention_mask: [batch_size, sequence_length] - position_ids: [batch_size, sequence_length] Returns: a data pass protocol containing \"reward\". Only the [EOS] position contains the reward. Other position should have zero reward. Note that this may change in the future if we use dense reward. So, we leave the interface for general case. - reward: [batch_size, sequence_length]. \"\"\" pass","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/reward_model/base.py#BasePPORewardModel.__init__","name":"__init__","kind":"method","path":"verl/workers/reward_model/base.py","span":"L25-L26","signature":"def __init__(self, config)","exported":false,"origin":"ast","body_hash":"c203a0217d1c1950415c427f1cb9c7d6085761d3f21b90b6d7c6e63a0643abbd","body_text":"def __init__(self, config): self.config = config","summary_state":"pending","summary":null,"crux":null,"owner":"BasePPORewardModel"},{"id":"verl/workers/reward_model/base.py#BasePPORewardModel.compute_reward","name":"compute_reward","kind":"method","path":"verl/workers/reward_model/base.py","span":"L29-L45","signature":"def compute_reward(self, data: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"4c4deda96c979a035f4f4e2779980a4cbed5a9d3c0c8f630873e173564ed9a13","body_text":"def compute_reward(self, data: DataProto) -> DataProto: \"\"\"Computing reward given input_ids. The transformers should output a tensor with shape [batch_size, sequence_length], and the value at [EOS] mask should be gathered. Args: data: must contain keys \"input_ids\", \"attention_mask\" and \"position_ids\". - input_ids: [batch_size, sequence_length] - attention_mask: [batch_size, sequence_length] - position_ids: [batch_size, sequence_length] Returns: a data pass protocol containing \"reward\". Only the [EOS] position contains the reward. Other position should have zero reward. Note that this may change in the future if we use dense reward. So, we leave the interface for general case. - reward: [batch_size, sequence_length]. \"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"BasePPORewardModel"}],"rawEdges":[{"source":"verl/workers/reward_model/base.py","relation":"imports","specifier":"abc","file":"verl/workers/reward_model/base.py"},{"source":"verl/workers/reward_model/base.py","relation":"imports","specifier":"verl","file":"verl/workers/reward_model/base.py"},{"source":"verl/workers/reward_model/base.py","relation":"contains","targetId":"verl/workers/reward_model/base.py#BasePPORewardModel","file":"verl/workers/reward_model/base.py"},{"source":"verl/workers/reward_model/base.py#BasePPORewardModel","relation":"extends","name":"ABC","file":"verl/workers/reward_model/base.py"},{"source":"verl/workers/reward_model/base.py#BasePPORewardModel","relation":"contains","targetId":"verl/workers/reward_model/base.py#BasePPORewardModel.__init__","file":"verl/workers/reward_model/base.py"},{"source":"verl/workers/reward_model/base.py#BasePPORewardModel","relation":"contains","targetId":"verl/workers/reward_model/base.py#BasePPORewardModel.compute_reward","file":"verl/workers/reward_model/base.py"}]},"verl/workers/reward_model/megatron/__init__.py":{"size":647,"mtimeMs":1789275830328.2031,"hash":"1a01e19bf170896603577bc8752239de2743685141424dcef64bb685cb2ffd96","nodes":[{"id":"verl/workers/reward_model/megatron/__init__.py","name":"__init__.py","kind":"file","path":"verl/workers/reward_model/megatron/__init__.py","span":"L1-L16","signature":null,"exported":true,"origin":"ast","body_hash":"1a01e19bf170896603577bc8752239de2743685141424dcef64bb685cb2ffd96","chars":647,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 .reward_model import MegatronRewardModel"}],"rawEdges":[{"source":"verl/workers/reward_model/megatron/__init__.py","relation":"imports","specifier":".reward_model","file":"verl/workers/reward_model/megatron/__init__.py"}]},"verl/workers/reward_model/megatron/reward_model.py":{"size":13316,"mtimeMs":1789275830328.2031,"hash":"e823efa4d1a943ad9f6eb79061d40aed46501528efd83c2fdeedd645b5d81d53","nodes":[{"id":"verl/workers/reward_model/megatron/reward_model.py","name":"reward_model.py","kind":"file","path":"verl/workers/reward_model/megatron/reward_model.py","span":"L1-L279","signature":null,"exported":true,"origin":"ast","body_hash":"e823efa4d1a943ad9f6eb79061d40aed46501528efd83c2fdeedd645b5d81d53","chars":13316,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Megatron Reward Model. \"\"\" from tensordict import TensorDict from functools import partial from verl import DataProto from verl.utils.torch_functional import logprobs_from_logits import torch import torch import torch.distributed from verl.utils.torch_functional import get_eos_mask, pad_sequence_to_length from verl.utils.megatron.pipeline_parallel import (compute_transformers_input_shapes, make_batch_generator) from verl import DataProto from verl.utils.torch_functional import logprobs_from_logits, broadcast_dict_tensor, split_dict_tensor_into_batches from verl.utils.torch_dtypes import PrecisionType from verl.workers.reward_model.base import BasePPORewardModel from verl.utils.megatron import sequence_parallel as sp_utils from megatron.core import parallel_state as mpu from megatron.core.pipeline_parallel import get_forward_backward_func"},{"id":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel","name":"MegatronRewardModel","kind":"class","path":"verl/workers/reward_model/megatron/reward_model.py","span":"L37-L278","signature":"class MegatronRewardModel(BasePPORewardModel)","exported":true,"origin":"ast","body_hash":"c1b14e25645d7baa9a6c33c7de2b6212bb06caa5627fad3641af8d17a4a6ef18","body_text":"class MegatronRewardModel(BasePPORewardModel): def __init__(self, config, model_config, reward_model_module: torch.nn.ModuleList, megatron_config, sft_tokenizer=None, rm_tokenizer=None): self.config = config self.reward_model_module = reward_model_module self.megatron_config = megatron_config self.model_config = model_config self.device = 'cuda' self.sft_tokenizer = sft_tokenizer self.rm_tokenizer = rm_tokenizer self.use_different_tokenizer = rm_tokenizer is not None if self.config.param_offload: self.offload_params_to_cpu() def re_encode_by_rm_tokenizer(self, data: DataProto) -> DataProto: assert self.use_different_tokenizer, 're-encode need rm tokenizer not be None!' # need to use rm tokenizer to re-generate input_ids, attention_mask and position_ids # 1. remove pad for each sequence # 2. decode by sft_tokenizer, remove sft system prompts # 3. encode by rm_tokenizer with rm system prompts, get rm_input_ids # 4. generate attention_mask and position_ids input_ids = data.batch['input_ids'] # (bs, seq_len) attention_mask = data.batch['attention_mask'] position_ids = data.batch['position_ids'] ori_values = {'input_ids': input_ids, 'attention_mask': attention_mask, 'position_ids': position_ids} ori_bs, ori_seqlen = input_ids.size(0), input_ids.size(1) input_ids_for_rm = [] attention_mask_for_rm = [] position_ids_for_rm = [] print_decode = True ori_seqlen = ori_seqlen + 128 for id, mask in zip(input_ids, attention_mask): # 1. remove pad for each sequence non_zero_indices = torch.nonzero(mask).view(-1) begin_pos, end_pos = non_zero_indices[0].item(), non_zero_indices[-1].item() valid_id = id[begin_pos:end_pos + 1] # 2. decode by sft_tokenizer, remove sft system prompts decode_result = self.sft_tokenizer.decode(valid_id) # workaround decode_with_rm_chat = decode_result.replace(\"<|user|>\\n\", \"[INST] \").replace( \"</s>\\n<|assistant|>\\n\", \" [/INST]\").replace(\"</s> \\n<|assistant|>\\n\", \" [/INST]\") + \"</s>\" print(f\"decode_with_rm_chat: {decode_with_rm_chat}\") if print_decode and torch.distributed.get_rank() == 0: # only print first decode result print(f'device {torch.cuda.current_device()}: sft decode result:\\n{decode_result}\\n \\ \\ndevice {torch.cuda.current_device()}: sft decode result with rm chat template:\\n{decode_with_rm_chat}\\n\\n' ) print_decode = False # 3. encode by rm_tokenizer rm_input_ids = self.rm_tokenizer(decode_with_rm_chat, return_tensors='pt')['input_ids'][0].to(input_ids.device) # 4. generate attention_mask and position_ids rm_attention_mask = torch.ones_like(rm_input_ids, device=input_ids.device) cur_seqlen = rm_input_ids.shape[-1] # NOTE(gh): the later reward compute will process the shape (bs, seqlen_pad_128) if cur_seqlen > ori_seqlen: print(f'warninig: rm encode seqlen {cur_seqlen} > sft encode seqlen {ori_seqlen}') rm_input_ids = rm_input_ids[:ori_seqlen] rm_attention_mask = rm_attention_mask[:ori_seqlen] else: # right padding rm_input_ids = pad_sequence_to_length(rm_input_ids, ori_seqlen, self.rm_tokenizer.pad_token_id) rm_attention_mask = pad_sequence_to_length(rm_attention_mask, ori_seqlen, 0) rm_position_ids = torch.arange(0, ori_seqlen, device=input_ids.device) input_ids_for_rm.append(torch.unsqueeze(rm_input_ids, dim=0)) attention_mask_for_rm.append(torch.unsqueeze(rm_attention_mask, dim=0)) position_ids_for_rm.append(torch.unsqueeze(rm_position_ids, dim=0)) input_ids_for_rm = torch.cat(input_ids_for_rm, dim=0) attention_mask_for_rm = torch.cat(attention_mask_for_rm, dim=0) position_ids_for_rm = torch.cat(position_ids_for_rm, dim=0) # (bs, seqlen) will not change, but input_ids, attention_mask and position_ids will change # NOTE(gh): need to replace into origin values after compute reward! data.batch['input_ids'] = input_ids_for_rm data.batch['attention_mask'] = attention_mask_for_rm data.batch['position_ids'] = position_ids_for_rm return data, ori_values @torch.no_grad() def compute_reward(self, data: DataProto) -> DataProto: if self.config.param_offload: self.load_params_to_cuda() if self.use_different_tokenizer: data, ori_values = self.re_encode_by_rm_tokenizer(data) input_ids = data.batch['input_ids'] # (bs, seq_len') attention_mask = data.batch['attention_mask'] position_ids = data.batch['position_ids'] responses = data.batch['responses'] batch_size = responses.size(0) response_length = responses.size(1) with torch.no_grad(): output = self.forward_batch(data) if mpu.is_pipeline_last_stage(ignore_virtual=True): logits = torch.cat([o['logits'] for o in output], dim=0) else: logits = torch.empty( (input_ids.shape[0], input_ids.shape[1]), dtype=torch.bfloat16, # TODO(sgm): check why is bfloat16 device=input_ids.device) # broadcast across pp ranks torch.distributed.broadcast(tensor=logits, src=mpu.get_pipeline_model_parallel_last_rank(), group=mpu.get_pipeline_model_parallel_group(), async_op=False) # (bs, seqlen', hidden_size) -> (bs, seqlen', 1) -> (bs, seqlen') token_level_rewards = logits # find the last token reward ends = attention_mask.cumsum(dim=-1).argmax(dim=-1).view(-1, 1) # (bs, ","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.__init__","name":"__init__","kind":"method","path":"verl/workers/reward_model/megatron/reward_model.py","span":"L39-L56","signature":"def __init__(self, config, model_config, reward_model_module: torch.nn.ModuleList, megatron_config, sft_tokenizer=None, rm_tokenizer=None)","exported":false,"origin":"ast","body_hash":"b8e87a9ff5b13674f39f86b8dff13b7f2f37a7c5d1cefb41acbb036f057587bb","body_text":"def __init__(self, config, model_config, reward_model_module: torch.nn.ModuleList, megatron_config, sft_tokenizer=None, rm_tokenizer=None): self.config = config self.reward_model_module = reward_model_module self.megatron_config = megatron_config self.model_config = model_config self.device = 'cuda' self.sft_tokenizer = sft_tokenizer self.rm_tokenizer = rm_tokenizer self.use_different_tokenizer = rm_tokenizer is not None if self.config.param_offload: self.offload_params_to_cpu()","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronRewardModel"},{"id":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","name":"re_encode_by_rm_tokenizer","kind":"method","path":"verl/workers/reward_model/megatron/reward_model.py","span":"L58-L123","signature":"def re_encode_by_rm_tokenizer(self, data: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"16912e40ed8b64195f0d485dcb4e0a83d45455ce0510359da0bafc053dc3dd82","body_text":"def re_encode_by_rm_tokenizer(self, data: DataProto) -> DataProto: assert self.use_different_tokenizer, 're-encode need rm tokenizer not be None!' # need to use rm tokenizer to re-generate input_ids, attention_mask and position_ids # 1. remove pad for each sequence # 2. decode by sft_tokenizer, remove sft system prompts # 3. encode by rm_tokenizer with rm system prompts, get rm_input_ids # 4. generate attention_mask and position_ids input_ids = data.batch['input_ids'] # (bs, seq_len) attention_mask = data.batch['attention_mask'] position_ids = data.batch['position_ids'] ori_values = {'input_ids': input_ids, 'attention_mask': attention_mask, 'position_ids': position_ids} ori_bs, ori_seqlen = input_ids.size(0), input_ids.size(1) input_ids_for_rm = [] attention_mask_for_rm = [] position_ids_for_rm = [] print_decode = True ori_seqlen = ori_seqlen + 128 for id, mask in zip(input_ids, attention_mask): # 1. remove pad for each sequence non_zero_indices = torch.nonzero(mask).view(-1) begin_pos, end_pos = non_zero_indices[0].item(), non_zero_indices[-1].item() valid_id = id[begin_pos:end_pos + 1] # 2. decode by sft_tokenizer, remove sft system prompts decode_result = self.sft_tokenizer.decode(valid_id) # workaround decode_with_rm_chat = decode_result.replace(\"<|user|>\\n\", \"[INST] \").replace( \"</s>\\n<|assistant|>\\n\", \" [/INST]\").replace(\"</s> \\n<|assistant|>\\n\", \" [/INST]\") + \"</s>\" print(f\"decode_with_rm_chat: {decode_with_rm_chat}\") if print_decode and torch.distributed.get_rank() == 0: # only print first decode result print(f'device {torch.cuda.current_device()}: sft decode result:\\n{decode_result}\\n \\ \\ndevice {torch.cuda.current_device()}: sft decode result with rm chat template:\\n{decode_with_rm_chat}\\n\\n' ) print_decode = False # 3. encode by rm_tokenizer rm_input_ids = self.rm_tokenizer(decode_with_rm_chat, return_tensors='pt')['input_ids'][0].to(input_ids.device) # 4. generate attention_mask and position_ids rm_attention_mask = torch.ones_like(rm_input_ids, device=input_ids.device) cur_seqlen = rm_input_ids.shape[-1] # NOTE(gh): the later reward compute will process the shape (bs, seqlen_pad_128) if cur_seqlen > ori_seqlen: print(f'warninig: rm encode seqlen {cur_seqlen} > sft encode seqlen {ori_seqlen}') rm_input_ids = rm_input_ids[:ori_seqlen] rm_attention_mask = rm_attention_mask[:ori_seqlen] else: # right padding rm_input_ids = pad_sequence_to_length(rm_input_ids, ori_seqlen, self.rm_tokenizer.pad_token_id) rm_attention_mask = pad_sequence_to_length(rm_attention_mask, ori_seqlen, 0) rm_position_ids = torch.arange(0, ori_seqlen, device=input_ids.device) input_ids_for_rm.append(torch.unsqueeze(rm_input_ids, dim=0)) attention_mask_for_rm.append(torch.unsqueeze(rm_attention_mask, dim=0)) position_ids_for_rm.append(torch.unsqueeze(rm_position_ids, dim=0)) input_ids_for_rm = torch.cat(input_ids_for_rm, dim=0) attention_mask_for_rm = torch.cat(attention_mask_for_rm, dim=0) position_ids_for_rm = torch.cat(position_ids_for_rm, dim=0) # (bs, seqlen) will not change, but input_ids, attention_mask and position_ids will change # NOTE(gh): need to replace into origin values after compute reward! data.batch['input_ids'] = input_ids_for_rm data.batch['attention_mask'] = attention_mask_for_rm data.batch['position_ids'] = position_ids_for_rm return data, ori_values","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronRewardModel"},{"id":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","name":"compute_reward","kind":"method","path":"verl/workers/reward_model/megatron/reward_model.py","span":"L126-L186","signature":"def compute_reward(self, data: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"76dc6fda0d17a0074b3bac801895c7f22b672643dcd78ac8e1fa30707c996a8e","body_text":"def compute_reward(self, data: DataProto) -> DataProto: if self.config.param_offload: self.load_params_to_cuda() if self.use_different_tokenizer: data, ori_values = self.re_encode_by_rm_tokenizer(data) input_ids = data.batch['input_ids'] # (bs, seq_len') attention_mask = data.batch['attention_mask'] position_ids = data.batch['position_ids'] responses = data.batch['responses'] batch_size = responses.size(0) response_length = responses.size(1) with torch.no_grad(): output = self.forward_batch(data) if mpu.is_pipeline_last_stage(ignore_virtual=True): logits = torch.cat([o['logits'] for o in output], dim=0) else: logits = torch.empty( (input_ids.shape[0], input_ids.shape[1]), dtype=torch.bfloat16, # TODO(sgm): check why is bfloat16 device=input_ids.device) # broadcast across pp ranks torch.distributed.broadcast(tensor=logits, src=mpu.get_pipeline_model_parallel_last_rank(), group=mpu.get_pipeline_model_parallel_group(), async_op=False) # (bs, seqlen', hidden_size) -> (bs, seqlen', 1) -> (bs, seqlen') token_level_rewards = logits # find the last token reward ends = attention_mask.cumsum(dim=-1).argmax(dim=-1).view(-1, 1) # (bs, 1) rewards = torch.gather(token_level_rewards, dim=1, index=ends) # (bs, 1) if self.use_different_tokenizer: data.batch.update(ori_values) input_ids = ori_values['input_ids'] attention_mask = ori_values['attention_mask'] position_ids = ori_values['position_ids'] token_level_rewards = rewards.expand(attention_mask.shape[0], attention_mask.shape[1]) # (bs, ori_seqlen) # assign last valid token reward to ori position eos_mask_idx = torch.argmax(position_ids * attention_mask, dim=-1) # (bs,) eos_mask = torch.zeros_like(attention_mask) eos_mask[torch.arange(batch_size), eos_mask_idx] = 1. token_level_rewards = token_level_rewards * eos_mask token_level_rewards = token_level_rewards[:, -response_length:] if self.config.param_offload: self.offload_params_to_cpu() else: # add empty cache after each compute torch.cuda.empty_cache() batch = TensorDict({'rm_scores': token_level_rewards}, batch_size=input_ids.shape[0]) return DataProto(batch=batch)","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronRewardModel"},{"id":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","name":"forward_batch","kind":"method","path":"verl/workers/reward_model/megatron/reward_model.py","span":"L188-L263","signature":"def forward_batch(self, data: DataProto)","exported":true,"origin":"ast","body_hash":"2ccf24d9701c58ec2048f5c6642f440e5d117a02c67d0b428a04e493295346e9","body_text":"def forward_batch(self, data: DataProto): \"\"\" We assume: - The model takes input: (input_ids, attention_mask, position_ids). No rmpad for the input - The communication shape is (total_nnz_pad_to_sp // tp_size, 1, hidden_size) if sequence parallel is enabled \"\"\" # broadcast from last pp rank to all other pp ranks # TODO: actually, we just need to control the sampling order. data.batch = data.batch.contiguous() broadcast_dict_tensor(data.batch, src=mpu.get_pipeline_model_parallel_last_rank(), group=mpu.get_pipeline_model_parallel_group()) # split into micro-batches if self.config is not None and 'ppo_micro_batch_size' in self.config: infer_batch_size = self.config.ppo_micro_batch_size else: infer_batch_size = data.batch.batch_size[0] data.batch['attention_mask'] = data.batch['attention_mask'].to(bool) batches = split_dict_tensor_into_batches(data.batch, batch_size=infer_batch_size) n_micro_batch = len(batches) seq_len = batches[0]['input_ids'].shape[1] # compute input shapes for pp stages input_shapes = compute_transformers_input_shapes( batches, meta_info={ 'sequence_parallel': self.megatron_config.sequence_parallel, 'hidden_size': self.model_config.hidden_size }) # compute input shapes for pp stages forward_backward_func = get_forward_backward_func() def loss_func(output): return 1., {'logits': output.logits} def forward_step(batch_iter, model): batch = next(batch_iter) input_ids = batch['input_ids'] attention_mask = batch['attention_mask'] position_ids = batch['position_ids'] output = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) return output, loss_func # batch should be a list of batches inside micro-batches batch_generator = make_batch_generator(batches, vpp_size=len(self.reward_model_module)) # TODO: we may use the new schedule instead # for flash-attn: (seq_len, batch_size, hidden_size) = (mbs*seq_len, 1, hidden_size) if mpu.get_pipeline_model_parallel_world_size() > 1: losses_reduced = forward_backward_func( forward_step_func=forward_step, data_iterator=batch_generator, model=self.reward_model_module, num_microbatches=n_micro_batch, input_shapes=input_shapes, # must set for flash-attn sequence packing seq_length=infer_batch_size * seq_len, # no use when input_shapes was set hidden_size=self.model_config.hidden_size, # no use when input_shapes was set micro_batch_size=1, # no use when input_shapes was set forward_only=True, ) else: losses_reduced = forward_backward_func( forward_step_func=forward_step, data_iterator=batch_generator, model=self.reward_model_module, num_microbatches=n_micro_batch, seq_length=infer_batch_size * seq_len, # in use for pp = 1 hidden_size=self.model_config.hidden_size, # in use for pp = 1 micro_batch_size=1, # in use for pp = 1 forward_only=True, ) # loss_reduces contains the stats returned from loss_func return losses_reduced","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronRewardModel"},{"id":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch.loss_func","name":"loss_func","kind":"function","path":"verl/workers/reward_model/megatron/reward_model.py","span":"L222-L223","signature":"def loss_func(output)","exported":true,"origin":"ast","body_hash":"f3a216b6d0a19fd180a1afe64c9487dc4c69cbf84f409689f3a545538bbbaa45","body_text":"def loss_func(output): return 1., {'logits': output.logits}","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch.forward_step","name":"forward_step","kind":"function","path":"verl/workers/reward_model/megatron/reward_model.py","span":"L225-L231","signature":"def forward_step(batch_iter, model)","exported":true,"origin":"ast","body_hash":"d44a40cfb38a0ec59f6a029648054ac6d6e9128475c9decd9662f2145bbc5faa","body_text":"def forward_step(batch_iter, model): batch = next(batch_iter) input_ids = batch['input_ids'] attention_mask = batch['attention_mask'] position_ids = batch['position_ids'] output = model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) return output, loss_func","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.offload_params_to_cpu","name":"offload_params_to_cpu","kind":"method","path":"verl/workers/reward_model/megatron/reward_model.py","span":"L265-L271","signature":"def offload_params_to_cpu(self)","exported":true,"origin":"ast","body_hash":"4455271e2192a01590ad75648561f8efe2f47818cc3c1516ac487a78b7710056","body_text":"def offload_params_to_cpu(self): if self.device == 'cuda': for reward_model_module in self.reward_model_module: for name, param in reward_model_module.named_parameters(): param.data = param.data.to('cpu', non_blocking=True) self.device = 'cpu' torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronRewardModel"},{"id":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.load_params_to_cuda","name":"load_params_to_cuda","kind":"method","path":"verl/workers/reward_model/megatron/reward_model.py","span":"L273-L278","signature":"def load_params_to_cuda(self)","exported":true,"origin":"ast","body_hash":"42a9a62d65760e78029f89c989410c547d7c8835a5797765ee4f58afe6474534","body_text":"def load_params_to_cuda(self): if self.device == 'cpu': for reward_model_module in self.reward_model_module: for name, param in reward_model_module.named_parameters(): param.data = param.data.to(torch.cuda.current_device(), non_blocking=True) self.device = 'cuda'","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronRewardModel"}],"rawEdges":[{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"tensordict","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"functools","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"verl","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"torch","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"torch","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"verl.utils.megatron.pipeline_parallel","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"verl","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"verl.utils.torch_dtypes","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"verl.workers.reward_model.base","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"verl.utils.megatron","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"megatron.core","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"imports","specifier":"megatron.core.pipeline_parallel","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py","relation":"contains","targetId":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel","relation":"extends","name":"BasePPORewardModel","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel","relation":"contains","targetId":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.__init__","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.__init__","relation":"calls","name":"offload_params_to_cpu","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py","recvType":"MegatronRewardModel"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel","relation":"contains","targetId":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"zip","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"view","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"nonzero","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"item","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"decode","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"replace","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"replace","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"replace","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"rm_tokenizer","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py","recvType":"MegatronRewardModel"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"ones_like","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"pad_sequence_to_length","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"pad_sequence_to_length","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"arange","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.re_encode_by_rm_tokenizer","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel","relation":"contains","targetId":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"load_params_to_cuda","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py","recvType":"MegatronRewardModel"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"re_encode_by_rm_tokenizer","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py","recvType":"MegatronRewardModel"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"forward_batch","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py","recvType":"MegatronRewardModel"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"is_pipeline_last_stage","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"empty","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"get_pipeline_model_parallel_last_rank","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"get_pipeline_model_parallel_group","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"view","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"argmax","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"cumsum","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"gather","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"update","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"expand","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"argmax","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"zeros_like","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"arange","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"offload_params_to_cpu","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py","recvType":"MegatronRewardModel"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.compute_reward","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel","relation":"contains","targetId":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"broadcast_dict_tensor","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"get_pipeline_model_parallel_last_rank","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"get_pipeline_model_parallel_group","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"split_dict_tensor_into_batches","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"compute_transformers_input_shapes","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"get_forward_backward_func","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"contains","targetId":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch.loss_func","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"contains","targetId":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch.forward_step","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch.forward_step","relation":"calls","name":"next","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch.forward_step","relation":"calls","name":"model","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"make_batch_generator","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"forward_backward_func","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.forward_batch","relation":"calls","name":"forward_backward_func","viaMember":false,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel","relation":"contains","targetId":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.offload_params_to_cpu","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.offload_params_to_cpu","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.offload_params_to_cpu","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.offload_params_to_cpu","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel","relation":"contains","targetId":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.load_params_to_cuda","file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.load_params_to_cuda","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.load_params_to_cuda","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"},{"source":"verl/workers/reward_model/megatron/reward_model.py#MegatronRewardModel.load_params_to_cuda","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/reward_model/megatron/reward_model.py"}]},"verl/workers/rollout/__init__.py":{"size":753,"mtimeMs":1789275830329.2031,"hash":"6dc2b3dd28631c76b826f545091bedc233111ad6bab00511a4fe40e6a1e4cf78","nodes":[{"id":"verl/workers/rollout/__init__.py","name":"__init__.py","kind":"file","path":"verl/workers/rollout/__init__.py","span":"L1-L20","signature":null,"exported":true,"origin":"ast","body_hash":"6dc2b3dd28631c76b826f545091bedc233111ad6bab00511a4fe40e6a1e4cf78","chars":753,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 from .naive import NaiveRollout from .hf_rollout import HFRollout __all__ = [\"BaseRollout\", \"NaiveRollout\", \"HFRollout\"]"}],"rawEdges":[{"source":"verl/workers/rollout/__init__.py","relation":"imports","specifier":".base","file":"verl/workers/rollout/__init__.py"},{"source":"verl/workers/rollout/__init__.py","relation":"imports","specifier":".naive","file":"verl/workers/rollout/__init__.py"},{"source":"verl/workers/rollout/__init__.py","relation":"imports","specifier":".hf_rollout","file":"verl/workers/rollout/__init__.py"}]},"verl/workers/rollout/base.py":{"size":1141,"mtimeMs":1789275830329.2031,"hash":"480b812f13b681735e18347220dc7462a7ae982fb61a46541dc713d4e2608d12","nodes":[{"id":"verl/workers/rollout/base.py","name":"base.py","kind":"file","path":"verl/workers/rollout/base.py","span":"L1-L38","signature":null,"exported":true,"origin":"ast","body_hash":"480b812f13b681735e18347220dc7462a7ae982fb61a46541dc713d4e2608d12","chars":1141,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 abc import ABC, abstractmethod from typing import Iterable, Union from verl import DataProto __all__ = ['BaseRollout']"},{"id":"verl/workers/rollout/base.py#BaseRollout","name":"BaseRollout","kind":"class","path":"verl/workers/rollout/base.py","span":"L23-L37","signature":"class BaseRollout(ABC)","exported":true,"origin":"ast","body_hash":"0bb29a174c27fb52aaef5a01f896b396b40cfa07ee8eb07ecdca34ce5c484d8d","body_text":"class BaseRollout(ABC): def __init__(self): \"\"\" Args: dataloader: an Iterable of TensorDict that consistently generates prompts. Note that the dataloader should handle when the training stops. \"\"\" super().__init__() @abstractmethod def generate_sequences(self, prompts: DataProto) -> DataProto: \"\"\"Generate sequences\"\"\" pass","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/rollout/base.py#BaseRollout.__init__","name":"__init__","kind":"method","path":"verl/workers/rollout/base.py","span":"L25-L32","signature":"def __init__(self)","exported":false,"origin":"ast","body_hash":"a1b0cdfe2c0b19340546ad0d89fe5f87e69ff58e9eecf4ab2a52ca701c0cc55b","body_text":"def __init__(self): \"\"\" Args: dataloader: an Iterable of TensorDict that consistently generates prompts. Note that the dataloader should handle when the training stops. \"\"\" super().__init__()","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRollout"},{"id":"verl/workers/rollout/base.py#BaseRollout.generate_sequences","name":"generate_sequences","kind":"method","path":"verl/workers/rollout/base.py","span":"L35-L37","signature":"def generate_sequences(self, prompts: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"a8d487dc760ca40ea6f9698b4a22b1114a5e3d594f1f75460d0ff457e6d6d051","body_text":"def generate_sequences(self, prompts: DataProto) -> DataProto: \"\"\"Generate sequences\"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"BaseRollout"}],"rawEdges":[{"source":"verl/workers/rollout/base.py","relation":"imports","specifier":"abc","file":"verl/workers/rollout/base.py"},{"source":"verl/workers/rollout/base.py","relation":"imports","specifier":"typing","file":"verl/workers/rollout/base.py"},{"source":"verl/workers/rollout/base.py","relation":"imports","specifier":"verl","file":"verl/workers/rollout/base.py"},{"source":"verl/workers/rollout/base.py","relation":"contains","targetId":"verl/workers/rollout/base.py#BaseRollout","file":"verl/workers/rollout/base.py"},{"source":"verl/workers/rollout/base.py#BaseRollout","relation":"extends","name":"ABC","file":"verl/workers/rollout/base.py"},{"source":"verl/workers/rollout/base.py#BaseRollout","relation":"contains","targetId":"verl/workers/rollout/base.py#BaseRollout.__init__","file":"verl/workers/rollout/base.py"},{"source":"verl/workers/rollout/base.py#BaseRollout.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/rollout/base.py"},{"source":"verl/workers/rollout/base.py#BaseRollout.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/rollout/base.py"},{"source":"verl/workers/rollout/base.py#BaseRollout","relation":"contains","targetId":"verl/workers/rollout/base.py#BaseRollout.generate_sequences","file":"verl/workers/rollout/base.py"}]},"verl/workers/rollout/hf_rollout.py":{"size":5867,"mtimeMs":1789275830329.2031,"hash":"38d84c1bcbe61fa560dc6ab3e5f2acb46a4561c1bb329cf344d6cccc781b1009","nodes":[{"id":"verl/workers/rollout/hf_rollout.py","name":"hf_rollout.py","kind":"file","path":"verl/workers/rollout/hf_rollout.py","span":"L1-L141","signature":null,"exported":true,"origin":"ast","body_hash":"38d84c1bcbe61fa560dc6ab3e5f2acb46a4561c1bb329cf344d6cccc781b1009","chars":5867,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 verl import DataProto from verl.utils.torch_functional import get_eos_mask from .base import BaseRollout from transformers import GenerationConfig __all__ = ['HFRollout']"},{"id":"verl/workers/rollout/hf_rollout.py#HFRollout","name":"HFRollout","kind":"class","path":"verl/workers/rollout/hf_rollout.py","span":"L35-L140","signature":"class HFRollout(BaseRollout)","exported":true,"origin":"ast","body_hash":"a28ff57aa2de02003823defe5fe8414e7eb09ee5dd6dfbc6d7ebf53f19a18a94","body_text":"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: 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'] pad_token_id = prompts.meta_info['pad_token_id'] batch_size = idx.size(0) prompt_length = idx.size(1) self.module.eval() param_ctx = contextlib.nullcontext() # make sampling args can be overriden by inputs do_sample = prompts.meta_info.get('do_sample', self.config.do_sample) 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 = prompts.meta_info.get('top_k', self.config.get('top_k', 0)) if top_k is None: top_k = 0 top_k = max(0, top_k) # to be compatible with vllm temperature = prompts.meta_info.get('temperature', self.config.temperature) generation_config = GenerationConfig(temperature=temperature, top_p=top_p, top_k=top_k) 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: with torch.autocast(device_type='cuda', dtype=torch.bfloat16): output = self.module.generate( input_ids=idx, attention_mask=attention_mask, do_sample=do_sample, max_new_tokens=response_length, # max_length=max_length, eos_token_id=eos_token_id, pad_token_id=pad_token_id, generation_config=generation_config, # renormalize_logits=True, 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 # 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=(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 prompt = seq[:, :prompt_length] # (bs, prompt_length) response = seq[:, prompt_length:] # (bs, 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(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_eos_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=batch_size) # empty cache before compute old_log_prob torch.cuda.empty_cache() self.module.train() return DataProto(batch=batch)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/rollout/hf_rollout.py#HFRollout.__init__","name":"__init__","kind":"method","path":"verl/workers/rollout/hf_rollout.py","span":"L37-L40","signature":"def __init__(self, module: nn.Module, config)","exported":false,"origin":"ast","body_hash":"938a898c776d914d6fff8fb9f563c0a8aab80bf9cd5b444d2a995ef8a1db1017","body_text":"def __init__(self, module: nn.Module, config): super().__init__() self.config = config self.module = module","summary_state":"pending","summary":null,"crux":null,"owner":"HFRollout"},{"id":"verl/workers/rollout/hf_rollout.py#HFRollout.generate_sequences","name":"generate_sequences","kind":"method","path":"verl/workers/rollout/hf_rollout.py","span":"L42-L48","signature":"def generate_sequences(self, prompts: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"8d1b7708714a9ec9961e4581ff95789c74f8a5d567f701f25c09b2e2d5f5c75b","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"HFRollout"},{"id":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","name":"_generate_minibatch","kind":"method","path":"verl/workers/rollout/hf_rollout.py","span":"L51-L140","signature":"def _generate_minibatch(self, prompts: DataProto) -> DataProto","exported":false,"origin":"ast","body_hash":"21b6ff7ae2d420c42606ff6a28381aafbf78d7c91af6a45a3511e87f5b646357","body_text":"def _generate_minibatch(self, prompts: DataProto) -> DataProto: 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'] pad_token_id = prompts.meta_info['pad_token_id'] batch_size = idx.size(0) prompt_length = idx.size(1) self.module.eval() param_ctx = contextlib.nullcontext() # make sampling args can be overriden by inputs do_sample = prompts.meta_info.get('do_sample', self.config.do_sample) 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 = prompts.meta_info.get('top_k', self.config.get('top_k', 0)) if top_k is None: top_k = 0 top_k = max(0, top_k) # to be compatible with vllm temperature = prompts.meta_info.get('temperature', self.config.temperature) generation_config = GenerationConfig(temperature=temperature, top_p=top_p, top_k=top_k) 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: with torch.autocast(device_type='cuda', dtype=torch.bfloat16): output = self.module.generate( input_ids=idx, attention_mask=attention_mask, do_sample=do_sample, max_new_tokens=response_length, # max_length=max_length, eos_token_id=eos_token_id, pad_token_id=pad_token_id, generation_config=generation_config, # renormalize_logits=True, 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 # 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=(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 prompt = seq[:, :prompt_length] # (bs, prompt_length) response = seq[:, prompt_length:] # (bs, 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(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_eos_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=batch_size) # empty cache before compute old_log_prob torch.cuda.empty_cache() self.module.train() return DataProto(batch=batch)","summary_state":"pending","summary":null,"crux":null,"owner":"HFRollout"}],"rawEdges":[{"source":"verl/workers/rollout/hf_rollout.py","relation":"imports","specifier":"contextlib","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py","relation":"imports","specifier":"torch","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py","relation":"imports","specifier":"tensordict","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py","relation":"imports","specifier":"torch","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py","relation":"imports","specifier":"torch.distributed.fsdp","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py","relation":"imports","specifier":"verl","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py","relation":"imports","specifier":".base","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py","relation":"imports","specifier":"transformers","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py","relation":"contains","targetId":"verl/workers/rollout/hf_rollout.py#HFRollout","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout","relation":"extends","name":"BaseRollout","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout","relation":"contains","targetId":"verl/workers/rollout/hf_rollout.py#HFRollout.__init__","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout","relation":"contains","targetId":"verl/workers/rollout/hf_rollout.py#HFRollout.generate_sequences","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout.generate_sequences","relation":"calls","name":"max","viaMember":false,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout.generate_sequences","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout.generate_sequences","relation":"calls","name":"chunk","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py","recvType":"DataProto"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout.generate_sequences","relation":"calls","name":"_generate_minibatch","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py","recvType":"HFRollout"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout.generate_sequences","relation":"calls","name":"concat","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout","relation":"contains","targetId":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"eval","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"nullcontext","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"max","viaMember":false,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"GenerationConfig","viaMember":false,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"isinstance","viaMember":false,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"summon_full_params","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"autocast","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"generate","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"ones","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"arange","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"repeat","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"get_eos_mask","viaMember":false,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"train","viaMember":true,"file":"verl/workers/rollout/hf_rollout.py"},{"source":"verl/workers/rollout/hf_rollout.py#HFRollout._generate_minibatch","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/workers/rollout/hf_rollout.py"}]},"verl/workers/rollout/naive/__init__.py":{"size":641,"mtimeMs":1789275830329.2031,"hash":"551983cdba422a02f03c8b9cce52face25963b511b9029e09247a16afa3c70fd","nodes":[{"id":"verl/workers/rollout/naive/__init__.py","name":"__init__.py","kind":"file","path":"verl/workers/rollout/naive/__init__.py","span":"L1-L16","signature":null,"exported":true,"origin":"ast","body_hash":"551983cdba422a02f03c8b9cce52face25963b511b9029e09247a16afa3c70fd","chars":641,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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"}],"rawEdges":[{"source":"verl/workers/rollout/naive/__init__.py","relation":"imports","specifier":".naive_rollout","file":"verl/workers/rollout/naive/__init__.py"}]},"verl/workers/rollout/naive/naive_rollout.py":{"size":4821,"mtimeMs":1789275830329.2031,"hash":"09dcb0d16d53b02d8cb665b21b4fc74f47ed479bdbe6b0292550d2b1b9a46d61","nodes":[{"id":"verl/workers/rollout/naive/naive_rollout.py","name":"naive_rollout.py","kind":"file","path":"verl/workers/rollout/naive/naive_rollout.py","span":"L1-L120","signature":null,"exported":true,"origin":"ast","body_hash":"09dcb0d16d53b02d8cb665b21b4fc74f47ed479bdbe6b0292550d2b1b9a46d61","chars":4821,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 \"\"\" from typing import Iterable, Union 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__ = ['NativeRollout']"},{"id":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout","name":"NaiveRollout","kind":"class","path":"verl/workers/rollout/naive/naive_rollout.py","span":"L36-L119","signature":"class NaiveRollout(BaseRollout)","exported":true,"origin":"ast","body_hash":"d033c70dbc91de7e89d9e0ee15df56d8a1a803581f641d089d46b1f39735885e","body_text":"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) prev_attention_mask = torch.logical_and(idx_next != eos_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)","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.__init__","name":"__init__","kind":"method","path":"verl/workers/rollout/naive/naive_rollout.py","span":"L38-L49","signature":"def __init__(self, module: nn.Module, config)","exported":false,"origin":"ast","body_hash":"5b26d9ea4261a1f3ebf130a225d89b053a49d225da50338111ff4ebcfac28369","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"NaiveRollout"},{"id":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","name":"generate_sequences","kind":"method","path":"verl/workers/rollout/naive/naive_rollout.py","span":"L52-L119","signature":"def generate_sequences(self, prompts: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"c391e003befed30bce23f0c306da44b94053b370033c05bb2a811edb31be628b","body_text":"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) prev_attention_mask = torch.logical_and(idx_next != eos_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)","summary_state":"pending","summary":null,"crux":null,"owner":"NaiveRollout"}],"rawEdges":[{"source":"verl/workers/rollout/naive/naive_rollout.py","relation":"imports","specifier":"typing","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py","relation":"imports","specifier":"torch","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py","relation":"imports","specifier":"tensordict","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py","relation":"imports","specifier":"torch","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py","relation":"imports","specifier":"verl","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py","relation":"imports","specifier":"..base","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py","relation":"contains","targetId":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout","relation":"extends","name":"BaseRollout","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout","relation":"contains","targetId":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.__init__","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout","relation":"contains","targetId":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"eval","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"ones","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"module","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py","recvType":"NaiveRollout"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"topk","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"min","viaMember":false,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"float","viaMember":false,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"softmax","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"multinomial","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"argmax","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"logical_and","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"bool","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"stack","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"logprobs_from_logits","viaMember":false,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"train","viaMember":true,"file":"verl/workers/rollout/naive/naive_rollout.py"},{"source":"verl/workers/rollout/naive/naive_rollout.py#NaiveRollout.generate_sequences","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/workers/rollout/naive/naive_rollout.py"}]},"verl/workers/rollout/tokenizer.py":{"size":5747,"mtimeMs":1789275830329.2031,"hash":"92fbaf5d1c93e2c0d2ec2e399af83be277ac744005466b3016331516708ed00b","nodes":[{"id":"verl/workers/rollout/tokenizer.py","name":"tokenizer.py","kind":"file","path":"verl/workers/rollout/tokenizer.py","span":"L1-L163","signature":null,"exported":true,"origin":"ast","body_hash":"92fbaf5d1c93e2c0d2ec2e399af83be277ac744005466b3016331516708ed00b","chars":5747,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 from typing import Dict, List, Union __all__ = ['HybridEngineBaseTokenizer']"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","name":"HybridEngineBaseTokenizer","kind":"class","path":"verl/workers/rollout/tokenizer.py","span":"L23-L162","signature":"class HybridEngineBaseTokenizer(ABC)","exported":true,"origin":"ast","body_hash":"ef4dbc43dcbeeaa3771984ea6a1bb64590819573f33a3ab9dbc0fe8d5f59b98f","body_text":"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(`'<unk>'`, `'<cls>'`, etc.) mapped to class attributes. \"\"\" pass @property @abstractmethod def all_special_tokens(self) -> List[str]: \"\"\" `List[str]`: A list of the unique special tokens (`'<unk>'`, `'<cls>'`, ..., 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: Union[int, List[int], \"np.ndarray\", \"torch.Tensor\", \"tf.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, tf.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: Union[int, List[int]], skip_special_tokens: bool = False) -> Union[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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.vocab_size","name":"vocab_size","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L28-L32","signature":"def vocab_size(self)","exported":true,"origin":"ast","body_hash":"3763bfe0e853185d33bd91be4bba3eb80d98ea087561b90dbb1fe0250ab31e88","body_text":"def vocab_size(self): \"\"\" `int`: Size of the base vocabulary (without the added tokens). \"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.pad_token_id","name":"pad_token_id","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L36-L40","signature":"def pad_token_id(self)","exported":true,"origin":"ast","body_hash":"2d07309ef776988ad38fa2fb3a0789e931e70b7725437da91d96aca544b81114","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.eos_token_id","name":"eos_token_id","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L44-L49","signature":"def eos_token_id(self)","exported":true,"origin":"ast","body_hash":"c84a77363c577bc350ba1d404909618bb8d33cbf2867b6e95d0adef3a2996bc9","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.all_special_ids","name":"all_special_ids","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L53-L57","signature":"def all_special_ids(self) -> List[int]","exported":true,"origin":"ast","body_hash":"44622c48dc64fa19fa8d213366b1901c053efc1608abb9ead883de58144c6edd","body_text":"def all_special_ids(self) -> List[int]: \"\"\" `List[int]`: List the ids of the special tokens(`'<unk>'`, `'<cls>'`, etc.) mapped to class attributes. \"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.all_special_tokens","name":"all_special_tokens","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L61-L67","signature":"def all_special_tokens(self) -> List[str]","exported":true,"origin":"ast","body_hash":"0afe0ed556c2d88131cc7a09c43db13d055d9af70f781b75fb469c7d6eb70ede","body_text":"def all_special_tokens(self) -> List[str]: \"\"\" `List[str]`: A list of the unique special tokens (`'<unk>'`, `'<cls>'`, ..., etc.). Convert tokens of `tokenizers.AddedToken` type to string. \"\"\" pass","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.encode","name":"encode","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L70-L83","signature":"def encode(self, text)","exported":true,"origin":"ast","body_hash":"6dd784fbefb8cc37419f33828d501941e9e64efb7a459b1c7e0b56bde8fd34d8","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.decode","name":"decode","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L86-L113","signature":"def decode( self, token_ids: Union[int, List[int], \"np.ndarray\", \"torch.Tensor\", \"tf.Tensor\"], skip_special_tokens: bool = False, clean_up_tokenization_spaces: bool = None, **kwargs, ) -> str","exported":true,"origin":"ast","body_hash":"9cae8bc84c54442f9fbcd954d2351a43f59686169b6d95e83f85b284bfe9968f","body_text":"def decode( self, token_ids: Union[int, List[int], \"np.ndarray\", \"torch.Tensor\", \"tf.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, tf.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","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.convert_ids_to_tokens","name":"convert_ids_to_tokens","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L116-L132","signature":"def convert_ids_to_tokens(self, ids: Union[int, List[int]], skip_special_tokens: bool = False) -> Union[str, List[str]]","exported":true,"origin":"ast","body_hash":"efcc765335f7dbbedbf8392a922c1223fe134ac84a2d97caa9aa221f9a164fe1","body_text":"def convert_ids_to_tokens(self, ids: Union[int, List[int]], skip_special_tokens: bool = False) -> Union[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","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.get_added_vocab","name":"get_added_vocab","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L135-L144","signature":"def get_added_vocab(self) -> Dict[str, int]","exported":true,"origin":"ast","body_hash":"5cd7dfc19c33240a1bc2392ee12b2b489159874198e19b5bba11e4ac2e9ce424","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.convert_tokens_to_string","name":"convert_tokens_to_string","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L147-L158","signature":"def convert_tokens_to_string(self, tokens: List[str]) -> str","exported":true,"origin":"ast","body_hash":"a0e0934445c028d437d89b0b47ef4a473e4a1282fbe4882f4670403b94033232","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"},{"id":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.is_fast","name":"is_fast","kind":"method","path":"verl/workers/rollout/tokenizer.py","span":"L161-L162","signature":"def is_fast(self)","exported":true,"origin":"ast","body_hash":"f3a56f7225d06df3ea38ff9e8c1daedf9c43a733bdd63b865c9ced9769bc1243","body_text":"def is_fast(self): return False","summary_state":"pending","summary":null,"crux":null,"owner":"HybridEngineBaseTokenizer"}],"rawEdges":[{"source":"verl/workers/rollout/tokenizer.py","relation":"imports","specifier":"abc","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py","relation":"imports","specifier":"typing","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"extends","name":"ABC","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.vocab_size","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.pad_token_id","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.eos_token_id","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.all_special_ids","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.all_special_tokens","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.encode","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.decode","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.convert_ids_to_tokens","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.get_added_vocab","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.convert_tokens_to_string","file":"verl/workers/rollout/tokenizer.py"},{"source":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer","relation":"contains","targetId":"verl/workers/rollout/tokenizer.py#HybridEngineBaseTokenizer.is_fast","file":"verl/workers/rollout/tokenizer.py"}]},"verl/workers/rollout/vllm_rollout/__init__.py":{"size":638,"mtimeMs":1789275830329.2031,"hash":"875ce539030328cfd2feeec2d531a12ca24175762177ba5b4db2256c8669e0f7","nodes":[{"id":"verl/workers/rollout/vllm_rollout/__init__.py","name":"__init__.py","kind":"file","path":"verl/workers/rollout/vllm_rollout/__init__.py","span":"L1-L15","signature":null,"exported":true,"origin":"ast","body_hash":"875ce539030328cfd2feeec2d531a12ca24175762177ba5b4db2256c8669e0f7","chars":638,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 .vllm_rollout import vLLMRollout"}],"rawEdges":[{"source":"verl/workers/rollout/vllm_rollout/__init__.py","relation":"imports","specifier":".vllm_rollout","file":"verl/workers/rollout/vllm_rollout/__init__.py"}]},"verl/workers/rollout/vllm_rollout/vllm_rollout.py":{"size":12028,"mtimeMs":1789275830329.2031,"hash":"4aca8567c90b651f2cd2d290f936b0b50bcc476c010f15e21c389ae7c0616dcb","nodes":[{"id":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","name":"vllm_rollout.py","kind":"file","path":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","span":"L1-L264","signature":null,"exported":true,"origin":"ast","body_hash":"4aca8567c90b651f2cd2d290f936b0b50bcc476c010f15e21c389ae7c0616dcb","chars":11862,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" from typing import List from contextlib import contextmanager from omegaconf import DictConfig import torch import torch.distributed from tensordict import TensorDict from torch import nn from verl import DataProto from verl.utils.torch_functional import get_eos_mask, pad_sequence_to_length from verl.workers.rollout.base import BaseRollout from verl.third_party.vllm import LLM, vllm_version from verl.third_party.vllm import parallel_state as vllm_ps from vllm import SamplingParams # TODO # 1. support pp in vllm # 2. passing tokenizer is not necessary? no encoding/decoding is happending here # 3. simplify init logics # NOTE(sgm): add for verl. We can optimize it by making the dataloader yield List[int] without padding."},{"id":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#_pre_process_inputs","name":"_pre_process_inputs","kind":"function","path":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","span":"L49-L62","signature":"def _pre_process_inputs(pad_token_id, prompt_token_ids: torch.Tensor) -> List[int]: # remove the left padding in the prompt token_id # pad_token_id = self.llm_engine.tokenizer.pad_token_id if self.llm_engine.tokenizer.pad_token_id is not None else self.llm_engine.tokenizer.eos_token_id","exported":false,"origin":"ast","body_hash":"a2f69b855d2de5c48422777aebe25196f6ffeeb516ef9d9a6c99f89352139118","body_text":"def _pre_process_inputs(pad_token_id, prompt_token_ids: torch.Tensor) -> List[int]: # remove the left padding in the prompt token_id # pad_token_id = self.llm_engine.tokenizer.pad_token_id if self.llm_engine.tokenizer.pad_token_id is not None else self.llm_engine.tokenizer.eos_token_id '''new from here''' non_pad_indices = torch.nonzero(prompt_token_ids != pad_token_id, as_tuple=False) # 添加边界检查：如果全是 padding，返回空列表而不是报错 if len(non_pad_indices) == 0: return [] '''new to here''' non_pad_index = torch.nonzero(prompt_token_ids != pad_token_id, as_tuple=False)[0][0] #token_ids = prompt_token_ids[non_pad_index:].tolist() token_ids = prompt_token_ids[non_pad_index:].long().tolist() # 使用 .long() 确保是整数类型 return token_ids","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout","name":"vLLMRollout","kind":"class","path":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","span":"L65-L263","signature":"class vLLMRollout(BaseRollout)","exported":true,"origin":"ast","body_hash":"b366d44562daf8902e5b4b91a7f93443ee02ccc67bb06f3de5f2e191d3d63533","body_text":"class vLLMRollout(BaseRollout): def __init__(self, actor_module: nn.Module, config: DictConfig, tokenizer, model_hf_config, **kwargs): \"\"\"A vLLM rollout. It requires the module is supported by the vllm. Args: module: module here follows huggingface APIs config: DictConfig tokenizer: the task/model tokenizer model_hf_config: the huggingface config to initiallize the generating model in vllm **kwargs: train_tp, for Megatron Backend to initialize hybrid engine (zero redundancy) process group \"\"\" super().__init__() self.config = config assert not (not config.enforce_eager and config.free_cache_engine), \\ \"disable CUDA graph (enforce_eager = False) if free cache engine\" tensor_parallel_size = self.config.get('tensor_model_parallel_size', 1) assert tensor_parallel_size <= torch.distributed.get_world_size(), \\ \"tensor parallel size should be less than or equal to the world size\" if kwargs.get('train_tp', None) is not None: # deployed with megatron import os os.environ['CUDA_TIMER_STREAM_KAFKA_ENABLE'] = '0' os.environ['MEGATRON_IMPORT_TIMERS'] = '0' train_tp = kwargs.get('train_tp', None) num_tp_per_train_tp = train_tp // tensor_parallel_size if vllm_version in ('0.4.2', '0.5.4', '0.6.3'): vllm_ps.initialize_parallel_state(tensor_model_parallel_size=tensor_parallel_size, num_tp_per_train_tp=num_tp_per_train_tp) assert model_hf_config.max_position_embeddings >= config.prompt_length + config.response_length, \\ \"model context length should be greater than total sequence length\" self.inference_engine = LLM(actor_module, tokenizer=tokenizer, model_hf_config=model_hf_config, tensor_parallel_size=tensor_parallel_size, dtype=config.dtype, enforce_eager=config.enforce_eager, gpu_memory_utilization=config.gpu_memory_utilization, skip_tokenizer_init=False, max_model_len=config.prompt_length + config.response_length, load_format=config.load_format) # Offload vllm model to reduce peak memory usage self.inference_engine.offload_model_weights() kwargs = dict( n=1, logprobs=1, # can be set to 0 and let actor to recompute max_tokens=config.response_length, ) # we may detokenize the result all together later if vllm_version in ('0.4.2', '0.5.4', '0.6.3'): kwargs['detokenize'] = False # supporting adding any sampling params from the config file for k in config.keys(): if hasattr(SamplingParams(), str(k)): kwargs[k] = config.get(k) print(f\"kwargs: {kwargs}\") self.sampling_params = SamplingParams(**kwargs) self.pad_token_id = tokenizer.pad_token_id @contextmanager def update_sampling_params(self, **kwargs): # update sampling params old_sampling_params_args = {} if kwargs: for key, value in kwargs.items(): if hasattr(self.sampling_params, key): old_value = getattr(self.sampling_params, key) old_sampling_params_args[key] = old_value setattr(self.sampling_params, key, value) yield # roll back to previous sampling params # if len(old_sampling_params_args): for key, value in old_sampling_params_args.items(): setattr(self.sampling_params, key, value) @torch.no_grad() def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: # rebuild vllm cache engine if self.config.free_cache_engine: self.inference_engine.init_cache_engine() idx = prompts.batch['input_ids'] # (bs, prompt_length) # left-padded attention_mask attention_mask = prompts.batch['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) idx_list = [] # parse idx from torch.Tensor to List[List[str]] '''new here''' # for i in range(batch_size): # idx_list.append(_pre_process_inputs(self.pad_token_id, idx[i])) for i in range(batch_size): token_ids = _pre_process_inputs(self.pad_token_id, idx[i]) # 添加调试信息和格式验证 if len(token_ids) == 0: print(f\"[WARNING] Empty token_ids for sample {i}, using minimal valid sequence\") # 使用一个有效的最小序列（比如只有 BOS token） token_ids = [self.inference_engine.llm_engine.tokenizer.bos_token_id] if hasattr( self.inference_engine.llm_engine.tokenizer, 'bos_token_id') else [1] idx_list.append(token_ids) # 添加调试信息 print(f\"[DEBUG] idx_list length: {len(idx_list)}\") print(f\"[DEBUG] idx_list[0] type: {type(idx_list[0])}, length: {len(idx_list[0]) if idx_list else 0}\") print(f\"[DEBUG] idx_list[0] sample: {idx_list[0][:10] if idx_list and len(idx_list[0]) > 0 else 'empty'}\") '''new to here''' do_sample = prompts.meta_info.get('do_sample', True) if not do_sample: kwargs = { 'best_of': 1, 'top_p': 1.0, 'top_k': -1, 'min_p': 0.0, 'temperature': 0, 'n': 1 # if greedy, only 1 response } # users can customize different sampling_params at different run with self.update_sampling_params(**kwargs): output = self.inference_engine.generate( prompts=None, # because we have already convert it to prompt token id sampling_params=self.sampling_params, prompt_token_ids=idx_list, use_tqdm=False) # TODO(sgm): disable logprob when recompute_log_prob is enable # if n = 1: (bs, response_length) ; if n > 1: (bs * n, response_length) response = output[0].to(idx.device) log_pro","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","name":"__init__","kind":"method","path":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","span":"L67-L131","signature":"def __init__(self, actor_module: nn.Module, config: DictConfig, tokenizer, model_hf_config, **kwargs)","exported":false,"origin":"ast","body_hash":"3ee4b029c5c2d473568eb31b1283fc8116fdfca2adcb20cf3d417752387038d9","body_text":"def __init__(self, actor_module: nn.Module, config: DictConfig, tokenizer, model_hf_config, **kwargs): \"\"\"A vLLM rollout. It requires the module is supported by the vllm. Args: module: module here follows huggingface APIs config: DictConfig tokenizer: the task/model tokenizer model_hf_config: the huggingface config to initiallize the generating model in vllm **kwargs: train_tp, for Megatron Backend to initialize hybrid engine (zero redundancy) process group \"\"\" super().__init__() self.config = config assert not (not config.enforce_eager and config.free_cache_engine), \\ \"disable CUDA graph (enforce_eager = False) if free cache engine\" tensor_parallel_size = self.config.get('tensor_model_parallel_size', 1) assert tensor_parallel_size <= torch.distributed.get_world_size(), \\ \"tensor parallel size should be less than or equal to the world size\" if kwargs.get('train_tp', None) is not None: # deployed with megatron import os os.environ['CUDA_TIMER_STREAM_KAFKA_ENABLE'] = '0' os.environ['MEGATRON_IMPORT_TIMERS'] = '0' train_tp = kwargs.get('train_tp', None) num_tp_per_train_tp = train_tp // tensor_parallel_size if vllm_version in ('0.4.2', '0.5.4', '0.6.3'): vllm_ps.initialize_parallel_state(tensor_model_parallel_size=tensor_parallel_size, num_tp_per_train_tp=num_tp_per_train_tp) assert model_hf_config.max_position_embeddings >= config.prompt_length + config.response_length, \\ \"model context length should be greater than total sequence length\" self.inference_engine = LLM(actor_module, tokenizer=tokenizer, model_hf_config=model_hf_config, tensor_parallel_size=tensor_parallel_size, dtype=config.dtype, enforce_eager=config.enforce_eager, gpu_memory_utilization=config.gpu_memory_utilization, skip_tokenizer_init=False, max_model_len=config.prompt_length + config.response_length, load_format=config.load_format) # Offload vllm model to reduce peak memory usage self.inference_engine.offload_model_weights() kwargs = dict( n=1, logprobs=1, # can be set to 0 and let actor to recompute max_tokens=config.response_length, ) # we may detokenize the result all together later if vllm_version in ('0.4.2', '0.5.4', '0.6.3'): kwargs['detokenize'] = False # supporting adding any sampling params from the config file for k in config.keys(): if hasattr(SamplingParams(), str(k)): kwargs[k] = config.get(k) print(f\"kwargs: {kwargs}\") self.sampling_params = SamplingParams(**kwargs) self.pad_token_id = tokenizer.pad_token_id","summary_state":"pending","summary":null,"crux":null,"owner":"vLLMRollout"},{"id":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.update_sampling_params","name":"update_sampling_params","kind":"method","path":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","span":"L134-L147","signature":"def update_sampling_params(self, **kwargs): # update sampling params","exported":true,"origin":"ast","body_hash":"ea0c4ff74d95491928936e36007f86f7f50b80ab022fe29addc7b11eff29b179","body_text":"def update_sampling_params(self, **kwargs): # update sampling params old_sampling_params_args = {} if kwargs: for key, value in kwargs.items(): if hasattr(self.sampling_params, key): old_value = getattr(self.sampling_params, key) old_sampling_params_args[key] = old_value setattr(self.sampling_params, key, value) yield # roll back to previous sampling params # if len(old_sampling_params_args): for key, value in old_sampling_params_args.items(): setattr(self.sampling_params, key, value)","summary_state":"pending","summary":null,"crux":null,"owner":"vLLMRollout"},{"id":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","name":"generate_sequences","kind":"method","path":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","span":"L150-L263","signature":"def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: # rebuild vllm cache engine","exported":true,"origin":"ast","body_hash":"477b7746564dc258e15f1a6a88284161a9ff32ddf7c403e4e95f075a6105043e","body_text":"def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: # rebuild vllm cache engine if self.config.free_cache_engine: self.inference_engine.init_cache_engine() idx = prompts.batch['input_ids'] # (bs, prompt_length) # left-padded attention_mask attention_mask = prompts.batch['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) idx_list = [] # parse idx from torch.Tensor to List[List[str]] '''new here''' # for i in range(batch_size): # idx_list.append(_pre_process_inputs(self.pad_token_id, idx[i])) for i in range(batch_size): token_ids = _pre_process_inputs(self.pad_token_id, idx[i]) # 添加调试信息和格式验证 if len(token_ids) == 0: print(f\"[WARNING] Empty token_ids for sample {i}, using minimal valid sequence\") # 使用一个有效的最小序列（比如只有 BOS token） token_ids = [self.inference_engine.llm_engine.tokenizer.bos_token_id] if hasattr( self.inference_engine.llm_engine.tokenizer, 'bos_token_id') else [1] idx_list.append(token_ids) # 添加调试信息 print(f\"[DEBUG] idx_list length: {len(idx_list)}\") print(f\"[DEBUG] idx_list[0] type: {type(idx_list[0])}, length: {len(idx_list[0]) if idx_list else 0}\") print(f\"[DEBUG] idx_list[0] sample: {idx_list[0][:10] if idx_list and len(idx_list[0]) > 0 else 'empty'}\") '''new to here''' do_sample = prompts.meta_info.get('do_sample', True) if not do_sample: kwargs = { 'best_of': 1, 'top_p': 1.0, 'top_k': -1, 'min_p': 0.0, 'temperature': 0, 'n': 1 # if greedy, only 1 response } # users can customize different sampling_params at different run with self.update_sampling_params(**kwargs): output = self.inference_engine.generate( prompts=None, # because we have already convert it to prompt token id sampling_params=self.sampling_params, prompt_token_ids=idx_list, use_tqdm=False) # TODO(sgm): disable logprob when recompute_log_prob is enable # if n = 1: (bs, response_length) ; if n > 1: (bs * n, response_length) response = output[0].to(idx.device) log_probs = output[1].to(idx.device) ### #关键修复：确保 response 是整数类型 if response.dtype != torch.long: response = response.long() # log_probs 保持浮点类型 if log_probs.dtype == torch.bfloat16: log_probs = log_probs.float() if idx.dtype != torch.long: idx = idx.long() ### if response.shape[1] < self.config.response_length: response = pad_sequence_to_length(response, self.config.response_length, self.pad_token_id) log_probs = pad_sequence_to_length(log_probs, self.config.response_length, self.pad_token_id) if self.config.n > 1 and do_sample: idx = idx.repeat_interleave(self.config.n, dim=0) attention_mask = attention_mask.repeat_interleave(self.config.n, dim=0) position_ids = position_ids.repeat_interleave(self.config.n, dim=0) batch_size = batch_size * self.config.n seq = torch.cat([idx, response], dim=-1) 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(batch_size, 1) # TODO(sgm): fix position_ids on right_pad # prompt: left pad + response: right pad # attention_mask: [0,0,0,0,1,1,1,1, | 1,1,1,0,0,0,0,0] # position_ids: [0,0,0,0,0,1,2,3, | 4,5,6,7,8,9,10,11] response_position_ids = position_ids[:, -1:] + delta_position_id position_ids = torch.cat([position_ids, response_position_ids], dim=-1) response_attention_mask = get_eos_mask(response_id=response, eos_token=eos_token_id, dtype=attention_mask.dtype) attention_mask = torch.cat((attention_mask, response_attention_mask), dim=-1) # all the tp ranks should contain the same data here. data in all ranks are valid batch = TensorDict( { 'prompts': idx, 'responses': response, 'input_ids': seq, # here input_ids become the whole sentences # 'old_log_probs': log_probs, # we will recompute old log prob with actor 'attention_mask': attention_mask, 'position_ids': position_ids }, batch_size=batch_size) # free vllm cache engine if self.config.free_cache_engine: self.inference_engine.free_cache_engine() return DataProto(batch=batch)","summary_state":"pending","summary":null,"crux":null,"owner":"vLLMRollout"}],"rawEdges":[{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"typing","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"contextlib","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"omegaconf","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"torch","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"tensordict","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"torch","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"verl","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"verl.workers.rollout.base","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"verl.third_party.vllm","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"verl.third_party.vllm","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"vllm","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"contains","targetId":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#_pre_process_inputs","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#_pre_process_inputs","relation":"calls","name":"nonzero","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#_pre_process_inputs","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#_pre_process_inputs","relation":"calls","name":"nonzero","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#_pre_process_inputs","relation":"calls","name":"tolist","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#_pre_process_inputs","relation":"calls","name":"long","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"contains","targetId":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout","relation":"extends","name":"BaseRollout","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout","relation":"contains","targetId":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"dict"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","relation":"imports","specifier":"os","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"dict"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"initialize_parallel_state","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"LLM","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"LLM"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"dict","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"keys","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"DictConfig"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"SamplingParams","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"str","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"DictConfig"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.__init__","relation":"calls","name":"SamplingParams","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout","relation":"contains","targetId":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.update_sampling_params","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.update_sampling_params","relation":"calls","name":"items","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.update_sampling_params","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.update_sampling_params","relation":"calls","name":"getattr","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.update_sampling_params","relation":"calls","name":"setattr","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.update_sampling_params","relation":"calls","name":"items","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.update_sampling_params","relation":"calls","name":"setattr","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout","relation":"calls","name":"no_grad","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout","relation":"contains","targetId":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"init_cache_engine","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"LLM"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"_pre_process_inputs","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"hasattr","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"type","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"update_sampling_params","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"vLLMRollout"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"generate","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"LLM"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"long","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"pad_sequence_to_length"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"float","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"pad_sequence_to_length"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"long","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"pad_sequence_to_length","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"pad_sequence_to_length","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"repeat_interleave","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"repeat_interleave","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"repeat_interleave","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"pad_sequence_to_length"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"arange","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"repeat","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"unsqueeze","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"get_eos_mask","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"TensorDict","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"free_cache_engine","viaMember":true,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py","recvType":"LLM"},{"source":"verl/workers/rollout/vllm_rollout/vllm_rollout.py#vLLMRollout.generate_sequences","relation":"calls","name":"DataProto","viaMember":false,"file":"verl/workers/rollout/vllm_rollout/vllm_rollout.py"}]},"verl/workers/sharding_manager/__init__.py":{"size":1169,"mtimeMs":1789275830329.2031,"hash":"32a39ab63010e701598f3c892357ede984a36a8e9013c78478b3f4a33e36ec21","nodes":[{"id":"verl/workers/sharding_manager/__init__.py","name":"__init__.py","kind":"file","path":"verl/workers/sharding_manager/__init__.py","span":"L1-L34","signature":null,"exported":true,"origin":"ast","body_hash":"32a39ab63010e701598f3c892357ede984a36a8e9013c78478b3f4a33e36ec21","chars":1169,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 verl.utils.import_utils import is_vllm_available, is_megatron_core_available from .base import BaseShardingManager from .fsdp_ulysses import FSDPUlyssesShardingManager AllGatherPPModel = None if is_megatron_core_available() and is_vllm_available(): from .megatron_vllm import AllGatherPPModel, MegatronVLLMShardingManager elif AllGatherPPModel is not None: pass else: AllGatherPPModel = None MegatronVLLMShardingManager = None if is_vllm_available(): from .fsdp_vllm import FSDPVLLMShardingManager else: FSDPVLLMShardingManager = None"}],"rawEdges":[{"source":"verl/workers/sharding_manager/__init__.py","relation":"imports","specifier":"verl.utils.import_utils","file":"verl/workers/sharding_manager/__init__.py"},{"source":"verl/workers/sharding_manager/__init__.py","relation":"imports","specifier":".base","file":"verl/workers/sharding_manager/__init__.py"},{"source":"verl/workers/sharding_manager/__init__.py","relation":"imports","specifier":".fsdp_ulysses","file":"verl/workers/sharding_manager/__init__.py"},{"source":"verl/workers/sharding_manager/__init__.py","relation":"calls","name":"is_megatron_core_available","viaMember":false,"file":"verl/workers/sharding_manager/__init__.py"},{"source":"verl/workers/sharding_manager/__init__.py","relation":"calls","name":"is_vllm_available","viaMember":false,"file":"verl/workers/sharding_manager/__init__.py"},{"source":"verl/workers/sharding_manager/__init__.py","relation":"imports","specifier":".megatron_vllm","file":"verl/workers/sharding_manager/__init__.py"},{"source":"verl/workers/sharding_manager/__init__.py","relation":"calls","name":"is_vllm_available","viaMember":false,"file":"verl/workers/sharding_manager/__init__.py"},{"source":"verl/workers/sharding_manager/__init__.py","relation":"imports","specifier":".fsdp_vllm","file":"verl/workers/sharding_manager/__init__.py"}]},"verl/workers/sharding_manager/base.py":{"size":982,"mtimeMs":1789275830330.2031,"hash":"d249884ee73ed51563443388163d57caaff198af11f60da44c70b85568c5127f","nodes":[{"id":"verl/workers/sharding_manager/base.py","name":"base.py","kind":"file","path":"verl/workers/sharding_manager/base.py","span":"L1-L34","signature":null,"exported":true,"origin":"ast","body_hash":"d249884ee73ed51563443388163d57caaff198af11f60da44c70b85568c5127f","chars":982,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Sharding manager to implement HybridEngine \"\"\" from verl import DataProto"},{"id":"verl/workers/sharding_manager/base.py#BaseShardingManager","name":"BaseShardingManager","kind":"class","path":"verl/workers/sharding_manager/base.py","span":"L21-L33","signature":"class BaseShardingManager","exported":true,"origin":"ast","body_hash":"e7315ea1f1753eda9f50225f1d41ab20efd00f80395e3b802f9433ed4923e1b0","body_text":"class BaseShardingManager: 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/sharding_manager/base.py#BaseShardingManager.__enter__","name":"__enter__","kind":"method","path":"verl/workers/sharding_manager/base.py","span":"L23-L24","signature":"def __enter__(self)","exported":false,"origin":"ast","body_hash":"e4a4082f5989cb33babba4c682814bb1e2f75e6381c629d401aa763c1aaf92a9","body_text":"def __enter__(self): pass","summary_state":"pending","summary":null,"crux":null,"owner":"BaseShardingManager"},{"id":"verl/workers/sharding_manager/base.py#BaseShardingManager.__exit__","name":"__exit__","kind":"method","path":"verl/workers/sharding_manager/base.py","span":"L26-L27","signature":"def __exit__(self, exc_type, exc_value, traceback)","exported":false,"origin":"ast","body_hash":"2fa00ae99dc27ec842b8ca24e542d9c989846583f491924b100fac2c031f6a71","body_text":"def __exit__(self, exc_type, exc_value, traceback): pass","summary_state":"pending","summary":null,"crux":null,"owner":"BaseShardingManager"},{"id":"verl/workers/sharding_manager/base.py#BaseShardingManager.preprocess_data","name":"preprocess_data","kind":"method","path":"verl/workers/sharding_manager/base.py","span":"L29-L30","signature":"def preprocess_data(self, data: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"76f6e1e029e6ab701c71671c546fad3d3978b229d2267706c8de653627c940f3","body_text":"def preprocess_data(self, data: DataProto) -> DataProto: return data","summary_state":"pending","summary":null,"crux":null,"owner":"BaseShardingManager"},{"id":"verl/workers/sharding_manager/base.py#BaseShardingManager.postprocess_data","name":"postprocess_data","kind":"method","path":"verl/workers/sharding_manager/base.py","span":"L32-L33","signature":"def postprocess_data(self, data: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"341da106b7585a37527445fff4f2e2c45c40dfc91d99fe0b372350be82515388","body_text":"def postprocess_data(self, data: DataProto) -> DataProto: return data","summary_state":"pending","summary":null,"crux":null,"owner":"BaseShardingManager"}],"rawEdges":[{"source":"verl/workers/sharding_manager/base.py","relation":"imports","specifier":"verl","file":"verl/workers/sharding_manager/base.py"},{"source":"verl/workers/sharding_manager/base.py","relation":"contains","targetId":"verl/workers/sharding_manager/base.py#BaseShardingManager","file":"verl/workers/sharding_manager/base.py"},{"source":"verl/workers/sharding_manager/base.py#BaseShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/base.py#BaseShardingManager.__enter__","file":"verl/workers/sharding_manager/base.py"},{"source":"verl/workers/sharding_manager/base.py#BaseShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/base.py#BaseShardingManager.__exit__","file":"verl/workers/sharding_manager/base.py"},{"source":"verl/workers/sharding_manager/base.py#BaseShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/base.py#BaseShardingManager.preprocess_data","file":"verl/workers/sharding_manager/base.py"},{"source":"verl/workers/sharding_manager/base.py#BaseShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/base.py#BaseShardingManager.postprocess_data","file":"verl/workers/sharding_manager/base.py"}]},"verl/workers/sharding_manager/fsdp_ulysses.py":{"size":3542,"mtimeMs":1789275830330.2031,"hash":"ec64073e9239d67d0dc7fa360013e779534e1bab8a6eb6991c6b9305aa573867","nodes":[{"id":"verl/workers/sharding_manager/fsdp_ulysses.py","name":"fsdp_ulysses.py","kind":"file","path":"verl/workers/sharding_manager/fsdp_ulysses.py","span":"L1-L88","signature":null,"exported":true,"origin":"ast","body_hash":"ec64073e9239d67d0dc7fa360013e779534e1bab8a6eb6991c6b9305aa573867","chars":3542,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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. \"\"\" Contains a resharding manager that binds weights from FSDP zero3 to XPerfGPT \"\"\" from typing import Optional from .base import BaseShardingManager import random from torch.distributed.device_mesh import DeviceMesh from verl.utils.torch_functional import allgather_dict_tensors from verl.utils.ulysses import set_ulysses_sequence_parallel_group, get_ulysses_sequence_parallel_group import numpy as np import torch import torch.distributed from verl import DataProto"},{"id":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager","name":"FSDPUlyssesShardingManager","kind":"class","path":"verl/workers/sharding_manager/fsdp_ulysses.py","span":"L33-L88","signature":"class FSDPUlyssesShardingManager(BaseShardingManager)","exported":true,"origin":"ast","body_hash":"2a26e4547c0ce5a7cbc425cbc720ca1db625e49ff038ebee13079e3d27bc548c","body_text":"class FSDPUlyssesShardingManager(BaseShardingManager): \"\"\" Sharding manager to support data resharding when using FSDP + Ulysses \"\"\" 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: # We have a global SP group # so we have to change to use model-specific sp group self.prev_sp_group = get_ulysses_sequence_parallel_group() set_ulysses_sequence_parallel_group(self.device_mesh['sp'].get_group()) # TODO: check how to set seed for each model def __exit__(self, exc_type, exc_value, traceback): # restore random states if self.device_mesh is not None: # revert to previous sp group set_ulysses_sequence_parallel_group(self.prev_sp_group) # TODO: check how to set seed for each model 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: sp_size = self.device_mesh['sp'].size() group = self.device_mesh['sp'].get_group() prev_device = data.batch.device data.batch = data.batch.cuda(device=torch.cuda.current_device()) data.batch = allgather_dict_tensors(data.batch.contiguous(), size=sp_size, group=group, dim=0) data.batch = data.batch.to(prev_device) # all gather non_tensor_batch all_non_tensor_batch = [None for _ in range(sp_size)] torch.distributed.all_gather_object(all_non_tensor_batch, data.non_tensor_batch, group=group) data.non_tensor_batch = { k: np.concatenate([d[k] for d in all_non_tensor_batch]) for k in data.non_tensor_batch } 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","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__init__","name":"__init__","kind":"method","path":"verl/workers/sharding_manager/fsdp_ulysses.py","span":"L38-L41","signature":"def __init__(self, device_mesh: DeviceMesh)","exported":false,"origin":"ast","body_hash":"af13162113c019abbd1152d19d1d90533a56c35fb68204028436e2c292eb17d2","body_text":"def __init__(self, device_mesh: DeviceMesh): super().__init__() self.device_mesh = device_mesh self.seed_offset = 12345","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPUlyssesShardingManager"},{"id":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__enter__","name":"__enter__","kind":"method","path":"verl/workers/sharding_manager/fsdp_ulysses.py","span":"L43-L49","signature":"def __enter__(self)","exported":false,"origin":"ast","body_hash":"71b7d1a755f3dc7fbfc57089464ef7b614b0935d134e8285ec11d683a6888b3f","body_text":"def __enter__(self): if self.device_mesh is not None: # We have a global SP group # so we have to change to use model-specific sp group self.prev_sp_group = get_ulysses_sequence_parallel_group() set_ulysses_sequence_parallel_group(self.device_mesh['sp'].get_group()) # TODO: check how to set seed for each model","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPUlyssesShardingManager"},{"id":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__exit__","name":"__exit__","kind":"method","path":"verl/workers/sharding_manager/fsdp_ulysses.py","span":"L51-L56","signature":"def __exit__(self, exc_type, exc_value, traceback): # restore random states","exported":false,"origin":"ast","body_hash":"643533cd250efc05e121171cd2c3016c60375f59024d3173672546900b098245","body_text":"def __exit__(self, exc_type, exc_value, traceback): # restore random states if self.device_mesh is not None: # revert to previous sp group set_ulysses_sequence_parallel_group(self.prev_sp_group) # TODO: check how to set seed for each model","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPUlyssesShardingManager"},{"id":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","name":"preprocess_data","kind":"method","path":"verl/workers/sharding_manager/fsdp_ulysses.py","span":"L58-L78","signature":"def preprocess_data(self, data: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"f3a65fca280f64b9ff3640063cb44eeedc9e6896d34a8397b814b67e24cac7e3","body_text":"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: sp_size = self.device_mesh['sp'].size() group = self.device_mesh['sp'].get_group() prev_device = data.batch.device data.batch = data.batch.cuda(device=torch.cuda.current_device()) data.batch = allgather_dict_tensors(data.batch.contiguous(), size=sp_size, group=group, dim=0) data.batch = data.batch.to(prev_device) # all gather non_tensor_batch all_non_tensor_batch = [None for _ in range(sp_size)] torch.distributed.all_gather_object(all_non_tensor_batch, data.non_tensor_batch, group=group) data.non_tensor_batch = { k: np.concatenate([d[k] for d in all_non_tensor_batch]) for k in data.non_tensor_batch } return data","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPUlyssesShardingManager"},{"id":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.postprocess_data","name":"postprocess_data","kind":"method","path":"verl/workers/sharding_manager/fsdp_ulysses.py","span":"L80-L88","signature":"def postprocess_data(self, data: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"7a746d97a2d6da051be0e5dc30a7e4b3117f7710c2e1c3ceab2c3ec7d4db7157","body_text":"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","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPUlyssesShardingManager"}],"rawEdges":[{"source":"verl/workers/sharding_manager/fsdp_ulysses.py","relation":"imports","specifier":"typing","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py","relation":"imports","specifier":".base","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py","relation":"imports","specifier":"random","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py","relation":"imports","specifier":"torch.distributed.device_mesh","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py","relation":"imports","specifier":"verl.utils.ulysses","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py","relation":"imports","specifier":"torch","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py","relation":"imports","specifier":"verl","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager","relation":"extends","name":"BaseShardingManager","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__init__","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__init__","relation":"calls","name":"__init__","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__init__","relation":"calls","name":"super","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__enter__","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__enter__","relation":"calls","name":"get_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__enter__","relation":"calls","name":"set_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__enter__","relation":"calls","name":"get_group","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__exit__","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.__exit__","relation":"calls","name":"set_ulysses_sequence_parallel_group","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","relation":"calls","name":"get_group","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","relation":"calls","name":"cuda","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","relation":"calls","name":"allgather_dict_tensors","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","relation":"calls","name":"all_gather_object","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.preprocess_data","relation":"calls","name":"concatenate","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.postprocess_data","file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.postprocess_data","relation":"calls","name":"size","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.postprocess_data","relation":"calls","name":"get_local_rank","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py"},{"source":"verl/workers/sharding_manager/fsdp_ulysses.py#FSDPUlyssesShardingManager.postprocess_data","relation":"calls","name":"chunk","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_ulysses.py","recvType":"DataProto"}]},"verl/workers/sharding_manager/fsdp_vllm.py":{"size":6186,"mtimeMs":1789275830330.2031,"hash":"c186b6c455d1e97c1b09c6fd8c80ea915d90f24e318b6fc377794e628583abd8","nodes":[{"id":"verl/workers/sharding_manager/fsdp_vllm.py","name":"fsdp_vllm.py","kind":"file","path":"verl/workers/sharding_manager/fsdp_vllm.py","span":"L1-L134","signature":null,"exported":true,"origin":"ast","body_hash":"c186b6c455d1e97c1b09c6fd8c80ea915d90f24e318b6fc377794e628583abd8","chars":6186,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 logging import torch from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP from torch.distributed.fsdp.api import ShardingStrategy, ShardedStateDictConfig, StateDictType, FullStateDictConfig from torch.distributed.device_mesh import DeviceMesh from verl.third_party.vllm import LLM from verl.third_party.vllm import parallel_state as vllm_ps from verl import DataProto from verl.utils.torch_functional import (broadcast_dict_tensor, allgather_dict_tensors) from verl.utils.debug import log_gpu_memory_usage from .base import BaseShardingManager logger = logging.getLogger(__file__) logger.setLevel(os.getenv('VERL_PPO_LOGGING_LEVEL', 'WARN'))"},{"id":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager","name":"FSDPVLLMShardingManager","kind":"class","path":"verl/workers/sharding_manager/fsdp_vllm.py","span":"L34-L133","signature":"class FSDPVLLMShardingManager(BaseShardingManager)","exported":true,"origin":"ast","body_hash":"ce00f12de93462843e586bd67d4466e297cdcded7fd5a6c4a9e597a35cab5756","body_text":"class FSDPVLLMShardingManager(BaseShardingManager): def __init__(self, module: FSDP, inference_engine: LLM, model_config, full_params: bool = False, device_mesh: DeviceMesh = None): self.module = module self.inference_engine = inference_engine self.model_config = model_config self.device_mesh = device_mesh # Full params self.full_params = full_params if full_params: FSDP.set_state_dict_type(self.module, state_dict_type=StateDictType.FULL_STATE_DICT, state_dict_config=FullStateDictConfig()) else: FSDP.set_state_dict_type(self.module, state_dict_type=StateDictType.SHARDED_STATE_DICT, state_dict_config=ShardedStateDictConfig()) # Note that torch_random_states may be different on each dp rank self.torch_random_states = torch.cuda.get_rng_state() # get a random rng states if self.device_mesh is not None: gen_dp_rank = self.device_mesh['dp'].get_local_rank() torch.cuda.manual_seed(gen_dp_rank + 1000) # make sure all tp ranks have the same random states self.gen_random_states = torch.cuda.get_rng_state() torch.cuda.set_rng_state(self.torch_random_states) else: self.gen_random_states = None def __enter__(self): log_gpu_memory_usage('Before state_dict() in sharding manager memory', logger=logger) params = self.module.state_dict() log_gpu_memory_usage('After state_dict() in sharding manager memory', logger=logger) # Copy, not share memory load_format = 'hf' if self.full_params else 'dtensor' self.inference_engine.sync_model_weights(params, load_format=load_format) log_gpu_memory_usage('After sync model weights in sharding manager', logger=logger) del params torch.cuda.empty_cache() log_gpu_memory_usage('After del state_dict and empty_cache in sharding manager', logger=logger) # TODO: offload FSDP model weights # self.module.cpu() # torch.cuda.empty_cache() # if torch.distributed.get_rank() == 0: # print(f'after model to cpu in sharding manager memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB') # important: need to manually set the random states of each tp to be identical. if self.device_mesh is not None: self.torch_random_states = torch.cuda.get_rng_state() torch.cuda.set_rng_state(self.gen_random_states) def __exit__(self, exc_type, exc_value, traceback): log_gpu_memory_usage('Before vllm offload in sharding manager', logger=logger) self.inference_engine.offload_model_weights() log_gpu_memory_usage('After vllm offload in sharding manager', logger=logger) # self.module.to('cuda') # if torch.distributed.get_rank() == 0: # print(f'after actor module to cuda in sharding manager memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB') self.module.train() # add empty cache after each compute torch.cuda.empty_cache() # restore random states if self.device_mesh is not None: self.gen_random_states = torch.cuda.get_rng_state() torch.cuda.set_rng_state(self.torch_random_states) def preprocess_data(self, data: DataProto) -> DataProto: # TODO: Current impl doesn't consider FSDP with torch micro-dp data.batch = allgather_dict_tensors(data.batch.contiguous(), size=vllm_ps.get_tensor_model_parallel_world_size(), group=vllm_ps.get_tensor_model_parallel_group(), dim=0) return data def postprocess_data(self, data: DataProto) -> DataProto: # TODO: Current impl doesn't consider FSDP with torch micro-dp broadcast_dict_tensor(data.batch, src=vllm_ps.get_tensor_model_parallel_src_rank(), group=vllm_ps.get_tensor_model_parallel_group()) dp_rank = torch.distributed.get_rank() dp_size = torch.distributed.get_world_size() # not consider torch micro-dp tp_size = vllm_ps.get_tensor_model_parallel_world_size() if tp_size > 1: # TODO: shall we build a micro_dp group for vllm when integrating with vLLM? local_prompts = data.chunk(chunks=tp_size) data = local_prompts[dp_rank % tp_size] return data","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","name":"__init__","kind":"method","path":"verl/workers/sharding_manager/fsdp_vllm.py","span":"L36-L67","signature":"def __init__(self, module: FSDP, inference_engine: LLM, model_config, full_params: bool = False, device_mesh: DeviceMesh = None)","exported":false,"origin":"ast","body_hash":"32577f4433ed2a3db6a9d112b8e464d2fca584df68e3935fc36c5363a3f02440","body_text":"def __init__(self, module: FSDP, inference_engine: LLM, model_config, full_params: bool = False, device_mesh: DeviceMesh = None): self.module = module self.inference_engine = inference_engine self.model_config = model_config self.device_mesh = device_mesh # Full params self.full_params = full_params if full_params: FSDP.set_state_dict_type(self.module, state_dict_type=StateDictType.FULL_STATE_DICT, state_dict_config=FullStateDictConfig()) else: FSDP.set_state_dict_type(self.module, state_dict_type=StateDictType.SHARDED_STATE_DICT, state_dict_config=ShardedStateDictConfig()) # Note that torch_random_states may be different on each dp rank self.torch_random_states = torch.cuda.get_rng_state() # get a random rng states if self.device_mesh is not None: gen_dp_rank = self.device_mesh['dp'].get_local_rank() torch.cuda.manual_seed(gen_dp_rank + 1000) # make sure all tp ranks have the same random states self.gen_random_states = torch.cuda.get_rng_state() torch.cuda.set_rng_state(self.torch_random_states) else: self.gen_random_states = None","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPVLLMShardingManager"},{"id":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","name":"__enter__","kind":"method","path":"verl/workers/sharding_manager/fsdp_vllm.py","span":"L69-L91","signature":"def __enter__(self)","exported":false,"origin":"ast","body_hash":"15feffbcdc55618bc089fbcb4de827d208cb334c76653634630847b488f49d35","body_text":"def __enter__(self): log_gpu_memory_usage('Before state_dict() in sharding manager memory', logger=logger) params = self.module.state_dict() log_gpu_memory_usage('After state_dict() in sharding manager memory', logger=logger) # Copy, not share memory load_format = 'hf' if self.full_params else 'dtensor' self.inference_engine.sync_model_weights(params, load_format=load_format) log_gpu_memory_usage('After sync model weights in sharding manager', logger=logger) del params torch.cuda.empty_cache() log_gpu_memory_usage('After del state_dict and empty_cache in sharding manager', logger=logger) # TODO: offload FSDP model weights # self.module.cpu() # torch.cuda.empty_cache() # if torch.distributed.get_rank() == 0: # print(f'after model to cpu in sharding manager memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB') # important: need to manually set the random states of each tp to be identical. if self.device_mesh is not None: self.torch_random_states = torch.cuda.get_rng_state() torch.cuda.set_rng_state(self.gen_random_states)","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPVLLMShardingManager"},{"id":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__exit__","name":"__exit__","kind":"method","path":"verl/workers/sharding_manager/fsdp_vllm.py","span":"L93-L110","signature":"def __exit__(self, exc_type, exc_value, traceback)","exported":false,"origin":"ast","body_hash":"ee6a8d29775e15efc158cc644ed5f9de0a21167f65b5dab575741bab4aeaeff3","body_text":"def __exit__(self, exc_type, exc_value, traceback): log_gpu_memory_usage('Before vllm offload in sharding manager', logger=logger) self.inference_engine.offload_model_weights() log_gpu_memory_usage('After vllm offload in sharding manager', logger=logger) # self.module.to('cuda') # if torch.distributed.get_rank() == 0: # print(f'after actor module to cuda in sharding manager memory allocated: {torch.cuda.memory_allocated() / 1e9}GB, reserved: {torch.cuda.memory_reserved() / 1e9}GB') self.module.train() # add empty cache after each compute torch.cuda.empty_cache() # restore random states if self.device_mesh is not None: self.gen_random_states = torch.cuda.get_rng_state() torch.cuda.set_rng_state(self.torch_random_states)","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPVLLMShardingManager"},{"id":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.preprocess_data","name":"preprocess_data","kind":"method","path":"verl/workers/sharding_manager/fsdp_vllm.py","span":"L112-L119","signature":"def preprocess_data(self, data: DataProto) -> DataProto: # TODO: Current impl doesn't consider FSDP with torch micro-dp","exported":true,"origin":"ast","body_hash":"7baddca806140ad14a3bc545ee5f2f05490a98aa4bd0549f08f6cd5d3a1e1238","body_text":"def preprocess_data(self, data: DataProto) -> DataProto: # TODO: Current impl doesn't consider FSDP with torch micro-dp data.batch = allgather_dict_tensors(data.batch.contiguous(), size=vllm_ps.get_tensor_model_parallel_world_size(), group=vllm_ps.get_tensor_model_parallel_group(), dim=0) return data","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPVLLMShardingManager"},{"id":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.postprocess_data","name":"postprocess_data","kind":"method","path":"verl/workers/sharding_manager/fsdp_vllm.py","span":"L121-L133","signature":"def postprocess_data(self, data: DataProto) -> DataProto: # TODO: Current impl doesn't consider FSDP with torch micro-dp","exported":true,"origin":"ast","body_hash":"12e949aae2e6d3782fb9fc110380931d8a753159f0278cb84c47df96c731460b","body_text":"def postprocess_data(self, data: DataProto) -> DataProto: # TODO: Current impl doesn't consider FSDP with torch micro-dp broadcast_dict_tensor(data.batch, src=vllm_ps.get_tensor_model_parallel_src_rank(), group=vllm_ps.get_tensor_model_parallel_group()) dp_rank = torch.distributed.get_rank() dp_size = torch.distributed.get_world_size() # not consider torch micro-dp tp_size = vllm_ps.get_tensor_model_parallel_world_size() if tp_size > 1: # TODO: shall we build a micro_dp group for vllm when integrating with vLLM? local_prompts = data.chunk(chunks=tp_size) data = local_prompts[dp_rank % tp_size] return data","summary_state":"pending","summary":null,"crux":null,"owner":"FSDPVLLMShardingManager"}],"rawEdges":[{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"os","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"logging","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"torch","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"torch.distributed.fsdp.fully_sharded_data_parallel","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"torch.distributed.fsdp.api","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"torch.distributed.device_mesh","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"verl.third_party.vllm","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"verl.third_party.vllm","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"verl","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":"verl.utils.debug","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"imports","specifier":".base","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"calls","name":"getLogger","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"calls","name":"setLevel","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"calls","name":"getenv","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager","relation":"extends","name":"BaseShardingManager","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","relation":"calls","name":"set_state_dict_type","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","relation":"calls","name":"FullStateDictConfig","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","relation":"calls","name":"set_state_dict_type","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","relation":"calls","name":"ShardedStateDictConfig","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","relation":"calls","name":"get_rng_state","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","relation":"calls","name":"get_local_rank","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","relation":"calls","name":"manual_seed","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","relation":"calls","name":"get_rng_state","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__init__","relation":"calls","name":"set_rng_state","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","relation":"calls","name":"state_dict","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","relation":"calls","name":"get_rng_state","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__enter__","relation":"calls","name":"set_rng_state","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__exit__","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__exit__","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__exit__","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__exit__","relation":"calls","name":"log_gpu_memory_usage","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__exit__","relation":"calls","name":"train","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__exit__","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__exit__","relation":"calls","name":"get_rng_state","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.__exit__","relation":"calls","name":"set_rng_state","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.preprocess_data","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.preprocess_data","relation":"calls","name":"allgather_dict_tensors","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.preprocess_data","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.preprocess_data","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.preprocess_data","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.postprocess_data","file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.postprocess_data","relation":"calls","name":"broadcast_dict_tensor","viaMember":false,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.postprocess_data","relation":"calls","name":"get_tensor_model_parallel_src_rank","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.postprocess_data","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.postprocess_data","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.postprocess_data","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.postprocess_data","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py"},{"source":"verl/workers/sharding_manager/fsdp_vllm.py#FSDPVLLMShardingManager.postprocess_data","relation":"calls","name":"chunk","viaMember":true,"file":"verl/workers/sharding_manager/fsdp_vllm.py","recvType":"DataProto"}]},"verl/workers/sharding_manager/megatron_vllm.py":{"size":18595,"mtimeMs":1789275830330.2031,"hash":"a96ff19985b83712f572dfb8bfe937b3d2c43d031f879c99bfc2fafc5d10eee8","nodes":[{"id":"verl/workers/sharding_manager/megatron_vllm.py","name":"megatron_vllm.py","kind":"file","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L1-L429","signature":null,"exported":true,"origin":"ast","body_hash":"a96ff19985b83712f572dfb8bfe937b3d2c43d031f879c99bfc2fafc5d10eee8","chars":18595,"summary_state":"pending","summary":null,"crux":null,"body_text":"# 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 file contains a Megatron style Hybrid Engine that shares the weights of the actor with the inference engine. \"\"\" import torch import torch.distributed as dist from torch import nn from megatron.core import parallel_state as mpu from megatron.core import DistributedDataParallel as LocalDDP from megatron.core.transformer.module import Float16Module from torch.nn.parallel.distributed import DistributedDataParallel as torchDDP from verl.utils.megatron_utils import get_model, unwrap_model from verl.utils.memory_buffer import ( build_memory_buffer, build_memory_reference_from_module, get_weight_buffer_meta_from_module, ) \"\"\" Megatron Hybrid Engine: - 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. \"\"\" from .base import BaseShardingManager import torch from torch import nn import torch.distributed from torch.distributed import new_group from verl import DataProto from verl.utils.torch_functional import (broadcast_dict_tensor, allgather_dict_tensors) import verl.utils.megatron.tensor_parallel as tp_utils from verl.third_party.vllm import parallel_state as vllm_ps from verl.third_party.vllm import LLM from verl.utils.model import normalize_pp_vpp_params # Micro Data parallel group. Micro data parallel group is additional dp group that origins from splitting training tp # into infer_tp and micro_tp. By default, we use order micro_dp - tp _MICRO_DATA_PARALLEL_GROUP = None \"\"\" Micro Data parallel group \"\"\""},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","name":"AllGatherPPModel","kind":"class","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L35-L208","signature":"class AllGatherPPModel","exported":true,"origin":"ast","body_hash":"1cbde619b8c982c8f122ba2dbd5caaf71affe4ca023ab5fcfbec886ddeed1a15","body_text":"class AllGatherPPModel: def __init__(self, model_provider) -> None: self._pp_group = mpu.get_pipeline_model_parallel_group() self._pp_rank = mpu.get_pipeline_model_parallel_rank() self._pp_size = mpu.get_pipeline_model_parallel_world_size() self._vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() self._model_chunk_size = self._vpp_size or 1 # each one holds a list of model_chunks in this pp stage self._pp_models = [None] * self.pp_size rank_list = list(range(self.pp_size)) # make current rank the last one to initialize rank_list[self.pp_rank], rank_list[-1] = rank_list[-1], rank_list[self.pp_rank] self._this_rank_models = None # store the parameter of each pp stage self.memory_buffers = [None] * self.pp_size for cur_pp_rank in rank_list: print( f'create pp model', f'torch allocated {torch.cuda.memory_allocated() / 1e9:.4f} GB, ' f'reserved {torch.cuda.memory_reserved() / 1e9:.4f} GB') # since the last initialized rank is the current pp rank, after init, the pp rank is still correct mpu.set_pipeline_model_parallel_rank(cur_pp_rank) if cur_pp_rank != self.pp_rank: models = get_model(model_provider, wrap_with_ddp=False) models = nn.ModuleList(models) assert len(models) == self._model_chunk_size, f\"{len(models)} != {self._model_chunk_size}\" self.pp_models[cur_pp_rank] = models else: # for regular model, we wrapped it with DDP models = get_model(model_provider) assert len(models) == self._model_chunk_size, f\"{len(models)} != {self._model_chunk_size}\" self._this_rank_models = nn.ModuleList(models) self.pp_models[cur_pp_rank] = nn.ModuleList(unwrap_model(models, (torchDDP, LocalDDP))) self._build_param_buffer(cur_pp_rank) self._build_param_references(cur_pp_rank, maintain_weight=cur_pp_rank == self.pp_rank) # TODO: after binding to the memory buffer, we can load the checkpoint here if cur_pp_rank != self.pp_rank: for model in self.pp_models[cur_pp_rank]: model.eval() self._offload_params_to_cpu(cur_pp_rank) def _build_param_buffer(self, pp_rank): \"\"\"Build the parameter buffer in each pp rank\"\"\" model = self.pp_models[pp_rank] weight_buffer_meta = get_weight_buffer_meta_from_module(model) self.memory_buffers[pp_rank] = build_memory_buffer(weight_buffer_meta) def _build_param_references(self, pp_rank, maintain_weight=False): model = self.pp_models[pp_rank] build_memory_reference_from_module(model, self.memory_buffers[pp_rank], maintain_weight=maintain_weight) def _load_params_to_cuda(self, pp_rank, to_empty=False): assert pp_rank != self.pp_rank, f\"unexpected to load current pp rank [{pp_rank}] back to cuda\" for buffer in self.memory_buffers[pp_rank].values(): if not to_empty: buffer.data = buffer.data.to(torch.cuda.current_device(), non_blocking=True) else: buffer.data = torch.empty_like(buffer.data, device='cuda') # rebuild reference after loading to CUDA self._build_param_references(pp_rank) def _offload_params_to_cpu(self, pp_rank, to_empty=False): assert pp_rank != self.pp_rank, f\"unexpected to offload current pp rank [{pp_rank}] to cpu\" for buffer in self.memory_buffers[pp_rank].values(): if not to_empty: # offload the whole memory buffer to CPU buffer.data = buffer.data.to('cpu', non_blocking=True) else: buffer.data = torch.empty_like(buffer.data, device='cpu') self._build_param_references(pp_rank) def load_params_to_cuda(self, to_empty=False): \"\"\"load all model params to cuda\"\"\" for cur_pp_rank in range(self.pp_size): if cur_pp_rank != self.pp_rank: self._load_params_to_cuda(cur_pp_rank, to_empty=to_empty) def allgather_params(self): \"\"\"allgather params of all pp ranks. Return a list of handles\"\"\" for cur_pp_rank in range(self.pp_size): global_src = dist.get_global_rank(group=self.pp_group, group_rank=cur_pp_rank) # NOTE(sgm): the async op may cause memory leakage of the memory_buffer/pp_models for memory_buffer in self.memory_buffers[cur_pp_rank].values(): dist.broadcast(tensor=memory_buffer.data, src=global_src, group=self.pp_group, async_op=False) def forward(self, *inputs, **kwargs): try: prev_output = None for cur_chunk_rank in range(self._model_chunk_size): if self._vpp_size: mpu.set_virtual_pipeline_model_parallel_rank(cur_chunk_rank) for cur_pp_rank in range(self.pp_size): mpu.set_pipeline_model_parallel_rank(cur_pp_rank) self.pp_models[cur_pp_rank][cur_chunk_rank].set_input_tensor(prev_output) ret = self.pp_models[cur_pp_rank][cur_chunk_rank](*inputs, **kwargs) self.pp_models[cur_pp_rank][cur_chunk_rank].set_input_tensor(None) prev_output = ret finally: if self._vpp_size: mpu.set_virtual_pipeline_model_parallel_rank(0) mpu.set_pipeline_model_parallel_rank(self.pp_rank) return ret def __call__(self, *inputs, **kwargs): return self.forward(*inputs, **kwargs) def eval(self): for model in self.pp_models[self.pp_rank]: model.eval() def train(self): for model in self.pp_models[self.pp_rank]: model.train() def offload_params_to_cpu(self, to_empty=False): \"\"\"offload params of models that are not of current pp rank to cpu\"\"\" for cur_pp_rank in range(self.pp_size): if cur_pp_rank != self.pp_rank:","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","name":"__init__","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L37-L80","signature":"def __init__(self, model_provider) -> None","exported":false,"origin":"ast","body_hash":"25592b5cd8fc407633b0cb0ea89e7ceb2a0e2cf33d6635d6dc95e3f47ae4049b","body_text":"def __init__(self, model_provider) -> None: self._pp_group = mpu.get_pipeline_model_parallel_group() self._pp_rank = mpu.get_pipeline_model_parallel_rank() self._pp_size = mpu.get_pipeline_model_parallel_world_size() self._vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() self._model_chunk_size = self._vpp_size or 1 # each one holds a list of model_chunks in this pp stage self._pp_models = [None] * self.pp_size rank_list = list(range(self.pp_size)) # make current rank the last one to initialize rank_list[self.pp_rank], rank_list[-1] = rank_list[-1], rank_list[self.pp_rank] self._this_rank_models = None # store the parameter of each pp stage self.memory_buffers = [None] * self.pp_size for cur_pp_rank in rank_list: print( f'create pp model', f'torch allocated {torch.cuda.memory_allocated() / 1e9:.4f} GB, ' f'reserved {torch.cuda.memory_reserved() / 1e9:.4f} GB') # since the last initialized rank is the current pp rank, after init, the pp rank is still correct mpu.set_pipeline_model_parallel_rank(cur_pp_rank) if cur_pp_rank != self.pp_rank: models = get_model(model_provider, wrap_with_ddp=False) models = nn.ModuleList(models) assert len(models) == self._model_chunk_size, f\"{len(models)} != {self._model_chunk_size}\" self.pp_models[cur_pp_rank] = models else: # for regular model, we wrapped it with DDP models = get_model(model_provider) assert len(models) == self._model_chunk_size, f\"{len(models)} != {self._model_chunk_size}\" self._this_rank_models = nn.ModuleList(models) self.pp_models[cur_pp_rank] = nn.ModuleList(unwrap_model(models, (torchDDP, LocalDDP))) self._build_param_buffer(cur_pp_rank) self._build_param_references(cur_pp_rank, maintain_weight=cur_pp_rank == self.pp_rank) # TODO: after binding to the memory buffer, we can load the checkpoint here if cur_pp_rank != self.pp_rank: for model in self.pp_models[cur_pp_rank]: model.eval() self._offload_params_to_cpu(cur_pp_rank)","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._build_param_buffer","name":"_build_param_buffer","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L82-L86","signature":"def _build_param_buffer(self, pp_rank)","exported":false,"origin":"ast","body_hash":"bff7f85ee99fc88e4d21caf48796dc0d78a83f679e4cbf74dc439a43a844f4f3","body_text":"def _build_param_buffer(self, pp_rank): \"\"\"Build the parameter buffer in each pp rank\"\"\" model = self.pp_models[pp_rank] weight_buffer_meta = get_weight_buffer_meta_from_module(model) self.memory_buffers[pp_rank] = build_memory_buffer(weight_buffer_meta)","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._build_param_references","name":"_build_param_references","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L88-L90","signature":"def _build_param_references(self, pp_rank, maintain_weight=False)","exported":false,"origin":"ast","body_hash":"fa9068b8d77ab80c172f6fe7e57e8d40c45e94d252debaa501e1fe5e8de4ed62","body_text":"def _build_param_references(self, pp_rank, maintain_weight=False): model = self.pp_models[pp_rank] build_memory_reference_from_module(model, self.memory_buffers[pp_rank], maintain_weight=maintain_weight)","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._load_params_to_cuda","name":"_load_params_to_cuda","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L92-L100","signature":"def _load_params_to_cuda(self, pp_rank, to_empty=False)","exported":false,"origin":"ast","body_hash":"b06aaf2c50370569248fe41ff3d86f49cfa0bb5d073208494a017350f369511c","body_text":"def _load_params_to_cuda(self, pp_rank, to_empty=False): assert pp_rank != self.pp_rank, f\"unexpected to load current pp rank [{pp_rank}] back to cuda\" for buffer in self.memory_buffers[pp_rank].values(): if not to_empty: buffer.data = buffer.data.to(torch.cuda.current_device(), non_blocking=True) else: buffer.data = torch.empty_like(buffer.data, device='cuda') # rebuild reference after loading to CUDA self._build_param_references(pp_rank)","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._offload_params_to_cpu","name":"_offload_params_to_cpu","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L102-L110","signature":"def _offload_params_to_cpu(self, pp_rank, to_empty=False)","exported":false,"origin":"ast","body_hash":"85df90a6ddc191947b23a9fcebb1bcbb205718f684b808233c913f3dbd21563a","body_text":"def _offload_params_to_cpu(self, pp_rank, to_empty=False): assert pp_rank != self.pp_rank, f\"unexpected to offload current pp rank [{pp_rank}] to cpu\" for buffer in self.memory_buffers[pp_rank].values(): if not to_empty: # offload the whole memory buffer to CPU buffer.data = buffer.data.to('cpu', non_blocking=True) else: buffer.data = torch.empty_like(buffer.data, device='cpu') self._build_param_references(pp_rank)","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.load_params_to_cuda","name":"load_params_to_cuda","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L112-L116","signature":"def load_params_to_cuda(self, to_empty=False)","exported":true,"origin":"ast","body_hash":"e90586b096909ffbb7ca506ec55229ab5085eedb72c9cc32838bf278eeca2b1e","body_text":"def load_params_to_cuda(self, to_empty=False): \"\"\"load all model params to cuda\"\"\" for cur_pp_rank in range(self.pp_size): if cur_pp_rank != self.pp_rank: self._load_params_to_cuda(cur_pp_rank, to_empty=to_empty)","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.allgather_params","name":"allgather_params","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L118-L125","signature":"def allgather_params(self)","exported":true,"origin":"ast","body_hash":"7c0923e5baee67a052d45fc85d34ea40a6dc5676cb159b8c92105d9c31c193cc","body_text":"def allgather_params(self): \"\"\"allgather params of all pp ranks. Return a list of handles\"\"\" for cur_pp_rank in range(self.pp_size): global_src = dist.get_global_rank(group=self.pp_group, group_rank=cur_pp_rank) # NOTE(sgm): the async op may cause memory leakage of the memory_buffer/pp_models for memory_buffer in self.memory_buffers[cur_pp_rank].values(): dist.broadcast(tensor=memory_buffer.data, src=global_src, group=self.pp_group, async_op=False)","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.forward","name":"forward","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L127-L144","signature":"def forward(self, *inputs, **kwargs)","exported":true,"origin":"ast","body_hash":"20c09979aac9ec48b6a682c264530e6c32b84f531ae3199a97ce14bb573ff662","body_text":"def forward(self, *inputs, **kwargs): try: prev_output = None for cur_chunk_rank in range(self._model_chunk_size): if self._vpp_size: mpu.set_virtual_pipeline_model_parallel_rank(cur_chunk_rank) for cur_pp_rank in range(self.pp_size): mpu.set_pipeline_model_parallel_rank(cur_pp_rank) self.pp_models[cur_pp_rank][cur_chunk_rank].set_input_tensor(prev_output) ret = self.pp_models[cur_pp_rank][cur_chunk_rank](*inputs, **kwargs) self.pp_models[cur_pp_rank][cur_chunk_rank].set_input_tensor(None) prev_output = ret finally: if self._vpp_size: mpu.set_virtual_pipeline_model_parallel_rank(0) mpu.set_pipeline_model_parallel_rank(self.pp_rank) return ret","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__call__","name":"__call__","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L146-L147","signature":"def __call__(self, *inputs, **kwargs)","exported":false,"origin":"ast","body_hash":"3f97bdb5c37a52010e89ae476d90fa0d3fc8e39bd02b5bda7066d1ca2cea7bc5","body_text":"def __call__(self, *inputs, **kwargs): return self.forward(*inputs, **kwargs)","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.eval","name":"eval","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L149-L151","signature":"def eval(self)","exported":true,"origin":"ast","body_hash":"c968584585449ead89bd7f1caf47f441fad996257b88f85b1dfee5b16fd0a177","body_text":"def eval(self): for model in self.pp_models[self.pp_rank]: model.eval()","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.train","name":"train","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L153-L155","signature":"def train(self)","exported":true,"origin":"ast","body_hash":"160506fb2be095a3baafec6071a17a7431f615d01f66404882ccf39258e24942","body_text":"def train(self): for model in self.pp_models[self.pp_rank]: model.train()","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.offload_params_to_cpu","name":"offload_params_to_cpu","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L157-L161","signature":"def offload_params_to_cpu(self, to_empty=False)","exported":true,"origin":"ast","body_hash":"3c07b15300c465fcddd4c93f0ccf6547644879d3add86170432a10400fccf13b","body_text":"def offload_params_to_cpu(self, to_empty=False): \"\"\"offload params of models that are not of current pp rank to cpu\"\"\" for cur_pp_rank in range(self.pp_size): if cur_pp_rank != self.pp_rank: self._offload_params_to_cpu(cur_pp_rank, to_empty=to_empty)","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.get_all_params","name":"get_all_params","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L163-L184","signature":"def get_all_params(self)","exported":true,"origin":"ast","body_hash":"2f627f1a1f1b7036c4dc9bc6cf3db5f9a5f72f0f927639b20c2f17a25d716b02","body_text":"def get_all_params(self): \"\"\"Get all the parameters of the models in all pp ranks Returns: params: List[List[Dict[str, Tensor]]]: a list of parameters in all pp, where each is a list of dict tensors of each model chunk \"\"\" params = [] for pp_rank in range(self.pp_size): params.append([]) for model_chunk_idx in range(len(self.pp_models[pp_rank])): params[pp_rank].append({}) pp_model = self.pp_models[pp_rank][model_chunk_idx] pp_model = unwrap_model(pp_model, ((torchDDP, LocalDDP, Float16Module))) # not use Float16Module for name, param in pp_model.named_parameters(): # NOTE(gh) workaround: should not get lora params for inference if 'lora' in name: continue params[pp_rank][model_chunk_idx][name] = param return params","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.update_this_rank_models","name":"update_this_rank_models","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L186-L188","signature":"def update_this_rank_models(self, new_models)","exported":true,"origin":"ast","body_hash":"5eb02bae13554c0550f079d618d7c2ce1cac145a2da7ef029ca03a8866814de2","body_text":"def update_this_rank_models(self, new_models): self._this_rank_models = new_models self._pp_models[self.pp_rank] = unwrap_model(new_models, (torchDDP, LocalDDP))","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.this_rank_models","name":"this_rank_models","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L191-L192","signature":"def this_rank_models(self)","exported":true,"origin":"ast","body_hash":"9ed9cc35aca98481f3b2f84fd0c5ac5d5e7688e79f9c8633b800acb85cfda317","body_text":"def this_rank_models(self): return self._this_rank_models","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.pp_size","name":"pp_size","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L195-L196","signature":"def pp_size(self)","exported":true,"origin":"ast","body_hash":"ccc8e638a09c29044a234d710fe6405dd416d5bc7ab62c03d11ac78127853921","body_text":"def pp_size(self): return self._pp_size","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.pp_rank","name":"pp_rank","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L199-L200","signature":"def pp_rank(self)","exported":true,"origin":"ast","body_hash":"0d35ee13e0caf0b14e7aac5cedfc3bf9423f890de09918d812eb96ea98613864","body_text":"def pp_rank(self): return self._pp_rank","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.pp_group","name":"pp_group","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L203-L204","signature":"def pp_group(self)","exported":true,"origin":"ast","body_hash":"68eb67bf4c2662ac8cab16396fdbb76463dd80ede92cfff62f61aa47eb6775af","body_text":"def pp_group(self): return self._pp_group","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.pp_models","name":"pp_models","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L207-L208","signature":"def pp_models(self)","exported":true,"origin":"ast","body_hash":"b1134e06c5e604fa2b48c5203653d56b7e6d3105e31af571e2615e916227d055","body_text":"def pp_models(self): return self._pp_models","summary_state":"pending","summary":null,"crux":null,"owner":"AllGatherPPModel"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager","name":"MegatronVLLMShardingManager","kind":"class","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L238-L410","signature":"class MegatronVLLMShardingManager(BaseShardingManager)","exported":true,"origin":"ast","body_hash":"cb7138e959092d8cf82cd80437debe4cf38800323d2cd58aa8340524bfe9dc4c","body_text":"class MegatronVLLMShardingManager(BaseShardingManager): def __init__(self, module: AllGatherPPModel, inference_engine: LLM, model_config, layer_name_mapping): self.module = module self.inference_engine = inference_engine self.model_config = model_config self.layer_name_mapping = layer_name_mapping # initialize micro_dp group for vllm inference global _MICRO_DATA_PARALLEL_GROUP world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() train_tensor_parallel_size = mpu.get_tensor_model_parallel_world_size() infer_tensor_parallel_size = vllm_ps.get_tensor_model_parallel_world_size() # TODO(sgm): this may not be true for FSDP -> vLLM assert infer_tensor_parallel_size <= train_tensor_parallel_size, \\ 'Not implemented for infer_tp > train_tp' assert train_tensor_parallel_size % infer_tensor_parallel_size == 0 micro_dp_size = train_tensor_parallel_size // infer_tensor_parallel_size num_micro_dp_groups = world_size // micro_dp_size assert _MICRO_DATA_PARALLEL_GROUP is None, (\"micro data parallel group is already initialized\") for i in range(num_micro_dp_groups): ranks = range(i * micro_dp_size, (i + 1) * micro_dp_size) group = new_group(ranks=ranks) if rank in ranks: _MICRO_DATA_PARALLEL_GROUP = group def default_tp_concat_fn(self, name, param, infer_params, model_config): \"\"\" name: name of the parameter param: training parameters infer_params (List[torch.Tensor]): a list of parameters all-gathered from micro_dp_group model_config: huggingface model_config TODO(zhangchi.usc1992): currently, the implementation is adhoc. We can move this function to the model definition so that it is model-agnostic. If the model doesn't implement this function, we can throw an error to force user disable TP HybridEngine. \"\"\" if self.layer_name_mapping.get(\"qkv_layer_name\") in name: # 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 model_config.num_attention_heads % model_config.num_key_value_heads == 0 num_q_per_kv = model_config.num_attention_heads // model_config.num_key_value_heads assert infer_params[0].shape[0] % (num_q_per_kv + 2) == 0 kv_size_per_tp = infer_params[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 infer_params: q, k, v = infer_param.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) infer_params = torch.cat((q, k, v), dim=0) elif self.layer_name_mapping.get(\"gate_proj_layer_name\") in name: # if the tensor is gate and proj gate_lst = [] up_lst = [] for infer_param in infer_params: 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) infer_params = torch.cat((gate, up), dim=0) else: # concat tensor infer_params = torch.cat(infer_params, dim=tp_utils.get_tensor_parallel_partition_dim(param)) return infer_params def _post_process_params(self, params): \"\"\" For each param, if it is a tp-splited param, we all-gather from micro_dp group. \"\"\" # here the params are in train tp format. we iterate params and all-gather # TODO(zhangchi.usc1992) We can consider copy non-tp weight to another infer buffer. # In this way, all the params in the original memory_buffers and can be offload. micro_dp_size = get_micro_data_parallel_world_size() micro_dp_group = get_micro_data_parallel_group() if micro_dp_size <= 1: return origin_params = {} for name in params.keys(): param = params[name] if tp_utils.is_tensor_parallel_param(param): # allocate a new tensor with proper size infer_params = [torch.empty_like(param) for _ in range(micro_dp_size)] torch.distributed.all_gather(infer_params, param, group=micro_dp_group) infer_params = self.default_tp_concat_fn(name, param, infer_params, self.model_config) # replace with original param params[name] = infer_params origin_params[name] = param return origin_params def __enter__(self): # create a new cuda space for parameters not in this pp rank self.module.load_params_to_cuda() # broadcast the parameters from pp rank to other ranks self.module.allgather_params() # obtain name to parameters in pp/vpp params = self.module.get_all_params() # bind the params to inference engine self.params = normalize_pp_vpp_params(params=params, num_hidden_layers=self.model_config.num_hidden_layers, layer_name='layers') self.origin_params = self._post_process_params(self.params) self.inference_engine.sync_model_weights(self.params, load_format='megatron') def __exit__(self, exc_type, exc_value, traceback): # offload parameters doesn't belong to this pp rank self.module.offload_params_to_cpu() # FIXME(sgm): the best practice is to delete the cuda tensor # rebind the model weights, can be any cpu tensor if get_micro_data_parallel_world_size() > 1: for name in self.params.keys(): self.params[name] = self.origin_params[name] # self.inference_engine.sync_m","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__init__","name":"__init__","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L240-L265","signature":"def __init__(self, module: AllGatherPPModel, inference_engine: LLM, model_config, layer_name_mapping)","exported":false,"origin":"ast","body_hash":"8fcbb42d498e0d3bc012e6443516e226c97abb0e164d56f4adbf9c3c25486014","body_text":"def __init__(self, module: AllGatherPPModel, inference_engine: LLM, model_config, layer_name_mapping): self.module = module self.inference_engine = inference_engine self.model_config = model_config self.layer_name_mapping = layer_name_mapping # initialize micro_dp group for vllm inference global _MICRO_DATA_PARALLEL_GROUP world_size = torch.distributed.get_world_size() rank = torch.distributed.get_rank() train_tensor_parallel_size = mpu.get_tensor_model_parallel_world_size() infer_tensor_parallel_size = vllm_ps.get_tensor_model_parallel_world_size() # TODO(sgm): this may not be true for FSDP -> vLLM assert infer_tensor_parallel_size <= train_tensor_parallel_size, \\ 'Not implemented for infer_tp > train_tp' assert train_tensor_parallel_size % infer_tensor_parallel_size == 0 micro_dp_size = train_tensor_parallel_size // infer_tensor_parallel_size num_micro_dp_groups = world_size // micro_dp_size assert _MICRO_DATA_PARALLEL_GROUP is None, (\"micro data parallel group is already initialized\") for i in range(num_micro_dp_groups): ranks = range(i * micro_dp_size, (i + 1) * micro_dp_size) group = new_group(ranks=ranks) if rank in ranks: _MICRO_DATA_PARALLEL_GROUP = group","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronVLLMShardingManager"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","name":"default_tp_concat_fn","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L267-L316","signature":"def default_tp_concat_fn(self, name, param, infer_params, model_config)","exported":true,"origin":"ast","body_hash":"0a0e902c54130d8a35ff923a039a915f1bb2e416e09dbe7dc1c54db881a6cd20","body_text":"def default_tp_concat_fn(self, name, param, infer_params, model_config): \"\"\" name: name of the parameter param: training parameters infer_params (List[torch.Tensor]): a list of parameters all-gathered from micro_dp_group model_config: huggingface model_config TODO(zhangchi.usc1992): currently, the implementation is adhoc. We can move this function to the model definition so that it is model-agnostic. If the model doesn't implement this function, we can throw an error to force user disable TP HybridEngine. \"\"\" if self.layer_name_mapping.get(\"qkv_layer_name\") in name: # 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 model_config.num_attention_heads % model_config.num_key_value_heads == 0 num_q_per_kv = model_config.num_attention_heads // model_config.num_key_value_heads assert infer_params[0].shape[0] % (num_q_per_kv + 2) == 0 kv_size_per_tp = infer_params[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 infer_params: q, k, v = infer_param.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) infer_params = torch.cat((q, k, v), dim=0) elif self.layer_name_mapping.get(\"gate_proj_layer_name\") in name: # if the tensor is gate and proj gate_lst = [] up_lst = [] for infer_param in infer_params: 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) infer_params = torch.cat((gate, up), dim=0) else: # concat tensor infer_params = torch.cat(infer_params, dim=tp_utils.get_tensor_parallel_partition_dim(param)) return infer_params","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronVLLMShardingManager"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager._post_process_params","name":"_post_process_params","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L318-L343","signature":"def _post_process_params(self, params)","exported":false,"origin":"ast","body_hash":"04a199679d75c48c5d30208a27796c38949b42d469b0c085f8ff8020d08e5393","body_text":"def _post_process_params(self, params): \"\"\" For each param, if it is a tp-splited param, we all-gather from micro_dp group. \"\"\" # here the params are in train tp format. we iterate params and all-gather # TODO(zhangchi.usc1992) We can consider copy non-tp weight to another infer buffer. # In this way, all the params in the original memory_buffers and can be offload. micro_dp_size = get_micro_data_parallel_world_size() micro_dp_group = get_micro_data_parallel_group() if micro_dp_size <= 1: return origin_params = {} for name in params.keys(): param = params[name] if tp_utils.is_tensor_parallel_param(param): # allocate a new tensor with proper size infer_params = [torch.empty_like(param) for _ in range(micro_dp_size)] torch.distributed.all_gather(infer_params, param, group=micro_dp_group) infer_params = self.default_tp_concat_fn(name, param, infer_params, self.model_config) # replace with original param params[name] = infer_params origin_params[name] = param return origin_params","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronVLLMShardingManager"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__enter__","name":"__enter__","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L345-L358","signature":"def __enter__(self): # create a new cuda space for parameters not in this pp rank","exported":false,"origin":"ast","body_hash":"6ed2d47d693a29548b5d0c804cd8a3ff84fe66f04b9cb235485a72c891e888c7","body_text":"def __enter__(self): # create a new cuda space for parameters not in this pp rank self.module.load_params_to_cuda() # broadcast the parameters from pp rank to other ranks self.module.allgather_params() # obtain name to parameters in pp/vpp params = self.module.get_all_params() # bind the params to inference engine self.params = normalize_pp_vpp_params(params=params, num_hidden_layers=self.model_config.num_hidden_layers, layer_name='layers') self.origin_params = self._post_process_params(self.params) self.inference_engine.sync_model_weights(self.params, load_format='megatron')","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronVLLMShardingManager"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__exit__","name":"__exit__","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L360-L376","signature":"def __exit__(self, exc_type, exc_value, traceback): # offload parameters doesn't belong to this pp rank","exported":false,"origin":"ast","body_hash":"5cbdaf4a363ca7999663b16e27d3d84b2284f510f4d37ded896d4f4b089bd7b0","body_text":"def __exit__(self, exc_type, exc_value, traceback): # offload parameters doesn't belong to this pp rank self.module.offload_params_to_cpu() # FIXME(sgm): the best practice is to delete the cuda tensor # rebind the model weights, can be any cpu tensor if get_micro_data_parallel_world_size() > 1: for name in self.params.keys(): self.params[name] = self.origin_params[name] # self.inference_engine.sync_model_weights(params) self.inference_engine.offload_model_weights() self.module.train() # add empty cache after each compute torch.cuda.empty_cache()","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronVLLMShardingManager"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.preprocess_data","name":"preprocess_data","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L378-L392","signature":"def preprocess_data(self, data: DataProto) -> DataProto: # prompts are identical for each training tp. We select for each inference tp","exported":true,"origin":"ast","body_hash":"3f8562d097a2b8ee46e749b442f2e76daef7e5a1e8743edbc18972faf7753f90","body_text":"def preprocess_data(self, data: DataProto) -> DataProto: # prompts are identical for each training tp. We select for each inference tp micro_dp_size = get_micro_data_parallel_world_size() micro_dp_rank = get_micro_data_parallel_rank() # broadcast from tp=0 to other tp ranks broadcast_dict_tensor(data.batch, src=mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group()) if micro_dp_size > 1: local_prompts = data.chunk(chunks=micro_dp_size) data = local_prompts[micro_dp_rank] return data","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronVLLMShardingManager"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","name":"postprocess_data","kind":"method","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L394-L410","signature":"def postprocess_data(self, data: DataProto) -> DataProto","exported":true,"origin":"ast","body_hash":"146827b2cff6d5ba0f86c5da325b2bc50d37e6acd3288569a71524afd35a1388","body_text":"def postprocess_data(self, data: DataProto) -> DataProto: meta_info = data.meta_info # all gather batch among micro-dp groups micro_dp_size = get_micro_data_parallel_world_size() if micro_dp_size > 1: data.batch = allgather_dict_tensors(data.batch.contiguous(), size=get_micro_data_parallel_world_size(), group=get_micro_data_parallel_group(), dim=0) # all gather batch among pp group if meta_info.get('allgather_pp_output', True): data.batch = allgather_dict_tensors(data.batch.contiguous(), size=mpu.get_pipeline_model_parallel_world_size(), group=mpu.get_pipeline_model_parallel_group(), dim=0) return data","summary_state":"pending","summary":null,"crux":null,"owner":"MegatronVLLMShardingManager"},{"id":"verl/workers/sharding_manager/megatron_vllm.py#get_micro_data_parallel_group","name":"get_micro_data_parallel_group","kind":"function","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L418-L420","signature":"def get_micro_data_parallel_group()","exported":true,"origin":"ast","body_hash":"2f7b380a27daa844595b0cfc44c14ac7b99261bb2450d06387b395f1dba78cf5","body_text":"def get_micro_data_parallel_group(): assert _MICRO_DATA_PARALLEL_GROUP is not None return _MICRO_DATA_PARALLEL_GROUP","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/sharding_manager/megatron_vllm.py#get_micro_data_parallel_world_size","name":"get_micro_data_parallel_world_size","kind":"function","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L423-L424","signature":"def get_micro_data_parallel_world_size()","exported":true,"origin":"ast","body_hash":"a1bfe7b553e575378267954148b439163e9ad2696ed12153bf9be879575bb9b8","body_text":"def get_micro_data_parallel_world_size(): return torch.distributed.get_world_size(group=get_micro_data_parallel_group())","summary_state":"pending","summary":null,"crux":null},{"id":"verl/workers/sharding_manager/megatron_vllm.py#get_micro_data_parallel_rank","name":"get_micro_data_parallel_rank","kind":"function","path":"verl/workers/sharding_manager/megatron_vllm.py","span":"L427-L428","signature":"def get_micro_data_parallel_rank()","exported":true,"origin":"ast","body_hash":"3557d8db20db94e469e45ac4f980e16bed562a5b5699954aa2b53229cc537cb7","body_text":"def get_micro_data_parallel_rank(): return torch.distributed.get_rank(group=get_micro_data_parallel_group())","summary_state":"pending","summary":null,"crux":null}],"rawEdges":[{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"torch","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"torch","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"megatron.core","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"megatron.core","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"megatron.core.transformer.module","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"torch.nn.parallel.distributed","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"verl.utils.megatron_utils","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"verl.utils.memory_buffer","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"get_pipeline_model_parallel_group","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"get_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"get_virtual_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"list","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"print","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"memory_allocated","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"memory_reserved","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"set_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"get_model","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"ModuleList","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"get_model","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"ModuleList","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"ModuleList","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"unwrap_model","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"_build_param_buffer","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"AllGatherPPModel"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"_build_param_references","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"AllGatherPPModel"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"eval","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__init__","relation":"calls","name":"_offload_params_to_cpu","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"AllGatherPPModel"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._build_param_buffer","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._build_param_buffer","relation":"calls","name":"get_weight_buffer_meta_from_module","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._build_param_buffer","relation":"calls","name":"build_memory_buffer","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._build_param_references","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._build_param_references","relation":"calls","name":"build_memory_reference_from_module","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._load_params_to_cuda","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._load_params_to_cuda","relation":"calls","name":"values","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._load_params_to_cuda","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._load_params_to_cuda","relation":"calls","name":"current_device","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._load_params_to_cuda","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._load_params_to_cuda","relation":"calls","name":"_build_param_references","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"AllGatherPPModel"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._offload_params_to_cpu","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._offload_params_to_cpu","relation":"calls","name":"values","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._offload_params_to_cpu","relation":"calls","name":"to","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._offload_params_to_cpu","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel._offload_params_to_cpu","relation":"calls","name":"_build_param_references","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"AllGatherPPModel"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.load_params_to_cuda","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.load_params_to_cuda","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.load_params_to_cuda","relation":"calls","name":"_load_params_to_cuda","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"AllGatherPPModel"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.allgather_params","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.allgather_params","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.allgather_params","relation":"calls","name":"get_global_rank","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.allgather_params","relation":"calls","name":"values","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.allgather_params","relation":"calls","name":"broadcast","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.forward","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.forward","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.forward","relation":"calls","name":"set_virtual_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.forward","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.forward","relation":"calls","name":"set_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.forward","relation":"calls","name":"set_input_tensor","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.forward","relation":"calls","name":"set_input_tensor","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.forward","relation":"calls","name":"set_virtual_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.forward","relation":"calls","name":"set_pipeline_model_parallel_rank","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__call__","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.__call__","relation":"calls","name":"forward","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"AllGatherPPModel"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.eval","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.eval","relation":"calls","name":"eval","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.train","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.train","relation":"calls","name":"train","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.offload_params_to_cpu","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.offload_params_to_cpu","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.offload_params_to_cpu","relation":"calls","name":"_offload_params_to_cpu","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"AllGatherPPModel"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.get_all_params","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.get_all_params","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.get_all_params","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.get_all_params","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.get_all_params","relation":"calls","name":"len","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.get_all_params","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.get_all_params","relation":"calls","name":"unwrap_model","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.get_all_params","relation":"calls","name":"named_parameters","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"unwrap_model"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.update_this_rank_models","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.update_this_rank_models","relation":"calls","name":"unwrap_model","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.this_rank_models","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.pp_size","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.pp_rank","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.pp_group","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#AllGatherPPModel.pp_models","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":".base","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"torch","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"torch","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"torch.distributed","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"verl","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"verl.utils.torch_functional","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"verl.third_party.vllm","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"verl.third_party.vllm","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"imports","specifier":"verl.utils.model","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager","relation":"extends","name":"BaseShardingManager","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__init__","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__init__","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__init__","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__init__","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__init__","relation":"calls","name":"get_tensor_model_parallel_world_size","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__init__","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__init__","relation":"calls","name":"new_group","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"split","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"chunk","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"append","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"cat","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.default_tp_concat_fn","relation":"calls","name":"get_tensor_parallel_partition_dim","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager._post_process_params","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager._post_process_params","relation":"calls","name":"get_micro_data_parallel_world_size","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager._post_process_params","relation":"calls","name":"get_micro_data_parallel_group","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager._post_process_params","relation":"calls","name":"keys","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager._post_process_params","relation":"calls","name":"is_tensor_parallel_param","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager._post_process_params","relation":"calls","name":"empty_like","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager._post_process_params","relation":"calls","name":"range","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager._post_process_params","relation":"calls","name":"all_gather","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager._post_process_params","relation":"calls","name":"default_tp_concat_fn","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"MegatronVLLMShardingManager"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__enter__","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__enter__","relation":"calls","name":"load_params_to_cuda","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__enter__","relation":"calls","name":"allgather_params","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__enter__","relation":"calls","name":"get_all_params","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__enter__","relation":"calls","name":"normalize_pp_vpp_params","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__enter__","relation":"calls","name":"_post_process_params","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"MegatronVLLMShardingManager"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__enter__","relation":"calls","name":"sync_model_weights","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__exit__","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__exit__","relation":"calls","name":"offload_params_to_cpu","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__exit__","relation":"calls","name":"get_micro_data_parallel_world_size","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__exit__","relation":"calls","name":"keys","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"normalize_pp_vpp_params"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__exit__","relation":"calls","name":"offload_model_weights","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__exit__","relation":"calls","name":"train","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.__exit__","relation":"calls","name":"empty_cache","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.preprocess_data","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.preprocess_data","relation":"calls","name":"get_micro_data_parallel_world_size","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.preprocess_data","relation":"calls","name":"get_micro_data_parallel_rank","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.preprocess_data","relation":"calls","name":"broadcast_dict_tensor","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.preprocess_data","relation":"calls","name":"get_tensor_model_parallel_src_rank","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.preprocess_data","relation":"calls","name":"get_tensor_model_parallel_group","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.preprocess_data","relation":"calls","name":"chunk","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py","recvType":"DataProto"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","relation":"calls","name":"get_micro_data_parallel_world_size","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","relation":"calls","name":"allgather_dict_tensors","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","relation":"calls","name":"get_micro_data_parallel_world_size","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","relation":"calls","name":"get_micro_data_parallel_group","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","relation":"calls","name":"get","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","relation":"calls","name":"allgather_dict_tensors","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","relation":"calls","name":"contiguous","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","relation":"calls","name":"get_pipeline_model_parallel_world_size","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#MegatronVLLMShardingManager.postprocess_data","relation":"calls","name":"get_pipeline_model_parallel_group","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#get_micro_data_parallel_group","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#get_micro_data_parallel_world_size","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#get_micro_data_parallel_world_size","relation":"calls","name":"get_world_size","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#get_micro_data_parallel_world_size","relation":"calls","name":"get_micro_data_parallel_group","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py","relation":"contains","targetId":"verl/workers/sharding_manager/megatron_vllm.py#get_micro_data_parallel_rank","file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#get_micro_data_parallel_rank","relation":"calls","name":"get_rank","viaMember":true,"file":"verl/workers/sharding_manager/megatron_vllm.py"},{"source":"verl/workers/sharding_manager/megatron_vllm.py#get_micro_data_parallel_rank","relation":"calls","name":"get_micro_data_parallel_group","viaMember":false,"file":"verl/workers/sharding_manager/megatron_vllm.py"}]}}}