diff --git a/Agent0/ENTERPRISE_GUIDE.md b/Agent0/ENTERPRISE_GUIDE.md new file mode 100644 index 0000000..d0fd609 --- /dev/null +++ b/Agent0/ENTERPRISE_GUIDE.md @@ -0,0 +1,64 @@ +# Enterprise Readiness Guide + +This guide upgrades the operational UX for **Agent0** deployments by documenting security integration, compliance expectations, lint/audit routines, and reproducible builds. + +## 1. UX / Operational Quality of Life +- **Standardized environment variables**: use a `.env` or secret manager so local and CI setups share the same configuration keys. +- **Clear paths**: keep all runtime artifacts under a single root (e.g., `$STORAGE_PATH`) to simplify cleanup and audits. +- **Runbook-first**: keep the primary workflow in a single script or Make target to reduce tribal knowledge. + +### Suggested environment variables +| Variable | Purpose | +| --- | --- | +| `STORAGE_PATH` | Central location for artifacts and checkpoints. | +| `HUGGINGFACENAME` | Hugging Face token or username. | +| `WANDB_API_KEY` | Weights & Biases API key. | +| `SANDBOX_API_URLS` | Comma-separated list of sandbox endpoints for tool execution. | + +## 2. Security Integration +- **Secrets management**: load credentials through your enterprise secret manager; avoid `.env` in production. +- **Network policy**: restrict outbound access from training workers to only model, logging, and sandbox endpoints. +- **Artifact integrity**: store checkpoints in immutable object storage with bucket versioning enabled. +- **Sandbox isolation**: treat the sandbox service as untrusted execution; use network isolation and per-request rate limiting. + +## 3. Compliance & Audit +- **Data lineage**: log dataset versions, question generation seeds, and filtering thresholds for every training run. +- **Model governance**: keep a manifest with model hash, base model ID, and training configuration. +- **Access control**: enforce RBAC on checkpoints, logs, and sandbox services. +- **Retention**: define retention policies for generated data and intermediate artifacts. + +## 4. Linting & Audit Checklist +Use these as baseline checks in CI (adjust for your environment). A `Makefile` target is provided for quick runs. + +- **Python linting**: `ruff` or `flake8` for style and static issues. +- **Type checks**: `mypy` for critical modules. +- **Dependency audit**: `pip-audit` or `safety` for known CVEs. +- **License scan**: `pip-licenses` to ensure dependency compliance. + +## 5. Build & Release Hygiene +- **Reproducible builds**: pin all dependencies in `requirements.txt` and use a lockfile for CI. +- **Immutable tags**: tag releases with model checkpoint hashes. +- **Container build**: prefer a single base image for all training and evaluation jobs to avoid drift. + +### Example CI sequence +```bash +python -m pip install -r requirements.txt +python -m pip install ruff mypy pip-audit pip-licenses +ruff check . +mypy . +pip-audit +pip-licenses --format=markdown +``` + +### Example local sequence +```bash +python -m pip install ruff mypy pip-audit pip-licenses +make lint +make audit +make build +``` + +## 6. Suggested Enhancements (Roadmap) +- Add a `Makefile` or `taskfile.yml` with standardized commands (`lint`, `audit`, `train`, `evaluate`). +- Add a `SECURITY.md` with responsible disclosure process and contact info. +- Add CI workflows for linting and dependency audits. diff --git a/Agent0/Makefile b/Agent0/Makefile new file mode 100644 index 0000000..b62c561 --- /dev/null +++ b/Agent0/Makefile @@ -0,0 +1,12 @@ +.PHONY: lint audit build + +lint: + python -m ruff check . + python -m mypy . + +audit: + python -m pip-audit + python -m pip-licenses --format=markdown + +build: + python -m pip install -r requirements.txt --dry-run diff --git a/Agent0/README.md b/Agent0/README.md index 1b32975..4b50b9d 100644 --- a/Agent0/README.md +++ b/Agent0/README.md @@ -119,4 +119,8 @@ If you find this work helpful, please consider citing our paper: author={Xia, Peng and Zeng, Kaide and Liu, Jiaqi and Qin, Can and Wu, Fang and Zhou, Yiyang and Xiong, Caiming and Yao, Huaxiu}, journal={arXiv preprint arXiv:2511.16043}, year={2025} -} \ No newline at end of file +} +``` + +## ๐Ÿข Enterprise Readiness +For security integration, compliance guidance, linting, and reproducible build recommendations, see [ENTERPRISE_GUIDE.md](./ENTERPRISE_GUIDE.md). diff --git a/Agent0/curriculum_train/examples/reward_function/curriculum_reward.py b/Agent0/curriculum_train/examples/reward_function/curriculum_reward.py index 6341a4d..5bb9578 100644 --- a/Agent0/curriculum_train/examples/reward_function/curriculum_reward.py +++ b/Agent0/curriculum_train/examples/reward_function/curriculum_reward.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,13 +27,17 @@ from sklearn.cluster import AgglomerativeClustering import numpy as np -STORAGE_PATH = os.getenv("STORAGE_PATH","") +STORAGE_PATH = os.getenv("STORAGE_PATH", "") + def _bleu_distance_matrix(sentences): n = len(sentences) dist = np.zeros((n, n)) smoother = SmoothingFunction().method1 - for i in tqdm(range(n), desc=" - Calculating BLEU distance matrix", leave=False): + for i in tqdm( + range(n), + desc=" - Calculating BLEU distance matrix", + leave=False): for j in range(i, n): if i == j: score = 1.0 @@ -44,13 +48,13 @@ def _bleu_distance_matrix(sentences): dist[i, j] = dist[j, i] = 1 - score return dist + def cluster_share_per_problem( - problems, - distance_threshold: float = 0.5, - linkage: str = "average"): + problems, distance_threshold: float = 0.5, linkage: str = "average" +): if not problems: return [] - print('start clustering') + print("start clustering") start_time = time.time() dist_mat = _bleu_distance_matrix(problems) @@ -58,10 +62,10 @@ def cluster_share_per_problem( n_clusters=None, distance_threshold=distance_threshold, metric="precomputed", - linkage=linkage + linkage=linkage, ) labels = clustering.fit_predict(dist_mat) - print(f'end clustering, time: {time.time() - start_time}') + print(f"end clustering, time: {time.time() - start_time}") total = len(problems) cluster_size = Counter(labels) cluster_ratio = {lab: sz / total for lab, sz in cluster_size.items()} @@ -69,41 +73,56 @@ def cluster_share_per_problem( proportions = [cluster_ratio[lab] for lab in labels] return proportions + def generate_temp_filename(prefix="temp", suffix=".json"): timestamp = int(time.time() * 1000) rand_part = random.randint(0, 99999) return f"{STORAGE_PATH}/temp_results/{prefix}_{timestamp}_{rand_part}{suffix}" + + def split_list(lst, n=4): k, m = divmod(len(lst), n) - return [lst[i*k + min(i, m):(i+1)*k + min(i+1, m)] for i in range(n)] + return [lst[i * k + min(i, m): (i + 1) * k + min(i + 1, m)] + for i in range(n)] + os.environ["NO_PROXY"] = "0.0.0.0,127.0.0.1" -def fetch(index,i): - response = requests.get(f"http://0.0.0.0:{5000+index}/hello?name={i}") + +def fetch(index, i): + response = requests.get(f"http://0.0.0.0:{5000 + index}/hello?name={i}") return True + def generate_results(data): - datas = split_list(data,4) - random_names = [generate_temp_filename(prefix=f"temp_{i}", suffix=".json") for i in range(4)] + datas = split_list(data, 4) + random_names = [ + generate_temp_filename( + prefix=f"temp_{i}", + suffix=".json") for i in range(4)] for i in range(4): - with open(random_names[i],'w') as f: - json.dump(datas[i],f,indent=4) + with open(random_names[i], "w") as f: + json.dump(datas[i], f, indent=4) final_results = [] with ThreadPoolExecutor(max_workers=4) as executor: - futures = [executor.submit(fetch, i,random_names[i]) for i in range(4)] + futures = [executor.submit(fetch, i, random_names[i]) + for i in range(4)] - for future in tqdm(as_completed(futures), total=len(futures), desc=" - Servers processing"): - future.result() # Simplified to just get the result + for future in tqdm( + as_completed(futures), + total=len(futures), + desc=" - Servers processing"): + future.result() # Simplified to just get the result for i in tqdm(range(4), desc=" - Reading result files", leave=False): - with open(random_names[i].replace('.json','_results.json'),'r') as f: + with open(random_names[i].replace(".json", "_results.json"), "r") as f: final_results.extend(json.load(f)) for i in range(4): - os.remove(random_names[i].replace('.json','_results.json')) + os.remove(random_names[i].replace(".json", "_results.json")) return final_results + def format_reward(predict: str) -> float: pattern = re.compile(r".*.*\\boxed\{.*\}.*", re.DOTALL) format_match = re.fullmatch(pattern, predict) @@ -114,7 +133,11 @@ def accuracy_reward(predict: str, ground_truth: str) -> float: answer = extract_boxed_content(predict) return 1.0 if grade_answer(answer, ground_truth) else 0.0 -def calculate_tool_reward(predict: str, weight: float = 0.05, cap: int = 4) -> float: + +def calculate_tool_reward( + predict: str, + weight: float = 0.05, + cap: int = 4) -> float: if not predict: return 0.0 @@ -125,28 +148,54 @@ def calculate_tool_reward(predict: str, weight: float = 0.05, cap: int = 4) -> f return capped_calls * weight -def compute_score(predicts: List[str], ground_truths: List[str], format_weight: float = 0.1, file_path: str = "") -> List[Dict[str, float]]: +def compute_score( + predicts: List[str], + ground_truths: List[str], + format_weight: float = 0.1, + file_path: str = "", +) -> List[Dict[str, float]]: results = [] - with open('test.json','w') as f: - json.dump(predicts,f,indent=4) + with open("test.json", "w") as f: + json.dump(predicts, f, indent=4) for i in tqdm(range(len(predicts)), desc=" - Parsing predictions"): - questions = re.findall(r"(.*?)", predicts[i], re.DOTALL) + questions = re.findall( + r"(.*?)", + predicts[i], + re.DOTALL) answers = extract_boxed_content(predicts[i]) if questions and answers: try: question = questions[-1].strip() answer = answers[-1].strip() results.append({"question": question, "answer": answer}) - except: + except BaseException: results.append({"question": "", "answer": ""}) else: results.append({"question": "", "answer": ""}) final_results = generate_results(results) - penalty = cluster_share_per_problem([result['question'] for result in final_results], distance_threshold=0.5) + penalty = cluster_share_per_problem( + [result["question"] for result in final_results], distance_threshold=0.5 + ) assert len(penalty) == len(final_results) scores = [] - for i in tqdm(range(len(final_results)), desc=" - Calculating final scores"): - final_score = (min(final_results[i]["score"],1-final_results[i]["score"]) if final_results[i]['question'] else -1)-penalty[i]+calculate_tool_reward(predicts[i]) - scores.append({"overall": final_score,"format": 1 if final_results[i]['question'] else 0,"accuracy": penalty[i],"tool_reward": calculate_tool_reward(predicts[i])}) - return scores \ No newline at end of file + for i in tqdm(range(len(final_results)), + desc=" - Calculating final scores"): + final_score = ( + ( + min(final_results[i]["score"], 1 - final_results[i]["score"]) + if final_results[i]["question"] + else -1 + ) + - penalty[i] + + calculate_tool_reward(predicts[i]) + ) + scores.append( + { + "overall": final_score, + "format": 1 if final_results[i]["question"] else 0, + "accuracy": penalty[i], + "tool_reward": calculate_tool_reward(predicts[i]), + } + ) + return scores diff --git a/Agent0/curriculum_train/examples/reward_function/math.py b/Agent0/curriculum_train/examples/reward_function/math.py index 1a8b675..8db20b4 100644 --- a/Agent0/curriculum_train/examples/reward_function/math.py +++ b/Agent0/curriculum_train/examples/reward_function/math.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -28,19 +28,24 @@ def accuracy_reward(predict: str, ground_truth: str) -> float: answer = extract_boxed_content(predict) try: return 1.0 if grade_answer(answer, ground_truth) else 0.0 - except: + except BaseException: return 0.0 -def compute_score(predicts: List[str], ground_truths: List[str], format_weight: float = 0.1) -> List[Dict[str, float]]: +def compute_score( + predicts: List[str], ground_truths: List[str], format_weight: float = 0.1 +) -> List[Dict[str, float]]: scores = [] for predict, ground_truth in zip(predicts, ground_truths): - predict = re.sub(r"\s*(<|>|/)\s*", r"\1", predict) # handle qwen2.5vl-32b format + predict = re.sub( + r"\s*(<|>|/)\s*", r"\1", predict + ) # handle qwen2.5vl-32b format format_score = format_reward(predict) accuracy_score = accuracy_reward(predict, ground_truth) scores.append( { - "overall": (1 - format_weight) * accuracy_score + format_weight * format_score, + "overall": (1 - format_weight) * accuracy_score + + format_weight * format_score, "format": format_score, "accuracy": accuracy_score, } diff --git a/Agent0/curriculum_train/examples/reward_function/r1v.py b/Agent0/curriculum_train/examples/reward_function/r1v.py index 204762f..8080396 100644 --- a/Agent0/curriculum_train/examples/reward_function/r1v.py +++ b/Agent0/curriculum_train/examples/reward_function/r1v.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,7 +19,9 @@ def format_reward(predict: str) -> float: - pattern = re.compile(r".*?\s*.*?", re.DOTALL) + pattern = re.compile( + r".*?\s*.*?", + re.DOTALL) format_match = re.fullmatch(pattern, predict) return 1.0 if format_match else 0.0 @@ -27,7 +29,8 @@ def format_reward(predict: str) -> float: def accuracy_reward(predict: str, ground_truth: str) -> float: try: content_match = re.search(r"(.*?)", predict) - given_answer = content_match.group(1).strip() if content_match else predict.strip() + given_answer = (content_match.group(1).strip() + if content_match else predict.strip()) if grade_answer(given_answer, ground_truth.strip()): return 1.0 @@ -37,11 +40,18 @@ def accuracy_reward(predict: str, ground_truth: str) -> float: return 0.0 -def compute_score(predict: str, ground_truth: str, format_weight: float = 0.5) -> Dict[str, float]: +def compute_score( + predict: str, ground_truth: str, format_weight: float = 0.5 +) -> Dict[str, float]: format_score = format_reward(predict) accuracy_score = accuracy_reward(predict, ground_truth) return { - "overall": (1 - format_weight) * accuracy_score + format_weight * format_score, + "overall": ( + 1 - + format_weight) * + accuracy_score + + format_weight * + format_score, "format": format_score, "accuracy": accuracy_score, } diff --git a/Agent0/curriculum_train/question_evaluate/evaluate.py b/Agent0/curriculum_train/question_evaluate/evaluate.py index b7106cc..43c80d9 100644 --- a/Agent0/curriculum_train/question_evaluate/evaluate.py +++ b/Agent0/curriculum_train/question_evaluate/evaluate.py @@ -1,6 +1,6 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- -''' +""" Description: This script evaluates generated answers against golden answers for a set of questions. It uses vLLM for efficient generation and a robust, timed grading mechanism to score the results. @@ -19,7 +19,7 @@ Example Usage (in a shell script): # This would run the script for GPU 0, with a specific model and save name. CUDA_VISIBLE_DEVICES=0 python evaluate.py --model "Qwen/Qwen3-4B-Base" --suffix 0 --save_name "my_experiment" & -''' +""" import json import vllm @@ -31,20 +31,44 @@ from mathruler.grader import extract_boxed_content, grade_answer # --- Argument Parsing --- -parser = argparse.ArgumentParser(description="Evaluate generated questions using vLLM.") -parser.add_argument("--model", type=str, default="Qwen/Qwen3-4B-Base", help="Path to the model in Hugging Face format.") -parser.add_argument("--num_samples", type=int, default=9, help="Number of candidate answers to generate per question (n).") -parser.add_argument("--suffix", type=str, default="0", help="A unique suffix for file naming, often the GPU index.") -parser.add_argument("--save_name", type=str, required=True, help="A base name for input and output files.") +parser = argparse.ArgumentParser( + description="Evaluate generated questions using vLLM.") +parser.add_argument( + "--model", + type=str, + default="Qwen/Qwen3-4B-Base", + help="Path to the model in Hugging Face format.", +) +parser.add_argument( + "--num_samples", + type=int, + default=9, + help="Number of candidate answers to generate per question (n).", +) +parser.add_argument( + "--suffix", + type=str, + default="0", + help="A unique suffix for file naming, often the GPU index.", +) +parser.add_argument( + "--save_name", + type=str, + required=True, + help="A base name for input and output files.", +) args = parser.parse_args() # --- Constants and Paths --- STORAGE_PATH = os.getenv("STORAGE_PATH", "") INPUT_FILE = f"{STORAGE_PATH}/generated_question/{args.save_name}_{args.suffix}.json" -OUTPUT_FILE = f"{STORAGE_PATH}/generated_question/{args.save_name}_{args.suffix}_results.json" +OUTPUT_FILE = ( + f"{STORAGE_PATH}/generated_question/{args.save_name}_{args.suffix}_results.json" +) + # --- Timeout-Protected Grading Function --- -@stopit.threading_timeoutable(default='TIMED_OUT') +@stopit.threading_timeoutable(default="TIMED_OUT") def grade_answer_with_timeout(res1, res2): """ Wraps the mathruler 'grade_answer' function with a timeout. @@ -53,6 +77,7 @@ def grade_answer_with_timeout(res1, res2): # The actual timeout value is passed as a keyword argument on each call. return grade_answer(res1, res2) + # --- Main Script Logic --- # 1. Load and Prepare Data @@ -67,7 +92,7 @@ def grade_answer_with_timeout(res1, res2): exit() # Filter data into questions that need processing -correct_data = [item for item in data if item.get('score') == 0] +correct_data = [item for item in data if item.get("score") == 0] if not correct_data: print(f"[{args.suffix}] No new questions to process (score=0). Exiting.") # Create an empty results file to signal completion @@ -99,14 +124,30 @@ def grade_answer_with_timeout(res1, res2): # 3. Generate Responses print(f"[{args.suffix}] Generating {args.num_samples} samples for each question...") -chats = [[{"role": "system", "content": "Please reason step by step, and put your final answer within \\boxed{}."},{"role": "user", "content": q}] for q in questions] +chats = [[{"role": "system", + "content": "Please reason step by step, and put your final answer within \\boxed{}.", + }, + {"role": "user", + "content": q}, + ] for q in questions] if tokenizer.chat_template: - prompts = [tokenizer.apply_chat_template(chat, tokenize=False, add_generation_prompt=True, add_special_tokens=True) for chat in chats] + prompts = [ + tokenizer.apply_chat_template( + chat, tokenize=False, add_generation_prompt=True, add_special_tokens=True + ) + for chat in chats + ] else: - prompts = ["system: " + chat[0]["content"] + '\n' + "user: " + chat[1]["content"] for chat in chats] - -responses = model.generate(prompts, sampling_params=sample_params, use_tqdm=True) + prompts = [ + "system: " + chat[0]["content"] + "\n" + "user: " + chat[1]["content"] + for chat in chats + ] + +responses = model.generate( + prompts, + sampling_params=sample_params, + use_tqdm=True) print(f"[{args.suffix}] Generation complete.") # 4. Process and Grade Responses @@ -115,11 +156,15 @@ def grade_answer_with_timeout(res1, res2): for response, golden_answer, question in zip(responses, answers, questions): try: # Extract the boxed content from all generated samples - results = [extract_boxed_content(output.text) for output in response.outputs] - results = [res for res in results if res] # Filter out None/empty results + results = [extract_boxed_content(output.text) + for output in response.outputs] + # Filter out None/empty results + results = [res for res in results if res] if not results: - print(f"[{args.suffix}] WARNING: No valid boxed answers found for question: '{question[:50]}...'") + print( + f"[{args.suffix}] WARNING: No valid boxed answers found for question: '{question[:50]}...'" + ) continue answer_counts = {} @@ -127,26 +172,33 @@ def grade_answer_with_timeout(res1, res2): matched = False for existing_answer in answer_counts: # OPTIMIZATION: Perform cheap string comparisons first. - if result == existing_answer or ('no ' in result.lower() and 'no ' in existing_answer.lower()): + if result == existing_answer or ( + "no " in result.lower() and "no " in existing_answer.lower()): answer_counts[existing_answer] += 1 matched = True break - + # If cheap checks fail, use the expensive, timed grader. # Check both directions (A vs B and B vs A). - match_1 = grade_answer_with_timeout(result, existing_answer, timeout=10) - if match_1 == 'TIMED_OUT': - print(f"[{args.suffix}] GRADER TIMEOUT on: '{result[:30]}...' vs '{existing_answer[:30]}...'") - continue # Skip to the next existing_answer - + match_1 = grade_answer_with_timeout( + result, existing_answer, timeout=10) + if match_1 == "TIMED_OUT": + print( + f"[{args.suffix}] GRADER TIMEOUT on: '{result[:30]}...' vs '{existing_answer[:30]}...'" + ) + continue # Skip to the next existing_answer + if match_1: answer_counts[existing_answer] += 1 matched = True break - match_2 = grade_answer_with_timeout(existing_answer, result, timeout=10) - if match_2 == 'TIMED_OUT': - print(f"[{args.suffix}] GRADER TIMEOUT on: '{existing_answer[:30]}...' vs '{result[:30]}...'") + match_2 = grade_answer_with_timeout( + existing_answer, result, timeout=10) + if match_2 == "TIMED_OUT": + print( + f"[{args.suffix}] GRADER TIMEOUT on: '{existing_answer[:30]}...' vs '{result[:30]}...'" + ) continue if match_2: @@ -166,23 +218,33 @@ def grade_answer_with_timeout(res1, res2): score = max_count / len(results) # Skip certain question types that are hard to grade automatically - if "่ฏๆ˜Ž" in question or 'box' in question.lower() or 'text' in majority_answer.lower(): + if ( + "่ฏๆ˜Ž" in question + or "box" in question.lower() + or "text" in majority_answer.lower() + ): continue - results_all.append({ - "question": question, - "answer": majority_answer, - "score": score, - 'results': results - }) + results_all.append( + { + "question": question, + "answer": majority_answer, + "score": score, + "results": results, + } + ) except Exception as e: - print(f"[{args.suffix}] CRITICAL ERROR processing question '{question[:50]}...': {e}") + print( + f"[{args.suffix}] CRITICAL ERROR processing question '{question[:50]}...': {e}" + ) continue # 5. Save Final Results -print(f"[{args.suffix}] Processed {len(results_all)} questions. Saving results to: {OUTPUT_FILE}") +print( + f"[{args.suffix}] Processed {len(results_all)} questions. Saving results to: {OUTPUT_FILE}" +) with open(OUTPUT_FILE, "w") as f: json.dump(results_all, f, indent=4) -print(f"[{args.suffix}] Script finished.") \ No newline at end of file +print(f"[{args.suffix}] Script finished.") diff --git a/Agent0/curriculum_train/question_evaluate/upload.py b/Agent0/curriculum_train/question_evaluate/upload.py index 95afd83..20142cc 100644 --- a/Agent0/curriculum_train/question_evaluate/upload.py +++ b/Agent0/curriculum_train/question_evaluate/upload.py @@ -16,32 +16,41 @@ parser = argparse.ArgumentParser() parser.add_argument("--max_score", type=float, default=0.7) parser.add_argument("--min_score", type=float, default=0.3) -parser.add_argument("--experiment_name", type=str, default="Qwen_Qwen3-4B-Base_all") +parser.add_argument( + "--experiment_name", + type=str, + default="Qwen_Qwen3-4B-Base_all") args = parser.parse_args() datas = [] for i in range(8): - file_path = f'{STORAGE_PATH}/generated_question/{args.experiment_name}_{i}_results.json' + file_path = ( + f"{STORAGE_PATH}/generated_question/{args.experiment_name}_{i}_results.json" + ) try: - with open(file_path, 'r') as f: + with open(file_path, "r") as f: data = json.load(f) datas.extend(data) except FileNotFoundError: - print(f"Warning: File {file_path} not found, skipping.", file=sys.stderr) + print( + f"Warning: File {file_path} not found, skipping.", + file=sys.stderr) continue print("Cleaning up temporary JSON files...", file=sys.stderr) for i in range(8): - file_path = f'{STORAGE_PATH}/generated_question/{args.experiment_name}_{i}_results.json' + file_path = ( + f"{STORAGE_PATH}/generated_question/{args.experiment_name}_{i}_results.json" + ) try: os.remove(file_path) except FileNotFoundError: pass filtered_datas = [ - {'problem': data['question'], 'answer': data['answer'], 'score': data['score']} + {"problem": data["question"], "answer": data["answer"], "score": data["score"]} for data in datas - if args.min_score <= data.get('score', 0) <= args.max_score and data.get('answer') + if args.min_score <= data.get("score", 0) <= args.max_score and data.get("answer") ] print(f"Filtered down to {len(filtered_datas)} samples.", file=sys.stderr) @@ -53,9 +62,9 @@ os.makedirs(save_dir, exist_ok=True) save_path = f"{save_dir}/train.parquet" - + train_dataset.to_parquet(save_path) - + print(save_path) else: - print("Warning: No data to save after filtering.", file=sys.stderr) \ No newline at end of file + print("Warning: No data to save after filtering.", file=sys.stderr) diff --git a/Agent0/curriculum_train/question_generate/question_generate.py b/Agent0/curriculum_train/question_generate/question_generate.py index dee5433..7ce5555 100644 --- a/Agent0/curriculum_train/question_generate/question_generate.py +++ b/Agent0/curriculum_train/question_generate/question_generate.py @@ -8,32 +8,35 @@ import json import regex as re import os + STORAGE_PATH = os.getenv("STORAGE_PATH") + def extract_boxed(text): results, i = [], 0 - prefix = r'\boxed{' + prefix = r"\boxed{" plen = len(prefix) while True: start = text.find(prefix, i) if start == -1: - break # no more \boxed{โ€ฆ} + break # no more \boxed{โ€ฆ} j = start + plen depth = 1 while j < len(text) and depth: - if text[j] == '{': + if text[j] == "{": depth += 1 - elif text[j] == '}': + elif text[j] == "}": depth -= 1 j += 1 - results.append(text[start + plen : j - 1]) + results.append(text[start + plen: j - 1]) i = j return results + def get_response_mask(response_ids, eos_token_id, dtype): batch_size, seq_len = response_ids.shape mask = torch.ones((batch_size, seq_len), dtype=dtype) @@ -44,6 +47,7 @@ def get_response_mask(response_ids, eos_token_id, dtype): break return mask + def main(args): tokenizer = AutoTokenizer.from_pretrained(args.model) if tokenizer.pad_token is None: @@ -62,8 +66,7 @@ def main(args): answer = answers[0] chat = [ { - "role": "system", - "content": ( + "role": "system", "content": ( "You are an expert competition-math problem setter.\n" "FIRST, in your private scratch-pad, think step-by-step to design a brand-new, non-trivial problem. " "The problem could come from any field of mathematics, including but not limited to algebra, geometry, number theory, combinatorics, prealgebra, probability, statistics, and calculus. " @@ -75,27 +78,20 @@ def main(args): "\n\n" r"\boxed{final_answer}" "\n\n" - "Do NOT output anything elseโ€”no explanations, no extra markup." - ) - }, - { - "role": "user", - "content": ( - "Generate one new, challenging reasoning question now. " - "Remember to format the output exactly as instructed." - ) - } - ] + "Do NOT output anything elseโ€”no explanations, no extra markup."), }, { + "role": "user", "content": ( + "Generate one new, challenging reasoning question now. " + "Remember to format the output exactly as instructed."), }, ] if tokenizer.chat_template: prompt = tokenizer.apply_chat_template( - chat, + chat, tokenize=False, - add_generation_prompt=True, - add_special_tokens=True - ) + add_generation_prompt=True, + add_special_tokens=True) else: - prompt = "system: " + chat[0]["content"] + '\n' + "user: " + chat[1]["content"] + prompt = "system: " + chat[0]["content"] + \ + "\n" + "user: " + chat[1]["content"] sample_params = vllm.SamplingParams( max_tokens=4096, temperature=1.0, @@ -104,31 +100,44 @@ def main(args): stop_token_ids=[tokenizer.eos_token_id], ) - completions: List[RequestOutput] = model.generate([prompt]*args.num_samples, sampling_params=sample_params) - results=[] + completions: List[RequestOutput] = model.generate( + [prompt] * args.num_samples, sampling_params=sample_params + ) + results = [] for completion in completions: response = completion.outputs[0].text try: - questions = re.findall(r"(.*?)", response, re.DOTALL) + questions = re.findall( + r"(.*?)", response, re.DOTALL) answers = extract_boxed(response) if questions and answers: question = questions[-1].strip() answer = answers[-1].strip() - results.append({"question": question, "answer": answer, "score": 0}) + results.append( + {"question": question, "answer": answer, "score": 0}) else: - results.append({"question": response, "answer": "", "score": -1}) - except: + results.append( + {"question": response, "answer": "", "score": -1}) + except Exception: results.append({"question": response, "answer": "", "score": -1}) - with open(f"{STORAGE_PATH}/generated_question/{args.save_name}_{args.suffix}.json", "w") as f: + with open( + f"{STORAGE_PATH}/generated_question/{args.save_name}_{args.suffix}.json", "w" + ) as f: json.dump(results, f, indent=4) + if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--model", type=str, default="Qwen/Qwen3-4B") - parser.add_argument("--num_samples", type=int, default=1250, help="Number of samples to generate") - parser.add_argument("--suffix", type=str, default="", help="Suffix to add to the output file") + parser.add_argument( + "--num_samples", + type=int, + default=1250, + help="Number of samples to generate") + parser.add_argument("--suffix", type=str, default="", + help="Suffix to add to the output file") parser.add_argument("--save_name", type=str, default="", help="") args = parser.parse_args() - main(args) \ No newline at end of file + main(args) diff --git a/Agent0/curriculum_train/requirements.txt b/Agent0/curriculum_train/requirements.txt index b63d664..fcdb8fe 100644 --- a/Agent0/curriculum_train/requirements.txt +++ b/Agent0/curriculum_train/requirements.txt @@ -38,7 +38,6 @@ fastapi==0.115.12 fastapi-cli==0.0.7 fastrlock==0.8.3 filelock==3.18.0 -flash_attn==2.7.4.post1 Flask==3.1.1 fonttools==4.58.2 frozenlist==1.7.0 diff --git a/Agent0/curriculum_train/scripts/model_merger.py b/Agent0/curriculum_train/scripts/model_merger.py index 4f4dd3d..53e2a19 100644 --- a/Agent0/curriculum_train/scripts/model_merger.py +++ b/Agent0/curriculum_train/scripts/model_merger.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -48,17 +48,31 @@ def upload_model_to_huggingface(local_path: str, remote_path: str): api = HfApi() api.create_repo(repo_id=remote_path, private=False, exist_ok=True) - api.upload_folder(repo_id=remote_path, folder_path=local_path, repo_type="model") + api.upload_folder( + repo_id=remote_path, + folder_path=local_path, + repo_type="model") if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--local_dir", required=True, type=str, help="The path for your saved model") - parser.add_argument("--hf_upload_path", default=False, type=str, help="The path of the huggingface repo to upload") + parser.add_argument( + "--local_dir", + required=True, + type=str, + help="The path for your saved model") + parser.add_argument( + "--hf_upload_path", + default=False, + type=str, + help="The path of the huggingface repo to upload", + ) args = parser.parse_args() local_dir: str = args.local_dir - assert not local_dir.endswith("huggingface"), "The local_dir should not end with huggingface." + assert not local_dir.endswith( + "huggingface" + ), "The local_dir should not end with huggingface." # copy rank zero to find the shape of (dp, fsdp) rank = 0 @@ -71,8 +85,13 @@ def upload_model_to_huggingface(local_path: str, remote_path: str): assert world_size, "No model file with the proper format." - rank0_weight_path = os.path.join(local_dir, f"model_world_size_{world_size}_rank_{rank}.pt") - state_dict = torch.load(rank0_weight_path, map_location="cpu", weights_only=False) + rank0_weight_path = os.path.join( + local_dir, f"model_world_size_{world_size}_rank_{rank}.pt" + ) + state_dict = torch.load( + rank0_weight_path, + map_location="cpu", + weights_only=False) pivot_key = sorted(state_dict.keys())[0] weight = state_dict[pivot_key] if isinstance(weight, DTensor): @@ -87,7 +106,10 @@ def upload_model_to_huggingface(local_path: str, remote_path: str): print(f"Got device mesh {mesh}, mesh_dim_names {mesh_dim_names}") - assert mesh_dim_names in (("fsdp",), ("ddp", "fsdp")), f"Unsupported mesh_dim_names {mesh_dim_names}." + assert mesh_dim_names in ( + ("fsdp",), + ("ddp", "fsdp"), + ), f"Unsupported mesh_dim_names {mesh_dim_names}." if "tp" in mesh_dim_names: # fsdp * tp @@ -104,8 +126,13 @@ def upload_model_to_huggingface(local_path: str, remote_path: str): model_state_dict_lst.extend([""] * (total_shards - 1)) def process_one_shard(rank, model_state_dict_lst): - model_path = os.path.join(local_dir, f"model_world_size_{world_size}_rank_{rank}.pt") - state_dict = torch.load(model_path, map_location="cpu", weights_only=False) + model_path = os.path.join( + local_dir, f"model_world_size_{world_size}_rank_{rank}.pt" + ) + state_dict = torch.load( + model_path, + map_location="cpu", + weights_only=False) model_state_dict_lst[rank] = state_dict return state_dict @@ -174,7 +201,9 @@ def process_one_shard(rank, model_state_dict_lst): raise NotImplementedError(f"Unknown architecture {architectures}.") with torch.device("meta"): - model: PreTrainedModel = AutoClass.from_config(config, torch_dtype=torch.bfloat16) + model: PreTrainedModel = AutoClass.from_config( + config, torch_dtype=torch.bfloat16 + ) assert isinstance(model, PreTrainedModel) model.to_empty(device="cpu") diff --git a/Agent0/curriculum_train/verl/__init__.py b/Agent0/curriculum_train/verl/__init__.py index cf49f90..28e6444 100644 --- a/Agent0/curriculum_train/verl/__init__.py +++ b/Agent0/curriculum_train/verl/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,6 +27,8 @@ if os.getenv("USE_MODELSCOPE_HUB", "0").lower() in ["true", "y", "1"]: # Patch hub to download models from modelscope to speed up. if not is_package_available("modelscope"): - raise ImportError("You are using the modelscope hub, please install modelscope by `pip install modelscope`.") + raise ImportError( + "You are using the modelscope hub, please install modelscope by `pip install modelscope`." + ) patch_hub() diff --git a/Agent0/curriculum_train/verl/protocol.py b/Agent0/curriculum_train/verl/protocol.py index 65d48be..a13f330 100644 --- a/Agent0/curriculum_train/verl/protocol.py +++ b/Agent0/curriculum_train/verl/protocol.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -45,7 +45,9 @@ __all__ = ["DataProto", "union_tensor_dict"] -def pad_dataproto_to_divisor(data: "DataProto", size_divisor: int) -> Tuple["DataProto", int]: +def pad_dataproto_to_divisor( + data: "DataProto", size_divisor: int +) -> Tuple["DataProto", int]: """Pad a DataProto to size divisible by size_divisor Args: @@ -81,15 +83,20 @@ def unpad_dataproto(data: "DataProto", pad_size: int) -> "DataProto": return data -def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict: +def union_tensor_dict( + tensor_dict1: TensorDict, + tensor_dict2: TensorDict) -> TensorDict: """Union two tensordicts.""" if tensor_dict1.batch_size != tensor_dict2.batch_size: raise ValueError( - f"Two tensor dict must have identical batch size. Got {tensor_dict1.batch_size} and {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 in tensor_dict1 and not torch.equal(tensor_dict1[key], tensor_dict2[key]): + if key in tensor_dict1 and not torch.equal( + tensor_dict1[key], tensor_dict2[key] + ): raise ValueError(f"Key already exists: {key}.") tensor_dict1[key] = tensor_dict2[key] @@ -97,7 +104,9 @@ def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> Ten return tensor_dict1 -def union_numpy_dict(tensor_dict1: Dict[str, NDArray], tensor_dict2: Dict[str, NDArray]) -> Dict[str, NDArray]: +def union_numpy_dict( + tensor_dict1: Dict[str, NDArray], tensor_dict2: Dict[str, NDArray] +) -> Dict[str, NDArray]: for key in tensor_dict2.keys(): if key in tensor_dict1: assert isinstance(tensor_dict2[key], np.ndarray) @@ -137,9 +146,13 @@ def fold_batch_dim(data: "DataProto", new_batch_size: int): tensor.auto_batch_size_(batch_dims=1) for key, value in non_tensor.items(): - non_tensor[key] = np.reshape(value, newshape=(new_batch_size, -1, *value.shape[1:])) + non_tensor[key] = np.reshape( + value, newshape=(new_batch_size, -1, *value.shape[1:]) + ) - return DataProto(batch=tensor, non_tensor_batch=non_tensor, meta_info=data.meta_info) + return DataProto( + batch=tensor, non_tensor_batch=non_tensor, meta_info=data.meta_info + ) def collate_fn(data_items: list["DataProtoItem"]): @@ -151,7 +164,11 @@ def collate_fn(data_items: list["DataProtoItem"]): batch = torch.stack(batch).contiguous() non_tensor_batch = batch_collate(non_tensor_batch) - non_tensor_batch = {key: np.array(value, dtype=object) for key, value in non_tensor_batch.items()} + non_tensor_batch = { + key: np.array( + value, + dtype=object) for key, + value in non_tensor_batch.items()} return DataProto(batch=batch, non_tensor_batch=non_tensor_batch) @@ -187,11 +204,19 @@ def __len__(self) -> int: else: return 0 - def __getitem__(self, item: Union[int, slice]) -> Union["DataProto", "DataProtoItem"]: + def __getitem__( + self, item: Union[int, slice] + ) -> Union["DataProto", "DataProtoItem"]: tensor_data = self.batch[item] - non_tensor_data = {key: value[item] for key, value in self.non_tensor_batch.items()} + non_tensor_data = { + key: value[item] for key, value in self.non_tensor_batch.items() + } return_type = DataProto if isinstance(item, slice) else DataProtoItem - return return_type(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info) + return return_type( + batch=tensor_data, + non_tensor_batch=non_tensor_data, + meta_info=self.meta_info, + ) def __getstate__(self) -> Tuple[bytes, Dict[str, NDArray], Dict[str, Any]]: buffer = io.BytesIO() @@ -203,10 +228,15 @@ def __getstate__(self) -> Tuple[bytes, Dict[str, NDArray], Dict[str, Any]]: buffer_bytes = buffer.getvalue() return buffer_bytes, self.non_tensor_batch, self.meta_info - def __setstate__(self, data: Tuple[bytes, Dict[str, NDArray], Dict[str, Any]]) -> None: + def __setstate__( + self, data: Tuple[bytes, Dict[str, NDArray], Dict[str, Any]] + ) -> None: batch_deserialized_bytes, non_tensor_batch, meta_info = data batch_deserialized = io.BytesIO(batch_deserialized_bytes) - batch = torch.load(batch_deserialized, weights_only=False, map_location="cpu") + batch = torch.load( + batch_deserialized, + weights_only=False, + map_location="cpu") self.batch = batch self.non_tensor_batch = non_tensor_batch self.meta_info = meta_info @@ -243,15 +273,20 @@ def check_consistency(self): 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" + assert len( + self.batch.batch_size) == 1, "only support num_batch_dims=1" 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." + 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, value in self.non_tensor_batch.items(): - assert len(value) == batch_size, f"key {key} length {len(value)} is not equal to bsz {batch_size}." + assert ( + len(value) == batch_size), f"key {key} length { + len(value)} is not equal to bsz {batch_size}." @classmethod def from_single_dict( @@ -268,7 +303,9 @@ def from_single_dict( else: raise ValueError(f"Unsupported type in data {type(value)}") - return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info) + return DataProto.from_dict( + tensors=tensors, non_tensors=non_tensors, meta_info=meta_info + ) @classmethod def from_dict( @@ -285,11 +322,14 @@ def from_dict( 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." + assert ( + num_batch_dims == 1 + ), "only support num_batch_dims=1 when non_tensors is not None." meta_info = meta_info or {} non_tensors = non_tensors or {} - assert isinstance(non_tensors, dict), "non_tensors should be a dictionary." + assert isinstance( + non_tensors, dict), "non_tensors should be a dictionary." # get and check batch size batch_size = None @@ -306,7 +346,10 @@ def from_dict( ) tensor_dict = TensorDict(source=tensors, batch_size=batch_size) - return cls(batch=tensor_dict, non_tensor_batch=non_tensors, meta_info=meta_info) + return cls( + batch=tensor_dict, + non_tensor_batch=non_tensors, + meta_info=meta_info) def to(self, device: torch.device) -> "DataProto": """move the batch to device @@ -347,7 +390,11 @@ def select( sub_batch = self.batch if non_tensor_batch_keys is not None: - non_tensor_batch = {k: v for k, v in self.non_tensor_batch.items() if k in non_tensor_batch_keys} + non_tensor_batch = { + k: v + for k, v in self.non_tensor_batch.items() + if k in non_tensor_batch_keys + } else: non_tensor_batch = self.non_tensor_batch @@ -355,14 +402,19 @@ def select( non_tensor_batch = copy.deepcopy(non_tensor_batch) if meta_info_keys is not None: - sub_meta_info = {k: v for k, v in self.meta_info.items() if k in meta_info_keys} + sub_meta_info = { + k: v for k, v in self.meta_info.items() if k 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) + return DataProto( + batch=sub_batch, + non_tensor_batch=non_tensor_batch, + meta_info=sub_meta_info) def pop( self, @@ -395,10 +447,14 @@ def pop( for key in 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) + return DataProto.from_dict( + tensors=tensors, non_tensors=non_tensors, meta_info=meta_info + ) def rename( - self, old_keys: Optional[Union[str, List[str]]] = None, new_keys: Optional[Union[str, List[str]]] = None + self, + old_keys: Optional[Union[str, List[str]]] = None, + new_keys: Optional[Union[str, List[str]]] = None, ) -> "DataProto": """ Note that this function only rename the key in the batch @@ -411,7 +467,9 @@ def validate_input(keys): elif isinstance(keys, list): pass else: - raise TypeError(f"keys must be a list or a string, but got {type(keys)}") + raise TypeError( + f"keys must be a list or a string, but got { + type(keys)}") return keys old_keys = validate_input(old_keys) @@ -419,8 +477,9 @@ def validate_input(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)}" - ) + 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)) @@ -440,12 +499,18 @@ def union(self, other: "DataProto") -> "DataProto": 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.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 def make_iterator( - self, mini_batch_size: int, epochs: int, seed: int = None, dataloader_kwargs: Dict[str, Any] = None + self, + mini_batch_size: int, + epochs: int, + seed: int = None, + dataloader_kwargs: Dict[str, Any] = 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. @@ -461,7 +526,9 @@ def make_iterator( 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" + 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 = {} @@ -474,7 +541,11 @@ def make_iterator( assert isinstance(dataloader_kwargs, Dict) train_dataloader = DataLoader( - dataset=self, batch_size=mini_batch_size, collate_fn=collate_fn, generator=generator, **dataloader_kwargs + dataset=self, + batch_size=mini_batch_size, + collate_fn=collate_fn, + generator=generator, + **dataloader_kwargs, ) def get_data(): @@ -494,9 +565,10 @@ def chunk(self, chunks: int) -> List["DataProto"]: 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}." - ) + 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: @@ -513,7 +585,11 @@ def chunk(self, chunks: int) -> List["DataProto"]: 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) + DataProto( + batch=batch_lst[i], + non_tensor_batch=non_tensor_batch_lst[i], + meta_info=self.meta_info, + ) ) return output @@ -543,7 +619,11 @@ def concat(data: List["DataProto"]) -> "DataProto": for key, value in non_tensor_batch.items(): non_tensor_batch[key] = np.concatenate(value, axis=0) - return DataProto(batch=new_batch, non_tensor_batch=non_tensor_batch, meta_info=data[0].meta_info) + return DataProto( + batch=new_batch, + non_tensor_batch=non_tensor_batch, + meta_info=data[0].meta_info, + ) def reorder(self, indices: torch.Tensor) -> None: """ @@ -551,9 +631,14 @@ def reorder(self, indices: torch.Tensor) -> None: """ indices_np = indices.detach().numpy() self.batch = self.batch[indices] - self.non_tensor_batch = {key: value[indices_np] for key, value in self.non_tensor_batch.items()} - - def repeat(self, repeat_times: int = 2, interleave: bool = True) -> "DataProto": + self.non_tensor_batch = { + key: value[indices_np] for key, + value in self.non_tensor_batch.items()} + + def repeat( + self, + repeat_times: int = 2, + interleave: bool = True) -> "DataProto": """ Repeat the batch data a specified number of times. @@ -568,12 +653,15 @@ def repeat(self, repeat_times: int = 2, interleave: bool = True) -> "DataProto": if interleave: # Interleave the data repeated_tensors = { - key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items() + 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:]) + key: tensor.unsqueeze(0) + .expand(repeat_times, *tensor.shape) + .reshape(-1, *tensor.shape[1:]) for key, tensor in self.batch.items() } @@ -587,9 +675,12 @@ def repeat(self, repeat_times: int = 2, interleave: bool = True) -> "DataProto": repeated_non_tensor_batch = {} for key, value in self.non_tensor_batch.items(): if interleave: - repeated_non_tensor_batch[key] = np.repeat(value, repeat_times, axis=0) + repeated_non_tensor_batch[key] = np.repeat( + value, repeat_times, axis=0) else: - repeated_non_tensor_batch[key] = np.tile(value, (repeat_times,) + (1,) * (value.ndim - 1)) + repeated_non_tensor_batch[key] = np.tile( + value, (repeat_times,) + (1,) * (value.ndim - 1) + ) return DataProto( batch=repeated_batch, @@ -631,7 +722,9 @@ 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 + 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 @@ -643,13 +736,17 @@ def get(self): outputs = self.collect_fn(outputs) # select dp, concat if self.dispatch_fn is not None: - outputs = self.dispatch_fn(outputs) # split in batch dim, select using dp + # split in batch dim, select using dp + outputs = self.dispatch_fn(outputs) return outputs def allgather_dict_tensors( - tensors: Union[Dict[str, torch.Tensor], TensorDict], size: int, group: ProcessGroup, dim: int = 0 + tensors: Union[Dict[str, torch.Tensor], TensorDict], + size: int, + group: ProcessGroup, + dim: int = 0, ) -> Union[Dict[str, torch.Tensor], TensorDict]: """ TODO: optimize this. @@ -668,22 +765,35 @@ def allgather_dict_tensors( for key in sorted_keys: value = tensors_as_dict[key] output[key] = [torch.empty_like(value) for _ in range(size)] - torch.distributed.all_gather(output[key], value, group=group, async_op=False) + torch.distributed.all_gather( + output[key], value, 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) + output = TensorDict(source=output, + batch_size=tensors.batch_size[0] * size) return output -def all_gather_data_proto(data: DataProto, size: int, group: ProcessGroup) -> None: - # Note that this is an inplace operator just like torch.distributed.all_gather +def all_gather_data_proto( + data: DataProto, + size: int, + group: ProcessGroup) -> None: + # Note that this is an inplace operator just like + # torch.distributed.all_gather prev_device = data.batch.device data.batch = data.batch.cuda(device=torch.cuda.current_device()) - data.batch = allgather_dict_tensors(data.batch.contiguous(), size=size, group=group, dim=0) + data.batch = allgather_dict_tensors( + data.batch.contiguous(), size=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(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} + 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 + } diff --git a/Agent0/curriculum_train/verl/single_controller/__init__.py b/Agent0/curriculum_train/verl/single_controller/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/curriculum_train/verl/single_controller/__init__.py +++ b/Agent0/curriculum_train/verl/single_controller/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/single_controller/base/__init__.py b/Agent0/curriculum_train/verl/single_controller/base/__init__.py index 46c9670..746bbca 100644 --- a/Agent0/curriculum_train/verl/single_controller/base/__init__.py +++ b/Agent0/curriculum_train/verl/single_controller/base/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/single_controller/base/decorator.py b/Agent0/curriculum_train/verl/single_controller/base/decorator.py index b0e85a3..9a69010 100644 --- a/Agent0/curriculum_train/verl/single_controller/base/decorator.py +++ b/Agent0/curriculum_train/verl/single_controller/base/decorator.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,7 +26,8 @@ from .worker_group import WorkerGroup -# here we add a magic number of avoid user-defined function already have this attribute +# here we add a magic number of avoid user-defined function already have +# this attribute MAGIC_ATTR = "attrs_3141562937" @@ -90,34 +91,53 @@ def _concat_data_proto_or_future(outputs: List[DataProto]) -> DataProto: def dispatch_dp_compute(worker_group: "WorkerGroup", *args, **kwargs): for arg in args: - assert isinstance(arg, (tuple, list)) and len(arg) == worker_group.world_size + assert isinstance(arg, (tuple, list)) and len( + arg) == worker_group.world_size for value in kwargs.values(): - assert isinstance(value, (tuple, list)) and len(value) == worker_group.world_size + assert (isinstance(value, (tuple, list)) + and len(value) == worker_group.world_size) return args, kwargs -def collect_dp_compute(worker_group: "WorkerGroup", outputs: List[DataProto]) -> List[DataProto]: +def collect_dp_compute( + worker_group: "WorkerGroup", outputs: List[DataProto] +) -> List[DataProto]: assert len(outputs) == worker_group.world_size return outputs -def dispatch_dp_compute_data_proto(worker_group: "WorkerGroup", *args, **kwargs): - splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(worker_group.world_size, *args, **kwargs) +def dispatch_dp_compute_data_proto( + worker_group: "WorkerGroup", + *args, + **kwargs): + splitted_args, splitted_kwargs = _split_args_kwargs_data_proto( + worker_group.world_size, *args, **kwargs + ) return splitted_args, splitted_kwargs -def dispatch_dp_compute_data_proto_with_func(worker_group: "WorkerGroup", *args, **kwargs): - assert type(args[0]) is 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 +def dispatch_dp_compute_data_proto_with_func( + worker_group: "WorkerGroup", *args, **kwargs +): + # NOTE: The first one args is a function! + assert type(args[0]) is FunctionType + 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 -def collect_dp_compute_data_proto(worker_group: "WorkerGroup", outputs: List[DataProto]) -> DataProto: +def collect_dp_compute_data_proto( + worker_group: "WorkerGroup", outputs: List[DataProto] +) -> DataProto: for output in outputs: - assert isinstance(output, (DataProto, ray.ObjectRef)), f"Expect a DataProto, but got {type(output)}" + assert isinstance( + output, (DataProto, ray.ObjectRef) + ), f"Expect a DataProto, but got {type(output)}" outputs = collect_dp_compute(worker_group, outputs) return _concat_data_proto_or_future(outputs) @@ -165,18 +185,26 @@ def get_predefined_execute_fn(execute_mode: Execute): return predefined_execute_mode_fn[execute_mode] -def _check_dispatch_mode(dispatch_mode: Union[Dispatch, Dict[Literal["dispatch_fn", "collect_fn"], FunctionType]]): - assert isinstance(dispatch_mode, (Dispatch, dict)), ( - f"dispatch_mode must be a Dispatch or a Dict. Got {dispatch_mode}" - ) +def _check_dispatch_mode( + dispatch_mode: Union[ + Dispatch, Dict[Literal["dispatch_fn", "collect_fn"], FunctionType] + ], +): + 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" + assert ( + key in dispatch_mode + ), f"key {key} should be in dispatch_mode if it is a dictionary" def _check_execute_mode(execute_mode: Execute): - assert isinstance(execute_mode, Execute), f"execute_mode must be a Execute. Got {execute_mode}" + assert isinstance( + execute_mode, Execute + ), f"execute_mode must be a Execute. Got {execute_mode}" def _materialize_futures(*args, **kwargs): @@ -195,7 +223,12 @@ def _materialize_futures(*args, **kwargs): return new_args, kwargs -def register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, blocking=True, materialize_futures=True): +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) @@ -206,7 +239,11 @@ def inner(*args, **kwargs): args, kwargs = _materialize_futures(*args, **kwargs) return func(*args, **kwargs) - attrs = {"dispatch_mode": dispatch_mode, "execute_mode": execute_mode, "blocking": blocking} + attrs = { + "dispatch_mode": dispatch_mode, + "execute_mode": execute_mode, + "blocking": blocking, + } setattr(inner, MAGIC_ATTR, attrs) return inner diff --git a/Agent0/curriculum_train/verl/single_controller/base/register_center/__init__.py b/Agent0/curriculum_train/verl/single_controller/base/register_center/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/curriculum_train/verl/single_controller/base/register_center/__init__.py +++ b/Agent0/curriculum_train/verl/single_controller/base/register_center/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/single_controller/base/register_center/ray.py b/Agent0/curriculum_train/verl/single_controller/base/register_center/ray.py index de7f702..0dbd906 100644 --- a/Agent0/curriculum_train/verl/single_controller/base/register_center/ray.py +++ b/Agent0/curriculum_train/verl/single_controller/base/register_center/ray.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/single_controller/base/worker.py b/Agent0/curriculum_train/verl/single_controller/base/worker.py index 9ecffca..3c74bba 100644 --- a/Agent0/curriculum_train/verl/single_controller/base/worker.py +++ b/Agent0/curriculum_train/verl/single_controller/base/worker.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -78,7 +78,10 @@ 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} + return { + f"_{key.lower()}": self._store.get(f"_{key.lower()}", None) + for key in WorkerMeta.keys + } # we assume that in each WorkerGroup, there is a Master Worker @@ -104,14 +107,22 @@ def __new__(cls, *args, **kwargs): rank = os.getenv("RANK", None) worker_group_prefix = os.getenv("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)) + # 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)}" + 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() @@ -119,18 +130,22 @@ def _configure_before_init(self, register_center_name: str, rank: int): "MASTER_ADDR": master_addr, "MASTER_PORT": master_port, } - self.register_center = create_worker_group_register_center(name=register_center_name, info=rank_zero_info) + 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 + # construct a meta from envrionment variable. Note that the import must + # be inside the class because it is executed remotely world_size = int(os.getenv("WORLD_SIZE")) rank = int(os.getenv("RANK")) self._rank = rank self._world_size = world_size if "AMD" in torch.cuda.get_device_name(): - os.environ["CUDA_VISIBLE_DEVICES"] = os.getenv("ROCR_VISIBLE_DEVICES") + os.environ["CUDA_VISIBLE_DEVICES"] = os.getenv( + "ROCR_VISIBLE_DEVICES") os.environ["LOCAL_RANK"] = os.getenv("RAY_LOCAL_RANK") cuda_visible_devices = os.getenv("LOCAL_RANK", "0") torch.cuda.set_device(int(cuda_visible_devices)) @@ -169,7 +184,9 @@ def _configure_with_meta(self, meta: WorkerMeta): os.environ[key] = str(val) os.environ["REDIS_STORE_SERVER_HOST"] = ( - str(self._master_addr).replace("[", "").replace("]", "") if self._master_addr else "" + str(self._master_addr).replace("[", "").replace("]", "") + if self._master_addr + else "" ) def get_master_addr_port(self): @@ -196,7 +213,8 @@ 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) + @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 diff --git a/Agent0/curriculum_train/verl/single_controller/base/worker_group.py b/Agent0/curriculum_train/verl/single_controller/base/worker_group.py index 8648fbf..7fede2a 100644 --- a/Agent0/curriculum_train/verl/single_controller/base/worker_group.py +++ b/Agent0/curriculum_train/verl/single_controller/base/worker_group.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -21,21 +21,30 @@ import time from typing import Any, Callable, Dict, List, Optional -from .decorator import MAGIC_ATTR, Dispatch, get_predefined_dispatch_fn, get_predefined_execute_fn +from .decorator import ( + MAGIC_ATTR, + Dispatch, + get_predefined_dispatch_fn, + get_predefined_execute_fn, +) class ResourcePool: """The resource pool with meta info such as world size.""" def __init__( - self, process_on_nodes: Optional[Any] = None, max_colocate_count: int = 10, n_gpus_per_node: int = 8 + self, + process_on_nodes: Optional[Any] = None, + max_colocate_count: int = 10, + n_gpus_per_node: int = 8, ) -> None: if process_on_nodes is None: process_on_nodes = [] self._store = process_on_nodes self.max_colocate_count = max_colocate_count - self.n_gpus_per_node = n_gpus_per_node # this is left for future huawei GPU that contains 16 GPUs per node + # this is left for future huawei GPU that contains 16 GPUs per node + self.n_gpus_per_node = n_gpus_per_node def add_node(self, process_count): self._store.append(process_count) @@ -53,12 +62,15 @@ def store(self): 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 + [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] # noqa: C416 + nested_local_rank_list = [ + [i for i in range(local_world_size)] for local_world_size in self._store + ] # noqa: C416 return [item for row in nested_local_rank_list for item in row] @@ -77,11 +89,15 @@ def __call__(self) -> Any: return self.cls(*self.args, **self.kwargs) -def check_workers_alive(workers: List, is_alive: Callable, gap_time: float = 1) -> None: +def check_workers_alive( + workers: List, + is_alive: Callable, + gap_time: float = 1) -> None: while True: for worker in workers: if not is_alive(worker): - logging.warning(f"Worker {worker} is not alive, sending signal to main thread") + logging.warning( + f"Worker {worker} is not alive, sending signal to main thread") signal.raise_signal(signal.SIGABRT) time.sleep(gap_time) @@ -108,22 +124,27 @@ def __init__(self, resource_pool: ResourcePool, **kwargs) -> None: self._checker_thread: threading.Thread = None def _is_worker_alive(self, worker): - raise NotImplementedError("WorkerGroup._is_worker_alive called, should be implemented in derived class.") + raise NotImplementedError( + "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] + 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 + # 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) + target=check_workers_alive, + args=(self._workers, self._is_worker_alive, every_n_seconds), ) self._checker_thread.start() @@ -138,16 +159,23 @@ def _bind_worker_method(self, user_defined_cls, func_generator): for method_name in dir(user_defined_cls): try: method = getattr(user_defined_cls, method_name) - assert callable(method), f"{method_name} in {user_defined_cls} is not callable" + assert callable( + method + ), f"{method_name} in {user_defined_cls} is not callable" except Exception: - # if it is a property, it will fail because Class doesn't have instance property + # 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, "attribute must contain dispatch_mode in its key" + assert isinstance( + attribute, Dict + ), f"attribute must be a dictionary. Got {type(attribute)}" + assert ( + "dispatch_mode" in attribute + ), "attribute must contain dispatch_mode in its key" dispatch_mode = attribute["dispatch_mode"] execute_mode = attribute["execute_mode"] @@ -156,7 +184,8 @@ def _bind_worker_method(self, user_defined_cls, func_generator): # get dispatch fn if isinstance(dispatch_mode, Dispatch): # get default dispatch fn - fn = get_predefined_dispatch_fn(dispatch_mode=dispatch_mode) + fn = get_predefined_dispatch_fn( + dispatch_mode=dispatch_mode) dispatch_fn = fn["dispatch_fn"] collect_fn = fn["collect_fn"] else: @@ -167,7 +196,8 @@ def _bind_worker_method(self, user_defined_cls, func_generator): collect_fn = dispatch_mode["collect_fn"] # get execute_fn_name - execute_mode = get_predefined_execute_fn(execute_mode=execute_mode) + execute_mode = get_predefined_execute_fn( + execute_mode=execute_mode) wg_execute_fn_name = execute_mode["execute_fn_name"] # get execute_fn from string diff --git a/Agent0/curriculum_train/verl/single_controller/ray/__init__.py b/Agent0/curriculum_train/verl/single_controller/ray/__init__.py index 25b3141..cf1eb36 100644 --- a/Agent0/curriculum_train/verl/single_controller/ray/__init__.py +++ b/Agent0/curriculum_train/verl/single_controller/ray/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,7 +12,17 @@ # See the License for the specific language governing permissions and # limitations under the License. -from .base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup, create_colocated_worker_cls +from .base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, + create_colocated_worker_cls, +) -__all__ = ["RayClassWithInitArgs", "RayResourcePool", "RayWorkerGroup", "create_colocated_worker_cls"] +__all__ = [ + "RayClassWithInitArgs", + "RayResourcePool", + "RayWorkerGroup", + "create_colocated_worker_cls", +] diff --git a/Agent0/curriculum_train/verl/single_controller/ray/base.py b/Agent0/curriculum_train/verl/single_controller/ray/base.py index 9827312..a7a8419 100644 --- a/Agent0/curriculum_train/verl/single_controller/ray/base.py +++ b/Agent0/curriculum_train/verl/single_controller/ray/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -25,7 +25,10 @@ from ray.experimental.state.api import get_actor from ray.util import list_named_actors from ray.util.placement_group import PlacementGroup, placement_group -from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy, PlacementGroupSchedulingStrategy +from ray.util.scheduling_strategies import ( + NodeAffinitySchedulingStrategy, + PlacementGroupSchedulingStrategy, +) from ..base import ClassWithInitArgs, ResourcePool, Worker, WorkerGroup from ..base.decorator import MAGIC_ATTR @@ -39,7 +42,13 @@ def get_random_string(length: int) -> str: return "".join(random.choice(letters_digits) for _ in range(length)) -def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking): +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) @@ -51,7 +60,8 @@ def func(*args, **kwargs): return func -def sort_placement_group_by_node_ip(pgs: List[PlacementGroup]) -> List[PlacementGroup]: +def sort_placement_group_by_node_ip( + pgs: List[PlacementGroup]) -> List[PlacementGroup]: """ Sort the placement groups by node ip, all bundles in a single placement group should be on the same node. @@ -61,7 +71,8 @@ def sort_placement_group_by_node_ip(pgs: List[PlacementGroup]) -> List[Placement With this function, if there's only one resource pool and there's no node change, RANK should be consistent across nodes in multiple ray jobs, even if the whole ray cluster is restarted. """ - node_ip = {node["NodeID"]: node["NodeManagerAddress"] for node in ray.nodes()} + node_ip = {node["NodeID"]: node["NodeManagerAddress"] + for node in ray.nodes()} pg_ip = {} for pg in pgs: specs = ray._private.state.state.placement_group_table(pg.id) @@ -88,17 +99,25 @@ def __init__( self.pgs = None self.detached = detached - def get_placement_groups(self, strategy: str = "STRICT_PACK", name: Optional[str] = None) -> List[PlacementGroup]: + def get_placement_groups( + self, strategy: str = "STRICT_PACK", name: Optional[str] = None + ) -> List[PlacementGroup]: 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])}:" + 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_colocate_count, "GPU": 1} if self.use_gpu else {"CPU": self.max_colocate_count} + ( + {"CPU": self.max_colocate_count, "GPU": 1} + if self.use_gpu + else {"CPU": self.max_colocate_count} + ) for _ in range(process_count) ] for process_count in self._store @@ -107,7 +126,12 @@ def get_placement_groups(self, strategy: str = "STRICT_PACK", name: Optional[str lifetime = "detached" if self.detached else None pgs = [ - placement_group(bundles=bundles, strategy=strategy, name=pg_name_prefix + str(idx), lifetime=lifetime) + placement_group( + bundles=bundles, + strategy=strategy, + name=pg_name_prefix + str(idx), + lifetime=lifetime, + ) for idx, bundles in enumerate(pg_scheme) ] @@ -118,7 +142,9 @@ def get_placement_groups(self, strategy: str = "STRICT_PACK", name: Optional[str def extract_pg_from_exist( - resource_pools: Dict[str, RayResourcePool], src_role_names: List[str], resource_pool: RayResourcePool + resource_pools: Dict[str, RayResourcePool], + src_role_names: List[str], + resource_pool: RayResourcePool, ) -> List[PlacementGroup]: src_pgs = [ pg @@ -127,33 +153,47 @@ def extract_pg_from_exist( 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) + 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" - ) + 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)] -def merge_resource_pool(rp1: RayResourcePool, rp2: RayResourcePool) -> RayResourcePool: +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_colocate_count == rp2.max_colocate_count, ( - "Both RayResourcePool must has the same max_colocate_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" + assert ( + rp1.max_colocate_count == rp2.max_colocate_count + ), "Both RayResourcePool must has the same max_colocate_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 = 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 @@ -182,15 +222,22 @@ def __call__( ) -> 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 + 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 + placement_group=placement_group, + placement_group_bundle_index=placement_group_bundle_idx, ) } options.update(self._options) @@ -221,7 +268,9 @@ def __init__( ) -> 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 + 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 @@ -231,15 +280,23 @@ def __init__( 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 + 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) + self._bind_worker_method( + self.ray_cls_with_init.cls, func_generator) def _is_worker_alive(self, worker: ActorHandle) -> bool: 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 + 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: List[str]) -> None: workers = [ray.get_actor(name=name) for name in worker_names] @@ -247,7 +304,11 @@ def _init_with_detached_workers(self, worker_names: List[str]) -> None: self._world_size = len(worker_names) def _init_with_resource_pool( - self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWithInitArgs, bin_pack: bool, detached: bool + self, + resource_pool: RayResourcePool, + ray_cls_with_init: RayClassWithInitArgs, + bin_pack: bool, + detached: bool, ): use_gpu = resource_pool.use_gpu @@ -264,11 +325,14 @@ def _init_with_resource_pool( rank = -1 local_world_size = resource_pool.store[0] for pg_idx, pg in enumerate(sort_placement_group_by_node_ip(pgs)): - assert local_world_size <= pg.bundle_count, f"when generating for {self.name_prefix}, for the " + 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 + # 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), @@ -282,18 +346,28 @@ def _init_with_resource_pool( env_vars["MASTER_PORT"] = self._master_port 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}) + 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" + # e.g. Worker_2:5 + name = f"{self.name_prefix}{cia_name}_{pg_idx}:{local_rank}" + + 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 + 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) @@ -301,16 +375,28 @@ def _init_with_resource_pool( if rank == 0: register_center_actor = None for _ in range(120): - if f"{self.name_prefix}_register_center" not in list_named_actors(): + 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") + 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)}" + 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"], ) - 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}") @@ -321,7 +407,10 @@ def worker_names(self): @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 + resource_pool=None, + ray_cls_with_init=ray_cls_with_init, + name_prefix=None, + worker_names=worker_names, ) return worker_group @@ -346,7 +435,8 @@ def _rebind_actor_methods(worker_group, actor_name): 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 + worker_names=self._worker_names, + ray_cls_with_init=self.ray_cls_with_init, ) _rebind_actor_methods(new_worker_group, prefix) @@ -354,7 +444,9 @@ def _rebind_actor_methods(worker_group, actor_name): return new_worker_group_dict def execute_rank_zero_sync(self, method_name: str, *args, **kwargs): - return ray.get(self.execute_rank_zero_async(method_name, *args, **kwargs)) + return ray.get( + self.execute_rank_zero_async( + method_name, *args, **kwargs)) def execute_rank_zero_async(self, method_name: str, *args, **kwargs): remote_call = getattr(self._workers[0], method_name) @@ -375,18 +467,28 @@ def execute_all_async(self, method_name: str, *args, **kwargs): # then we will send each element in the list to the corresponding 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()): + 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)) + result.append( + remote_call.remote( + *sliced_args, + **sliced_kwargs)) return result - return [getattr(worker, method_name).remote(*args, **kwargs) for worker in self._workers] + return [ + getattr(worker, method_name).remote(*args, **kwargs) + for worker in self._workers + ] @property def master_address(self): @@ -419,9 +521,12 @@ def _bind_workers_method_to_parent(cls, key, user_defined_cls): 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" + assert callable( + method + ), f"{method_name} in {user_defined_cls} is not callable" except Exception: - # if it is a property, it will fail because Class doesn't have instance property + # if it is a property, it will fail because Class doesn't have + # instance property continue if hasattr(method, MAGIC_ATTR): @@ -429,7 +534,9 @@ def _bind_workers_method_to_parent(cls, key, user_defined_cls): def generate_function(name): def func(self, *args, **kwargs): # dispatch to the actual worker - return getattr(self.worker_dict[key], name)(*args, **kwargs) + return getattr( + self.worker_dict[key], name)( + *args, **kwargs) return func @@ -462,9 +569,9 @@ def create_colocated_worker_cls(class_dict: dict[str, RayClassWithInitArgs]): if worker_cls is 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" - ) + 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} @@ -480,7 +587,8 @@ def __init__(self): # 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", {}) + *init_args_dict[key].get("args", ()), + **init_args_dict[key].get("kwargs", {}), ) # now monkey-patch the methods from inner class to WorkerDict diff --git a/Agent0/curriculum_train/verl/trainer/__init__.py b/Agent0/curriculum_train/verl/trainer/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/curriculum_train/verl/trainer/__init__.py +++ b/Agent0/curriculum_train/verl/trainer/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/trainer/config.py b/Agent0/curriculum_train/verl/trainer/config.py index ef2852d..6995b87 100644 --- a/Agent0/curriculum_train/verl/trainer/config.py +++ b/Agent0/curriculum_train/verl/trainer/config.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -52,7 +52,8 @@ class DataConfig: def post_init(self): if self.format_prompt is not None: - if os.path.exists(self.format_prompt): # ray job uses absolute path + if os.path.exists( + self.format_prompt): # ray job uses absolute path self.format_prompt = os.path.abspath(self.format_prompt) else: self.format_prompt = None @@ -72,6 +73,7 @@ class AlgorithmConfig: kl_target: float = 0.0 mock_data: str = "" + @dataclass class TrainerConfig: total_epochs: int = 10 @@ -93,11 +95,16 @@ class TrainerConfig: def post_init(self): if self.save_checkpoint_path is None: - self.save_checkpoint_path = os.path.join("checkpoints", self.project_name, self.experiment_name) + self.save_checkpoint_path = os.path.join( + "checkpoints", self.project_name, self.experiment_name + ) - self.save_checkpoint_path = os.path.abspath(self.save_checkpoint_path) # ray job uses absolute path + self.save_checkpoint_path = os.path.abspath( + self.save_checkpoint_path + ) # ray job uses absolute path if self.load_checkpoint_path is not None: - self.load_checkpoint_path = os.path.abspath(self.load_checkpoint_path) + self.load_checkpoint_path = os.path.abspath( + self.load_checkpoint_path) @dataclass @@ -110,7 +117,9 @@ class PPOConfig: def post_init(self): self.worker.rollout.prompt_length = self.data.max_prompt_length self.worker.rollout.response_length = self.data.max_response_length - self.worker.rollout.trust_remote_code = self.worker.actor.model.trust_remote_code + self.worker.rollout.trust_remote_code = ( + self.worker.actor.model.trust_remote_code + ) self.worker.actor.disable_kl = self.algorithm.disable_kl self.worker.actor.use_kl_loss = self.algorithm.use_kl_loss self.worker.actor.kl_penalty = self.algorithm.kl_penalty diff --git a/Agent0/curriculum_train/verl/trainer/core_algos.py b/Agent0/curriculum_train/verl/trainer/core_algos.py index 86f9410..24d6764 100644 --- a/Agent0/curriculum_train/verl/trainer/core_algos.py +++ b/Agent0/curriculum_train/verl/trainer/core_algos.py @@ -1,5 +1,5 @@ # Copyright 2022 The HuggingFace Team -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -46,7 +46,8 @@ def update(self, current_kl: float, n_steps: int) -> None: class AdaptiveKLController(KLController): """Adaptive KL controller described in: https://arxiv.org/pdf/1909.08593.pdf - Copied from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/utils.py#L54""" + Copied from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/utils.py#L54 + """ def __init__(self, init_kl_coef: float, target_kl: float, horizon: float): self.kl_coef = init_kl_coef @@ -63,7 +64,8 @@ def update(self, current_kl: float, n_steps: int) -> None: class FixedKLController(KLController): """Fixed KL controller. - Copeid from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/utils.py#L72""" + Copeid from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/utils.py#L72 + """ def __init__(self, init_kl_coef: float): self.kl_coef = init_kl_coef @@ -77,7 +79,9 @@ def get_kl_controller(algorithm_config: "AlgorithmConfig") -> KLController: if algorithm_config.kl_type == "fixed": kl_ctrl = FixedKLController(init_kl_coef=algorithm_config.kl_coef) elif algorithm_config.kl_type == "adaptive": - assert algorithm_config.kl_horizon > 0, f"horizon must be larger than 0. Got {algorithm_config.kl_horizon}." + assert ( + algorithm_config.kl_horizon > 0 + ), f"horizon must be larger than 0. Got {algorithm_config.kl_horizon}." kl_ctrl = AdaptiveKLController( init_kl_coef=algorithm_config.kl_coef, target_kl=algorithm_config.kl_target, @@ -133,10 +137,14 @@ def compute_gae_advantage_return( return advantages, returns -# NOTE(sgm): this implementation only consider outcome supervision, where the reward is a scalar. +# NOTE(sgm): this implementation only consider outcome supervision, where +# the reward is a scalar. @torch.no_grad() def compute_grpo_outcome_advantage( - token_level_rewards: torch.Tensor, response_mask: torch.Tensor, index: torch.Tensor, eps: float = 1e-6 + token_level_rewards: torch.Tensor, + response_mask: torch.Tensor, + index: torch.Tensor, + eps: float = 1e-6, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Compute advantage for GRPO, operating only on Outcome reward @@ -251,7 +259,9 @@ def compute_reinforce_plus_plus_outcome_advantage( @torch.no_grad() def compute_remax_outcome_advantage( - token_level_rewards: torch.Tensor, reward_baselines: torch.Tensor, response_mask: torch.Tensor + token_level_rewards: torch.Tensor, + reward_baselines: torch.Tensor, + response_mask: torch.Tensor, ) -> Tuple[torch.Tensor, torch.Tensor]: """ Compute advantage for ReMax, operating only on Outcome reward @@ -333,18 +343,29 @@ def compute_policy_loss( # see: https://github.com/pytorch/pytorch/issues/10729 ratio = torch.exp(negative_approx_kl) clipped_ratio = torch.exp( - torch.clamp(negative_approx_kl, np.log(1.0 - clip_ratio_low), np.log(1.0 + clip_ratio_high)) + torch.clamp( + negative_approx_kl, + np.log(1.0 - clip_ratio_low), + np.log(1.0 + clip_ratio_high), + ) ) pg_loss = -advantages * ratio pg_loss2 = -advantages * clipped_ratio pg_loss3 = -advantages * clip_ratio_dual - clipped_pg_loss_higher = torch.max(pg_loss, pg_loss2) # clip if pg_loss < pg_loss2 + clipped_pg_loss_higher = torch.max( + pg_loss, pg_loss2) # clip if pg_loss < pg_loss2 pg_clipfrac_higher = (pg_loss < pg_loss2).float() - clipped_pg_loss_lower = torch.min(clipped_pg_loss_higher, pg_loss3) # clip if pg_loss > pg_loss3 and adv < 0 - final_pg_loss = torch.where(advantages < 0, clipped_pg_loss_lower, clipped_pg_loss_higher) - pg_clipfrac_lower = (clipped_pg_loss_higher > pg_loss3).float() * (advantages < 0).float() + clipped_pg_loss_lower = torch.min( + clipped_pg_loss_higher, pg_loss3 + ) # clip if pg_loss > pg_loss3 and adv < 0 + final_pg_loss = torch.where( + advantages < 0, clipped_pg_loss_lower, clipped_pg_loss_higher + ) + pg_clipfrac_lower = (clipped_pg_loss_higher > pg_loss3).float() * ( + advantages < 0 + ).float() final_pg_loss = VF.masked_mean(final_pg_loss, response_mask) pg_clipfrac_higher = VF.masked_mean(pg_clipfrac_higher, response_mask) @@ -383,15 +404,22 @@ def compute_value_loss( The ratio of vf being clipped """ - vpredclipped = torch.clamp(vpreds, values - cliprange_value, values + cliprange_value) + vpredclipped = torch.clamp( + vpreds, values - cliprange_value, values + cliprange_value + ) vf_loss1 = torch.square(vpreds - returns) vf_loss2 = torch.square(vpredclipped - returns) - vf_loss = 0.5 * VF.masked_mean(torch.max(vf_loss1, vf_loss2), action_mask) # clip if vf_loss1 < vf_loss2 + vf_loss = 0.5 * VF.masked_mean( + torch.max(vf_loss1, vf_loss2), action_mask + ) # clip if vf_loss1 < vf_loss2 vf_clipfrac = VF.masked_mean((vf_loss1 < vf_loss2).float(), action_mask) return vf_loss, vf_clipfrac -def compute_kl(log_probs: torch.FloatTensor, ref_log_probs: torch.FloatTensor, kl_penalty: str) -> torch.Tensor: +def compute_kl( + log_probs: torch.FloatTensor, + ref_log_probs: torch.FloatTensor, + kl_penalty: str) -> torch.Tensor: """Compute KL divergence given log_probs and ref_log_probs. Adapted from https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L1150 @@ -423,6 +451,8 @@ def compute_kl(log_probs: torch.FloatTensor, ref_log_probs: torch.FloatTensor, k return torch.clamp(kld, min=-10, max=10) if kl_penalty == "full": - return F.kl_div(ref_log_probs, log_probs, log_target=True, reduction="none").sum(-1) + return F.kl_div( + ref_log_probs, log_probs, log_target=True, reduction="none" + ).sum(-1) raise NotImplementedError(f"Unknown KL penalty: {kl_penalty}.") diff --git a/Agent0/curriculum_train/verl/trainer/data_loader.py b/Agent0/curriculum_train/verl/trainer/data_loader.py index cb6881b..5b3a486 100644 --- a/Agent0/curriculum_train/verl/trainer/data_loader.py +++ b/Agent0/curriculum_train/verl/trainer/data_loader.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,7 +23,11 @@ from .config import DataConfig -def create_dataloader(config: DataConfig, tokenizer: PreTrainedTokenizer, processor: Optional[ProcessorMixin]) -> None: +def create_dataloader( + config: DataConfig, + tokenizer: PreTrainedTokenizer, + processor: Optional[ProcessorMixin], +) -> None: train_dataset = RLHFDataset( data_path=config.train_files, tokenizer=tokenizer, @@ -42,7 +46,9 @@ def create_dataloader(config: DataConfig, tokenizer: PreTrainedTokenizer, proces if config.shuffle: train_dataloader_generator = torch.Generator() train_dataloader_generator.manual_seed(config.seed) - sampler = RandomSampler(data_source=train_dataset, generator=train_dataloader_generator) + sampler = RandomSampler( + data_source=train_dataset, generator=train_dataloader_generator + ) else: sampler = SequentialSampler(data_source=train_dataset) @@ -72,7 +78,9 @@ def create_dataloader(config: DataConfig, tokenizer: PreTrainedTokenizer, proces ) val_dataloader = StatefulDataLoader( dataset=val_dataset, - batch_size=len(val_dataset) if config.val_batch_size == -1 else config.val_batch_size, + batch_size=( + len(val_dataset) if config.val_batch_size == - + 1 else config.val_batch_size), shuffle=False, num_workers=8, collate_fn=collate_fn, diff --git a/Agent0/curriculum_train/verl/trainer/main.py b/Agent0/curriculum_train/verl/trainer/main.py index 2c552bd..467e20d 100644 --- a/Agent0/curriculum_train/verl/trainer/main.py +++ b/Agent0/curriculum_train/verl/trainer/main.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -58,27 +58,38 @@ def run(self, config: PPOConfig): } global_pool_id = "global_pool" resource_pool_spec = { - global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + 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, } - resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, mapping=mapping + ) if config.worker.reward.reward_type == "sequential": RewardManager = SequentialFunctionRewardManager elif config.worker.reward.reward_type == "batch": RewardManager = BatchFunctionRewardManager else: - raise NotImplementedError(f"Unknown reward type {config.worker.reward.reward_type}.") + raise NotImplementedError( + f"Unknown reward type {config.worker.reward.reward_type}." + ) - RemoteRewardManager = ray.remote(RewardManager).options(num_cpus=config.worker.reward.num_cpus) + RemoteRewardManager = ray.remote(RewardManager).options( + num_cpus=config.worker.reward.num_cpus + ) reward_fn = RemoteRewardManager.remote(config.worker.reward, tokenizer) - val_reward_fn = RemoteRewardManager.remote(config.worker.reward, tokenizer) + val_reward_fn = RemoteRewardManager.remote( + config.worker.reward, tokenizer) - train_dataloader, val_dataloader = create_dataloader(config.data, tokenizer, processor) + train_dataloader, val_dataloader = create_dataloader( + config.data, tokenizer, processor + ) trainer = RayPPOTrainer( config=config, @@ -99,10 +110,10 @@ def run(self, config: PPOConfig): def main(): cli_args = OmegaConf.from_cli() default_config = OmegaConf.structured(PPOConfig()) - with open('tokens.json', 'r') as f: + with open("tokens.json", "r") as f: tokens = json.load(f) - os.environ['HF_TOKEN'] = tokens['huggingface'] - os.environ['WANDB_API_KEY'] = tokens['wandb'] + os.environ["HF_TOKEN"] = tokens["huggingface"] + os.environ["WANDB_API_KEY"] = tokens["wandb"] if hasattr(cli_args, "config"): config_path = cli_args.pop("config", None) file_config = OmegaConf.load(config_path) @@ -123,7 +134,7 @@ def main(): "PYTHONUNBUFFERED": "1", } } - ray.init(runtime_env=runtime_env,num_cpus=16) + ray.init(runtime_env=runtime_env, num_cpus=16) runner = Runner.remote() ray.get(runner.run.remote(ppo_config)) diff --git a/Agent0/curriculum_train/verl/trainer/metrics.py b/Agent0/curriculum_train/verl/trainer/metrics.py index 02cd233..309b457 100644 --- a/Agent0/curriculum_train/verl/trainer/metrics.py +++ b/Agent0/curriculum_train/verl/trainer/metrics.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,7 +24,8 @@ def reduce_metrics(metrics: Dict[str, List[Any]]) -> Dict[str, Any]: return {key: np.mean(value) for key, value in metrics.items()} -def compute_data_metrics(batch: DataProto, use_critic: bool = False) -> Dict[str, Any]: +def compute_data_metrics( + batch: DataProto, use_critic: bool = False) -> Dict[str, Any]: sequence_score = batch.batch["token_level_scores"].sum(-1) sequence_reward = batch.batch["token_level_rewards"].sum(-1) @@ -33,8 +34,10 @@ def compute_data_metrics(batch: DataProto, use_critic: bool = False) -> Dict[str max_response_length = batch.batch["responses"].size(-1) - prompt_mask = batch.batch["attention_mask"][:, :-max_response_length].bool() - response_mask = batch.batch["attention_mask"][:, -max_response_length:].bool() + 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) prompt_length = prompt_mask.sum(-1).float() @@ -73,7 +76,9 @@ def compute_data_metrics(batch: DataProto, use_critic: bool = False) -> Dict[str "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(), + "critic/vf_explained_var": (1.0 - return_diff_var / (return_var + 1e-5)) + .detach() + .item(), } if use_critic else {} @@ -82,35 +87,50 @@ def compute_data_metrics(batch: DataProto, use_critic: bool = False) -> Dict[str "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()) + "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(), + "prompt_length/clip_ratio": torch.mean( + torch.eq(prompt_length, max_prompt_length).float() + ) + .detach() + .item(), } return metrics -def compute_timing_metrics(batch: DataProto, timing_raw: Dict[str, float]) -> Dict[str, Any]: +def compute_timing_metrics( + batch: DataProto, timing_raw: Dict[str, float] +) -> Dict[str, Any]: num_response_tokens = torch.sum(batch.batch["response_mask"]).item() num_overall_tokens = sum(batch.meta_info["global_token_num"]) num_tokens_of_section = { **dict.fromkeys(["gen", "reward"], num_response_tokens), - **dict.fromkeys(["ref", "old", "values", "adv", "update_critic", "update_actor"], num_overall_tokens), + **dict.fromkeys( + ["ref", "old", "values", "adv", "update_critic", "update_actor"], + num_overall_tokens, + ), } 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] + 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()) }, } -def compute_throughout_metrics(batch: DataProto, timing_raw: Dict[str, float], num_gpus: int) -> Dict[str, Any]: +def compute_throughout_metrics( + batch: DataProto, timing_raw: Dict[str, float], num_gpus: int +) -> Dict[str, Any]: total_num_tokens = sum(batch.meta_info["global_token_num"]) time = timing_raw["step"] return { diff --git a/Agent0/curriculum_train/verl/trainer/ray_trainer.py b/Agent0/curriculum_train/verl/trainer/ray_trainer.py index 0ba89d3..bc47eda 100644 --- a/Agent0/curriculum_train/verl/trainer/ray_trainer.py +++ b/Agent0/curriculum_train/verl/trainer/ray_trainer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,18 +33,30 @@ from ..protocol import DataProto, pad_dataproto_to_divisor, unpad_dataproto from ..single_controller.base import Worker -from ..single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from ..single_controller.ray import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) from ..single_controller.ray.base import create_colocated_worker_cls from ..utils import torch_functional as VF from ..utils.checkpoint import CHECKPOINT_TRACKER, remove_obsolete_ckpt from ..utils.logger import Tracker from ..utils.py_functional import convert_dict_to_str, timer -from ..utils.seqlen_balancing import get_seqlen_balanced_partitions, log_seqlen_unbalance +from ..utils.seqlen_balancing import ( + get_seqlen_balanced_partitions, + log_seqlen_unbalance, +) from ..workers.fsdp_workers import FSDPWorker from ..workers.reward import FunctionRewardManager from . import core_algos from .config import PPOConfig -from .metrics import compute_data_metrics, compute_throughout_metrics, compute_timing_metrics, reduce_metrics +from .metrics import ( + compute_data_metrics, + compute_throughout_metrics, + compute_timing_metrics, + reduce_metrics, +) class Role(IntEnum): @@ -81,15 +93,20 @@ class ResourcePoolManager: resource_pool_spec: dict[str, list[int]] mapping: dict[Role, str] - resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict) + 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 + # 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 + 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 @@ -101,37 +118,58 @@ def get_resource_pool(self, role: Role) -> RayResourcePool: def get_num_gpus(self) -> int: """Get the number of gpus in this cluster.""" - return sum([n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes]) + return sum( + [ + n_gpus + for process_on_nodes in self.resource_pool_spec.values() + for n_gpus in process_on_nodes + ] + ) def _check_resource_available(self): """Check if the resource pool can be satisfied in this ray cluster.""" gpus_available = ray.available_resources().get("GPU", 0) gpus_required = self.get_num_gpus() if gpus_available < gpus_required: - raise ValueError(f"Total available GPUs {gpus_available} is less than total desired GPUs {gpus_required}.") + raise ValueError( + f"Total available GPUs {gpus_available} is less than total desired GPUs {gpus_required}.") -def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.KLController, kl_penalty="kl"): +def apply_kl_penalty( + data: DataProto, kl_ctrl: core_algos.KLController, kl_penalty="kl" +): token_level_scores = data.batch["token_level_scores"] batch_size = data.batch.batch_size[0] response_mask = data.batch["response_mask"] # compute kl between ref_policy and current policy - kld = core_algos.compute_kl(data.batch["old_log_probs"], data.batch["ref_log_probs"], kl_penalty=kl_penalty) + kld = core_algos.compute_kl( + data.batch["old_log_probs"], + data.batch["ref_log_probs"], + kl_penalty=kl_penalty) kld = kld * response_mask # (batch_size, response_length) - data.batch["token_level_rewards"] = token_level_scores - kl_ctrl.kl_coef * kld + data.batch["token_level_rewards"] = token_level_scores - \ + kl_ctrl.kl_coef * kld - current_kl = VF.masked_mean(kld, mask=response_mask, dim=-1) # average over sequence + current_kl = VF.masked_mean( + kld, mask=response_mask, dim=-1 + ) # average over sequence current_kl = torch.mean(current_kl, dim=0).item() metrics = {"critic/kl": current_kl, "critic/kl_coef": kl_ctrl.kl_coef} - # According to https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L880 + # According to + # https://github.com/huggingface/trl/blob/v0.11.0/trl/trainer/ppo_trainer.py#L880 kl_ctrl.update(current_kl=current_kl, n_steps=batch_size) return data, metrics -def compute_advantage(data: DataProto, adv_estimator: AdvantageEstimator, gamma: float = 1.0, lam: float = 1.0): +def compute_advantage( + data: DataProto, + adv_estimator: AdvantageEstimator, + gamma: float = 1.0, + lam: float = 1.0, +): token_level_rewards = data.batch["token_level_rewards"] response_mask = data.batch["response_mask"] index = data.non_tensor_batch["uid"] @@ -141,18 +179,21 @@ def compute_advantage(data: DataProto, adv_estimator: AdvantageEstimator, gamma: token_level_rewards, values, response_mask, gamma, lam ) elif adv_estimator == AdvantageEstimator.GRPO: - advantages, returns = core_algos.compute_grpo_outcome_advantage(token_level_rewards, response_mask, index) + advantages, returns = core_algos.compute_grpo_outcome_advantage( + token_level_rewards, response_mask, index + ) elif adv_estimator == AdvantageEstimator.REINFORCE_PLUS_PLUS: advantages, returns = core_algos.compute_reinforce_plus_plus_outcome_advantage( - token_level_rewards, response_mask, gamma - ) + token_level_rewards, response_mask, gamma) elif adv_estimator == AdvantageEstimator.REMAX: reward_baselines = data.batch["reward_baselines"] advantages, returns = core_algos.compute_remax_outcome_advantage( token_level_rewards, reward_baselines, response_mask ) elif adv_estimator == AdvantageEstimator.RLOO: - advantages, returns = core_algos.compute_rloo_outcome_advantage(token_level_rewards, response_mask, index) + advantages, returns = core_algos.compute_rloo_outcome_advantage( + token_level_rewards, response_mask, index + ) else: raise NotImplementedError @@ -189,9 +230,9 @@ def __init__( self.hybrid_engine = config.worker.hybrid_engine if self.hybrid_engine: - assert Role.ActorRollout in role_worker_mapping, ( - f"ActorRollout should be included in {role_worker_mapping.keys()}." - ) + assert ( + Role.ActorRollout in role_worker_mapping + ), f"ActorRollout should be included in {role_worker_mapping.keys()}." else: raise NotImplementedError @@ -207,7 +248,9 @@ def __init__( else: self.use_reference_policy = False self.kl_ctrl = core_algos.FixedKLController(init_kl_coef=0.0) - print("KL is disabled, no KL metrics will be logged. Please set `kl_coef=0` to log KL metrics.") + print( + "KL is disabled, no KL metrics will be logged. Please set `kl_coef=0` to log KL metrics." + ) if config.algorithm.adv_estimator == AdvantageEstimator.GAE: self.use_critic = True @@ -215,10 +258,14 @@ def __init__( self.use_critic = False if config.algorithm.adv_estimator not in list(AdvantageEstimator): - raise NotImplementedError(f"Unknown advantage estimator: {config.algorithm.adv_estimator}.") + raise NotImplementedError( + f"Unknown advantage estimator: { + config.algorithm.adv_estimator}.") if config.data.rollout_batch_size % config.worker.actor.global_batch_size != 0: - raise ValueError("Rollout batch size must be divisible by actor global batch size.") + raise ValueError( + "Rollout batch size must be divisible by actor global batch size." + ) if ( config.data.rollout_batch_size * config.worker.rollout.n @@ -228,8 +275,11 @@ def __init__( ) if self.use_critic: - if config.data.rollout_batch_size % config.worker.critic.global_batch_size != 0: - raise ValueError("Rollout batch size must be divisible by critic global batch size.") + if (config.data.rollout_batch_size % + config.worker.critic.global_batch_size != 0): + raise ValueError( + "Rollout batch size must be divisible by critic global batch size." + ) if ( config.data.rollout_batch_size * config.worker.rollout.n @@ -239,22 +289,30 @@ def __init__( ) if ( - config.algorithm.adv_estimator in (AdvantageEstimator.GRPO, AdvantageEstimator.RLOO) + config.algorithm.adv_estimator + in (AdvantageEstimator.GRPO, AdvantageEstimator.RLOO) and config.worker.rollout.n == 1 ): - raise ValueError("GRPO and RLOO algorithm need `config.worker.rollout.n > 1`.") + raise ValueError( + "GRPO and RLOO algorithm need `config.worker.rollout.n > 1`." + ) if config.trainer.max_steps is not None: self.training_steps = config.trainer.max_steps else: - self.training_steps = len(train_dataloader) * config.trainer.total_epochs + self.training_steps = len( + train_dataloader) * config.trainer.total_epochs config.worker.actor.optim.training_steps = self.training_steps config.worker.critic.optim.training_steps = self.training_steps print(f"Total training steps: {self.training_steps}") def _maybe_log_val_generations( - self, inputs: List[str], outputs: List[str], labels: List[str], scores: List[float] + self, + inputs: List[str], + outputs: List[str], + labels: List[str], + scores: List[float], ) -> None: """Log a table of validation samples""" if self.config.trainer.val_generations_to_log <= 0: @@ -280,13 +338,21 @@ def _validate(self) -> Dict[str, Any]: test_batch = DataProto.from_single_dict(batch_dict) # Store original inputs input_ids = test_batch.batch["input_ids"] - input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids] + input_texts = [ + self.tokenizer.decode(ids, skip_special_tokens=True) + for ids in input_ids + ] sample_inputs.extend(input_texts) if "multi_modal_data" in test_batch.non_tensor_batch.keys(): test_gen_batch = test_batch.pop( - batch_keys=["input_ids", "attention_mask", "position_ids"], - non_tensor_batch_keys=["raw_prompt_ids", "multi_modal_data"], + batch_keys=[ + "input_ids", + "attention_mask", + "position_ids"], + non_tensor_batch_keys=[ + "raw_prompt_ids", + "multi_modal_data"], ) else: test_gen_batch = test_batch.pop( @@ -295,23 +361,37 @@ def _validate(self) -> Dict[str, Any]: ) test_gen_batch.meta_info = self.config.worker.rollout.val_override_config - test_gen_batch.meta_info.update({ - "min_pixels": self.config.data.min_pixels, - "max_pixels": self.config.data.max_pixels, - }) - test_gen_batch, pad_size = pad_dataproto_to_divisor(test_gen_batch, self.actor_rollout_wg.world_size) - test_output_gen_batch = self.actor_rollout_wg.generate_sequences(test_gen_batch) - test_output_gen_batch = unpad_dataproto(test_output_gen_batch, pad_size=pad_size) + test_gen_batch.meta_info.update( + { + "min_pixels": self.config.data.min_pixels, + "max_pixels": self.config.data.max_pixels, + } + ) + test_gen_batch, pad_size = pad_dataproto_to_divisor( + test_gen_batch, self.actor_rollout_wg.world_size + ) + test_output_gen_batch = self.actor_rollout_wg.generate_sequences( + test_gen_batch + ) + test_output_gen_batch = unpad_dataproto( + test_output_gen_batch, pad_size=pad_size + ) # Store generated outputs output_ids = test_output_gen_batch.batch["responses"] - output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids] + output_texts = [ + self.tokenizer.decode(ids, skip_special_tokens=True) + for ids in output_ids + ] sample_outputs.extend(output_texts) - sample_labels.extend(test_batch.non_tensor_batch["ground_truth"].tolist()) + sample_labels.extend( + test_batch.non_tensor_batch["ground_truth"].tolist()) test_batch = test_batch.union(test_output_gen_batch) # evaluate using reward_function - reward_tensor, reward_metrics = ray.get(self.val_reward_fn.compute_reward.remote(test_batch)) + reward_tensor, reward_metrics = ray.get( + self.val_reward_fn.compute_reward.remote(test_batch) + ) # Store scores scores = reward_tensor.sum(-1).cpu().tolist() @@ -321,63 +401,92 @@ def _validate(self) -> Dict[str, Any]: for key, value in reward_metrics.items(): reward_metrics_lst[key].extend(value) - self._maybe_log_val_generations(sample_inputs, sample_outputs, sample_labels, sample_scores) - reward_score = torch.cat(reward_tensor_lst, dim=0).sum(-1).mean().item() - val_reward_metrics = {f"val/{key}_reward": value for key, value in reduce_metrics(reward_metrics_lst).items()} + self._maybe_log_val_generations( + sample_inputs, sample_outputs, sample_labels, sample_scores + ) + reward_score = torch.cat( + reward_tensor_lst, dim=0).sum(-1).mean().item() + val_reward_metrics = { + f"val/{key}_reward": value + for key, value in reduce_metrics(reward_metrics_lst).items() + } return {"val/reward_score": reward_score, **val_reward_metrics} def init_workers(self) -> None: """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()} + 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) + 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.worker, role="actor_rollout" + cls=self.role_worker_mapping[Role.ActorRollout], + config=self.config.worker, + role="actor_rollout", ) - self.resource_pool_to_cls[resource_pool]["actor_rollout"] = actor_rollout_cls + self.resource_pool_to_cls[resource_pool][ + "actor_rollout" + ] = actor_rollout_cls else: raise NotImplementedError # create critic if self.use_critic: - resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic) + resource_pool = self.resource_pool_manager.get_resource_pool( + Role.Critic) critic_cls = RayClassWithInitArgs( - cls=self.role_worker_mapping[Role.Critic], config=self.config.worker, role="critic" + cls=self.role_worker_mapping[Role.Critic], + config=self.config.worker, + role="critic", ) self.resource_pool_to_cls[resource_pool]["critic"] = critic_cls # create reference policy if needed if self.use_reference_policy: - resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy) + resource_pool = self.resource_pool_manager.get_resource_pool( + Role.RefPolicy) ref_policy_cls = RayClassWithInitArgs( - self.role_worker_mapping[Role.RefPolicy], config=self.config.worker, role="ref" + self.role_worker_mapping[Role.RefPolicy], + config=self.config.worker, + 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_reward_model: # we create a RM here - resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) + resource_pool = self.resource_pool_manager.get_resource_pool( + Role.RewardModel + ) rm_cls = RayClassWithInitArgs( - cls=self.role_worker_mapping[Role.RewardModel], config=self.config.worker, role="reward" + cls=self.role_worker_mapping[Role.RewardModel], + config=self.config.worker, + role="reward", ) 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. + # See + # https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb + # for more information. all_wg: Dict[str, FSDPWorker] = {} 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) + 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 + # 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: @@ -392,16 +501,22 @@ def init_workers(self) -> None: 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 + # 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() def _save_checkpoint(self) -> None: # path: {save_checkpoint_path}/global_step_{global_step}/{actor,critic} remove_obsolete_ckpt( - self.config.trainer.save_checkpoint_path, self.global_step, self.config.trainer.save_limit + self.config.trainer.save_checkpoint_path, + self.global_step, + self.config.trainer.save_limit, ) - folder_path = os.path.join(self.config.trainer.save_checkpoint_path, f"global_step_{self.global_step}") + folder_path = os.path.join( + self.config.trainer.save_checkpoint_path, + f"global_step_{ + self.global_step}") actor_path = os.path.join(folder_path, "actor") self.actor_rollout_wg.save_checkpoint(actor_path) @@ -413,7 +528,9 @@ def _save_checkpoint(self) -> None: dataloader_state_dict = self.train_dataloader.state_dict() torch.save(dataloader_state_dict, dataloader_path) - last_global_step_path = os.path.join(self.config.trainer.save_checkpoint_path, CHECKPOINT_TRACKER) + last_global_step_path = os.path.join( + self.config.trainer.save_checkpoint_path, CHECKPOINT_TRACKER + ) with open(last_global_step_path, "w") as f: f.write(str(self.global_step)) @@ -421,38 +538,65 @@ def _load_checkpoint(self) -> None: if self.config.trainer.load_checkpoint_path is None: return - if "global_step_" not in self.config.trainer.load_checkpoint_path.strip(os.path.sep).split(os.path.sep)[-1]: - raise ValueError("`load_checkpoint_path` should end with `global_step_*`.") - - print(f"Load from checkpoint: {self.config.trainer.load_checkpoint_path}.") - self.global_step = int(self.config.trainer.load_checkpoint_path.strip(os.path.sep).split("global_step_")[-1]) - actor_path = os.path.join(self.config.trainer.load_checkpoint_path, "actor") + if ("global_step_" not in self.config.trainer.load_checkpoint_path.strip( + os.path.sep).split(os.path.sep)[-1]): + raise ValueError( + "`load_checkpoint_path` should end with `global_step_*`.") + + print( + f"Load from checkpoint: { + self.config.trainer.load_checkpoint_path}.") + self.global_step = int( + self.config.trainer.load_checkpoint_path.strip(os.path.sep).split( + "global_step_" + )[-1] + ) + actor_path = os.path.join( + self.config.trainer.load_checkpoint_path, "actor") self.actor_rollout_wg.load_checkpoint(actor_path) if self.use_critic: - critic_path = os.path.join(self.config.trainer.load_checkpoint_path, "critic") + critic_path = os.path.join( + self.config.trainer.load_checkpoint_path, "critic" + ) self.critic_wg.load_checkpoint(critic_path) - dataloader_path = os.path.join(self.config.trainer.load_checkpoint_path, "dataloader.pt") + dataloader_path = os.path.join( + self.config.trainer.load_checkpoint_path, "dataloader.pt" + ) if os.path.exists(dataloader_path): - dataloader_state_dict = torch.load(dataloader_path, weights_only=False) + dataloader_state_dict = torch.load( + dataloader_path, weights_only=False) self.train_dataloader.load_state_dict(dataloader_state_dict) else: - print(f"No dataloader state found at {dataloader_path}, will start from scratch.") + print( + f"No dataloader state found at {dataloader_path}, will start from scratch.") - def _balance_batch(self, batch: DataProto, metrics: Dict[str, Any], logging_prefix: str = "global_seqlen") -> None: + def _balance_batch( + self, + batch: DataProto, + metrics: Dict[str, Any], + logging_prefix: str = "global_seqlen", + ) -> None: """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 = batch.batch["attention_mask"].view(batch_size, -1).sum(-1).tolist() # (train_batch_size,) + global_seqlen_lst = ( + batch.batch["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]) + # 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 + seqlen_list=global_seqlen_lst, + partitions=global_partition_lst, + prefix=logging_prefix, ) metrics.update(global_balance_stats) @@ -462,7 +606,9 @@ def fit(self): 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. """ - self.logger = Tracker(loggers=self.config.trainer.logger, config=self.config.to_dict()) + self.logger = Tracker( + loggers=self.config.trainer.logger, config=self.config.to_dict() + ) self.global_step = 0 val_metrics: Optional[Dict[str, Any]] = None @@ -477,8 +623,12 @@ def fit(self): if self.config.trainer.val_only: return - for _ in tqdm(range(self.config.trainer.total_epochs), desc="Epoch", position=0): - for batch_dict in tqdm(self.train_dataloader, desc="Running step", position=1): + for _ in tqdm( + range(self.config.trainer.total_epochs), desc="Epoch", position=0 + ): + for batch_dict in tqdm( + self.train_dataloader, desc="Running step", position=1 + ): self.global_step += 1 if self.global_step > self.training_steps: break @@ -489,67 +639,96 @@ def fit(self): # pop those keys for generation if "multi_modal_data" in batch.non_tensor_batch.keys(): gen_batch = batch.pop( - batch_keys=["input_ids", "attention_mask", "position_ids"], - non_tensor_batch_keys=["raw_prompt_ids", "multi_modal_data"], + batch_keys=[ + "input_ids", + "attention_mask", + "position_ids"], + non_tensor_batch_keys=[ + "raw_prompt_ids", + "multi_modal_data"], + ) + gen_batch.meta_info.update( + { + "min_pixels": self.config.data.min_pixels, + "max_pixels": self.config.data.max_pixels, + } ) - gen_batch.meta_info.update({ - "min_pixels": self.config.data.min_pixels, - "max_pixels": self.config.data.max_pixels, - }) else: gen_batch = batch.pop( - batch_keys=["input_ids", "attention_mask", "position_ids"], + batch_keys=[ + "input_ids", + "attention_mask", + "position_ids"], non_tensor_batch_keys=["raw_prompt_ids"], ) with timer("step", timing_raw): # generate a batch with timer("gen", timing_raw): # wg: worker group - gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) + gen_batch_output = self.actor_rollout_wg.generate_sequences( + gen_batch) if self.config.algorithm.adv_estimator == "remax": with timer("gen_max", timing_raw): gen_baseline_batch = deepcopy(gen_batch) gen_baseline_batch.meta_info["temperature"] = 0 gen_baseline_batch.meta_info["n"] = 1 - gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch) + gen_baseline_output = ( + self.actor_rollout_wg.generate_sequences( + gen_baseline_batch + ) + ) batch = batch.union(gen_baseline_output) - reward_baseline_tensor, _ = ray.get(self.reward_fn.compute_reward.remote(batch)) - reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1) - - batch.pop(batch_keys=list(gen_baseline_output.batch.keys())) + reward_baseline_tensor, _ = ray.get( + self.reward_fn.compute_reward.remote(batch) + ) + reward_baseline_tensor = reward_baseline_tensor.sum( + dim=-1) + + batch.pop( + batch_keys=list( + gen_baseline_output.batch.keys())) batch.batch["reward_baselines"] = reward_baseline_tensor del gen_baseline_batch, gen_baseline_output batch.non_tensor_batch["uid"] = np.array( - [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object + [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.worker.rollout.n, interleave=True) + batch = batch.repeat( + repeat_times=self.config.worker.rollout.n, + interleave=True) batch = batch.union(gen_batch_output) # 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 + # 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_token_num"] = torch.sum( + batch.batch["attention_mask"], dim=-1 + ).tolist() # compute reward with timer("reward", timing_raw): - reward_ref = self.reward_fn.compute_reward.remote(batch) + reward_ref = self.reward_fn.compute_reward.remote( + batch) # recompute old_log_probs with timer("old", timing_raw): - old_log_probs = self.actor_rollout_wg.compute_log_probs(batch) + old_log_probs = self.actor_rollout_wg.compute_log_probs( + batch) batch = batch.union(old_log_probs) # compute ref_log_probs if self.use_reference_policy: with timer("ref", timing_raw): - ref_log_probs = self.ref_policy_wg.compute_ref_log_probs(batch) + ref_log_probs = self.ref_policy_wg.compute_ref_log_probs( + batch) batch = batch.union(ref_log_probs) # compute values @@ -562,16 +741,25 @@ def fit(self): # get token level scores reward_tensor, reward_metrics = ray.get(reward_ref) batch.batch["token_level_scores"] = reward_tensor - reward_metrics = {f"reward/{k}": v for k, v in reduce_metrics(reward_metrics).items()} + reward_metrics = { + f"reward/{k}": v + for k, v in reduce_metrics(reward_metrics).items() + } metrics.update(reward_metrics) # apply kl penalty if available - if not self.config.algorithm.use_kl_loss and self.use_reference_policy: + if ( + not self.config.algorithm.use_kl_loss + and self.use_reference_policy + ): # apply kl penalty to reward - batch, kl_metrics = apply_kl_penalty(batch, self.kl_ctrl, self.config.algorithm.kl_penalty) + batch, kl_metrics = apply_kl_penalty( + batch, self.kl_ctrl, self.config.algorithm.kl_penalty) metrics.update(kl_metrics) else: - batch.batch["token_level_rewards"] = batch.batch["token_level_scores"] + batch.batch["token_level_rewards"] = batch.batch[ + "token_level_scores" + ] # compute advantages, executed on the driver process batch = compute_advantage( @@ -586,15 +774,18 @@ def fit(self): with timer("update_critic", timing_raw): critic_output = self.critic_wg.update_critic(batch) - critic_metrics = reduce_metrics(critic_output.non_tensor_batch) + critic_metrics = reduce_metrics( + critic_output.non_tensor_batch) metrics.update(critic_metrics) # update actor if self.config.trainer.critic_warmup <= self.global_step: with timer("update_actor", timing_raw): - actor_output = self.actor_rollout_wg.update_actor(batch) + actor_output = self.actor_rollout_wg.update_actor( + batch) - actor_metrics = reduce_metrics(actor_output.non_tensor_batch) + actor_metrics = reduce_metrics( + actor_output.non_tensor_batch) metrics.update(actor_metrics) # validate @@ -608,15 +799,25 @@ def fit(self): metrics.update(val_metrics) - if self.config.trainer.save_freq > 0 and self.global_step % self.config.trainer.save_freq == 0: + if (self.config.trainer.save_freq > 0 and self.global_step % + self.config.trainer.save_freq == 0): with timer("save_checkpoint", timing_raw): self._save_checkpoint() # collect metrics num_gpus = self.resource_pool_manager.get_num_gpus() - metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic)) - metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) - metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, num_gpus=num_gpus)) + metrics.update( + compute_data_metrics( + batch=batch, + use_critic=self.use_critic)) + metrics.update( + compute_timing_metrics(batch=batch, timing_raw=timing_raw) + ) + metrics.update( + compute_throughout_metrics( + batch=batch, timing_raw=timing_raw, num_gpus=num_gpus + ) + ) self.logger.log(data=metrics, step=self.global_step) @@ -630,7 +831,12 @@ def fit(self): val_metrics = self._validate() self.logger.log(data=val_metrics, step=self.global_step) - print(f"Final validation metrics: {convert_dict_to_str(val_metrics)}") + print( + f"Final validation metrics: { + convert_dict_to_str(val_metrics)}") - if self.config.trainer.save_freq <= 0 or self.global_step % self.config.trainer.save_freq != 0: + if ( + self.config.trainer.save_freq <= 0 + or self.global_step % self.config.trainer.save_freq != 0 + ): self._save_checkpoint() diff --git a/Agent0/curriculum_train/verl/utils/__init__.py b/Agent0/curriculum_train/verl/utils/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/curriculum_train/verl/utils/__init__.py +++ b/Agent0/curriculum_train/verl/utils/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/utils/checkpoint/__init__.py b/Agent0/curriculum_train/verl/utils/checkpoint/__init__.py index de1a2fc..4fc90fa 100644 --- a/Agent0/curriculum_train/verl/utils/checkpoint/__init__.py +++ b/Agent0/curriculum_train/verl/utils/checkpoint/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/utils/checkpoint/checkpoint_manager.py b/Agent0/curriculum_train/verl/utils/checkpoint/checkpoint_manager.py index 749b60c..6653e3a 100644 --- a/Agent0/curriculum_train/verl/utils/checkpoint/checkpoint_manager.py +++ b/Agent0/curriculum_train/verl/utils/checkpoint/checkpoint_manager.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -85,7 +85,9 @@ def local_mkdir(path: str) -> str: os.makedirs(path, exist_ok=True) except Exception as e: print(f"Warning: Failed to acquire lock for {path}: {e}") - os.makedirs(path, exist_ok=True) # even if the lock is not acquired, try to create the directory + os.makedirs( + path, exist_ok=True + ) # even if the lock is not acquired, try to create the directory return path @@ -107,7 +109,9 @@ def load_rng_state(rng_state: Dict[str, Any]): random.setstate(rng_state["random"]) -def find_latest_ckpt_path(path: Optional[str] = None, directory_format: str = "global_step_{}") -> Optional[str]: +def find_latest_ckpt_path( + path: Optional[str] = None, directory_format: str = "global_step_{}" +) -> Optional[str]: if path is None: return None @@ -135,7 +139,12 @@ def get_checkpoint_tracker_filename(root_path: str) -> str: return os.path.join(root_path, CHECKPOINT_TRACKER) -def remove_obsolete_ckpt(path: str, global_step: int, save_limit: int = -1, directory_format: str = "global_step_{}"): +def remove_obsolete_ckpt( + path: str, + global_step: int, + save_limit: int = -1, + directory_format: str = "global_step_{}", +): """ Remove the obsolete checkpoints that exceed the save_limit. """ @@ -154,7 +163,7 @@ def remove_obsolete_ckpt(path: str, global_step: int, save_limit: int = -1, dire ckpt_folders.append((step, folder)) ckpt_folders.sort(reverse=True) - for _, folder in ckpt_folders[save_limit - 1 :]: + for _, folder in ckpt_folders[save_limit - 1:]: folder_path = os.path.join(path, folder) shutil.rmtree(folder_path, ignore_errors=True) print(f"Removed obsolete checkpoint: {folder_path}") diff --git a/Agent0/curriculum_train/verl/utils/checkpoint/fsdp_checkpoint_manager.py b/Agent0/curriculum_train/verl/utils/checkpoint/fsdp_checkpoint_manager.py index 1318bfe..1eda3bf 100644 --- a/Agent0/curriculum_train/verl/utils/checkpoint/fsdp_checkpoint_manager.py +++ b/Agent0/curriculum_train/verl/utils/checkpoint/fsdp_checkpoint_manager.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,7 +17,11 @@ import torch import torch.distributed as dist -from torch.distributed.checkpoint.state_dict import StateDictOptions, get_state_dict, set_state_dict +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_state_dict, + set_state_dict, +) from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from transformers import PreTrainedModel, PreTrainedTokenizer, ProcessorMixin @@ -53,12 +57,24 @@ def load_checkpoint(self, path: Optional[str] = None): return # every rank download its own checkpoint - model_path = os.path.join(path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt") - optim_path = os.path.join(path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt") - extra_path = os.path.join(path, f"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt") - print(f"[rank-{self.rank}]: Loading model from {os.path.abspath(model_path)}.") - print(f"[rank-{self.rank}]: Loading optimizer from {os.path.abspath(optim_path)}.") - print(f"[rank-{self.rank}]: Loading extra_state from {os.path.abspath(extra_path)}.") + model_path = os.path.join( + path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt" + ) + optim_path = os.path.join( + path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt" + ) + extra_path = os.path.join( + path, f"extra_state_world_size_{ + self.world_size}_rank_{ + self.rank}.pt") + print( + f"[rank-{self.rank}]: Loading model from {os.path.abspath(model_path)}.") + print( + f"[rank-{self.rank}]: Loading optimizer from {os.path.abspath(optim_path)}." + ) + print( + f"[rank-{self.rank}]: Loading extra_state from {os.path.abspath(extra_path)}." + ) model_state_dict = torch.load(model_path, weights_only=False) optim_state_dict = torch.load(optim_path, weights_only=False) extra_state_dict = torch.load(extra_path, weights_only=False) @@ -83,18 +99,31 @@ def save_checkpoint(self, path: str): # every rank will save its own model and optim shard state_dict_options = StateDictOptions(cpu_offload=True) - model_state_dict, optim_state_dict = get_state_dict(self.model, self.optimizer, options=state_dict_options) + model_state_dict, optim_state_dict = get_state_dict( + self.model, self.optimizer, options=state_dict_options + ) extra_state_dict = { "lr_scheduler": self.lr_scheduler.state_dict(), "rng": self.get_rng_state(), } - model_path = os.path.join(path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt") - optim_path = os.path.join(path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt") - extra_path = os.path.join(path, f"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt") - - print(f"[rank-{self.rank}]: Saving model to {os.path.abspath(model_path)}.") - print(f"[rank-{self.rank}]: Saving optimizer to {os.path.abspath(optim_path)}.") - print(f"[rank-{self.rank}]: Saving extra_state to {os.path.abspath(extra_path)}.") + model_path = os.path.join( + path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt" + ) + optim_path = os.path.join( + path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt" + ) + extra_path = os.path.join( + path, f"extra_state_world_size_{ + self.world_size}_rank_{ + self.rank}.pt") + + print( + f"[rank-{self.rank}]: Saving model to {os.path.abspath(model_path)}.") + print( + f"[rank-{self.rank}]: Saving optimizer to {os.path.abspath(optim_path)}.") + print( + f"[rank-{self.rank}]: Saving extra_state to {os.path.abspath(extra_path)}." + ) torch.save(model_state_dict, model_path) torch.save(optim_state_dict, optim_path) torch.save(extra_state_dict, extra_path) @@ -107,7 +136,8 @@ def save_checkpoint(self, path: str): os.makedirs(hf_path, exist_ok=True) assert isinstance(self.model._fsdp_wrapped_module, PreTrainedModel) self.model._fsdp_wrapped_module.config.save_pretrained(hf_path) - self.model._fsdp_wrapped_module.generation_config.save_pretrained(hf_path) + self.model._fsdp_wrapped_module.generation_config.save_pretrained( + hf_path) self.processing_class.save_pretrained(hf_path) dist.barrier() diff --git a/Agent0/curriculum_train/verl/utils/code_executor.py b/Agent0/curriculum_train/verl/utils/code_executor.py index 82b9c67..d3758ed 100644 --- a/Agent0/curriculum_train/verl/utils/code_executor.py +++ b/Agent0/curriculum_train/verl/utils/code_executor.py @@ -2,7 +2,8 @@ import json import re -SANDBOX_API_URL = 'http://172.22.1.105:8080/run_code' +SANDBOX_API_URL = "http://172.22.1.105:8080/run_code" + def execute_code_in_sandbox(code: str) -> str: """ @@ -14,15 +15,12 @@ def execute_code_in_sandbox(code: str) -> str: Returns: ๆ‰ง่กŒ็ป“ๆžœ๏ผˆstdout๏ผ‰๏ผŒๅฆ‚ๆžœๅ‡บ้”™ๅˆ™่ฟ”ๅ›ž้”™่ฏฏไฟกๆฏใ€‚ """ - payload = { - "code": code, - "language": "python" - } - headers = { - 'Content-Type': 'application/json' - } - - response = requests.post(SANDBOX_API_URL, headers=headers, data=json.dumps(payload), timeout=10) + payload = {"code": code, "language": "python"} + headers = {"Content-Type": "application/json"} + + response = requests.post( + SANDBOX_API_URL, headers=headers, data=json.dumps(payload), timeout=10 + ) response.raise_for_status() result = response.json() @@ -32,18 +30,21 @@ def execute_code_in_sandbox(code: str) -> str: if run_info.get("status") == "Finished": return run_info.get("stdout", "") else: - return f"Execution failed with status: {run_info.get('status')}\nStderr: {run_info.get('stderr', '')}" + return f"Execution failed with status: { + run_info.get('status')}\nStderr: { + run_info.get( + 'stderr', '')}" else: return f"{result}" -if __name__ == '__main__': +if __name__ == "__main__": hello_world_code = 'print("Hello, world!")' print(f"Executing code:\n---\n{hello_world_code}\n---") output = execute_code_in_sandbox(hello_world_code) print(f"Result:\n---\n{output}\n---") - error_code = 'print(1 / 0)' + error_code = "print(1 / 0)" print(f"Executing code with error:\n---\n{error_code}\n---") output = execute_code_in_sandbox(error_code) - print(f"Result:\n---\n{output}\n---") \ No newline at end of file + print(f"Result:\n---\n{output}\n---") diff --git a/Agent0/curriculum_train/verl/utils/dataset.py b/Agent0/curriculum_train/verl/utils/dataset.py index 0002e53..c246e86 100644 --- a/Agent0/curriculum_train/verl/utils/dataset.py +++ b/Agent0/curriculum_train/verl/utils/dataset.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,6 +32,8 @@ import json import random + + def collate_fn(features: List[Dict[str, Any]]) -> Dict[str, Any]: tensors = defaultdict(list) non_tensors = defaultdict(list) @@ -51,8 +53,9 @@ def collate_fn(features: List[Dict[str, Any]]) -> Dict[str, Any]: return {**tensors, **non_tensors} - -def process_image(image: Union[Dict[str, Any], ImageObject, str], min_pixels: int, max_pixels: int) -> ImageObject: +def process_image( + image: Union[Dict[str, Any], ImageObject, str], min_pixels: int, max_pixels: int +) -> ImageObject: if isinstance(image, str): image = Image.open(image) elif isinstance(image, dict): @@ -62,12 +65,16 @@ def process_image(image: Union[Dict[str, Any], ImageObject, str], min_pixels: in if (image.width * image.height) > max_pixels: resize_factor = math.sqrt(max_pixels / (image.width * image.height)) - width, height = int(image.width * resize_factor), int(image.height * resize_factor) + width, height = int(image.width * resize_factor), int( + image.height * resize_factor + ) image = image.resize((width, height)) if (image.width * image.height) < min_pixels: resize_factor = math.sqrt(min_pixels / (image.width * image.height)) - width, height = int(image.width * resize_factor), int(image.height * resize_factor) + width, height = int(image.width * resize_factor), int( + image.height * resize_factor + ) image = image.resize((width, height)) if image.mode != "RGB": @@ -113,10 +120,13 @@ def __init__( data_split = "train" if os.path.isdir(data_path): - # when we use dataset builder, we should always refer to the train split - self.dataset = load_dataset("parquet", data_dir=data_path, split="train") + # when we use dataset builder, we should always refer to the train + # split + self.dataset = load_dataset( + "parquet", data_dir=data_path, split="train") elif os.path.isfile(data_path): - self.dataset = load_dataset("parquet", data_files=data_path, split="train") + self.dataset = load_dataset( + "parquet", data_files=data_path, split="train") else: # load remote dataset from huggingface hub self.dataset = load_dataset(data_path, split=data_split) @@ -128,11 +138,16 @@ def __init__( if "questioner_format_with_persona" in self.format_prompt: print("load personas") - personas_dataset = load_dataset("proj-persona/PersonaHub", "math", split="train") - self.personas = [item['input persona'] for item in personas_dataset] + personas_dataset = load_dataset( + "proj-persona/PersonaHub", "math", split="train" + ) + self.personas = [item["input persona"] + for item in personas_dataset] # self.personas = self.personas.select(range(100)) if self.filter_overlong_prompts: - self.dataset = self.dataset.filter(self._filter_overlong_prompts, desc="Filtering overlong prompts") + self.dataset = self.dataset.filter( + self._filter_overlong_prompts, + desc="Filtering overlong prompts") def _build_messages(self, example: Dict[str, Any]) -> List[Dict[str, Any]]: prompt_str: str = example[self.prompt_key] @@ -140,10 +155,10 @@ def _build_messages(self, example: Dict[str, Any]) -> List[Dict[str, Any]]: print("load personas") return [ { - "role": "system", - "content": ( - f"You are {random.choice(self.personas)}.\n" - "FIRST, in your private scratch-pad, think step-by-step to design a brand-new, non-trivial problem. " + "role": "system", "content": ( + f"You are { + random.choice( + self.personas)}.\n" "FIRST, in your private scratch-pad, think step-by-step to design a brand-new, non-trivial problem. " "The problem could come from any field of mathematics, including but not limited to algebra, geometry, number theory, combinatorics, prealgebra, probability, statistics, and calculus. " "Aim for a difficulty such that fewer than 30 % of advanced high-school students could solve it. " "Avoid re-using textbook clichรฉs or famous contest problems.\n" @@ -153,23 +168,15 @@ def _build_messages(self, example: Dict[str, Any]) -> List[Dict[str, Any]]: "\n\n" r"\boxed{final_answer}" "\n\n" - "Do NOT output anything elseโ€”no explanations, no extra markup." - ) - }, - { - "role": "user", - "content": ( - "Generate one new, challenging reasoning question now. " - "Remember to format the output exactly as instructed." - ) - } - ] + "Do NOT output anything elseโ€”no explanations, no extra markup."), }, { + "role": "user", "content": ( + "Generate one new, challenging reasoning question now. " + "Remember to format the output exactly as instructed."), }, ] if "questioner_format" in self.format_prompt: # print('detected questioner_format') return [ { - "role": "system", - "content": ( + "role": "system", "content": ( "You are an expert competition-math problem setter.\n" "FIRST, in your private scratch-pad, think step-by-step to design a brand-new, non-trivial problem. " "The problem could come from any field of mathematics, including but not limited to algebra, geometry, number theory, combinatorics, prealgebra, probability, statistics, and calculus. " @@ -181,32 +188,22 @@ def _build_messages(self, example: Dict[str, Any]) -> List[Dict[str, Any]]: "\n\n" r"\boxed{final_answer}" "\n\n" - "Do NOT output anything elseโ€”no explanations, no extra markup." - ) - }, - { - "role": "user", - "content": ( + "Do NOT output anything elseโ€”no explanations, no extra markup."), }, { + "role": "user", "content": ( "Generate one new, challenging reasoning question now. " - "Remember to format the output exactly as instructed." - ) - } - ] + "Remember to format the output exactly as instructed."), }, ] if "solver_format" in self.format_prompt: return [ { - "role": "system", - "content": r"Please reason step by step, and put your final answer within \boxed{}." + "role": "system", + "content": r"Please reason step by step, and put your final answer within \boxed{}.", }, - { - "role": "user", - "content": prompt_str - } - ] + {"role": "user", "content": prompt_str}, + ] if self.format_prompt: format_prompt = Template(self.format_prompt.strip()) prompt_str = format_prompt.render(content=prompt_str) - + if self.image_key in example: # https://huggingface.co/docs/transformers/en/tasks/image_text_to_text content_list = [] @@ -223,16 +220,29 @@ def _build_messages(self, example: Dict[str, Any]) -> List[Dict[str, Any]]: def _filter_overlong_prompts(self, example: Dict[str, Any]) -> bool: messages = self._build_messages(example) - processing_class = self.processor if self.processor is not None else self.tokenizer + processing_class = ( + self.processor if self.processor is not None else self.tokenizer + ) if self.tokenizer.chat_template: return ( - len(processing_class.apply_chat_template(messages, add_generation_prompt=True)) <= self.max_prompt_length + len( + processing_class.apply_chat_template( + messages, add_generation_prompt=True + ) + ) + <= self.max_prompt_length ) else: return ( - len("system: " + messages[0]["content"] + '\n' + "user: " + messages[1]["content"]) <= self.max_prompt_length + len( + "system: " + + messages[0]["content"] + + "\n" + + "user: " + + messages[1]["content"] + ) + <= self.max_prompt_length ) - def __len__(self): return len(self.dataset) @@ -242,26 +252,46 @@ def __getitem__(self, index): messages = self._build_messages(example) if self.image_key in example: - prompt = self.processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + prompt = self.processor.apply_chat_template( + messages, add_generation_prompt=True, tokenize=False + ) raw_image_data = example.pop(self.image_key) images = [ - process_image(image, min_pixels=self.min_pixels, max_pixels=self.max_pixels) + process_image( + image, min_pixels=self.min_pixels, max_pixels=self.max_pixels + ) for image in raw_image_data ] - model_inputs = self.processor(images, [prompt], add_special_tokens=False, return_tensors="pt") + model_inputs = self.processor( + images, [prompt], add_special_tokens=False, return_tensors="pt" + ) input_ids = model_inputs.pop("input_ids")[0] attention_mask = model_inputs.pop("attention_mask")[0] example["multi_modal_data"] = {"image": raw_image_data} else: if self.tokenizer.chat_template: - prompt = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + prompt = self.tokenizer.apply_chat_template( + messages, add_generation_prompt=True, tokenize=False + ) else: - prompt = "system: " + messages[0]["content"] + '\n' + "user: " + messages[1]["content"] - model_inputs = self.tokenizer([prompt], add_special_tokens=False, return_tensors="pt") + prompt = ( + "system: " + + messages[0]["content"] + + "\n" + + "user: " + + messages[1]["content"] + ) + model_inputs = self.tokenizer( + [prompt], add_special_tokens=False, return_tensors="pt" + ) input_ids = model_inputs.pop("input_ids")[0] attention_mask = model_inputs.pop("attention_mask")[0] - if self.processor is not None and self.processor.image_processor.__class__.__name__ == "Qwen2VLImageProcessor": + if ( + self.processor is not None + and self.processor.image_processor.__class__.__name__ + == "Qwen2VLImageProcessor" + ): # qwen2vl mrope position_ids = get_rope_index( self.processor, @@ -270,7 +300,9 @@ def __getitem__(self, index): attention_mask=attention_mask, ) # (3, seq_length) else: - position_ids = torch.clip(attention_mask.cumsum(dim=0) - 1, min=0, max=None) # (seq_length,) + position_ids = torch.clip( + attention_mask.cumsum(dim=0) - 1, min=0, max=None + ) # (seq_length,) input_ids, attention_mask, position_ids = VF.postprocess_data( input_ids=input_ids, @@ -281,14 +313,18 @@ def __getitem__(self, index): left_pad=True, truncation=self.truncation, ) - raw_prompt_ids = self.tokenizer.encode(prompt, add_special_tokens=False) + raw_prompt_ids = self.tokenizer.encode( + prompt, add_special_tokens=False) if len(raw_prompt_ids) > self.max_prompt_length: if self.truncation == "left": - raw_prompt_ids = raw_prompt_ids[-self.max_prompt_length :] + raw_prompt_ids = raw_prompt_ids[-self.max_prompt_length:] elif self.truncation == "right": raw_prompt_ids = raw_prompt_ids[: self.max_prompt_length] elif self.truncation == "error": - raise RuntimeError(f"Prompt length {len(raw_prompt_ids)} is longer than {self.max_prompt_length}.") + raise RuntimeError( + f"Prompt length { + len(raw_prompt_ids)} is longer than { + self.max_prompt_length}.") example["input_ids"] = input_ids example["attention_mask"] = attention_mask diff --git a/Agent0/curriculum_train/verl/utils/flops_counter.py b/Agent0/curriculum_train/verl/utils/flops_counter.py index dee7623..0672ff9 100644 --- a/Agent0/curriculum_train/verl/utils/flops_counter.py +++ b/Agent0/curriculum_train/verl/utils/flops_counter.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -66,7 +66,9 @@ class FlopsCounter: def __init__(self, config: "LlamaConfig"): if config.model_type not in VALID_MODLE_TYPE: - print(f"Only support {VALID_MODLE_TYPE}, but got {config.model_type}. MFU will always be zero.") + print( + f"Only support {VALID_MODLE_TYPE}, but got { + config.model_type}. MFU will always be zero.") self.estimate_func = { "llama": self._estimate_llama_flops, @@ -76,10 +78,14 @@ def __init__(self, config: "LlamaConfig"): } self.config = config - def _estimate_unknown_flops(self, tokens_sum: int, batch_seqlens: List[int], delta_time: float) -> float: + def _estimate_unknown_flops( + self, tokens_sum: int, batch_seqlens: List[int], delta_time: float + ) -> float: return 0 - def _estimate_llama_flops(self, tokens_sum: int, batch_seqlens: List[int], delta_time: float) -> float: + def _estimate_llama_flops( + self, tokens_sum: int, batch_seqlens: List[int], delta_time: float + ) -> float: hidden_size = self.config.hidden_size vocab_size = self.config.vocab_size num_hidden_layers = self.config.num_hidden_layers @@ -95,10 +101,13 @@ def _estimate_llama_flops(self, tokens_sum: int, batch_seqlens: List[int], delta # 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) + 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 + 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 @@ -107,14 +116,21 @@ def _estimate_llama_flops(self, tokens_sum: int, batch_seqlens: List[int], delta 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 + 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: List[int], delta_time: float) -> Tuple[float, float]: + def estimate_flops( + self, batch_seqlens: List[int], delta_time: float + ) -> Tuple[float, float]: """ Estimate the FLOPS based on the number of valid tokens in the current batch and the time taken. @@ -127,7 +143,9 @@ def estimate_flops(self, batch_seqlens: List[int], delta_time: float) -> Tuple[f 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) + 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 diff --git a/Agent0/curriculum_train/verl/utils/fsdp_utils.py b/Agent0/curriculum_train/verl/utils/fsdp_utils.py index 1ca563a..f8abe7d 100644 --- a/Agent0/curriculum_train/verl/utils/fsdp_utils.py +++ b/Agent0/curriculum_train/verl/utils/fsdp_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,23 +27,31 @@ from transformers.trainer_pt_utils import get_module_class_from_name -def get_init_fn(model: nn.Module, device: Union[str, torch.device]) -> Callable[[nn.Module], None]: +def get_init_fn( + model: nn.Module, device: Union[str, torch.device] +) -> Callable[[nn.Module], None]: param_occurrence = defaultdict(int) for _, param in model.named_parameters(remove_duplicate=False): param_occurrence[param] += 1 - duplicated_params = {param for param in param_occurrence.keys() if param_occurrence[param] > 1} + duplicated_params = { + param for param in param_occurrence.keys() if param_occurrence[param] > 1} materialized_params = {} def init_fn(module: nn.Module): for name, param in module.named_parameters(recurse=False): if param in duplicated_params: module._parameters[name] = materialized_params.setdefault( - param, nn.Parameter(torch.empty_like(param.data, device=device), requires_grad=param.requires_grad) + param, + nn.Parameter( + torch.empty_like(param.data, device=device), + requires_grad=param.requires_grad, + ), ) else: module._parameters[name] = nn.Parameter( - torch.empty_like(param.data, device=device), requires_grad=param.requires_grad + torch.empty_like(param.data, device=device), + requires_grad=param.requires_grad, ) return init_fn @@ -63,7 +71,8 @@ def get_fsdp_wrap_policy(model: PreTrainedModel): else: transformer_cls_to_wrap.add(transformer_cls) - return partial(transformer_auto_wrap_policy, transformer_layer_cls=transformer_cls_to_wrap) + return partial(transformer_auto_wrap_policy, + transformer_layer_cls=transformer_cls_to_wrap) @torch.no_grad() diff --git a/Agent0/curriculum_train/verl/utils/logger/__init__.py b/Agent0/curriculum_train/verl/utils/logger/__init__.py index 557c477..c67f0ff 100644 --- a/Agent0/curriculum_train/verl/utils/logger/__init__.py +++ b/Agent0/curriculum_train/verl/utils/logger/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/utils/logger/gen_logger.py b/Agent0/curriculum_train/verl/utils/logger/gen_logger.py index b62cde6..af59882 100644 --- a/Agent0/curriculum_train/verl/utils/logger/gen_logger.py +++ b/Agent0/curriculum_train/verl/utils/logger/gen_logger.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,22 +31,34 @@ @dataclass class GenerationLogger(ABC): @abstractmethod - def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None: ... + def log(self, samples: List[Tuple[str, str, + str, float]], step: int) -> None: ... @dataclass class ConsoleGenerationLogger(GenerationLogger): - def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None: + def log(self, samples: List[Tuple[str, str, + str, float]], step: int) -> None: for inp, out, lab, score in samples: - print(f"[prompt] {inp}\n[output] {out}\n[ground_truth] {lab}\n[score] {score}\n") + print( + f"[prompt] {inp}\n[output] {out}\n[ground_truth] {lab}\n[score] {score}\n") @dataclass class WandbGenerationLogger(GenerationLogger): - def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None: + def log(self, samples: List[Tuple[str, str, + str, float]], step: int) -> None: # Create column names for all samples columns = ["step"] + sum( - [[f"input_{i + 1}", f"output_{i + 1}", f"label_{i + 1}", f"score_{i + 1}"] for i in range(len(samples))], + [ + [ + f"input_{i + 1}", + f"output_{i + 1}", + f"label_{i + 1}", + f"score_{i + 1}", + ] + for i in range(len(samples)) + ], [], ) @@ -55,8 +67,11 @@ def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None: self.validation_table = wandb.Table(columns=columns) # Create a new table with same columns and existing data - # Workaround for https://github.com/wandb/wandb/issues/2981#issuecomment-1997445737 - new_table = wandb.Table(columns=columns, data=self.validation_table.data) + # Workaround for + # https://github.com/wandb/wandb/issues/2981#issuecomment-1997445737 + new_table = wandb.Table( + columns=columns, + data=self.validation_table.data) # Add new row with all data row_data = [step] @@ -70,13 +85,23 @@ def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None: @dataclass class SwanlabGenerationLogger(GenerationLogger): - def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None: + def log(self, samples: List[Tuple[str, str, + str, float]], step: int) -> None: swanlab_text_list = [] for i, sample in enumerate(samples): row_text = "\n\n---\n\n".join( - (f"input: {sample[0]}", f"output: {sample[1]}", f"label: {sample[2]}", f"score: {sample[3]}") + ( + f"input: {sample[0]}", + f"output: {sample[1]}", + f"label: {sample[2]}", + f"score: {sample[3]}", + ) ) - swanlab_text_list.append(swanlab.Text(row_text, caption=f"sample {i + 1}")) + swanlab_text_list.append( + swanlab.Text( + row_text, + caption=f"sample { + i + 1}")) swanlab.log({"val/generations": swanlab_text_list}, step=step) @@ -97,6 +122,7 @@ def __init__(self, loggers: List[str]): if logger in GEN_LOGGERS: self.loggers.append(GEN_LOGGERS[logger]()) - def log(self, samples: List[Tuple[str, str, str, float]], step: int) -> None: + def log(self, samples: List[Tuple[str, str, + str, float]], step: int) -> None: for logger in self.loggers: logger.log(samples, step) diff --git a/Agent0/curriculum_train/verl/utils/logger/logger.py b/Agent0/curriculum_train/verl/utils/logger/logger.py index a29fb50..381737f 100644 --- a/Agent0/curriculum_train/verl/utils/logger/logger.py +++ b/Agent0/curriculum_train/verl/utils/logger/logger.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -21,7 +21,12 @@ import torch -from ..py_functional import convert_dict_to_str, flatten_dict, is_package_available, unflatten_dict +from ..py_functional import ( + convert_dict_to_str, + flatten_dict, + is_package_available, + unflatten_dict, +) from .gen_logger import AggregateGenerationsLogger @@ -105,7 +110,10 @@ def __init__(self, config: Dict[str, Any]) -> None: else: config_dict[key] = str(value) - self.writer.add_hparams(hparam_dict=config_dict, metric_dict={"placeholder": 0}) + self.writer.add_hparams( + hparam_dict=config_dict, + metric_dict={ + "placeholder": 0}) def log(self, data: Dict[str, Any], step: int) -> None: for key, value in data.items(): @@ -140,7 +148,11 @@ def finish(self) -> None: class Tracker: - def __init__(self, loggers: Union[str, List[str]] = "console", config: Optional[Dict[str, Any]] = None): + def __init__( + self, + loggers: Union[str, List[str]] = "console", + config: Optional[Dict[str, Any]] = None, + ): if isinstance(loggers, str): loggers = [loggers] @@ -157,7 +169,9 @@ def log(self, data: Dict[str, Any], step: int) -> None: for logger in self.loggers: logger.log(data=data, step=step) - def log_generation(self, samples: List[Tuple[str, str, str, float]], step: int) -> None: + def log_generation( + self, samples: List[Tuple[str, str, str, float]], step: int + ) -> None: self.gen_logger.log(samples, step) def __del__(self): diff --git a/Agent0/curriculum_train/verl/utils/model_utils.py b/Agent0/curriculum_train/verl/utils/model_utils.py index 71d4fe2..4b53c68 100644 --- a/Agent0/curriculum_train/verl/utils/model_utils.py +++ b/Agent0/curriculum_train/verl/utils/model_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,10 +32,14 @@ def print_gpu_memory_usage(prefix: str = "GPU memory usage") -> None: """Report the current GPU VRAM usage.""" if is_rank0(): free_mem, total_mem = torch.cuda.mem_get_info() - print(f"{prefix}: {(total_mem - free_mem) / (1024**3):.2f} GB / {total_mem / (1024**3):.2f} GB.") + print(f"{prefix}: {(total_mem - + free_mem) / + (1024**3):.2f} GB / {total_mem / + (1024**3):.2f} GB.") -def _get_model_size(model: nn.Module, scale: str = "auto") -> Tuple[float, str]: +def _get_model_size( + model: nn.Module, scale: str = "auto") -> Tuple[float, str]: """Compute the model size.""" n_params = sum(p.numel() for p in model.parameters()) diff --git a/Agent0/curriculum_train/verl/utils/py_functional.py b/Agent0/curriculum_train/verl/utils/py_functional.py index 1a9ed3c..08478b6 100644 --- a/Agent0/curriculum_train/verl/utils/py_functional.py +++ b/Agent0/curriculum_train/verl/utils/py_functional.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,7 +32,8 @@ def is_sci_notation(number: float) -> bool: return bool(pattern.match(str(number))) -def float_representer(dumper: Dumper, number: Union[float, np.float32, np.float64]): +def float_representer( + dumper: Dumper, number: Union[float, np.float32, np.float64]): if is_sci_notation(number): value = str(number) if "." not in value and "e" in value: @@ -53,18 +54,22 @@ def is_package_available(name: str) -> bool: return importlib.util.find_spec(name) is not None -def union_two_dict(dict1: Dict[str, Any], dict2: Dict[str, Any]) -> Dict[str, Any]: +def union_two_dict(dict1: Dict[str, Any], + dict2: Dict[str, Any]) -> Dict[str, Any]: """Union two dict. Will throw an error if there is an item not the same object with the same key.""" for key in dict2.keys(): if key in dict1: - assert dict1[key] == dict2[key], f"{key} in dict1 and dict2 are not the same object" + assert ( + dict1[key] == dict2[key] + ), f"{key} in dict1 and dict2 are not the same object" dict1[key] = dict2[key] return dict1 -def append_to_dict(data: Dict[str, List[Any]], new_data: Dict[str, Any]) -> None: +def append_to_dict(data: Dict[str, List[Any]], + new_data: Dict[str, Any]) -> None: """Append dict to a dict of list.""" for key, val in new_data.items(): if key not in data: @@ -89,7 +94,9 @@ def unflatten_dict(data: Dict[str, Any], sep: str = "/") -> Dict[str, Any]: return unflattened -def flatten_dict(data: Dict[str, Any], parent_key: str = "", sep: str = "/") -> Dict[str, Any]: +def flatten_dict( + data: Dict[str, Any], parent_key: str = "", sep: str = "/" +) -> Dict[str, Any]: flattened = {} for key, value in data.items(): new_key = parent_key + sep + key if parent_key else key diff --git a/Agent0/curriculum_train/verl/utils/seqlen_balancing.py b/Agent0/curriculum_train/verl/utils/seqlen_balancing.py index 5889784..14ea0c7 100644 --- a/Agent0/curriculum_train/verl/utils/seqlen_balancing.py +++ b/Agent0/curriculum_train/verl/utils/seqlen_balancing.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -94,12 +94,18 @@ def __repr__(self) -> str: return repr_str -def karmarkar_karp(seqlen_list: List[int], k_partitions: int, equal_size: bool): +def karmarkar_karp( + seqlen_list: List[int], + k_partitions: int, + equal_size: bool): # see: https://en.wikipedia.org/wiki/Largest_differencing_method - sorted_seqlen_list = sorted([(seqlen, i) for i, seqlen in enumerate(seqlen_list)]) + sorted_seqlen_list = sorted([(seqlen, i) + for i, seqlen in enumerate(seqlen_list)]) states_pq: List[State] = [] if equal_size: - assert len(seqlen_list) % k_partitions == 0, f"{len(seqlen_list)} % {k_partitions} != 0" + 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): @@ -108,7 +114,10 @@ def karmarkar_karp(seqlen_list: List[int], k_partitions: int, equal_size: bool): 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)) + heapq.heappush( + states_pq, State( + items=[ + (idx, seqlen)], k=k_partitions)) while len(states_pq) > 1: state0 = heapq.heappop(states_pq) @@ -121,15 +130,19 @@ def karmarkar_karp(seqlen_list: List[int], k_partitions: int, equal_size: bool): 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)}" - ) + assert len(partition) * k_partitions == len( + seqlen_list + ), f"{len(partition)} * {k_partitions} != {len(seqlen_list)}" return partitions -def greedy_partition(seqlen_list: List[int], k_partitions: int, equal_size: bool): +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)] + 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: @@ -141,13 +154,15 @@ def greedy_partition(seqlen_list: List[int], k_partitions: int, equal_size: bool 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)}" - ) + assert len(partition) * k_partitions == len( + seqlen_list + ), f"{len(partition)} * {k_partitions} != {len(seqlen_list)}" return partitions -def get_seqlen_balanced_partitions(seqlen_list: List[int], k_partitions: int, equal_size: bool): +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: @@ -163,10 +178,13 @@ def get_seqlen_balanced_partitions(seqlen_list: List[int], k_partitions: int, eq 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}]" + 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}" + assert len(partitions) == k_partitions, f"{ + len(partitions)} != {k_partitions}" seen_idx = set() sorted_partitions = [None] * k_partitions for i, partition in enumerate(partitions): @@ -177,11 +195,15 @@ def _check_and_sort_partitions(partitions): 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) + partitions = karmarkar_karp( + seqlen_list=seqlen_list, + k_partitions=k_partitions, + equal_size=equal_size) return _check_and_sort_partitions(partitions) -def log_seqlen_unbalance(seqlen_list: List[int], partitions: List[List[int]], prefix): +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 @@ -190,7 +212,7 @@ def log_seqlen_unbalance(seqlen_list: List[int], partitions: List[List[int]], pr 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]) + 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: @@ -225,29 +247,35 @@ def rearrange_micro_batches(batch: TensorDict, max_token_len, dp_group=None): """ # 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=}" - ) + 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) + 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_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.append(batch[idx: idx + 1]) curr_micro_batch = torch.cat(curr_micro_batch) micro_batches.append(curr_micro_batch) diff --git a/Agent0/curriculum_train/verl/utils/tokenizer.py b/Agent0/curriculum_train/verl/utils/tokenizer.py index b339e2a..39dc646 100644 --- a/Agent0/curriculum_train/verl/utils/tokenizer.py +++ b/Agent0/curriculum_train/verl/utils/tokenizer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,10 +15,17 @@ from typing import Optional -from transformers import AutoProcessor, AutoTokenizer, PreTrainedTokenizer, ProcessorMixin +from transformers import ( + AutoProcessor, + AutoTokenizer, + PreTrainedTokenizer, + ProcessorMixin, +) -def get_tokenizer(model_path: str, override_chat_template: Optional[str] = None, **kwargs) -> PreTrainedTokenizer: +def get_tokenizer( + model_path: str, override_chat_template: Optional[str] = None, **kwargs +) -> PreTrainedTokenizer: """Create a huggingface pretrained tokenizer.""" tokenizer = AutoTokenizer.from_pretrained(model_path, **kwargs) if override_chat_template is not None: @@ -27,7 +34,9 @@ def get_tokenizer(model_path: str, override_chat_template: Optional[str] = None, if tokenizer.bos_token == "" and tokenizer.eos_token == "": # the EOS token in gemma2 & gemma3 is ambiguious, which may worsen RL performance. # https://huggingface.co/google/gemma-2-2b-it/commit/17a01657f5c87135bcdd0ec7abb4b2dece04408a - print("Found gemma model. Set eos_token and eos_token_id to and 107.") + print( + "Found gemma model. Set eos_token and eos_token_id to and 107." + ) tokenizer.eos_token = "" if tokenizer.pad_token_id is None: @@ -37,7 +46,9 @@ def get_tokenizer(model_path: str, override_chat_template: Optional[str] = None, return tokenizer -def get_processor(model_path: str, override_chat_template: Optional[str] = None, **kwargs) -> Optional[ProcessorMixin]: +def get_processor( + model_path: str, override_chat_template: Optional[str] = None, **kwargs +) -> Optional[ProcessorMixin]: """Create a huggingface pretrained processor.""" processor = AutoProcessor.from_pretrained(model_path, **kwargs) if override_chat_template is not None: diff --git a/Agent0/curriculum_train/verl/utils/torch_dtypes.py b/Agent0/curriculum_train/verl/utils/torch_dtypes.py index e50c5c3..65a0be9 100644 --- a/Agent0/curriculum_train/verl/utils/torch_dtypes.py +++ b/Agent0/curriculum_train/verl/utils/torch_dtypes.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/utils/torch_functional.py b/Agent0/curriculum_train/verl/utils/torch_functional.py index 0bf926e..44b2090 100644 --- a/Agent0/curriculum_train/verl/utils/torch_functional.py +++ b/Agent0/curriculum_train/verl/utils/torch_functional.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Copyright Meta Platforms, Inc. and affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -35,7 +35,9 @@ @torch.compiler.disable() -def log_probs_from_logits_flash_attn(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: +def log_probs_from_logits_flash_attn( + logits: torch.Tensor, labels: torch.Tensor +) -> torch.Tensor: output = cross_entropy_loss(logits, labels, inplace_backward=True) if not isinstance(output, tuple): raise ValueError( @@ -45,7 +47,9 @@ def log_probs_from_logits_flash_attn(logits: torch.Tensor, labels: torch.Tensor) return -output[0] -def log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: +def log_probs_from_logits( + logits: torch.Tensor, + labels: torch.Tensor) -> torch.Tensor: """Compute log probs on the label ids given logits. We may use torch compile to speed up computing. @@ -69,12 +73,18 @@ def log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.T return output.view(*batch_dim) -def masked_mean(values: torch.Tensor, mask: torch.Tensor, dim: int = None, eps: float = 1e-8) -> torch.Tensor: +def masked_mean( + values: torch.Tensor, + mask: torch.Tensor, + dim: int = None, + eps: float = 1e-8) -> torch.Tensor: """Compute mean of tensor with a masked values.""" return (values * mask).sum(dim=dim) / (mask.sum(dim=dim) + eps) -def masked_var(values: torch.Tensor, mask: torch.Tensor, unbiased: bool = True) -> torch.Tensor: +def masked_var( + values: torch.Tensor, mask: torch.Tensor, unbiased: bool = True +) -> torch.Tensor: """Compute variance of tensor with masked values.""" mean = masked_mean(values, mask) centered_values = values - mean @@ -82,7 +92,9 @@ def masked_var(values: torch.Tensor, mask: torch.Tensor, unbiased: bool = True) if unbiased: mask_sum = mask.sum() if mask_sum <= 1: - print("The sum of the mask is less than one, which can cause a division by zero.") + print( + "The sum of the mask is less than one, which can cause a division by zero." + ) return variance bessel_correction = mask_sum / (mask_sum - 1) @@ -91,14 +103,18 @@ def masked_var(values: torch.Tensor, mask: torch.Tensor, unbiased: bool = True) return variance -def masked_whiten(values: torch.Tensor, mask: torch.Tensor, eps: float = 1e-8) -> torch.Tensor: +def masked_whiten( + values: torch.Tensor, mask: torch.Tensor, eps: float = 1e-8 +) -> torch.Tensor: """Whiten values with masked values.""" mean, var = masked_mean(values, mask), masked_var(values, mask) return (values - mean) * torch.rsqrt(var + eps) def get_response_mask( - response_ids: torch.Tensor, eos_token_id: Union[int, List[int]] = 2, dtype: torch.dtype = torch.long + response_ids: torch.Tensor, + eos_token_id: Union[int, List[int]] = 2, + dtype: torch.dtype = torch.long, ): """Get the mask for the response ids, the mask will be 0 after the first eos token. @@ -132,22 +148,35 @@ def pad_2d_list_to_length( else: target_length = max_response_length - padded_response = [tuple(sub_list) + (pad_token_id,) * (target_length - len(sub_list)) for sub_list in response] + padded_response = [ + tuple(sub_list) + (pad_token_id,) * (target_length - len(sub_list)) + for sub_list in response + ] tensor = torch.tensor(padded_response) return tensor def pad_sequence_to_length( - tensor: torch.Tensor, max_seq_len: int, pad_token_id: int, left_pad: bool = False -) -> torch.Tensor: + tensor: torch.Tensor, + max_seq_len: int, + pad_token_id: int, + left_pad: bool = False) -> torch.Tensor: """Pad a nD tensors in the last dim to max_seq_len.""" if tensor.size(-1) >= max_seq_len: return tensor pad_shape = list(tensor.shape) pad_shape[-1] = max_seq_len - tensor.size(-1) - pad_tensor = torch.full(pad_shape, fill_value=pad_token_id, dtype=tensor.dtype, device=tensor.device) - return torch.cat((pad_tensor, tensor), dim=-1) if left_pad else torch.cat((tensor, pad_tensor), dim=-1) + pad_tensor = torch.full( + pad_shape, + fill_value=pad_token_id, + dtype=tensor.dtype, + device=tensor.device) + return ( + torch.cat((pad_tensor, tensor), dim=-1) + if left_pad + else torch.cat((tensor, pad_tensor), dim=-1) + ) def postprocess_data( @@ -164,12 +193,21 @@ def postprocess_data( seq_length = len(input_ids) if seq_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 + 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 - ) - position_ids = pad_sequence_to_length(position_ids, max_seq_len=max_length, pad_token_id=0, left_pad=left_pad) + attention_mask, + max_seq_len=max_length, + pad_token_id=0, + left_pad=left_pad) + position_ids = pad_sequence_to_length( + position_ids, + max_seq_len=max_length, + pad_token_id=0, + left_pad=left_pad) elif seq_length > max_length: if truncation == "left": # actually, left truncation may not be reasonable input_ids = input_ids[..., -max_length:] @@ -180,9 +218,11 @@ def postprocess_data( attention_mask = attention_mask[..., :max_length] position_ids = position_ids[..., :max_length] elif truncation == "error": - raise RuntimeError(f"Input sequence length {seq_length} is longer than max length {max_length}.") + raise RuntimeError( + f"Input sequence length {seq_length} is longer than max length {max_length}.") else: - raise NotImplementedError(f"Unknown truncation method {truncation}.") + raise NotImplementedError( + f"Unknown truncation method {truncation}.") return input_ids, attention_mask, position_ids @@ -282,14 +322,18 @@ def step(self, closure=None): momentum_dtype = PrecisionType.to_dtype(group["momentum_dtype"]) variance_dtype = PrecisionType.to_dtype(group["variance_dtype"]) - compensation_buffer_dtype = PrecisionType.to_dtype(group["compensation_buffer_dtype"]) + compensation_buffer_dtype = PrecisionType.to_dtype( + group["compensation_buffer_dtype"] + ) for p in group["params"]: assert isinstance(p, torch.Tensor) # lint if p.grad is None: continue if p.grad.is_sparse: - raise RuntimeError("AnyPrecisionAdamW does not support sparse gradients.") + raise RuntimeError( + "AnyPrecisionAdamW does not support sparse gradients." + ) state = self.state[p] # State initialization @@ -297,14 +341,18 @@ def step(self, closure=None): state["step"] = torch.tensor(0.0) # momentum - EMA of gradient values - state["exp_avg"] = torch.zeros_like(p, dtype=momentum_dtype) + state["exp_avg"] = torch.zeros_like( + p, dtype=momentum_dtype) # variance uncentered - EMA of squared gradient values - state["exp_avg_sq"] = torch.zeros_like(p, dtype=variance_dtype) + state["exp_avg_sq"] = torch.zeros_like( + p, dtype=variance_dtype) # optional Kahan summation - accumulated error tracker if use_kahan_summation: - state["compensation"] = torch.zeros_like(p, dtype=compensation_buffer_dtype) + state["compensation"] = torch.zeros_like( + p, dtype=compensation_buffer_dtype + ) # Main processing # update the steps for each param group update @@ -318,18 +366,28 @@ def step(self, closure=None): if weight_decay: # weight decay, AdamW style p.data.mul_(1 - lr * weight_decay) - exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) # update momentum - exp_avg_sq.mul_(beta2).addcmul_(grad, grad, value=1 - beta2) # update uncentered variance + exp_avg.mul_(beta1).add_( + grad, alpha=1 - beta1) # update momentum + exp_avg_sq.mul_(beta2).addcmul_( + grad, grad, value=1 - beta2 + ) # update uncentered variance bias_correction1 = 1 - beta1**step # adjust using bias1 step_size = lr / bias_correction1 - denom_correction = (1 - beta2**step) ** 0.5 # adjust using bias2 and avoids math import - centered_variance = (exp_avg_sq.sqrt() / denom_correction).add_(eps, alpha=1) + denom_correction = ( + 1 - beta2**step + ) ** 0.5 # adjust using bias2 and avoids math import + centered_variance = ( + exp_avg_sq.sqrt() / + denom_correction).add_( + eps, + alpha=1) if use_kahan_summation: # lr update to compensation compensation = state["compensation"] - compensation.addcdiv_(exp_avg, centered_variance, value=-step_size) + compensation.addcdiv_( + exp_avg, centered_variance, value=-step_size) # update weights with compensation (Kahan summation) # save error back to compensation for next iteration @@ -337,4 +395,5 @@ def step(self, closure=None): p.data.add_(compensation) compensation.add_(temp_buffer.sub_(p.data)) else: # usual AdamW updates - p.data.addcdiv_(exp_avg, centered_variance, value=-step_size) + p.data.addcdiv_( + exp_avg, centered_variance, value=-step_size) diff --git a/Agent0/curriculum_train/verl/utils/ulysses.py b/Agent0/curriculum_train/verl/utils/ulysses.py index 18e07b4..f582d1b 100644 --- a/Agent0/curriculum_train/verl/utils/ulysses.py +++ b/Agent0/curriculum_train/verl/utils/ulysses.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -44,7 +44,8 @@ def get_ulysses_sequence_parallel_group() -> Optional[dist.ProcessGroup]: return _ULYSSES_SEQUENCE_PARALLEL_GROUP -def get_ulysses_sequence_parallel_world_size(group: ProcessGroup = None) -> int: +def get_ulysses_sequence_parallel_world_size( + group: ProcessGroup = None) -> int: """ Get ulysses sequence parallel world size. """ @@ -84,7 +85,9 @@ def gather_seq_scatter_heads( return x -def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int, group: ProcessGroup = None) -> Tensor: +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: @@ -115,7 +118,9 @@ def _unpad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor: return x[slc] -def slice_input_tensor(x: Tensor, dim: int, padding: bool = True, group: ProcessGroup = None) -> Tensor: +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() @@ -140,9 +145,17 @@ def all_to_all_tensor( ): 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) + 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(): @@ -153,13 +166,23 @@ def wait(): return torch.cat(output_list, dim=gather_dim).contiguous() -def all_gather_tensor(local_tensor: Tensor, group: Optional[dist.ProcessGroup] = None, async_op: bool = False): +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) + 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 @@ -177,17 +200,26 @@ def forward( 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) + 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]: + 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() + 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), + all_to_all_tensor( + input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group, False + ), None, None, None, @@ -230,7 +262,9 @@ def backward(ctx: Any, grad_output: Tensor) -> Any: grad_output = grad_output * ctx.sp_world_size return ( None, - grad_output.split(ctx.part_size, dim=ctx.gather_dim)[ctx.sp_rank].contiguous(), + grad_output.split(ctx.part_size, dim=ctx.gather_dim)[ + ctx.sp_rank + ].contiguous(), None, None, None, @@ -252,7 +286,9 @@ def gather_outputs_and_unpad( 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" + 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) @@ -260,7 +296,9 @@ def gather_outputs_and_unpad( def ulysses_pad_and_slice_inputs( - input_ids_rmpad: torch.Tensor, position_ids_rmpad: Optional[torch.Tensor] = None, sp_size: int = 1 + 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 @@ -289,10 +327,15 @@ def ulysses_pad_and_slice_inputs( _, 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) + 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) + 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 diff --git a/Agent0/curriculum_train/verl/workers/__init__.py b/Agent0/curriculum_train/verl/workers/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/curriculum_train/verl/workers/__init__.py +++ b/Agent0/curriculum_train/verl/workers/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/workers/actor/__init__.py b/Agent0/curriculum_train/verl/workers/actor/__init__.py index 7472ab3..5aae76b 100644 --- a/Agent0/curriculum_train/verl/workers/actor/__init__.py +++ b/Agent0/curriculum_train/verl/workers/actor/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/workers/actor/base.py b/Agent0/curriculum_train/verl/workers/actor/base.py index bd264ca..fe3826c 100644 --- a/Agent0/curriculum_train/verl/workers/actor/base.py +++ b/Agent0/curriculum_train/verl/workers/actor/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/workers/actor/config.py b/Agent0/curriculum_train/verl/workers/actor/config.py index e792bc4..acc6990 100644 --- a/Agent0/curriculum_train/verl/workers/actor/config.py +++ b/Agent0/curriculum_train/verl/workers/actor/config.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,10 +33,13 @@ def post_init(self): if self.tokenizer_path is None: self.tokenizer_path = self.model_path - if self.model_path is not None and os.path.exists(self.model_path): # ray job uses absolute path + if self.model_path is not None and os.path.exists( + self.model_path + ): # ray job uses absolute path self.model_path = os.path.abspath(self.model_path) - if self.tokenizer_path is not None and os.path.exists(self.tokenizer_path): + if self.tokenizer_path is not None and os.path.exists( + self.tokenizer_path): self.tokenizer_path = os.path.abspath(self.tokenizer_path) @@ -104,7 +107,8 @@ class RefConfig: fsdp: FSDPConfig = field(default_factory=FSDPConfig) offload: OffloadConfig = field(default_factory=OffloadConfig) """auto keys""" - micro_batch_size_per_device_for_experience: int = field(default=-1, init=False) + micro_batch_size_per_device_for_experience: int = field( + default=-1, init=False) padding_free: bool = field(default=False, init=False) ulysses_sequence_parallel_size: int = field(default=1, init=False) use_torch_compile: bool = field(default=True, init=False) diff --git a/Agent0/curriculum_train/verl/workers/actor/dp_actor.py b/Agent0/curriculum_train/verl/workers/actor/dp_actor.py index 6b771ba..e38648b 100644 --- a/Agent0/curriculum_train/verl/workers/actor/dp_actor.py +++ b/Agent0/curriculum_train/verl/workers/actor/dp_actor.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,7 +24,11 @@ from ray.experimental.tqdm_ray import tqdm from torch import nn from torch.distributed.fsdp import FullyShardedDataParallel as FSDP -from transformers.modeling_flash_attention_utils import index_first_axis, pad_input, unpad_input +from transformers.modeling_flash_attention_utils import ( + index_first_axis, + pad_input, + unpad_input, +) from ...protocol import DataProto from ...trainer import core_algos @@ -53,11 +57,15 @@ def __init__( self.actor_module = actor_module self.actor_optimizer = actor_optimizer if config.use_torch_compile: - self.log_probs_from_logits = torch.compile(VF.log_probs_from_logits, dynamic=True) + self.log_probs_from_logits = torch.compile( + VF.log_probs_from_logits, dynamic=True + ) else: self.log_probs_from_logits = VF.log_probs_from_logits - def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor], temperature: float) -> torch.Tensor: + def _forward_micro_batch( + self, micro_batch: Dict[str, torch.Tensor], temperature: float + ) -> torch.Tensor: """ Returns: log_probs: # (bs, response_len) @@ -69,7 +77,9 @@ def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor], temperature responses = micro_batch["responses"] response_length = responses.size(-1) if position_ids.dim() == 3: # qwen2vl mrope - position_ids = position_ids.transpose(0, 1) # (bsz, 3, seqlen) -> (3, bsz, seqlen) + position_ids = position_ids.transpose( + 0, 1 + ) # (bsz, 3, seqlen) -> (3, bsz, seqlen) multi_modal_inputs = {} if "multi_modal_inputs" in micro_batch: @@ -87,28 +97,42 @@ def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor], temperature # unpad the position_ids to align the rotary if position_ids.dim() == 3: position_ids_rmpad = ( - index_first_axis(rearrange(position_ids, "c b s ... -> (b s) c ..."), indices) - .transpose(0, 1) - .unsqueeze(1) - ) # (3, bsz, seqlen) -> (3, 1, bsz * seqlen) + index_first_axis( + rearrange( + position_ids, + "c b s ... -> (b s) c ..."), + indices) .transpose( + 0, + 1) .unsqueeze(1)) # (3, bsz, seqlen) -> (3, 1, bsz * seqlen) else: position_ids_rmpad = index_first_axis( - rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices + 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) + 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.config.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.config.ulysses_sequence_parallel_size + input_ids_rmpad, position_ids_rmpad, pad_size = ( + ulysses_pad_and_slice_inputs( + input_ids_rmpad, + position_ids_rmpad, + sp_size=self.config.ulysses_sequence_parallel_size, + ) ) input_ids_rmpad_rolled, _, _ = ulysses_pad_and_slice_inputs( - input_ids_rmpad_rolled, None, self.config.ulysses_sequence_parallel_size + input_ids_rmpad_rolled, + None, + self.config.ulysses_sequence_parallel_size, ) - input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0) # ((total_nnz / sp) + pad) + 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( @@ -121,18 +145,27 @@ def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor], temperature logits_rmpad = output.logits.squeeze(0) # (total_nnz, vocab_size) logits_rmpad.div_(temperature) # ((total_nnz / sp) + pad) - log_probs = self.log_probs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled) + log_probs = self.log_probs_from_logits( + logits=logits_rmpad, labels=input_ids_rmpad_rolled + ) # gather log_prob if sp > 1 if self.config.ulysses_sequence_parallel_size > 1: # gather and unpad for the ulysses sp - log_probs = gather_outputs_and_unpad(log_probs, gather_dim=0, unpad_dim=0, padding_size=pad_size) + log_probs = gather_outputs_and_unpad( + log_probs, gather_dim=0, unpad_dim=0, padding_size=pad_size + ) # pad back to (bsz, seqlen) full_log_probs = pad_input( - hidden_states=log_probs.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen + hidden_states=log_probs.unsqueeze(-1), + indices=indices, + batch=batch_size, + seqlen=seqlen, ) - log_probs = full_log_probs.squeeze(-1)[:, -response_length - 1 : -1] # (bsz, response_length) + log_probs = full_log_probs.squeeze(-1)[ + :, -response_length - 1: -1 + ] # (bsz, response_length) else: output = self.actor_module( input_ids=input_ids, @@ -143,16 +176,23 @@ def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor], temperature ) logits: torch.Tensor = output.logits logits.div_(temperature) - logits = logits[:, -response_length - 1 : -1, :] # (bsz, response_length, vocab_size) - log_probs = self.log_probs_from_logits(logits, responses) # (bsz, response_length) + logits = logits[ + :, -response_length - 1: -1, : + ] # (bsz, response_length, vocab_size) + log_probs = self.log_probs_from_logits( + logits, responses + ) # (bsz, response_length) return log_probs def _optimizer_step(self) -> torch.Tensor: if isinstance(self.actor_module, FSDP): - grad_norm = self.actor_module.clip_grad_norm_(self.config.max_grad_norm) + grad_norm = self.actor_module.clip_grad_norm_( + self.config.max_grad_norm) else: - grad_norm = nn.utils.clip_grad_norm_(self.actor_module.parameters(), max_norm=self.config.max_grad_norm) + grad_norm = nn.utils.clip_grad_norm_( + self.actor_module.parameters(), + max_norm=self.config.max_grad_norm) if not torch.isfinite(grad_norm): print("Gradient norm is not finite. Skip update.") @@ -184,7 +224,11 @@ def compute_log_prob(self, data: DataProto) -> torch.Tensor: self.actor_module.eval() temperature = data.meta_info["temperature"] - select_keys = ["responses", "input_ids", "attention_mask", "position_ids"] + select_keys = [ + "responses", + "input_ids", + "attention_mask", + "position_ids"] if "multi_modal_inputs" in data.non_tensor_batch.keys(): non_tensor_select_keys = ["multi_modal_inputs"] else: @@ -195,11 +239,17 @@ def compute_log_prob(self, data: DataProto) -> torch.Tensor: ) log_probs_lst = [] if self.rank == 0: - micro_batches = tqdm(micro_batches, desc="Compute log probs", position=2) + micro_batches = tqdm( + micro_batches, + desc="Compute log probs", + position=2) for micro_batch in micro_batches: - model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch} - log_probs = self._forward_micro_batch(model_inputs, temperature=temperature) + model_inputs = { + **micro_batch.batch, + **micro_batch.non_tensor_batch} + log_probs = self._forward_micro_batch( + model_inputs, temperature=temperature) log_probs_lst.append(log_probs) log_probs = torch.concat(log_probs_lst, dim=0) @@ -208,8 +258,17 @@ def compute_log_prob(self, data: DataProto) -> torch.Tensor: def update_policy(self, data: DataProto) -> Dict[str, Any]: self.actor_module.train() - 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"] + 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.use_kl_loss and not self.config.disable_kl: select_keys.append("ref_log_probs") @@ -220,23 +279,35 @@ def update_policy(self, data: DataProto) -> Dict[str, Any]: # Split to make minibatch iterator for updating the actor # See PPO paper for details. https://arxiv.org/abs/1707.06347 - mini_batches = data.select(select_keys, non_tensor_select_keys).split(self.config.global_batch_size_per_device) + mini_batches = data.select(select_keys, non_tensor_select_keys).split( + self.config.global_batch_size_per_device + ) metrics = defaultdict(list) for _ in range(self.config.ppo_epochs): if self.rank == 0: - mini_batches = tqdm(mini_batches, desc="Train mini-batches", position=2) + mini_batches = tqdm( + mini_batches, + desc="Train mini-batches", + position=2) for mini_batch in mini_batches: gradient_accumulation = ( - self.config.global_batch_size_per_device // self.config.micro_batch_size_per_device_for_update + self.config.global_batch_size_per_device + // self.config.micro_batch_size_per_device_for_update + ) + micro_batches = mini_batch.split( + self.config.micro_batch_size_per_device_for_update ) - micro_batches = mini_batch.split(self.config.micro_batch_size_per_device_for_update) if self.rank == 0: - micro_batches = tqdm(micro_batches, desc="Update policy", position=3) + micro_batches = tqdm( + micro_batches, desc="Update policy", position=3 + ) for micro_batch in micro_batches: - model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch} + model_inputs = { + **micro_batch.batch, + **micro_batch.non_tensor_batch} responses = model_inputs["responses"] response_length = responses.size(1) attention_mask = model_inputs["attention_mask"] @@ -245,17 +316,23 @@ def update_policy(self, data: DataProto) -> Dict[str, Any]: advantages = model_inputs["advantages"] # all return: (bsz, response_length) - log_probs = self._forward_micro_batch(model_inputs, temperature=temperature) - entropy_loss = -VF.masked_mean(log_probs, response_mask) # estimator of entropy loss - - pg_loss, pg_clipfrac_higher, pg_clipfrac_lower, ppo_kl = core_algos.compute_policy_loss( - old_log_probs=old_log_probs, - log_probs=log_probs, - advantages=advantages, - response_mask=response_mask, - clip_ratio_low=self.config.clip_ratio_low, - clip_ratio_high=self.config.clip_ratio_high, - clip_ratio_dual=self.config.clip_ratio_dual, + log_probs = self._forward_micro_batch( + model_inputs, temperature=temperature + ) + entropy_loss = -VF.masked_mean( + log_probs, response_mask + ) # estimator of entropy loss + + pg_loss, pg_clipfrac_higher, pg_clipfrac_lower, ppo_kl = ( + core_algos.compute_policy_loss( + old_log_probs=old_log_probs, + log_probs=log_probs, + advantages=advantages, + response_mask=response_mask, + clip_ratio_low=self.config.clip_ratio_low, + clip_ratio_high=self.config.clip_ratio_high, + clip_ratio_dual=self.config.clip_ratio_dual, + ) ) if "ref_log_probs" in model_inputs: ref_log_probs = model_inputs["ref_log_probs"] @@ -283,6 +360,7 @@ def update_policy(self, data: DataProto) -> Dict[str, Any]: append_to_dict(metrics, batch_metrics) grad_norm = self._optimizer_step() - append_to_dict(metrics, {"actor/grad_norm": grad_norm.detach().item()}) + append_to_dict(metrics, + {"actor/grad_norm": grad_norm.detach().item()}) return metrics diff --git a/Agent0/curriculum_train/verl/workers/config.py b/Agent0/curriculum_train/verl/workers/config.py index ba21b0e..561accd 100644 --- a/Agent0/curriculum_train/verl/workers/config.py +++ b/Agent0/curriculum_train/verl/workers/config.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -46,7 +46,11 @@ class WorkerConfig: rollout: RolloutConfig = field(default_factory=RolloutConfig) def post_init(self): - self.ref.micro_batch_size_per_device_for_experience = self.actor.micro_batch_size_per_device_for_experience + self.ref.micro_batch_size_per_device_for_experience = ( + self.actor.micro_batch_size_per_device_for_experience + ) self.ref.padding_free = self.actor.padding_free - self.ref.ulysses_sequence_parallel_size = self.actor.ulysses_sequence_parallel_size + self.ref.ulysses_sequence_parallel_size = ( + self.actor.ulysses_sequence_parallel_size + ) self.ref.use_torch_compile = self.actor.use_torch_compile diff --git a/Agent0/curriculum_train/verl/workers/critic/__init__.py b/Agent0/curriculum_train/verl/workers/critic/__init__.py index 0bc8ee4..bc0c535 100644 --- a/Agent0/curriculum_train/verl/workers/critic/__init__.py +++ b/Agent0/curriculum_train/verl/workers/critic/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/workers/critic/base.py b/Agent0/curriculum_train/verl/workers/critic/base.py index 3d54146..951576d 100644 --- a/Agent0/curriculum_train/verl/workers/critic/base.py +++ b/Agent0/curriculum_train/verl/workers/critic/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/workers/critic/config.py b/Agent0/curriculum_train/verl/workers/critic/config.py index d18d2f0..733243d 100644 --- a/Agent0/curriculum_train/verl/workers/critic/config.py +++ b/Agent0/curriculum_train/verl/workers/critic/config.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/workers/critic/dp_critic.py b/Agent0/curriculum_train/verl/workers/critic/dp_critic.py index 013c8e5..3ba4afd 100644 --- a/Agent0/curriculum_train/verl/workers/critic/dp_critic.py +++ b/Agent0/curriculum_train/verl/workers/critic/dp_critic.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -34,7 +34,12 @@ try: - from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input + from flash_attn.bert_padding import ( + index_first_axis, + pad_input, + rearrange, + unpad_input, + ) except ImportError: pass @@ -43,13 +48,20 @@ class DataParallelPPOCritic(BasePPOCritic): - def __init__(self, config: CriticConfig, critic_module: nn.Module, critic_optimizer: torch.optim.Optimizer): + def __init__( + self, + config: CriticConfig, + critic_module: nn.Module, + critic_optimizer: torch.optim.Optimizer, + ): super().__init__(config) self.rank = int(os.getenv("RANK", "0")) self.critic_module = critic_module self.critic_optimizer = critic_optimizer - def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor]) -> torch.Tensor: + def _forward_micro_batch( + self, micro_batch: Dict[str, torch.Tensor] + ) -> torch.Tensor: input_ids = micro_batch["input_ids"] batch_size, seqlen = input_ids.shape attention_mask = micro_batch["attention_mask"] @@ -57,7 +69,9 @@ def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor]) -> torch.Te responses = micro_batch["responses"] response_length = responses.size(-1) if position_ids.dim() == 3: # qwen2vl mrope - position_ids = position_ids.transpose(0, 1) # (bsz, 3, seqlen) -> (3, bsz, seqlen) + position_ids = position_ids.transpose( + 0, 1 + ) # (bsz, 3, seqlen) -> (3, bsz, seqlen) multi_modal_inputs = {} if "multi_modal_inputs" in micro_batch: @@ -75,19 +89,27 @@ def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor]) -> torch.Te # unpad the position_ids to align the rotary if position_ids.dim() == 3: position_ids_rmpad = ( - index_first_axis(rearrange(position_ids, "c b s ... -> (b s) c ..."), indices) - .transpose(0, 1) - .unsqueeze(1) - ) # (3, bsz, seqlen) -> (3, 1, bsz * seqlen) + index_first_axis( + rearrange( + position_ids, + "c b s ... -> (b s) c ..."), + indices) .transpose( + 0, + 1) .unsqueeze(1)) # (3, bsz, seqlen) -> (3, 1, bsz * seqlen) else: position_ids_rmpad = index_first_axis( - rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices + rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), + indices, ).transpose(0, 1) # pad and slice the inputs if sp > 1 if self.config.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.config.ulysses_sequence_parallel_size + input_ids_rmpad, position_ids_rmpad, pad_size = ( + ulysses_pad_and_slice_inputs( + input_ids_rmpad, + position_ids_rmpad, + sp_size=self.config.ulysses_sequence_parallel_size, + ) ) # only pass input_ids and position_ids to enable flash_attn_varlen @@ -103,11 +125,14 @@ def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor]) -> torch.Te # gather output if sp > 1 if self.config.ulysses_sequence_parallel_size > 1: - values_rmpad = gather_outputs_and_unpad(values_rmpad, gather_dim=0, unpad_dim=0, padding_size=pad_size) + values_rmpad = gather_outputs_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_size, seqlen=seqlen).squeeze(-1) - values = values[:, -response_length - 1 : -1] + values = pad_input( + values_rmpad, indices=indices, batch=batch_size, seqlen=seqlen + ).squeeze(-1) + values = values[:, -response_length - 1: -1] else: output = self.critic_module( input_ids=input_ids, @@ -117,17 +142,20 @@ def _forward_micro_batch(self, micro_batch: Dict[str, torch.Tensor]) -> torch.Te use_cache=False, ) values: torch.Tensor = output.logits - values = values[:, -response_length - 1 : -1].squeeze(-1) # (bsz, response_length, vocab_size) + values = values[:, -response_length - 1: -1].squeeze( + -1 + ) # (bsz, response_length, vocab_size) return values def _optimizer_step(self) -> torch.Tensor: if isinstance(self.critic_module, FSDP): - grad_norm = self.critic_module.clip_grad_norm_(self.config.max_grad_norm) + grad_norm = self.critic_module.clip_grad_norm_( + self.config.max_grad_norm) else: grad_norm = torch.nn.utils.clip_grad_norm_( - self.critic_module.parameters(), max_norm=self.config.max_grad_norm - ) + self.critic_module.parameters(), + max_norm=self.config.max_grad_norm) if not torch.isfinite(grad_norm): print("Gradient norm is not finite. Skip update.") @@ -141,7 +169,11 @@ def _optimizer_step(self) -> torch.Tensor: def compute_values(self, data: DataProto) -> torch.Tensor: self.critic_module.eval() - select_keys = ["responses", "input_ids", "attention_mask", "position_ids"] + select_keys = [ + "responses", + "input_ids", + "attention_mask", + "position_ids"] if "multi_modal_inputs" in data.non_tensor_batch.keys(): non_tensor_select_keys = ["multi_modal_inputs"] else: @@ -152,10 +184,15 @@ def compute_values(self, data: DataProto) -> torch.Tensor: ) values_lst = [] if self.rank == 0: - micro_batches = tqdm(micro_batches, desc="Compute values", position=2) + micro_batches = tqdm( + micro_batches, + desc="Compute values", + position=2) for micro_batch in micro_batches: - model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch} + model_inputs = { + **micro_batch.batch, + **micro_batch.non_tensor_batch} values = self._forward_micro_batch(model_inputs) values_lst.append(values) @@ -163,13 +200,20 @@ def compute_values(self, data: DataProto) -> torch.Tensor: responses = data.batch["responses"] attention_mask = data.batch["attention_mask"] response_length = responses.size(1) - values = values * attention_mask[:, -response_length - 1 : -1] + values = values * attention_mask[:, -response_length - 1: -1] return values def update_critic(self, data: DataProto) -> Dict[str, Any]: self.critic_module.train() - select_keys = ["input_ids", "responses", "attention_mask", "position_ids", "values", "returns"] + select_keys = [ + "input_ids", + "responses", + "attention_mask", + "position_ids", + "values", + "returns", + ] if "multi_modal_inputs" in data.non_tensor_batch.keys(): non_tensor_select_keys = ["multi_modal_inputs"] else: @@ -177,29 +221,43 @@ def update_critic(self, data: DataProto) -> Dict[str, Any]: # Split to make minibatch iterator for updating the actor # See PPO paper for details. https://arxiv.org/abs/1707.06347 - mini_batches = data.select(select_keys, non_tensor_select_keys).split(self.config.global_batch_size_per_device) + mini_batches = data.select(select_keys, non_tensor_select_keys).split( + self.config.global_batch_size_per_device + ) metrics = defaultdict(list) for _ in range(self.config.ppo_epochs): if self.rank == 0: - mini_batches = tqdm(mini_batches, desc="Train mini-batches", position=2) + mini_batches = tqdm( + mini_batches, + desc="Train mini-batches", + position=2) for mini_batch in mini_batches: gradient_accumulation = ( - self.config.global_batch_size_per_device // self.config.micro_batch_size_per_device_for_update + self.config.global_batch_size_per_device + // self.config.micro_batch_size_per_device_for_update + ) + micro_batches = mini_batch.split( + self.config.micro_batch_size_per_device_for_update ) - micro_batches = mini_batch.split(self.config.micro_batch_size_per_device_for_update) if self.rank == 0: - micro_batches = tqdm(micro_batches, desc="Update critic", position=3) + micro_batches = tqdm( + micro_batches, desc="Update critic", position=3 + ) for micro_batch in micro_batches: - model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch} + model_inputs = { + **micro_batch.batch, + **micro_batch.non_tensor_batch} responses = model_inputs["responses"] attention_mask = model_inputs["attention_mask"] values = model_inputs["values"] returns = model_inputs["returns"] response_length = responses.size(1) - action_mask = attention_mask[:, -response_length - 1 : -1] # shift left for value computation + action_mask = attention_mask[ + :, -response_length - 1: -1 + ] # shift left for value computation vpreds = self._forward_micro_batch(model_inputs) vf_loss, vf_clipfrac = core_algos.compute_value_loss( @@ -215,11 +273,14 @@ def update_critic(self, data: DataProto) -> Dict[str, Any]: batch_metrics = { "critic/vf_loss": vf_loss.detach().item(), "critic/vf_clipfrac": vf_clipfrac.detach().item(), - "critic/vpred_mean": VF.masked_mean(vpreds, action_mask).detach().item(), + "critic/vpred_mean": VF.masked_mean( + vpreds, + action_mask) .detach() .item(), } append_to_dict(metrics, batch_metrics) grad_norm = self._optimizer_step() - append_to_dict(metrics, {"critic/grad_norm": grad_norm.detach().item()}) + append_to_dict(metrics, + {"critic/grad_norm": grad_norm.detach().item()}) return metrics diff --git a/Agent0/curriculum_train/verl/workers/fsdp_workers.py b/Agent0/curriculum_train/verl/workers/fsdp_workers.py index 17c65a9..a973212 100644 --- a/Agent0/curriculum_train/verl/workers/fsdp_workers.py +++ b/Agent0/curriculum_train/verl/workers/fsdp_workers.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -55,19 +55,34 @@ from ..utils.model_utils import print_gpu_memory_usage, print_model_size from ..utils.tokenizer import get_processor, get_tokenizer from ..utils.torch_dtypes import PrecisionType -from ..utils.torch_functional import AnyPrecisionAdamW, get_constant_schedule_with_warmup -from .config import ActorConfig, CriticConfig, FSDPConfig, ModelConfig, OptimConfig, RefConfig, WorkerConfig +from ..utils.torch_functional import ( + AnyPrecisionAdamW, + get_constant_schedule_with_warmup, +) +from .config import ( + ActorConfig, + CriticConfig, + FSDPConfig, + ModelConfig, + OptimConfig, + RefConfig, + WorkerConfig, +) from .rollout import vLLMRollout from .sharding_manager import FSDPVLLMShardingManager from .sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager class FSDPWorker(Worker): - def __init__( - self, - config: WorkerConfig, - role: Literal["actor", "critic", "rollout", "ref", "actor_rollout", "actor_rollout_ref"], - ): + def __init__(self, + config: WorkerConfig, + role: Literal["actor", + "critic", + "rollout", + "ref", + "actor_rollout", + "actor_rollout_ref"], + ): super().__init__() self.config = config self.role = role @@ -79,9 +94,14 @@ def __init__( torch.backends.cuda.matmul.allow_tf32 = False torch.backends.cuda.matmul.allow_bf16_reduced_precision_reduction = False - self._is_actor = self.role in ["actor", "actor_rollout", "actor_rollout_ref"] + self._is_actor = self.role in [ + "actor", "actor_rollout", "actor_rollout_ref"] self._is_critic = self.role == "critic" - self._is_rollout = self.role in ["rollout", "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._cache = {} @@ -95,20 +115,28 @@ def __init__( self._use_param_offload = self.config.critic.offload.offload_params self._use_optimizer_offload = self.config.critic.offload.offload_optimizer self._init_config(self.config.critic, "critic") - elif self._is_ref: # NOTE: it seems that manual offload is slower than FSDP offload + elif ( + self._is_ref + ): # NOTE: it seems that manual offload is slower than FSDP offload self._use_param_offload = self.config.ref.offload.offload_params self._init_config(self.config.ref, "ref") def _init_config( - self, config: Union[ActorConfig, CriticConfig, RefConfig], role: Literal["actor", "critic", "ref"] + self, + config: Union[ActorConfig, CriticConfig, RefConfig], + role: Literal["actor", "critic", "ref"], ): world_size = dist.get_world_size() fsdp_size = config.fsdp.fsdp_size if fsdp_size <= 0 or fsdp_size >= world_size: - self.device_mesh = init_device_mesh("cuda", mesh_shape=(world_size,), mesh_dim_names=("fsdp",)) + self.device_mesh = init_device_mesh( + "cuda", mesh_shape=(world_size,), mesh_dim_names=("fsdp",) + ) else: # hsdp self.device_mesh = init_device_mesh( - "cuda", mesh_shape=(world_size // fsdp_size, fsdp_size), mesh_dim_names=("ddp", "fsdp") + "cuda", + mesh_shape=(world_size // fsdp_size, fsdp_size), + mesh_dim_names=("ddp", "fsdp"), ) if config.ulysses_sequence_parallel_size > 1: @@ -123,29 +151,43 @@ def _init_config( else: self.ulysses_device_mesh = None - self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) + self.ulysses_sharding_manager = FSDPUlyssesShardingManager( + self.ulysses_device_mesh + ) if not hasattr(config, "global_batch_size"): # ref model return if self.config.rollout.n > 1: config.global_batch_size *= self.config.rollout.n - self.print_rank0(f"{role} will use global batch size {config.global_batch_size}.") + self.print_rank0( + f"{role} will use global batch size { + config.global_batch_size}.") config.global_batch_size_per_device = ( - config.global_batch_size // self.device_mesh.size() * config.ulysses_sequence_parallel_size + config.global_batch_size + // self.device_mesh.size() + * config.ulysses_sequence_parallel_size ) if config.global_batch_size_per_device == 0: - raise ValueError(f"{role} global batch size * ulysses size must be larger than num gpus.") + raise ValueError( + f"{role} global batch size * ulysses size must be larger than num gpus.") - if config.global_batch_size_per_device % config.micro_batch_size_per_device_for_update != 0: - raise ValueError(f"{role} global batch size per device must be divisible by the micro batch size.") + if ( + config.global_batch_size_per_device + % config.micro_batch_size_per_device_for_update + != 0 + ): + raise ValueError( + f"{role} global batch size per device must be divisible by the micro batch size.") if ( config.fsdp.enable_cpu_offload - and config.global_batch_size_per_device != config.micro_batch_size_per_device_for_update + and config.global_batch_size_per_device + != config.micro_batch_size_per_device_for_update ): - raise ValueError(f"{role} cannot use FSDP's CPU offload when gradient accumulation is enabled.") + raise ValueError( + f"{role} cannot use FSDP's CPU offload when gradient accumulation is enabled.") def _build_model_optimizer( self, @@ -174,9 +216,13 @@ def _build_model_optimizer( ) try: - self.generation_config = GenerationConfig.from_pretrained(model_config.model_path) + self.generation_config = GenerationConfig.from_pretrained( + model_config.model_path + ) except Exception: - self.generation_config = GenerationConfig.from_model_config(self.model_config) + self.generation_config = GenerationConfig.from_model_config( + self.model_config + ) self.print_rank0(f"Model config: {self.model_config}") @@ -185,7 +231,8 @@ def _build_model_optimizer( self.print_rank0("Ulysses patch applied!") if fsdp_config.torch_dtype is None: - torch_dtype = torch.float32 if self._is_actor or self._is_critic else torch.bfloat16 + torch_dtype = ( + torch.float32 if self._is_actor or self._is_critic else torch.bfloat16) else: torch_dtype = PrecisionType.to_dtype(fsdp_config.torch_dtype) @@ -196,11 +243,12 @@ def _build_model_optimizer( else: auto_class = AutoModelForCausalLM - if (not fsdp_config.enable_rank0_init) or self.device_mesh.get_local_rank("fsdp") == 0: + if (not fsdp_config.enable_rank0_init) or self.device_mesh.get_local_rank( + "fsdp") == 0: model = auto_class.from_pretrained( model_config.model_path, config=self.model_config, - torch_dtype='bfloat16', + torch_dtype="bfloat16", attn_implementation="flash_attention_2", device_map="cpu" if fsdp_config.enable_rank0_init else "cuda", low_cpu_mem_usage=True, @@ -210,7 +258,7 @@ def _build_model_optimizer( with no_init_weights(), init_empty_weights(): model = auto_class.from_config( self.model_config, - torch_dtype='bfloat16', + torch_dtype="bfloat16", attn_implementation="flash_attention_2", trust_remote_code=model_config.trust_remote_code, ) @@ -219,7 +267,9 @@ def _build_model_optimizer( model.tie_weights() # avoid hanging model = model.to(torch_dtype) if model_config.enable_gradient_checkpointing: - model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) if not (self._is_actor or self._is_critic): model.requires_grad_(False) @@ -261,7 +311,9 @@ def _build_model_optimizer( if fsdp_config.enable_rank0_init: sync_module_states = True - param_init_fn = get_init_fn(model, device="cuda") if self.rank != 0 else None + param_init_fn = ( + get_init_fn(model, device="cuda") if self.rank != 0 else None + ) else: sync_module_states = False param_init_fn = None @@ -284,7 +336,9 @@ def _build_model_optimizer( if self._is_actor or self._is_critic: if optim_config.strategy == "adamw": self.optimizer = torch.optim.AdamW( - filter(lambda p: p.requires_grad, self.fsdp_module.parameters()), + filter( + lambda p: p.requires_grad, + self.fsdp_module.parameters()), lr=optim_config.lr, betas=optim_config.betas, weight_decay=optim_config.weight_decay, @@ -292,15 +346,21 @@ def _build_model_optimizer( ) elif optim_config.strategy == "adamw_bf16": self.optimizer = AnyPrecisionAdamW( - filter(lambda p: p.requires_grad, self.fsdp_module.parameters()), + filter( + lambda p: p.requires_grad, + self.fsdp_module.parameters()), lr=optim_config.lr, betas=optim_config.betas, weight_decay=optim_config.weight_decay, ) else: - raise NotImplementedError(f"Optimizer {optim_config.strategy} not supported.") + raise NotImplementedError( + f"Optimizer {optim_config.strategy} not supported." + ) - num_warmup_steps = int(optim_config.lr_warmup_ratio * optim_config.training_steps) + num_warmup_steps = int( + optim_config.lr_warmup_ratio * optim_config.training_steps + ) self.lr_scheduler = get_constant_schedule_with_warmup( optimizer=self.optimizer, num_warmup_steps=num_warmup_steps ) @@ -311,10 +371,13 @@ def _build_model_optimizer( def _build_rollout(self) -> None: tp_size = self.config.rollout.tensor_parallel_size dp_size = self.world_size // tp_size - assert self.world_size % tp_size == 0, ( - f"rollout world size: {self.world_size} is not divisible by tp size: {tp_size}" + assert ( + self.world_size % + tp_size == 0), f"rollout world size: { + self.world_size} is not divisible by tp size: {tp_size}" + rollout_device_mesh = init_device_mesh( + "cuda", mesh_shape=(dp_size, tp_size), mesh_dim_names=("dp", "tp") ) - rollout_device_mesh = init_device_mesh("cuda", mesh_shape=(dp_size, tp_size), mesh_dim_names=("dp", "tp")) self.rollout = vLLMRollout( model_path=self.config.actor.model.model_path, config=self.config.rollout, @@ -359,11 +422,13 @@ def init_model(self): ) if self._use_param_offload: offload_fsdp_model(self.fsdp_module) - print_gpu_memory_usage(f"After offload {role} model during init") + print_gpu_memory_usage( + f"After offload {role} model during init") if self._use_optimizer_offload: offload_fsdp_optimizer(optimizer=self.optimizer) - print_gpu_memory_usage(f"After offload {role} optimizer during init") + print_gpu_memory_usage( + f"After offload {role} optimizer during init") if self._is_actor: from .actor.dp_actor import DataParallelPPOActor # lazy import @@ -400,7 +465,8 @@ def init_model(self): model=self.fsdp_module, optimizer=self.optimizer, lr_scheduler=self.lr_scheduler, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), ) @register(dispatch_mode=Dispatch.ONE_TO_ALL) @@ -431,15 +497,18 @@ def preprocess_multi_modal_data(self, data: DataProto): # inplace load & process image data min_pixels = data.meta_info["min_pixels"] max_pixels = data.meta_info["max_pixels"] - multi_modal_data_copy = deepcopy(data.non_tensor_batch["multi_modal_data"]) + multi_modal_data_copy = deepcopy( + data.non_tensor_batch["multi_modal_data"]) processed_images = [] for multi_modal_data in multi_modal_data_copy: processed_per_query_images = [] - for image in multi_modal_data['image']: + for image in multi_modal_data["image"]: processed_per_query_images.append( - process_image(image, min_pixels=min_pixels, max_pixels=max_pixels) - ) + process_image( + image, + min_pixels=min_pixels, + max_pixels=max_pixels)) processed_images.append(processed_per_query_images) # Note: Using the alternative (commented) code below to process images can lead to subtle resize issues: @@ -454,17 +523,24 @@ def preprocess_multi_modal_data(self, data: DataProto): # for j, image in enumerate(per_query_images): # images[i][j] = process_image(image, min_pixels=min_pixels, max_pixels=max_pixels) - multi_modal_inputs = np.array([ - dict(self.processor.image_processor(images=per_query_images, videos=None)) - for per_query_images in processed_images - ], dtype=object) + multi_modal_inputs = np.array( + [ + dict( + self.processor.image_processor(images=per_query_images, videos=None) + ) + for per_query_images in processed_images + ], + dtype=object, + ) data.non_tensor_batch["multi_modal_inputs"] = multi_modal_inputs @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) def update_actor(self, data: DataProto): assert self._is_actor if "multi_modal_inputs" in self._cache: - data.non_tensor_batch['multi_modal_inputs'] = deepcopy(self._cache['multi_modal_inputs']) + data.non_tensor_batch["multi_modal_inputs"] = deepcopy( + self._cache["multi_modal_inputs"] + ) elif "multi_modal_data" in data.non_tensor_batch: self.preprocess_multi_modal_data(data) @@ -483,26 +559,35 @@ def update_actor(self, data: DataProto): 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) + estimated_flops, promised_flops = self.flops_counter.estimate_flops( + global_num_tokens, delta_time) metrics["perf/mfu_actor"] = ( - estimated_flops * self.config.actor.ppo_epochs / (promised_flops * self.world_size) + estimated_flops + * self.config.actor.ppo_epochs + / (promised_flops * self.world_size) ) metrics["perf/max_memory_allocated_gb"] = ( - torch.cuda.max_memory_allocated() - self.rollout_sharding_manager.freed_bytes + torch.cuda.max_memory_allocated() + - self.rollout_sharding_manager.freed_bytes ) / (1024**3) metrics["perf/max_memory_reserved_gb"] = ( - torch.cuda.max_memory_reserved() - self.rollout_sharding_manager.freed_bytes + torch.cuda.max_memory_reserved() + - self.rollout_sharding_manager.freed_bytes ) / (1024**3) - metrics["perf/cpu_memory_used_gb"] = psutil.virtual_memory().used / (1024**3) + metrics["perf/cpu_memory_used_gb"] = psutil.virtual_memory().used / ( + 1024**3 + ) self.lr_scheduler.step() lr = self.lr_scheduler.get_last_lr()[0] metrics["actor/lr"] = lr - # Metrics should be in non_tensor_batch instead of meta_info, as DataProto not concat meta_info. + # Metrics should be in non_tensor_batch instead of meta_info, as + # DataProto not concat meta_info. output = DataProto( non_tensor_batch={ - key: np.array([value] if np.isscalar(value) else value) for key, value in metrics.items() + key: np.array([value] if np.isscalar(value) else value) + for key, value in metrics.items() } ) @@ -523,12 +608,16 @@ def generate_sequences(self, prompts: DataProto): load_fsdp_model(self.fsdp_module) meta_info = { - "eos_token_id": self.generation_config.eos_token_id - if self.generation_config is not None - else self.tokenizer.eos_token_id, - "pad_token_id": self.generation_config.pad_token_id - if self.generation_config is not None - else self.tokenizer.pad_token_id, + "eos_token_id": ( + self.generation_config.eos_token_id + if self.generation_config is not None + else self.tokenizer.eos_token_id + ), + "pad_token_id": ( + self.generation_config.pad_token_id + if self.generation_config is not None + else self.tokenizer.pad_token_id + ), } prompts.meta_info.update(meta_info) with self.rollout_sharding_manager: @@ -544,13 +633,19 @@ def generate_sequences(self, prompts: DataProto): # load image data cached_multi_modal_data = None if "multi_modal_data" in prompts.non_tensor_batch: - cached_multi_modal_data = deepcopy(prompts.non_tensor_batch["multi_modal_data"]) - min_pixels = prompts.meta_info['min_pixels'] - max_pixels = prompts.meta_info['max_pixels'] + cached_multi_modal_data = deepcopy( + prompts.non_tensor_batch["multi_modal_data"] + ) + min_pixels = prompts.meta_info["min_pixels"] + max_pixels = prompts.meta_info["max_pixels"] processed_images = [] - for i, multi_modal_data in enumerate(prompts.non_tensor_batch["multi_modal_data"]): + for i, multi_modal_data in enumerate( + prompts.non_tensor_batch["multi_modal_data"] + ): for j, image in enumerate(multi_modal_data["image"]): - multi_modal_data['image'][j] = process_image(image, min_pixels=min_pixels, max_pixels=max_pixels) + multi_modal_data["image"][j] = process_image( + image, min_pixels=min_pixels, max_pixels=max_pixels + ) processed_images.append(multi_modal_data) prompts.non_tensor_batch["multi_modal_data"] = processed_images @@ -562,7 +657,9 @@ def generate_sequences(self, prompts: DataProto): output.non_tensor_batch["multi_modal_data"] = cached_multi_modal_data if sampling_n > 1: output.non_tensor_batch["multi_modal_data"] = np.repeat( - output.non_tensor_batch["multi_modal_data"], repeats=sampling_n, axis=0, + output.non_tensor_batch["multi_modal_data"], + repeats=sampling_n, + axis=0, ) output = self.rollout_sharding_manager.postprocess_data(output) @@ -577,7 +674,9 @@ def compute_log_probs(self, data: DataProto): if "multi_modal_data" in data.non_tensor_batch: self.preprocess_multi_modal_data(data) # create cache for multi_modal_inputs - self._cache['multi_modal_inputs'] = deepcopy(data.non_tensor_batch['multi_modal_inputs']) + self._cache["multi_modal_inputs"] = deepcopy( + data.non_tensor_batch["multi_modal_inputs"] + ) data = data.to(torch.cuda.current_device()) if self._use_param_offload: @@ -590,7 +689,8 @@ def compute_log_probs(self, data: DataProto): data = self.ulysses_sharding_manager.preprocess_data(data) output = self.actor.compute_log_prob(data=data) output = DataProto.from_dict( - tensors={"old_log_probs": output}, meta_info={"temperature": self.config.rollout.temperature} + tensors={"old_log_probs": output}, + meta_info={"temperature": self.config.rollout.temperature}, ) output = self.ulysses_sharding_manager.postprocess_data(output) @@ -611,7 +711,9 @@ def compute_ref_log_probs(self, data: DataProto): # not in the ref_policy's or critic's caches. assert self._is_ref if "multi_modal_inputs" in self._cache: - data.non_tensor_batch['multi_modal_inputs'] = deepcopy(self._cache['multi_modal_inputs']) + data.non_tensor_batch["multi_modal_inputs"] = deepcopy( + self._cache["multi_modal_inputs"] + ) elif "multi_modal_data" in data.non_tensor_batch: self.preprocess_multi_modal_data(data) @@ -643,7 +745,9 @@ def compute_values(self, data: DataProto): # The `self._cache` is empty here since cached `multi_modal_inputs` is only saved in the actor's _cache, # not in the ref_policy's or critic's caches. if "multi_modal_inputs" in self._cache: - data.non_tensor_batch['multi_modal_inputs'] = deepcopy(self._cache['multi_modal_inputs']) + data.non_tensor_batch["multi_modal_inputs"] = deepcopy( + self._cache["multi_modal_inputs"] + ) elif "multi_modal_data" in data.non_tensor_batch: self.preprocess_multi_modal_data(data) @@ -655,7 +759,8 @@ def compute_values(self, data: DataProto): 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 = self.ulysses_sharding_manager.postprocess_data( + data=output) if self._use_param_offload: offload_fsdp_model(self.fsdp_module) @@ -668,7 +773,9 @@ def update_critic(self, data: DataProto): # The `self._cache` is empty here since cached `multi_modal_inputs` is only saved in the actor's _cache, # not in the ref_policy's or critic's caches. if "multi_modal_inputs" in self._cache: - data.non_tensor_batch['multi_modal_inputs'] = deepcopy(self._cache['multi_modal_inputs']) + data.non_tensor_batch["multi_modal_inputs"] = deepcopy( + self._cache["multi_modal_inputs"] + ) elif "multi_modal_data" not in data.non_tensor_batch: self.preprocess_multi_modal_data(data) @@ -686,19 +793,24 @@ def update_critic(self, data: DataProto): 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) + estimated_flops, promised_flops = self.flops_counter.estimate_flops( + global_num_tokens, delta_time) metrics["perf/mfu_critic"] = ( - estimated_flops * self.config.actor.ppo_epochs / (promised_flops * self.world_size) + estimated_flops + * self.config.actor.ppo_epochs + / (promised_flops * self.world_size) ) self.lr_scheduler.step() lr = self.lr_scheduler.get_last_lr()[0] metrics["critic/lr"] = lr - # Metrics should be in non_tensor_batch instead of meta_info, as DataProto not concat meta_info. + # Metrics should be in non_tensor_batch instead of meta_info, as + # DataProto not concat meta_info. output = DataProto( non_tensor_batch={ - metric: np.array([value] if np.isscalar(value) else value) for metric, value in metrics.items() + metric: np.array([value] if np.isscalar(value) else value) + for metric, value in metrics.items() } ) diff --git a/Agent0/curriculum_train/verl/workers/reward/__init__.py b/Agent0/curriculum_train/verl/workers/reward/__init__.py index 9d476f6..d9227ec 100644 --- a/Agent0/curriculum_train/verl/workers/reward/__init__.py +++ b/Agent0/curriculum_train/verl/workers/reward/__init__.py @@ -13,7 +13,16 @@ # limitations under the License. from .config import RewardConfig -from .function import BatchFunctionRewardManager, FunctionRewardManager, SequentialFunctionRewardManager +from .function import ( + BatchFunctionRewardManager, + FunctionRewardManager, + SequentialFunctionRewardManager, +) -__all__ = ["BatchFunctionRewardManager", "FunctionRewardManager", "RewardConfig", "SequentialFunctionRewardManager"] +__all__ = [ + "BatchFunctionRewardManager", + "FunctionRewardManager", + "RewardConfig", + "SequentialFunctionRewardManager", +] diff --git a/Agent0/curriculum_train/verl/workers/reward/config.py b/Agent0/curriculum_train/verl/workers/reward/config.py index 7e11bdb..18f1581 100644 --- a/Agent0/curriculum_train/verl/workers/reward/config.py +++ b/Agent0/curriculum_train/verl/workers/reward/config.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,13 +31,18 @@ class RewardConfig: reward_function_name: Optional[str] = field(default=None, init=False) def post_init(self): - if self.reward_function is not None: # support custom reward function, e.g., ./math.py:main + if ( + self.reward_function is not None + ): # support custom reward function, e.g., ./math.py:main if ":" not in self.reward_function: self.reward_function_name = "main" else: - self.reward_function, self.reward_function_name = self.reward_function.rsplit(":", maxsplit=1) + self.reward_function, self.reward_function_name = ( + self.reward_function.rsplit(":", maxsplit=1) + ) - if os.path.exists(self.reward_function): # ray job uses absolute path + if os.path.exists( + self.reward_function): # ray job uses absolute path self.reward_function = os.path.abspath(self.reward_function) else: self.reward_function = None diff --git a/Agent0/curriculum_train/verl/workers/reward/function.py b/Agent0/curriculum_train/verl/workers/reward/function.py index a7af022..eaabc66 100644 --- a/Agent0/curriculum_train/verl/workers/reward/function.py +++ b/Agent0/curriculum_train/verl/workers/reward/function.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -46,9 +46,13 @@ def __init__(self, config: RewardConfig, tokenizer: PreTrainedTokenizer): raise ValueError("Reward function is not provided.") if not os.path.exists(config.reward_function): - raise FileNotFoundError(f"Reward function file {config.reward_function} not found.") + raise FileNotFoundError( + f"Reward function file {config.reward_function} not found." + ) - spec = importlib.util.spec_from_file_location("custom_reward_fn", config.reward_function) + spec = importlib.util.spec_from_file_location( + "custom_reward_fn", config.reward_function + ) module = importlib.util.module_from_spec(spec) try: sys.modules["custom_reward_fn"] = module @@ -57,16 +61,23 @@ def __init__(self, config: RewardConfig, tokenizer: PreTrainedTokenizer): raise RuntimeError(f"Failed to load reward function: {e}") if not hasattr(module, config.reward_function_name): - raise AttributeError(f"Module {module} does not have function {config.reward_function_name}.") + raise AttributeError( + f"Module {module} does not have function { + config.reward_function_name}.") reward_fn = getattr(module, config.reward_function_name) - print(f"Using reward function `{config.reward_function_name}` from `{config.reward_function}`.") + print( + f"Using reward function `{ + config.reward_function_name}` from `{ + config.reward_function}`.") self.reward_fn = partial(reward_fn, **config.reward_function_kwargs) self.config = config self.tokenizer = tokenizer @abstractmethod - def compute_reward(self, data: DataProto) -> Tuple[torch.Tensor, Dict[str, List[float]]]: + def compute_reward( + self, data: DataProto + ) -> Tuple[torch.Tensor, Dict[str, List[float]]]: """Compute reward for a batch of data.""" ... @@ -74,16 +85,18 @@ def compute_reward(self, data: DataProto) -> Tuple[torch.Tensor, Dict[str, List[ class SequentialFunctionRewardManager(FunctionRewardManager): reward_fn: SequentialRewardFunction - def compute_reward(self, data: DataProto) -> Tuple[torch.Tensor, Dict[str, List[float]]]: - reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) + def compute_reward( + self, data: DataProto + ) -> Tuple[torch.Tensor, Dict[str, List[float]]]: + reward_tensor = torch.zeros_like( + data.batch["responses"], dtype=torch.float32) reward_metrics = defaultdict(list) response_ids = data.batch["responses"] response_length = data.batch["response_mask"].sum(dim=-1) for i in range(len(data)): valid_response_ids = response_ids[i][: response_length[i]] response_str = self.tokenizer.decode( - valid_response_ids, skip_special_tokens=self.config.skip_special_tokens - ) + valid_response_ids, skip_special_tokens=self.config.skip_special_tokens) ground_truth = data.non_tensor_batch["ground_truth"][i] score = self.reward_fn(response_str, ground_truth) @@ -97,19 +110,25 @@ def compute_reward(self, data: DataProto) -> Tuple[torch.Tensor, Dict[str, List[ class BatchFunctionRewardManager(FunctionRewardManager): reward_fn: BatchRewardFunction - def compute_reward(self, data: DataProto) -> Tuple[torch.Tensor, Dict[str, List[float]]]: + def compute_reward( + self, data: DataProto + ) -> Tuple[torch.Tensor, Dict[str, List[float]]]: response_str, ground_truth = [], [] response_ids = data.batch["responses"] response_length = data.batch["response_mask"].sum(dim=-1) for i in range(len(data)): valid_response_ids = response_ids[i][: response_length[i]] response_str.append( - self.tokenizer.decode(valid_response_ids, skip_special_tokens=self.config.skip_special_tokens) + self.tokenizer.decode( + valid_response_ids, + skip_special_tokens=self.config.skip_special_tokens, + ) ) ground_truth.append(data.non_tensor_batch["ground_truth"][i]) scores = self.reward_fn(response_str, ground_truth) - reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) + reward_tensor = torch.zeros_like( + data.batch["responses"], dtype=torch.float32) reward_metrics = defaultdict(list) for i, score in enumerate(scores): reward_tensor[i, response_length[i] - 1] = score["overall"] diff --git a/Agent0/curriculum_train/verl/workers/rollout/__init__.py b/Agent0/curriculum_train/verl/workers/rollout/__init__.py index 89cbcf0..9bd2aea 100644 --- a/Agent0/curriculum_train/verl/workers/rollout/__init__.py +++ b/Agent0/curriculum_train/verl/workers/rollout/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/workers/rollout/base.py b/Agent0/curriculum_train/verl/workers/rollout/base.py index 0a07eee..c985574 100644 --- a/Agent0/curriculum_train/verl/workers/rollout/base.py +++ b/Agent0/curriculum_train/verl/workers/rollout/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/workers/rollout/config.py b/Agent0/curriculum_train/verl/workers/rollout/config.py index e4c96ca..5237340 100644 --- a/Agent0/curriculum_train/verl/workers/rollout/config.py +++ b/Agent0/curriculum_train/verl/workers/rollout/config.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/workers/rollout/vllm_rollout_spmd.py b/Agent0/curriculum_train/verl/workers/rollout/vllm_rollout_spmd.py index 13cb4d7..521b013 100644 --- a/Agent0/curriculum_train/verl/workers/rollout/vllm_rollout_spmd.py +++ b/Agent0/curriculum_train/verl/workers/rollout/vllm_rollout_spmd.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,24 +31,36 @@ from .config import RolloutConfig import traceback -def _repeat_interleave(value: Union[torch.Tensor, np.ndarray], repeats: int) -> Union[torch.Tensor, List[Any]]: + + +def _repeat_interleave( + value: Union[torch.Tensor, np.ndarray], repeats: int +) -> Union[torch.Tensor, List[Any]]: if isinstance(value, torch.Tensor): return value.repeat_interleave(repeats, dim=0) else: return np.repeat(value, repeats, axis=0) -def _get_logit_bias(model_path: str, trust_remote_code: bool) -> Optional[Dict[int, float]]: +def _get_logit_bias( + model_path: str, trust_remote_code: bool +) -> Optional[Dict[int, float]]: processor = get_processor(model_path, trust_remote_code=trust_remote_code) if processor is not None and hasattr(processor, "image_token"): - image_token_id = processor.tokenizer.convert_tokens_to_ids(processor.image_token) + image_token_id = processor.tokenizer.convert_tokens_to_ids( + processor.image_token + ) return {image_token_id: -100} else: return None class vLLMRollout(BaseRollout): - def __init__(self, model_path: str, config: RolloutConfig, tokenizer: PreTrainedTokenizer): + def __init__( + self, + model_path: str, + config: RolloutConfig, + tokenizer: PreTrainedTokenizer): """A vLLM rollout. It requires the module is supported by the vllm. Args: @@ -61,14 +73,21 @@ def __init__(self, model_path: str, config: RolloutConfig, tokenizer: PreTrained self.config = config self.pad_token_id = tokenizer.pad_token_id if config.tensor_parallel_size > torch.distributed.get_world_size(): - raise ValueError("Tensor parallelism size should be less than world size.") - - if config.max_num_batched_tokens < config.prompt_length + config.response_length: - raise ValueError("max_num_batched_tokens should be greater than prompt_length + response_length.") + raise ValueError( + "Tensor parallelism size should be less than world size.") + + if ( + config.max_num_batched_tokens + < config.prompt_length + config.response_length + ): + raise ValueError( + "max_num_batched_tokens should be greater than prompt_length + response_length." + ) engine_kwargs = {} if config.limit_images: - engine_kwargs["limit_mm_per_prompt"] = {"image": config.limit_images} + engine_kwargs["limit_mm_per_prompt"] = { + "image": config.limit_images} self.inference_engine = LLM( model=model_path, @@ -77,7 +96,8 @@ def __init__(self, model_path: str, config: RolloutConfig, tokenizer: PreTrained load_format="dummy", dtype=PrecisionType.to_str(PrecisionType.to_dtype(config.dtype)), seed=config.seed, - max_model_len=config.max_model_len or config.prompt_length + config.response_length, + max_model_len=config.max_model_len + or config.prompt_length + config.response_length, distributed_executor_backend="external_launcher", tensor_parallel_size=config.tensor_parallel_size, gpu_memory_utilization=config.gpu_memory_utilization, @@ -97,11 +117,13 @@ def __init__(self, model_path: str, config: RolloutConfig, tokenizer: PreTrained sampling_kwargs = { "max_tokens": config.response_length, "detokenize": False, - "logit_bias": _get_logit_bias(model_path, trust_remote_code=config.trust_remote_code), + "logit_bias": _get_logit_bias( + model_path, trust_remote_code=config.trust_remote_code + ), } default_sampling_params = SamplingParams() for key in config.to_dict().keys(): - if key == 'seed': + if key == "seed": continue if hasattr(default_sampling_params, key): sampling_kwargs[key] = getattr(config, key) @@ -131,7 +153,8 @@ def generate_sequences(self, prompts: DataProto) -> DataProto: # traceback.print_stack() # exit() # left-padded attention_mask - input_ids: torch.Tensor = prompts.batch["input_ids"] # (bs, prompt_length) + # (bs, prompt_length) + input_ids: torch.Tensor = prompts.batch["input_ids"] attention_mask: torch.Tensor = prompts.batch["attention_mask"] position_ids: torch.Tensor = prompts.batch["position_ids"] eos_token_id: int = prompts.meta_info["eos_token_id"] @@ -144,36 +167,59 @@ def generate_sequences(self, prompts: DataProto) -> DataProto: if "multi_modal_data" in non_tensor_batch: vllm_inputs = [] for raw_prompt_ids, multi_modal_data in zip( - non_tensor_batch.pop("raw_prompt_ids"), non_tensor_batch.pop("multi_modal_data") + non_tensor_batch.pop("raw_prompt_ids"), + non_tensor_batch.pop("multi_modal_data"), ): - vllm_inputs.append({"prompt_token_ids": list(raw_prompt_ids), "multi_modal_data": multi_modal_data}) + vllm_inputs.append( + { + "prompt_token_ids": list(raw_prompt_ids), + "multi_modal_data": multi_modal_data, + } + ) else: vllm_inputs = [ - {"prompt_token_ids": list(raw_prompt_ids)} for raw_prompt_ids in non_tensor_batch.pop("raw_prompt_ids") + {"prompt_token_ids": list(raw_prompt_ids)} + for raw_prompt_ids in non_tensor_batch.pop("raw_prompt_ids") ] # users can customize different sampling_params at different run with self.update_sampling_params(**prompts.meta_info): completions: List[RequestOutput] = self.inference_engine.generate( - prompts=vllm_inputs, sampling_params=self.sampling_params, use_tqdm=False + prompts=vllm_inputs, + sampling_params=self.sampling_params, + use_tqdm=False, ) - response_ids = [output.token_ids for completion in completions for output in completion.outputs] + response_ids = [ + output.token_ids + for completion in completions + for output in completion.outputs + ] response_ids = VF.pad_2d_list_to_length( - response_ids, self.pad_token_id, max_length=self.config.response_length - ).to(input_ids.device) + response_ids, + self.pad_token_id, + max_length=self.config.response_length).to( + input_ids.device) if self.sampling_params.n > 1: batch_size = batch_size * self.sampling_params.n - input_ids = _repeat_interleave(input_ids, self.sampling_params.n) - attention_mask = _repeat_interleave(attention_mask, self.sampling_params.n) - position_ids = _repeat_interleave(position_ids, self.sampling_params.n) + input_ids = _repeat_interleave( + input_ids, self.sampling_params.n) + attention_mask = _repeat_interleave( + attention_mask, self.sampling_params.n + ) + position_ids = _repeat_interleave( + position_ids, self.sampling_params.n) sequence_ids = torch.cat([input_ids, response_ids], dim=-1) response_length = response_ids.size(1) - delta_position_id = torch.arange(1, response_length + 1, device=position_ids.device) - delta_position_id = delta_position_id.view(1, -1).expand(batch_size, -1) + delta_position_id = torch.arange( + 1, response_length + 1, device=position_ids.device + ) + delta_position_id = delta_position_id.view( + 1, -1).expand(batch_size, -1) if position_ids.dim() == 3: # qwen2vl mrope - delta_position_id = delta_position_id.view(batch_size, 1, -1).expand(batch_size, 3, -1) + delta_position_id = delta_position_id.view( + batch_size, 1, -1).expand(batch_size, 3, -1) # prompt: left pad + response: right pad # attention_mask: [0,0,0,0,1,1,1,1 | 1,1,1,0,0,0,0,0] @@ -181,10 +227,13 @@ def generate_sequences(self, prompts: DataProto) -> DataProto: response_position_ids = position_ids[..., -1:] + delta_position_id position_ids = torch.cat([position_ids, response_position_ids], dim=-1) response_mask = VF.get_response_mask( - response_ids=response_ids, eos_token_id=eos_token_id, dtype=attention_mask.dtype + response_ids=response_ids, + eos_token_id=eos_token_id, + dtype=attention_mask.dtype, ) attention_mask = torch.cat((attention_mask, response_mask), dim=-1) - # all the tp ranks should contain the same data here. data in all ranks are valid + # all the tp ranks should contain the same data here. data in all ranks + # are valid batch = TensorDict( { "prompts": input_ids, diff --git a/Agent0/curriculum_train/verl/workers/sharding_manager/__init__.py b/Agent0/curriculum_train/verl/workers/sharding_manager/__init__.py index 88eaee4..fda0477 100644 --- a/Agent0/curriculum_train/verl/workers/sharding_manager/__init__.py +++ b/Agent0/curriculum_train/verl/workers/sharding_manager/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,4 +18,8 @@ from .fsdp_vllm import FSDPVLLMShardingManager -__all__ = ["BaseShardingManager", "FSDPUlyssesShardingManager", "FSDPVLLMShardingManager"] +__all__ = [ + "BaseShardingManager", + "FSDPUlyssesShardingManager", + "FSDPVLLMShardingManager", +] diff --git a/Agent0/curriculum_train/verl/workers/sharding_manager/base.py b/Agent0/curriculum_train/verl/workers/sharding_manager/base.py index dc29756..6ce3197 100644 --- a/Agent0/curriculum_train/verl/workers/sharding_manager/base.py +++ b/Agent0/curriculum_train/verl/workers/sharding_manager/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/curriculum_train/verl/workers/sharding_manager/fsdp_ulysses.py b/Agent0/curriculum_train/verl/workers/sharding_manager/fsdp_ulysses.py index c2ce5b9..5b322e0 100644 --- a/Agent0/curriculum_train/verl/workers/sharding_manager/fsdp_ulysses.py +++ b/Agent0/curriculum_train/verl/workers/sharding_manager/fsdp_ulysses.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,7 +18,10 @@ from torch.distributed.device_mesh import DeviceMesh from ...protocol import DataProto, all_gather_data_proto -from ...utils.ulysses import get_ulysses_sequence_parallel_group, set_ulysses_sequence_parallel_group +from ...utils.ulysses import ( + get_ulysses_sequence_parallel_group, + set_ulysses_sequence_parallel_group, +) from .base import BaseShardingManager @@ -34,7 +37,8 @@ def __init__(self, device_mesh: DeviceMesh): def __enter__(self): if self.device_mesh is not None: self.prev_sp_group = get_ulysses_sequence_parallel_group() - set_ulysses_sequence_parallel_group(self.device_mesh["sp"].get_group()) + set_ulysses_sequence_parallel_group( + self.device_mesh["sp"].get_group()) def __exit__(self, exc_type, exc_value, traceback): if self.device_mesh is not None: diff --git a/Agent0/curriculum_train/verl/workers/sharding_manager/fsdp_vllm.py b/Agent0/curriculum_train/verl/workers/sharding_manager/fsdp_vllm.py index 11f1090..103eaca 100644 --- a/Agent0/curriculum_train/verl/workers/sharding_manager/fsdp_vllm.py +++ b/Agent0/curriculum_train/verl/workers/sharding_manager/fsdp_vllm.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -21,7 +21,9 @@ from torch.distributed._tensor import DTensor from torch.distributed.checkpoint.state_dict import get_model_state_dict from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp.fully_sharded_data_parallel import ( + FullyShardedDataParallel as FSDP, +) from transformers import PreTrainedModel from vllm import LLM from vllm.distributed import parallel_state as vllm_ps @@ -55,20 +57,31 @@ def __init__( self.torch_random_states = torch.cuda.get_rng_state() # get a random rng states 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 + 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) - def _rename_weight_keys(self, actor_weights: Dict[str, Union[torch.Tensor, DTensor]], model: PreTrainedModel): - # convert state dict keys: https://github.com/huggingface/transformers/pull/38385 + def _rename_weight_keys( + self, + actor_weights: Dict[str, Union[torch.Tensor, DTensor]], + model: PreTrainedModel, + ): + # convert state dict keys: + # https://github.com/huggingface/transformers/pull/38385 if not hasattr(model, "_checkpoint_conversion_mapping"): return actor_weights - reverse_key_mapping = {v: k for k, v in model._checkpoint_conversion_mapping.items()} + reverse_key_mapping = { + v: k for k, v in model._checkpoint_conversion_mapping.items() + } original_weights = {} for key, value in actor_weights.items(): for pattern, replacement in reverse_key_mapping.items(): - replacement = replacement.lstrip("^") # strip off un-needed chars and patterns + replacement = replacement.lstrip( + "^" + ) # strip off un-needed chars and patterns replacement = re.sub(r"\(.*\)", "", replacement) key, n_replace = re.subn(pattern, replacement, key) # Early exit of the loop @@ -92,30 +105,40 @@ def __enter__(self): # to speed up memory allocations. # # pytorch: https://pytorch.org/docs/stable/notes/cuda.html#memory-management - # vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/device_allocator/cumem.py#L103 + # vllm: + # https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/device_allocator/cumem.py#L103 torch.cuda.empty_cache() print_gpu_memory_usage("Before state_dict() in sharding manager") actor_weights = get_model_state_dict(self.module) - actor_weights = self._rename_weight_keys(actor_weights, self.module._fsdp_wrapped_module) + actor_weights = self._rename_weight_keys( + actor_weights, self.module._fsdp_wrapped_module + ) print_gpu_memory_usage("After state_dict() in sharding manager") - if "tags" in inspect.signature(self.inference_engine.wake_up).parameters: + if "tags" in inspect.signature( + self.inference_engine.wake_up).parameters: self.inference_engine.wake_up(tags=["weights"]) else: self.inference_engine.wake_up() - model = self.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner.model + model = ( + self.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner.model + ) model.load_weights(self._make_weight_iterator(actor_weights)) print_gpu_memory_usage("After sync model weights in sharding manager") del actor_weights torch.cuda.empty_cache() - if "tags" in inspect.signature(self.inference_engine.wake_up).parameters: + if "tags" in inspect.signature( + self.inference_engine.wake_up).parameters: self.inference_engine.wake_up(tags=["kv_cache"]) - print_gpu_memory_usage("After del state_dict and empty_cache in sharding manager") - # important: need to manually set the random states of each tp to be identical. + print_gpu_memory_usage( + "After del state_dict and empty_cache in sharding manager" + ) + # 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) diff --git a/Agent0/curriculum_train/vllm_service_init/start_vllm_server_tool.py b/Agent0/curriculum_train/vllm_service_init/start_vllm_server_tool.py index 888960b..c998653 100644 --- a/Agent0/curriculum_train/vllm_service_init/start_vllm_server_tool.py +++ b/Agent0/curriculum_train/vllm_service_init/start_vllm_server_tool.py @@ -1,7 +1,7 @@ #!/usr/bin/env python # -*- coding: utf-8 -*- -''' -This script enhances the LLM's problem-solving capabilities by integrating a code execution tool. +""" +This script enhances the LLM's problem-solving capabilities by integrating a code execution tool. It processes each question through a multi-turn conversational approach, allowing the model to generate, execute, and reason based on code output. The generation process for each of the 10 candidates is now a stateful, iterative loop. @@ -13,7 +13,7 @@ # 3. Run the server python your_server_file_name.py --port 5000 --model_path Qwen/Qwen3-4B-Base -''' +""" from flask import Flask, request, jsonify import vllm @@ -34,15 +34,16 @@ # ---------------------------- Code Execution Tool --------------------------- # SANDBOX_API_URLS = [ - 'IP1:PORT1/run_code', - 'IP2:PORT2/run_code', - 'IP3:PORT3/run_code', - 'IP4:PORT4/run_code' + "IP1:PORT1/run_code", + "IP2:PORT2/run_code", + "IP3:PORT3/run_code", + "IP4:PORT4/run_code", ] api_counter_lock = threading.Lock() api_counter = 0 + def execute_code_in_sandbox(code: str) -> str: """ Calls an external sandbox API to execute Python code, with load balancing. @@ -54,8 +55,10 @@ def execute_code_in_sandbox(code: str) -> str: try: payload = {"code": code, "language": "python"} - headers = {'Content-Type': 'application/json'} - response = requests.post(target_url, headers=headers, data=json.dumps(payload), timeout=20) + headers = {"Content-Type": "application/json"} + response = requests.post( + target_url, headers=headers, data=json.dumps(payload), timeout=20 + ) response.raise_for_status() result = response.json() @@ -65,8 +68,9 @@ def execute_code_in_sandbox(code: str) -> str: stdout = run_info.get("stdout", "") return stdout if stdout else "[No output]" else: - stderr = run_info.get('stderr', '') - return f"Execution failed with status: {run_info.get('status')}\nStderr: {stderr}" + stderr = run_info.get("stderr", "") + return f"Execution failed with status: { + run_info.get('status')}\nStderr: {stderr}" else: return f"API Error: {result}" except Exception as e: @@ -76,14 +80,18 @@ def execute_code_in_sandbox(code: str) -> str: # ---------------------------- Initial Setup --------------------------------- # parser = argparse.ArgumentParser() -parser.add_argument('--port', type=str, default='5000') -parser.add_argument('--model_path', type=str, default='Qwen/Qwen3-4B-Base') -parser.add_argument('--gpu_mem_util', type=float, default=0.8, - help='The maximum GPU memory utilization fraction for vLLM.') +parser.add_argument("--port", type=str, default="5000") +parser.add_argument("--model_path", type=str, default="Qwen/Qwen3-4B-Base") +parser.add_argument( + "--gpu_mem_util", + type=float, + default=0.8, + help="The maximum GPU memory utilization fraction for vLLM.", +) args = parser.parse_args() -print('[init] Loading model...') +print("[init] Loading model...") tokenizer = AutoTokenizer.from_pretrained(args.model_path) model = vllm.LLM( model=args.model_path, @@ -96,7 +104,7 @@ def execute_code_in_sandbox(code: str) -> str: temperature=0.7, top_p=0.9, n=1, - stop_token_ids=[tokenizer.eos_token_id] + stop_token_ids=[tokenizer.eos_token_id], ) SYSTEM_PROMPT = ( @@ -108,15 +116,15 @@ def execute_code_in_sandbox(code: str) -> str: "Code Format:\n" "Each code snippet is wrapped between ```. You need to use print() to output intermediate results.\n" "Answer Format:\n" - "The last part of your response should be: \\boxed{...}" -) + "The last part of your response should be: \\boxed{...}") # ---------------------------- GPU Idle Worker ------------------- # stop_event = threading.Event() pause_event = threading.Event() + def gpu_idle_worker(): - print('[idle_worker] GPU idle worker started.') + print("[idle_worker] GPU idle worker started.") running = True while not stop_event.is_set(): if pause_event.is_set(): @@ -128,31 +136,44 @@ def gpu_idle_worker(): if not running: running = True try: - a = torch.rand((2000, 2000), dtype=torch.float32, device='cuda') - b = torch.rand((2000, 2000), dtype=torch.float32, device='cuda') + a = torch.rand((2000, 2000), dtype=torch.float32, device="cuda") + b = torch.rand((2000, 2000), dtype=torch.float32, device="cuda") torch.matmul(a, b) torch.cuda.synchronize() except RuntimeError: time.sleep(1) - print('[idle_worker] GPU idle worker stopped.') + print("[idle_worker] GPU idle worker stopped.") + idle_thread = threading.Thread(target=gpu_idle_worker, daemon=True) idle_thread.start() + # ---------------------------- Core Logic (Refactored) ----------------------- # -@stopit.threading_timeoutable(default='TIMED_OUT') +@stopit.threading_timeoutable(default="TIMED_OUT") def grade_answer_with_timeout(res1, res2): return grade_answer(res1, res2) + sandbox_executor = ThreadPoolExecutor(max_workers=64) -def generate_with_tool_use(question: str, num_candidates: int = 10, max_turns: int = 4): + +def generate_with_tool_use( + question: str, + num_candidates: int = 10, + max_turns: int = 4): """ Generates answers using a multi-turn conversation loop (up to max_turns). Handles code execution and history updates dynamically. """ # Initialize conversation history for all candidates - conversations = [[{'role': 'system', 'content': SYSTEM_PROMPT}, {'role': 'user', 'content': question}] for _ in range(num_candidates)] + conversations = [ + [ + {"role": "system", "content": SYSTEM_PROMPT}, + {"role": "user", "content": question}, + ] + for _ in range(num_candidates) + ] final_assistant_messages = [""] * num_candidates active_indices = list(range(num_candidates)) @@ -161,10 +182,18 @@ def generate_with_tool_use(question: str, num_candidates: int = 10, max_turns: i break # Prepare prompts only for active candidates - prompts = [tokenizer.apply_chat_template(conversations[i], tokenize=False, add_generation_prompt=True) for i in active_indices] - + prompts = [ + tokenizer.apply_chat_template( + conversations[i], tokenize=False, add_generation_prompt=True + ) + for i in active_indices + ] + # Batch generate - responses = model.generate(prompts, sampling_params_single_turn, use_tqdm=False) + responses = model.generate( + prompts, + sampling_params_single_turn, + use_tqdm=False) tasks_to_run = [] indices_with_code = set() @@ -173,30 +202,40 @@ def generate_with_tool_use(question: str, num_candidates: int = 10, max_turns: i for i, response in enumerate(responses): original_index = active_indices[i] model_output = response.outputs[0].text.strip() - + # Clean up potential incomplete code blocks code_block_start_tag = "```python" code_block_end_tag = "```" start_index = model_output.find(code_block_start_tag) if start_index != -1: - end_index = model_output.find(code_block_end_tag, start_index + len(code_block_start_tag)) + end_index = model_output.find( + code_block_end_tag, start_index + len(code_block_start_tag) + ) if end_index != -1: - model_output = model_output[:end_index + len(code_block_end_tag)] - + model_output = model_output[: end_index + + len(code_block_end_tag)] + # Update history - conversations[original_index].append({'role': 'assistant', 'content': model_output}) + conversations[original_index].append( + {"role": "assistant", "content": model_output} + ) # Check for Code - code_match = re.search(r"```python\n(.*?)\n```", model_output, re.DOTALL) - + code_match = re.search( + r"```python\n(.*?)\n```", + model_output, + re.DOTALL) + # Check for Boxed Answer - has_boxed = r'\boxed' in model_output + has_boxed = r"\boxed" in model_output if code_match and not has_boxed: # Found code, no final answer yet -> Queue for execution code_to_run = (code_match.group(1) or "").strip() if code_to_run: - future = sandbox_executor.submit(execute_code_in_sandbox, code_to_run) + future = sandbox_executor.submit( + execute_code_in_sandbox, code_to_run + ) tasks_to_run.append((future, original_index)) indices_with_code.add(original_index) else: @@ -206,8 +245,9 @@ def generate_with_tool_use(question: str, num_candidates: int = 10, max_turns: i # Found answer -> Mark as finished final_assistant_messages[original_index] = model_output else: - # Pure text reasoning -> Will continue to next turn if logic requires, - # or strictly speaking, we keep it active to allow further reasoning. + # Pure text reasoning -> Will continue to next turn if logic requires, + # or strictly speaking, we keep it active to allow further + # reasoning. pass # Step 2: Collect Sandbox Results @@ -222,64 +262,74 @@ def generate_with_tool_use(question: str, num_candidates: int = 10, max_turns: i next_active_indices = [] for i, response in enumerate(responses): original_index = active_indices[i] - + # If we already found a boxed answer, this candidate is done. if final_assistant_messages[original_index]: continue - + # If it had code, append result and keep active if original_index in indices_with_code: - exec_result = results_map.get(original_index, "Result not found.") + exec_result = results_map.get( + original_index, "Result not found.") tool_feedback = f"Code execution result: {exec_result}" - conversations[original_index].append({'role': 'user', 'content': tool_feedback}) + conversations[original_index].append( + {"role": "user", "content": tool_feedback} + ) next_active_indices.append(original_index) - + # If it was just text (and no boxed), we keep it active for the next turn # (assuming it needs more steps), unless it was the last turn. else: next_active_indices.append(original_index) - + active_indices = next_active_indices - # Fill in any candidates that didn't finish with \boxed with their last output + # Fill in any candidates that didn't finish with \boxed with their last + # output for i in range(num_candidates): if not final_assistant_messages[i]: # Use the last assistant message as the best effort result # Traverse backwards to find the last assistant message for msg in reversed(conversations[i]): - if msg['role'] == 'assistant': - final_assistant_messages[i] = msg['content'] + if msg["role"] == "assistant": + final_assistant_messages[i] = msg["content"] break - + return final_assistant_messages def consolidate_and_grade(question, golden_answer, assistant_messages): - '''Consolidates and grades LLM outputs for a single question.''' + """Consolidates and grades LLM outputs for a single question.""" results = [extract_boxed_content(msg) for msg in assistant_messages] - + answer_counts = {} for res in results: - if not res: continue + if not res: + continue matched = False - + for exist_ans in list(answer_counts.keys()): - if res == exist_ans or ('no ' in res.lower() and 'no ' in exist_ans.lower()): + if res == exist_ans or ( + "no " in res.lower() and "no " in exist_ans.lower() + ): answer_counts[exist_ans] += 1 matched = True break - + try: is_match = False - match_result_1 = grade_answer_with_timeout(res, exist_ans, timeout=20) - if match_result_1 and match_result_1 != 'TIMED_OUT': + match_result_1 = grade_answer_with_timeout( + res, exist_ans, timeout=20) + if match_result_1 and match_result_1 != "TIMED_OUT": is_match = True if not is_match: - match_result_2 = grade_answer_with_timeout(exist_ans, res, timeout=20) - if match_result_2 and match_result_2 != 'TIMED_OUT': + match_result_2 = grade_answer_with_timeout( + exist_ans, res, timeout=20 + ) + if match_result_2 and match_result_2 != "TIMED_OUT": is_match = True - + if is_match: answer_counts[exist_ans] += 1 matched = True @@ -287,12 +337,12 @@ def consolidate_and_grade(question, golden_answer, assistant_messages): except Exception: continue - + if not matched: answer_counts[res] = 1 if not answer_counts: - majority_ans, max_count = '', 0 + majority_ans, max_count = "", 0 else: majority_ans = max(answer_counts, key=answer_counts.get) max_count = answer_counts[majority_ans] @@ -300,66 +350,91 @@ def consolidate_and_grade(question, golden_answer, assistant_messages): score = max_count / len(assistant_messages) if assistant_messages else 0.0 return { - 'question': question, - 'answer': majority_ans, - 'score': score if grade_answer(majority_ans, golden_answer) and score > 0.1 else 0, - 'all_outputs': assistant_messages, - 'extracted_results': results + "question": question, + "answer": majority_ans, + "score": ( + score if grade_answer( + majority_ans, + golden_answer) and score > 0.1 else 0), + "all_outputs": assistant_messages, + "extracted_results": results, } + # ---------------------------- Flask Application --------------------------- # app = Flask(__name__) -@app.route('/hello', methods=['GET']) + +@app.route("/hello", methods=["GET"]) def hello(): pause_event.set() torch.cuda.synchronize() - name = request.args.get('name', 'None') - - with open(name, 'r') as f: + name = request.args.get("name", "None") + + with open(name, "r") as f: data = json.load(f) os.remove(name) - questions = [item.get('question', '') for item in data] - answers = [item.get('answer', '') for item in data] + questions = [item.get("question", "") for item in data] + answers = [item.get("answer", "") for item in data] results_all = [] - + # Using TQDM for clean progress visualization - progress_bar = tqdm(zip(questions, answers), total=len(questions), desc=f"Processing {os.path.basename(name)}") - + progress_bar = tqdm( + zip(questions, answers), + total=len(questions), + desc=f"Processing {os.path.basename(name)}", + ) + for q, a in progress_bar: try: if q and a: # Multi-turn generation - final_assistant_messages = generate_with_tool_use(q, max_turns=4) - + final_assistant_messages = generate_with_tool_use( + q, max_turns=4) + # Consolidate and Grade item = consolidate_and_grade(q, a, final_assistant_messages) results_all.append(item) else: - results_all.append({'question': q, 'answer': a, 'score': -1, 'all_outputs': [], 'extracted_results': []}) + results_all.append( + { + "question": q, + "answer": a, + "score": -1, + "all_outputs": [], + "extracted_results": [], + } + ) except Exception as e: # Only printing critical errors to not mess up TQDM too much - print(f'\n[server] Error processing question: {str(e)}') - results_all.append({ - 'question': q, 'answer': a, 'score': -1, 'error': f'unhandled exception: {str(e)}' - }) - - out_path = name.replace('.json', '_results.json') - with open(out_path, 'w') as f: + print(f"\n[server] Error processing question: {str(e)}") + results_all.append( + { + "question": q, + "answer": a, + "score": -1, + "error": f"unhandled exception: {str(e)}", + } + ) + + out_path = name.replace(".json", "_results.json") + with open(out_path, "w") as f: json.dump(results_all, f, indent=4) pause_event.clear() - return jsonify({'message': f'Processed {name}, results saved to {out_path}.'}) + return jsonify( + {"message": f"Processed {name}, results saved to {out_path}."}) + # ------------------------- Main Application Entrypoint --------------------------- # -if __name__ == '__main__': +if __name__ == "__main__": try: - app.run(host='127.0.0.1', port=int(args.port), threaded=True) + app.run(host="127.0.0.1", port=int(args.port), threaded=True) finally: stop_event.set() if idle_thread.is_alive(): idle_thread.join() - print('[main] Application shutdown complete.') \ No newline at end of file + print("[main] Application shutdown complete.") diff --git a/Agent0/executor_train/LICENSE b/Agent0/executor_train/LICENSE index 297adda..c73360f 100644 --- a/Agent0/executor_train/LICENSE +++ b/Agent0/executor_train/LICENSE @@ -1,6 +1,6 @@ MIT License -Copyright (c) 2025 TIGER Lab +Copyright (c) 2025-2026 TIGER Lab Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal diff --git a/Agent0/executor_train/eval_service/app.py b/Agent0/executor_train/eval_service/app.py index 63b347a..c2c1126 100644 --- a/Agent0/executor_train/eval_service/app.py +++ b/Agent0/executor_train/eval_service/app.py @@ -16,32 +16,33 @@ # Set up logging logging.basicConfig( level=logging.ERROR, - format='%(asctime)s - %(name)s - %(levelname)s - %(message)s', - handlers=[ - logging.FileHandler("error_log.txt"), - logging.StreamHandler() - ] + format="%(asctime)s - %(name)s - %(levelname)s - %(message)s", + handlers=[logging.FileHandler("error_log.txt"), logging.StreamHandler()], ) logger = logging.getLogger(__name__) -def create_app(server_config: ServerConfig, model_config: ModelConfig, tool_config: ToolConfig) -> FastAPI: + +def create_app( + server_config: ServerConfig, + model_config: ModelConfig, + tool_config: ToolConfig) -> FastAPI: """ Create and configure the FastAPI application - + Args: server_config: Server configuration object model_config: Model configuration object tool_config: Tool configuration object - + Returns: Configured FastAPI application instance """ app = FastAPI( title="LLM Code Tool Service", description="Large language model code tool calling service compatible with OpenAI API", - version="1.0.0" + version="1.0.0", ) - + # Add CORS middleware to allow cross-origin requests app.add_middleware( CORSMiddleware, @@ -50,18 +51,21 @@ def create_app(server_config: ServerConfig, model_config: ModelConfig, tool_conf allow_methods=["*"], allow_headers=["*"], ) - + # Set debug mode based on environment - if hasattr(server_config, "environment") and server_config.environment == "development": + if ( + hasattr(server_config, "environment") + and server_config.environment == "development" + ): app.debug = True - + # Initialize the model service model_service = ModelService(model_config, tool_config) model_service.load_model() - + # Store service in application state app.state.model_service = model_service - + # Add middleware for global exception handling @app.middleware("http") async def log_exceptions(request: Request, call_next): @@ -71,74 +75,94 @@ async def log_exceptions(request: Request, call_next): error_details = traceback.format_exc() logger.error(f"Unhandled exception: {str(e)}\n{error_details}") raise - + @app.post("/completions") async def chat_completions(request: Request): """ Chat completion API endpoint compatible with OpenAI - + Processes chat messages and returns model-generated responses with tool calling capabilities """ try: request_body = await request.json() - logger.debug(f"Received completions request: {json.dumps(request_body)}") + logger.debug( + f"Received completions request: { + json.dumps(request_body)}") response = await app.state.model_service.completions_async(request_body) return response except Exception as e: error_details = traceback.format_exc() - logger.error(f"Error in completions endpoint: {str(e)}\n{error_details}") - raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") - + logger.error( + f"Error in completions endpoint: { + str(e)}\n{error_details}") + raise HTTPException( + status_code=500, detail=f"Internal server error: {str(e)}" + ) + @app.post("/chat/completions") async def completions(request: Request): """ Chat completion API endpoint compatible with OpenAI - + Processes chat messages and returns model-generated responses with tool calling capabilities """ try: request_body = await request.json() - logger.debug(f"Received chat completions request: {json.dumps(request_body)}") - response = await app.state.model_service.chat_completions_async(request_body) + logger.debug( + f"Received chat completions request: { + json.dumps(request_body)}") + response = await app.state.model_service.chat_completions_async( + request_body + ) return response except Exception as e: error_details = traceback.format_exc() - logger.error(f"Error in chat completions endpoint: {str(e)}\n{error_details}") - raise HTTPException(status_code=500, detail=f"Internal server error: {str(e)}") - + logger.error( + f"Error in chat completions endpoint: { + str(e)}\n{error_details}") + raise HTTPException( + status_code=500, detail=f"Internal server error: {str(e)}" + ) + @app.get("/health") async def health_check(): """Health check endpoint to verify service availability""" return {"status": "healthy"} - + return app + async def main_async(): # Set up command line argument parsing hf_parser = HfArgumentParser((ServerConfig, ModelConfig, ToolConfig)) - server_config, model_config, tool_config = hf_parser.parse_args_into_dataclasses() + server_config, model_config, tool_config = hf_parser.parse_args_into_dataclasses() tool_config.post_init() - + # Create and run the application app = create_app(server_config, model_config, tool_config) - + # Configure and start the server with enhanced logging config = uvicorn.Config( - app, - host=server_config.host, - port=server_config.port, - log_level=server_config.log_level, # Changed from "error" to "debug" for better visibility - ws_max_queue=server_config.ws_max_queue, - workers=server_config.workers*model_config.num_models, + app, + host=server_config.host, + port=server_config.port, + log_level=server_config.log_level, + # Changed from "error" to "debug" for better visibility + ws_max_queue=server_config.ws_max_queue, + workers=server_config.workers * model_config.num_models, access_log=True, - timeout_keep_alive=server_config.timeout_keep_alive # Added keep-alive timeout setting + # Added keep-alive timeout setting + timeout_keep_alive=server_config.timeout_keep_alive, ) server = uvicorn.Server(config) await server.serve() + def main(): import asyncio + asyncio.run(main_async()) + if __name__ == "__main__": - main() \ No newline at end of file + main() diff --git a/Agent0/executor_train/eval_service/config.py b/Agent0/executor_train/eval_service/config.py index 64f9a2c..9db0e10 100644 --- a/Agent0/executor_train/eval_service/config.py +++ b/Agent0/executor_train/eval_service/config.py @@ -2,6 +2,7 @@ from typing import Optional, List, Dict, Any, Union from dataclasses import dataclass + @dataclass class ModelConfig: model: str @@ -10,18 +11,25 @@ class ModelConfig: trust_remote_code: bool = True num_models: int = 1 max_model_len: int = 4096 + + @dataclass class ToolConfig: tool_server_url: str = "http://localhost:30150/get_observation" max_turns: int = 5 # max generation turns - truncate_obs_side: str = "left" # "left" or "right", which side to truncate when the observation is too long + truncate_obs_side: str = ( + "left" # "left" or "right", which side to truncate when the observation is too long + ) action_stop_tokens: str = None max_obs_length: int = 512 # maximum length of observation - enable_mtrl: bool=False - mtrl_sep: str=None # "\n<|im_start|>system\n{obs}<|im_end|>\n<|im_start|>assistant\n" - turn_end_token: str="<|im_end|>" - min_turns: int=0 - + enable_mtrl: bool = False + mtrl_sep: str = ( + # "\n<|im_start|>system\n{obs}<|im_end|>\n<|im_start|>assistant\n" + None + ) + turn_end_token: str = "<|im_end|>" + min_turns: int = 0 + def post_init(self): """ Post-initialization processing for ToolConfig (will not call automatically) @@ -30,15 +38,20 @@ def post_init(self): if isinstance(self.action_stop_tokens, str): if os.path.exists(self.action_stop_tokens): with open(self.action_stop_tokens, "r") as f: - self.action_stop_tokens = f.read().split(',') + self.action_stop_tokens = f.read().split(",") else: - self.action_stop_tokens = self.action_stop_tokens.split(',') - self.action_stop_tokens = [token.strip('\n ') for token in self.action_stop_tokens] - self.action_stop_tokens = [token for token in self.action_stop_tokens if token] + self.action_stop_tokens = self.action_stop_tokens.split(",") + self.action_stop_tokens = [ + token.strip("\n ") for token in self.action_stop_tokens + ] + self.action_stop_tokens = [ + token for token in self.action_stop_tokens if token + ] else: self.action_stop_tokens = None print(f"using action_stop_tokens: {self.action_stop_tokens}") + @dataclass class ServerConfig: host: str = "0.0.0.0" @@ -46,4 +59,4 @@ class ServerConfig: workers: int = 32 ws_max_queue: int = 1000 log_level: str = "error" - timeout_keep_alive: int = 60 \ No newline at end of file + timeout_keep_alive: int = 60 diff --git a/Agent0/executor_train/eval_service/model_service.py b/Agent0/executor_train/eval_service/model_service.py index 1d35cb1..2497b94 100644 --- a/Agent0/executor_train/eval_service/model_service.py +++ b/Agent0/executor_train/eval_service/model_service.py @@ -18,9 +18,10 @@ # other C0 control characters except common whitespace). CONTROL_CHAR_RE = re.compile( # this matches U+0000 through U+001F, excluding tab(09), LF(0A), CR(0D) - r'[\x00-\x08\x0B\x0C\x0E-\x1F]' + r"[\x00-\x08\x0B\x0C\x0E-\x1F]" ) + def sanitize_request(obj: Any) -> Any: """ Recursively walk through obj and: @@ -30,18 +31,20 @@ def sanitize_request(obj: Any) -> Any: - Leave other types untouched """ if isinstance(obj, dict): - return {sanitize_request(key): sanitize_request(val) for key, val in obj.items()} + return {sanitize_request(key): sanitize_request(val) + for key, val in obj.items()} elif isinstance(obj, (list, tuple)): return type(obj)(sanitize_request(item) for item in obj) elif isinstance(obj, str): # strip NUL (\x00) and other C0 control chars - return CONTROL_CHAR_RE.sub('', obj) + return CONTROL_CHAR_RE.sub("", obj) else: return obj - + + class ModelService: """verl-tool model inference service""" - + def __init__(self, model_config: ModelConfig, tool_config: ToolConfig): """initialize model service""" self.model_config = model_config @@ -52,10 +55,18 @@ def __init__(self, model_config: ModelConfig, tool_config: ToolConfig): self.encode_lock = asyncio.Lock() if self.tool_config.mtrl_sep is None: messages = [{"role": "system", "content": "{obs}"}] - self.tool_config.mtrl_sep = "\n" + self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + self.tool_config.mtrl_sep = "\n" + self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) # self.tool_config.mtrl_sep = self.tool_config.mtrl_sep.replace("system", "user") - - def call_tool_server(self, trajectory_ids: List[str], actions: List[str], finish: List[bool], **kwargs: Dict[str, List[Any]]) -> Dict[str, Any]: + + def call_tool_server( + self, + trajectory_ids: List[str], + actions: List[str], + finish: List[bool], + **kwargs: Dict[str, List[Any]], + ) -> Dict[str, Any]: """querying the tool server for the observation and done flag""" server_url = self.tool_config.tool_server_url # prepare payload @@ -63,23 +74,32 @@ def call_tool_server(self, trajectory_ids: List[str], actions: List[str], finish "trajectory_ids": trajectory_ids, "actions": actions, "finish": finish, - **kwargs + **kwargs, } try: data = sanitize_request(data) response = requests.post(server_url, json=data) response.raise_for_status() result = response.json() - return result + return result except Exception as e: print(f"Error calling tool server: {str(e)}") return { - "observations": [f"Error calling tool server: {str(e)}" for _ in range(len(trajectory_ids))], + "observations": [ + f"Error calling tool server: {str(e)}" + for _ in range(len(trajectory_ids)) + ], "dones": [True for _ in range(len(trajectory_ids))], - "valids": [False for _ in range(len(trajectory_ids))] + "valids": [False for _ in range(len(trajectory_ids))], } - - async def call_tool_server_async(self, trajectory_ids: List[str], actions: List[str], finish: List[bool], **kwargs: Dict[str, List[Any]]) -> Dict[str, Any]: + + async def call_tool_server_async( + self, + trajectory_ids: List[str], + actions: List[str], + finish: List[bool], + **kwargs: Dict[str, List[Any]], + ) -> Dict[str, Any]: """querying the tool server for the observation and done flag using aiohttp""" server_url = self.tool_config.tool_server_url # prepare payload @@ -87,13 +107,13 @@ async def call_tool_server_async(self, trajectory_ids: List[str], actions: List[ "trajectory_ids": trajectory_ids, "actions": actions, "finish": finish, - **kwargs + **kwargs, } - + # Create aiohttp session if it doesn't exist if self.session is None: self.session = aiohttp.ClientSession() - + try: data = sanitize_request(data) async with self.session.post(server_url, json=data) as response: @@ -103,78 +123,117 @@ async def call_tool_server_async(self, trajectory_ids: List[str], actions: List[ except Exception as e: print(f"Error calling tool server: {str(e)}") return { - "observations": [f"Error calling tool server: {str(e)}" for _ in range(len(trajectory_ids))], + "observations": [ + f"Error calling tool server: {str(e)}" + for _ in range(len(trajectory_ids)) + ], "dones": [True for _ in range(len(trajectory_ids))], - "valids": [False for _ in range(len(trajectory_ids))] + "valids": [False for _ in range(len(trajectory_ids))], } - - async def post_process_observations(self, next_obs: List[str], dones: List[bool], valid_action: List[bool], finishs: List[bool]): + + async def post_process_observations( + self, + next_obs: List[str], + dones: List[bool], + valid_action: List[bool], + finishs: List[bool], + ): """Process observations using the tokenizer with proper async locks""" - next_obs = [obs if not done else "" for obs, done in zip(next_obs, dones)] + next_obs = [ + obs if not done else "" for obs, + done in zip( + next_obs, + dones)] async with self.encode_lock: mtrl_sep = self.tool_config.mtrl_sep - if self.tool_config.truncate_obs_side == 'left': + if self.tool_config.truncate_obs_side == "left": next_obs_ids = self.tokenizer( next_obs, - padding='longest', - return_tensors='pt', + padding="longest", + return_tensors="pt", add_special_tokens=False, # Prevents adding special tokens - padding_side='left', - )['input_ids'].to(torch.int64) + padding_side="left", + )["input_ids"].to(torch.int64) if next_obs_ids.shape[1] > self.tool_config.max_obs_length: - print(f"[WARNING] OBSERVATION TOO LONG, CONSIDER CHANGING YOUR CONFIG, {next_obs_ids.shape[1]} & {self.tool_config.max_obs_length}") - next_obs_ids = next_obs_ids[:, -self.tool_config.max_obs_length:] - elif self.tool_config.truncate_obs_side == 'right': + print( + f"[WARNING] OBSERVATION TOO LONG, CONSIDER CHANGING YOUR CONFIG, { + next_obs_ids.shape[1]} & { + self.tool_config.max_obs_length}") + next_obs_ids = next_obs_ids[:, - + self.tool_config.max_obs_length:] + elif self.tool_config.truncate_obs_side == "right": next_obs_ids = self.tokenizer( next_obs, - padding='longest', - return_tensors='pt', + padding="longest", + return_tensors="pt", add_special_tokens=False, # Prevents adding special tokens - padding_side='right', - )['input_ids'].to(torch.int64) + padding_side="right", + )["input_ids"].to(torch.int64) if next_obs_ids.shape[1] > self.tool_config.max_obs_length: - print(f"[WARNING] OBSERVATION TOO LONG, CONSIDER CHANGING YOUR CONFIG, {next_obs_ids.shape[1]} & {self.tool_config.max_obs_length}") - next_obs_ids = next_obs_ids[:, :self.tool_config.max_obs_length] + print( + f"[WARNING] OBSERVATION TOO LONG, CONSIDER CHANGING YOUR CONFIG, { + next_obs_ids.shape[1]} & { + self.tool_config.max_obs_length}") + next_obs_ids = next_obs_ids[:, + : self.tool_config.max_obs_length] else: - raise ValueError(f"Invalid truncate_obs_side: {self.tool_config.truncate_obs_side}") + raise ValueError( + f"Invalid truncate_obs_side: { + self.tool_config.truncate_obs_side}") if self.tool_config.enable_mtrl: next_obs = self.tokenizer.batch_decode( - next_obs_ids, - skip_special_tokens=True + next_obs_ids, skip_special_tokens=True ) processed_next_obs = [] for i in range(len(next_obs)): if finishs[i] or dones[i]: # do action is false - assert next_obs[i] == "", f"next_obs should be empty when finishs is True, but got {next_obs[i]}" + assert ( + next_obs[i] == ""), f"next_obs should be empty when finishs is True, but got { + next_obs[i]}" processed_next_obs.append("") elif valid_action[i]: - processed_next_obs.append(mtrl_sep.format(obs=next_obs[i])) + processed_next_obs.append( + mtrl_sep.format(obs=next_obs[i])) else: - processed_next_obs.append(mtrl_sep.format(obs="Your action is not valid, please check the format and try again." + next_obs[i])) + processed_next_obs.append( + mtrl_sep.format( + obs="Your action is not valid, please check the format and try again." + + next_obs[i])) next_obs = processed_next_obs next_obs_ids = self.tokenizer( next_obs, - padding='longest', - return_tensors='pt', + padding="longest", + return_tensors="pt", add_special_tokens=False, # Prevents adding special tokens - )['input_ids'].to(torch.int64) + )["input_ids"].to(torch.int64) next_obs = self.tokenizer.batch_decode( next_obs_ids, skip_special_tokens=True, ) return next_obs - - async def _postprocess_responses(self, outputs: torch.Tensor, action_step: int) -> torch.Tensor: + + async def _postprocess_responses( + self, outputs: torch.Tensor, action_step: int + ) -> torch.Tensor: """Process responses to stop at python operation or answer operation.""" - active_responses = [outputs.choices[i].text for i in range(len(outputs.choices))] - active_finish_reasons = [outputs.choices[i].finish_reason for i in range(len(outputs.choices))] - + active_responses = [ + outputs.choices[i].text for i in range(len(outputs.choices)) + ] + active_finish_reasons = [ + outputs.choices[i].finish_reason for i in range(len(outputs.choices)) + ] + finishes = [] for i in range(len(active_responses)): finish = True - if active_finish_reasons[i] == "stop" and outputs.choices[i].stop_reason is not None: - active_responses[i] = active_responses[i] + outputs.choices[i].stop_reason + if ( + active_finish_reasons[i] == "stop" + and outputs.choices[i].stop_reason is not None + ): + active_responses[i] = ( + active_responses[i] + outputs.choices[i].stop_reason + ) if self.tool_config.enable_mtrl: active_responses[i] += self.tool_config.turn_end_token finish = False @@ -187,32 +246,53 @@ async def _postprocess_responses(self, outputs: torch.Tensor, action_step: int) active_responses[i] += self.tool_config.turn_end_token finishes.append(finish) return active_responses, finishes, active_finish_reasons - + def load_model(self): """load the model using VLLM backend""" print(f"Loading Model using VLLM: {self.model_config.model}...") # start a VLLM server using vllm.serve - vllm_args = [f"--{k.replace('_', '-')}" for k in self.model_config.__dict__.keys() if k not in ["model", "api_key", "num_models", "host", "port"]] + vllm_args = [ + f"--{k.replace('_', '-')}" + for k in self.model_config.__dict__.keys() + if k not in ["model", "api_key", "num_models", "host", "port"] + ] vllm_args = [] for k, v in self.model_config.__dict__.items(): if k not in ["model", "api_key", "num_models", "host", "port"]: - vllm_args.append(f"--{k.replace('_', '-')}") - if not isinstance(v, bool): - vllm_args.append(str(v)) - + vllm_args.append(f"--{k.replace('_', '-')}") + if not isinstance(v, bool): + vllm_args.append(str(v)) + host = "0.0.0.0" num_models = self.model_config.num_models ports = random.sample(range(8000, 9000), num_models) self.vllm_processes = [] - gpu_ids = os.environ.get("CUDA_VISIBLE_DEVICES", ",".join([str(i) for i in range(torch.cuda.device_count())])).split(",") + gpu_ids = os.environ.get( + "CUDA_VISIBLE_DEVICES", + ",".join([str(i) for i in range(torch.cuda.device_count())]), + ).split(",") tensor_parallel_size = self.model_config.tensor_parallel_size - gpu_ids_per_model = [gpu_ids[i:i+tensor_parallel_size] for i in range(0, len(gpu_ids), tensor_parallel_size)] - assert len(gpu_ids) >= num_models * tensor_parallel_size, f"Not enough GPUs available: {len(gpu_ids)} < {num_models * tensor_parallel_size}" + gpu_ids_per_model = [ + gpu_ids[i: i + tensor_parallel_size] + for i in range(0, len(gpu_ids), tensor_parallel_size) + ] + assert ( + len(gpu_ids) >= num_models * tensor_parallel_size + ), f"Not enough GPUs available: {len(gpu_ids)} < {num_models * tensor_parallel_size}" for i in range(num_models): cmd = [ - "vllm", "serve", self.model_config.model, "--api-key", "token-abc123", - "--host", host, "--port", str(ports[i]), - "--disable-uvicorn-access-log", "--disable-log-stats", "--disable-log-requests" + "vllm", + "serve", + self.model_config.model, + "--api-key", + "token-abc123", + "--host", + host, + "--port", + str(ports[i]), + "--disable-uvicorn-access-log", + "--disable-log-stats", + "--disable-log-requests", ] + vllm_args env = os.environ.copy() env["CUDA_VISIBLE_DEVICES"] = ",".join(gpu_ids_per_model[i]) @@ -220,9 +300,12 @@ def load_model(self): vllm_process = subprocess.Popen(cmd, env=env) self.vllm_processes.append(vllm_process) self.clients = [ - openai.Client(api_key="token-abc123", base_url=f"http://{host}:{ports[i]}/v1") for i in range(num_models) + openai.Client( + api_key="token-abc123", base_url=f"http://{host}:{ports[i]}/v1" + ) + for i in range(num_models) ] - + # Wait for the service to start (poll the health endpoint) max_retries = 60 retry_interval = 10 @@ -239,66 +322,77 @@ def load_model(self): # print(f"vLLM instance model-{j} at {host}:{ports[j]} is not ready yet: {str(e)}") continue if all(vllm_model_status): - print(f"โœ… vLLM service started successfully with model: {self.model_config.model}") - return + print( + f"โœ… vLLM service started successfully with model: { + self.model_config.model}") + return else: time.sleep(retry_interval) - + # If we get here, the service failed to start print("Failed to start one or more vLLM services. Check vLLM logs.") for process in self.vllm_processes: stderr = process.stderr.read() print(f"vLLM stderr: {stderr}") process.terminate() - + raise RuntimeError("Failed to start vLLM services") - - async def send_request(self, client, prompts: List[str], model:str, sampling_params: dict) -> str: + + async def send_request( + self, client, prompts: List[str], model: str, sampling_params: dict + ) -> str: # Send the request using the client sampling_params = sampling_params.copy() # Use the async encode method to get tokens async with self.encode_lock: - prompt_lens = [len(self.tokenizer.encode(prompt)) for prompt in prompts] + prompt_lens = [len(self.tokenizer.encode(prompt)) + for prompt in prompts] max_prompt_tokens = max(prompt_lens) - - sampling_params['max_tokens'] = min(max(self.model_config.max_model_len - max_prompt_tokens, 0), sampling_params['max_tokens']) + + sampling_params["max_tokens"] = min( + max(self.model_config.max_model_len - max_prompt_tokens, 0), + sampling_params["max_tokens"], + ) # print(f"Sending request to {client.base_url} with sampling params: {sampling_params}") - + # Run the API call in an executor to not block the event loop response = await asyncio.get_event_loop().run_in_executor( None, lambda: client.completions.create( - model=model, - prompt=prompts, - echo=False, - stream=False, - **sampling_params - ) + model=model, prompt=prompts, echo=False, stream=False, **sampling_params + ), ) return response - - async def generate_with_tools(self, prompts: List[str], sampling_params: dict) -> Tuple[List[str], List[str]]: + + async def generate_with_tools( + self, prompts: List[str], sampling_params: dict + ) -> Tuple[List[str], List[str]]: """ Generate text with tool calls in a multi-turn loop. - + Args: prompts: Initial prompts for generation sampling_params: Sampling parameters for the model - + Returns: Tuple of (full_responses, finish_reasons) """ - client = random.choice(self.clients) # ensure the same trajectory uses the same client for prefix caching - assert sampling_params.get("n", 1) <= 1, "n > 1 is not supported yet for tool generation" + client = random.choice( + self.clients + ) # ensure the same trajectory uses the same client for prefix caching + assert ( + sampling_params.get("n", 1) <= 1 + ), "n > 1 is not supported yet for tool generation" contexts = prompts final_responses = ["" for _ in range(len(prompts))] traj_ids = [str(uuid.uuid4()) for _ in range(len(prompts))] active_masks = [True for _ in range(len(prompts))] finish_reasons = [None for _ in range(len(prompts))] model = self.model_config.model - - # keep trying to generate the response until reached the tool-calling limit - for action_step in range(self.tool_config.max_turns+1): + + # keep trying to generate the response until reached the tool-calling + # limit + for action_step in range(self.tool_config.max_turns + 1): # print(f"Action step: {action_step}/{self.tool_config.max_turns}") if action_step == self.tool_config.max_turns: # last turn, don't stop by action stop tokens @@ -306,43 +400,49 @@ async def generate_with_tools(self, prompts: List[str], sampling_params: dict) - for action_stop_token in self.tool_config.action_stop_tokens: if action_stop_token in sampling_params["stop"]: sampling_params["stop"].remove(action_stop_token) - - active_traj_ids = [traj_ids[i] for i in range(len(traj_ids)) if active_masks[i]] - active_contexts = [contexts[i] for i in range(len(contexts)) if active_masks[i]] + + active_traj_ids = [ + traj_ids[i] for i in range(len(traj_ids)) if active_masks[i] + ] + active_contexts = [ + contexts[i] for i in range(len(contexts)) if active_masks[i] + ] if len(active_contexts) == 0: break - + # send request asynchronously outputs = await self.send_request( - client, - active_contexts, - model, - sampling_params + client, active_contexts, model, sampling_params ) - active_responses, finishes, active_finish_reasons = await self._postprocess_responses(outputs, action_step) - + active_responses, finishes, active_finish_reasons = ( + await self._postprocess_responses(outputs, action_step) + ) + # Use async tool server call if possible - if hasattr(self, 'call_tool_server_async'): + if hasattr(self, "call_tool_server_async"): tool_responses = await self.call_tool_server_async( - active_traj_ids, - active_responses, - finishes + active_traj_ids, active_responses, finishes ) else: # Fallback to sync version but run in executor tool_responses = await asyncio.get_event_loop().run_in_executor( - None, + None, self.call_tool_server, active_traj_ids, active_responses, - finishes + finishes, ) - + # print(f"Active observations (preprocess): {tool_responses['observations']}") - observations = await self.post_process_observations(tool_responses["observations"], tool_responses["dones"], tool_responses["valids"], finishes) + observations = await self.post_process_observations( + tool_responses["observations"], + tool_responses["dones"], + tool_responses["valids"], + finishes, + ) dones = tool_responses["dones"] valids = tool_responses["valids"] - + # print(f"Active step: {action_step}") # print(f"Active responses: {active_responses}") # print(f"Active observations: {observations}") @@ -354,51 +454,64 @@ async def generate_with_tools(self, prompts: List[str], sampling_params: dict) - active_idx = 0 for i in range(len(contexts)): if active_masks[i]: - contexts[i] += active_responses[active_idx] + observations[active_idx] - final_responses[i] += active_responses[active_idx] + observations[active_idx] + contexts[i] += ( + active_responses[active_idx] + observations[active_idx] + ) + final_responses[i] += ( + active_responses[active_idx] + observations[active_idx] + ) finish_reasons[i] = active_finish_reasons[active_idx] active_masks[i] = not dones[active_idx] active_idx += 1 - + return final_responses, finish_reasons - - async def chat_completions_async(self, body: Dict[str, Any]) -> Dict[str, Any]: + + async def chat_completions_async( + self, body: Dict[str, Any]) -> Dict[str, Any]: """process API request and generate response""" # print(f"Received request: {body}") - + if "messages" not in body or not body["messages"]: raise ValueError("No messages found in the request.") - if not 'user' in [message["role"] for message in body["messages"]]: + if "user" not in [message["role"] for message in body["messages"]]: raise ValueError("No user message found in the request.") - - assert body["model"] == self.model_config.model, f"model mismatch: {body['model']} != {self.model_config.model}" - + + assert ( + body["model"] == self.model_config.model + ), f"model mismatch: {body['model']} != {self.model_config.model}" + async with self.encode_lock: - prompt = self.tokenizer.apply_chat_template(body['messages'], - add_generation_prompt=True, - tokenize=False) - if body.get('n', 1) > 1: + prompt = self.tokenizer.apply_chat_template( + body["messages"], add_generation_prompt=True, tokenize=False + ) + if body.get("n", 1) > 1: prompts = [prompt for _ in range(body["n"])] else: prompts = [prompt] sampling_params = { "temperature": body.get("temperature", 1.0), - "max_tokens": body.get("max_tokens", body.get("max_completion_tokens", 512)), + "max_tokens": body.get( + "max_tokens", body.get("max_completion_tokens", 512) + ), "top_p": body.get("top_p", 1.0), - "stop": list(set(body.get("stop", []) + self.tool_config.action_stop_tokens)), + "stop": list( + set(body.get("stop", []) + self.tool_config.action_stop_tokens) + ), } # print(f"Sampling params: {sampling_params}") - all_responses, finish_reasons = await self.generate_with_tools(prompts, sampling_params) - + all_responses, finish_reasons = await self.generate_with_tools( + prompts, sampling_params + ) + async with self.encode_lock: prompt_tokens = len(self.tokenizer.encode(prompt)) completion_tokens = 0 for response in all_responses: completion_tokens += len(self.tokenizer.encode(response)) total_tokens = prompt_tokens + completion_tokens - + # format the response into OpenAI-compliant format return { "id": f"chatcmpl-{str(uuid.uuid4())}", @@ -412,49 +525,58 @@ async def chat_completions_async(self, body: Dict[str, Any]) -> Dict[str, Any]: "role": "assistant", "content": all_responses[i], }, - "finish_reason": finish_reasons[i] - } for i in range(len(all_responses)) + "finish_reason": finish_reasons[i], + } + for i in range(len(all_responses)) ], "usage": { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, - "total_tokens": total_tokens - } + "total_tokens": total_tokens, + }, } - + def chat_completions(self, body: Dict[str, Any]) -> Dict[str, Any]: """Synchronous wrapper for chat_completions""" return asyncio.run(self.chat_completions_async(body)) - + async def completions_async(self, body: Dict[str, Any]) -> Dict[str, Any]: """process API request and generate response async""" # print(f"Received request: {body}") - if 'prompt' not in body: + if "prompt" not in body: raise ValueError("No prompt found in the request.") - assert body["model"] == self.model_config.model, f"model mismatch: {body['model']} != {self.model_config.model}" - prompt = body['prompt'] + assert ( + body["model"] == self.model_config.model + ), f"model mismatch: {body['model']} != {self.model_config.model}" + prompt = body["prompt"] - if body.get('n', 1) > 1: + if body.get("n", 1) > 1: prompts = [prompt for _ in range(body["n"])] else: prompts = [prompt] sampling_params = { "temperature": body.get("temperature", 1.0), - "max_tokens": body.get("max_tokens", body.get("max_completion_tokens", 512)), + "max_tokens": body.get( + "max_tokens", body.get("max_completion_tokens", 512) + ), "top_p": body.get("top_p", 1.0), - "stop": list(set(body.get("stop", []) + self.tool_config.action_stop_tokens)), + "stop": list( + set(body.get("stop", []) + self.tool_config.action_stop_tokens) + ), } - all_responses, finish_reasons = await self.generate_with_tools(prompts, sampling_params) - + all_responses, finish_reasons = await self.generate_with_tools( + prompts, sampling_params + ) + async with self.encode_lock: prompt_tokens = len(self.tokenizer.encode(prompt)) completion_tokens = 0 for response in all_responses: completion_tokens += len(self.tokenizer.encode(response)) total_tokens = prompt_tokens + completion_tokens - + # format the response into OpenAI-compliant format return { "id": f"chatcmpl-{str(uuid.uuid4())}", @@ -465,27 +587,28 @@ async def completions_async(self, body: Dict[str, Any]) -> Dict[str, Any]: { "index": i, "text": all_responses[i], - "finish_reason": finish_reasons[i] - } for i in range(len(all_responses)) + "finish_reason": finish_reasons[i], + } + for i in range(len(all_responses)) ], "usage": { "prompt_tokens": prompt_tokens, "completion_tokens": completion_tokens, - "total_tokens": total_tokens - } + "total_tokens": total_tokens, + }, } - + def completions(self, body: Dict[str, Any]) -> Dict[str, Any]: """Synchronous wrapper for completions_async""" return asyncio.run(self.completions_async(body)) - + async def close(self): """Close any resources (like HTTP sessions and processes) when shutting down""" # Close HTTP session if self.session: await self.session.close() self.session = None - + # Terminate all VLLM processes for process in self.vllm_processes: if process: @@ -494,14 +617,15 @@ async def close(self): process.wait(timeout=5) except subprocess.TimeoutExpired: process.kill() - + self.vllm_processes = [] self.clients = [] - + def __del__(self): """Destructor to ensure resources are cleaned up""" try: asyncio.run(self.close()) except RuntimeError: - # Handle "Event loop is closed" error that can happen during shutdown + # Handle "Event loop is closed" error that can happen during + # shutdown pass diff --git a/Agent0/executor_train/eval_service/test/test_api.py b/Agent0/executor_train/eval_service/test/test_api.py index fff6835..9108a14 100644 --- a/Agent0/executor_train/eval_service/test/test_api.py +++ b/Agent0/executor_train/eval_service/test/test_api.py @@ -2,20 +2,23 @@ from openai import OpenAI from transformers import AutoTokenizer + def main( model_name: str, base_url: str, test_task: str = "math", - test_type: str = "chat_completion", # or "completion" + test_type: str = "chat_completion", # or "completion" api_key: str = "sk-proj-1234567890", temperature: float = 0.0, max_tokens: int = 2048, top_p: float = 1.0, n: int = 1, ): - client = OpenAI(api_key=api_key, base_url=base_url) # Replace with your local server address + client = OpenAI( + api_key=api_key, base_url=base_url + ) # Replace with your local server address tokenizer = AutoTokenizer.from_pretrained(model_name) - + # get test_task if test_task == "math": print("Testing math task...") @@ -23,33 +26,25 @@ def main( math_problem = "Convert the point $(0,3)$ in rectangular coordinates to polar coordinates. Enter your answer in the form $(r,\\theta),$ where $r > 0$ and $0 \\le \\theta < 2 \\pi.$" chat_messages = [ - { - "role": "system", - "content": system_prompt - }, - { - "role": "user", - "content": math_problem - } + {"role": "system", "content": system_prompt}, + {"role": "user", "content": math_problem}, ] - prompt = tokenizer.apply_chat_template(chat_messages, tokenize=False, add_generation_prompt=True) + prompt = tokenizer.apply_chat_template( + chat_messages, tokenize=False, add_generation_prompt=True + ) else: raise ValueError(f"Unknown test task: {test_task}") - - + if test_type == "chat_completion": - print(f"Testing {test_task} with {test_type} on model {model_name} at {base_url}", flush=True) + print( + f"Testing {test_task} with {test_type} on model {model_name} at {base_url}", + flush=True, + ) completion = client.chat.completions.create( model=model_name, messages=[ - { - "role": "system", - "content": system_prompt - }, - { - "role": "user", - "content": math_problem - } + {"role": "system", "content": system_prompt}, + {"role": "user", "content": math_problem}, ], temperature=temperature, max_tokens=max_tokens, @@ -58,18 +53,17 @@ def main( ) print(completion.choices[0].message.content) elif test_type == "completion": - print(f"Testing {test_task} with {test_type} on model {model_name} at {base_url}", flush=True) + print( + f"Testing {test_task} with {test_type} on model {model_name} at {base_url}", + flush=True, + ) chat_messages = [ - { - "role": "system", - "content": system_prompt - }, - { - "role": "user", - "content": math_problem - } + {"role": "system", "content": system_prompt}, + {"role": "user", "content": math_problem}, ] - prompt = tokenizer.apply_chat_template(chat_messages, tokenize=False, add_generation_prompt=True) + prompt = tokenizer.apply_chat_template( + chat_messages, tokenize=False, add_generation_prompt=True + ) completion = client.completions.create( model=model_name, prompt=prompt, @@ -82,8 +76,10 @@ def main( else: raise ValueError(f"Unknown test type: {test_type}") + if __name__ == "__main__": import fire + fire.Fire(main) """ diff --git a/Agent0/executor_train/eval_service/test/test_api_mp.py b/Agent0/executor_train/eval_service/test/test_api_mp.py index 313b571..671b2e2 100644 --- a/Agent0/executor_train/eval_service/test/test_api_mp.py +++ b/Agent0/executor_train/eval_service/test/test_api_mp.py @@ -11,89 +11,102 @@ # Different variations of the math problem to simulate diverse requests math_problems = [ math_problem, - math_problem.replace("9-kilometer", "10-kilometer").replace("4 hours", "5 hours").replace("2 hours and 24 minutes", "3 hours"), - math_problem.replace("9-kilometer", "8-kilometer").replace("4 hours", "3 hours").replace("2 hours and 24 minutes", "1 hour and 48 minutes"), + math_problem.replace("9-kilometer", "10-kilometer") + .replace("4 hours", "5 hours") + .replace("2 hours and 24 minutes", "3 hours"), + math_problem.replace("9-kilometer", "8-kilometer") + .replace("4 hours", "3 hours") + .replace("2 hours and 24 minutes", "1 hour and 48 minutes"), math_problem.replace("s+\\frac{1}{2}", "s+\\frac{2}{3}"), - math_problem.replace("s+\\frac{1}{2}", "s+1") + math_problem.replace("s+\\frac{1}{2}", "s+1"), ] + async def send_request(client, problem_text, request_id): """Send a single request and measure the time it takes""" start_time = time.time() print(f"Starting request {request_id}...") - + try: completion = await client.chat.completions.create( model="GAIR/ToRL-1.5B", messages=[ - { - "role": "system", - "content": system_prompt - }, - { - "role": "user", - "content": problem_text - } + {"role": "system", "content": system_prompt}, + {"role": "user", "content": problem_text}, ], temperature=0, max_tokens=2048, top_p=1, n=1, ) - + end_time = time.time() - print(f"Request {request_id} completed in {end_time - start_time:.2f} seconds") - + print( + f"Request {request_id} completed in { + end_time - + start_time:.2f} seconds") + # Print a shortened version of the response for verification response_content = completion.choices[0].message.content - print(f"Request {request_id} response (truncated): {response_content}...\n") - + print( + f"Request {request_id} response (truncated): {response_content}...\n") + return { "request_id": request_id, "duration": end_time - start_time, - "response": response_content + "response": response_content, } except Exception as e: end_time = time.time() - print(f"Request {request_id} failed after {end_time - start_time:.2f} seconds: {str(e)}") + print( + f"Request {request_id} failed after { + end_time - + start_time:.2f} seconds: { + str(e)}") return { "request_id": request_id, "duration": end_time - start_time, - "error": str(e) + "error": str(e), } + async def run_concurrent_test(num_concurrent=5, num_total=10): """Run multiple concurrent requests to test server performance""" - client = AsyncOpenAI(api_key="sk-proj-1234567890", base_url="http://0.0.0.0:5000") - - print(f"Starting concurrent test with {num_concurrent} concurrent requests, {num_total} total requests") + client = AsyncOpenAI( + api_key="sk-proj-1234567890", + base_url="http://0.0.0.0:5000") + + print( + f"Starting concurrent test with {num_concurrent} concurrent requests, {num_total} total requests" + ) start_time = time.time() - + # Create tasks for all requests tasks = [] for i in range(num_total): problem = math_problems[i % len(math_problems)] - tasks.append(send_request(client, problem, i+1)) - + tasks.append(send_request(client, problem, i + 1)) + # Run requests in batches of num_concurrent results = [] for i in range(0, len(tasks), num_concurrent): - batch = tasks[i:i+num_concurrent] + batch = tasks[i: i + num_concurrent] batch_results = await asyncio.gather(*batch) results.extend(batch_results) - + end_time = time.time() total_duration = end_time - start_time - + # Calculate statistics successful_requests = [r for r in results if "error" not in r] failed_requests = [r for r in results if "error" in r] - + if successful_requests: - avg_request_time = sum(r["duration"] for r in successful_requests) / len(successful_requests) + avg_request_time = sum( + r["duration"] for r in successful_requests) / len(successful_requests) else: avg_request_time = 0 - + # Print summary print("\n===== TEST RESULTS =====") print(f"Total test duration: {total_duration:.2f} seconds") @@ -102,46 +115,53 @@ async def run_concurrent_test(num_concurrent=5, num_total=10): print(f"Failed requests: {len(failed_requests)}") print(f"Average request time: {avg_request_time:.2f} seconds") print(f"Requests per second: {num_total / total_duration:.2f}") - + if failed_requests: print("\nFailed requests:") for req in failed_requests: print(f" Request {req['request_id']}: {req['error']}") + async def sequential_test_for_comparison(num_requests=5): """Run sequential requests as a baseline for comparison""" - client = AsyncOpenAI(api_key="sk-proj-1234567890", base_url="http://0.0.0.0:5000") - - print(f"\nStarting sequential test with {num_requests} requests for comparison") + client = AsyncOpenAI( + api_key="sk-proj-1234567890", + base_url="http://0.0.0.0:5000") + + print( + f"\nStarting sequential test with {num_requests} requests for comparison") start_time = time.time() - + results = [] for i in range(num_requests): problem = math_problems[i % len(math_problems)] - result = await send_request(client, problem, f"seq-{i+1}") + result = await send_request(client, problem, f"seq-{i + 1}") results.append(result) - + end_time = time.time() total_duration = end_time - start_time - + # Calculate statistics successful_requests = [r for r in results if "error" not in r] - + if successful_requests: - avg_request_time = sum(r["duration"] for r in successful_requests) / len(successful_requests) + avg_request_time = sum( + r["duration"] for r in successful_requests) / len(successful_requests) else: avg_request_time = 0 - + # Print summary print("\n===== SEQUENTIAL TEST RESULTS =====") print(f"Total test duration: {total_duration:.2f} seconds") print(f"Average request time: {avg_request_time:.2f} seconds") print(f"Requests per second: {num_requests / total_duration:.2f}") + async def main(): # Run both tests await run_concurrent_test(num_concurrent=3, num_total=6) await sequential_test_for_comparison(num_requests=3) + if __name__ == "__main__": - asyncio.run(main()) \ No newline at end of file + asyncio.run(main()) diff --git a/Agent0/executor_train/scripts/visualize_entropy.py b/Agent0/executor_train/scripts/visualize_entropy.py index da3cb31..703a31c 100644 --- a/Agent0/executor_train/scripts/visualize_entropy.py +++ b/Agent0/executor_train/scripts/visualize_entropy.py @@ -10,10 +10,13 @@ from tqdm import tqdm from collections import defaultdict -def plot_entropy_bar(entropy, labels, title="Token Entropy", save_path="entropy_plot.png"): + +def plot_entropy_bar( + entropy, labels, title="Token Entropy", save_path="entropy_plot.png" +): """ Plot the token entropy with color highlighting based on masks and background shading. - + Args: entropy (list): List of entropy values corresponding to each token. labels (List[str]): List of labels for the tokens, e.g., "prompt", "action" or "obs". @@ -22,28 +25,46 @@ def plot_entropy_bar(entropy, labels, title="Token Entropy", save_path="entropy_ """ # Color map for distinguishing between the parts color_map = {"prompt": "green", "action": "red", "obs": "blue"} - + plt.figure(figsize=(15 + len(entropy) * 0.01, 4)) clipped_entropy = np.clip(entropy, 0, 10) token_indices = np.arange(len(entropy)) # Initialize to hold color and label settings token_colors = [color_map.get(label, "gray") for label in labels] - alpha_values = [0.6 if label == "prompt" else 0.9 for label in labels] # Lighter for prompts, darker for actions and obs - + alpha_values = [ + 0.6 if label == "prompt" else 0.9 for label in labels + ] # Lighter for prompts, darker for actions and obs + # Plot background color for each section last_idx = 0 last_label = labels[0] for i in range(len(labels)): if labels[i] != last_label: - plt.axvspan(last_idx, i - 1, color=color_map[last_label], alpha=0.1, label=f"{last_label.capitalize()} Background") + plt.axvspan( + last_idx, + i - 1, + color=color_map[last_label], + alpha=0.1, + label=f"{last_label.capitalize()} Background", + ) last_idx = i last_label = labels[i] - plt.axvspan(last_idx, len(labels) - 1, color=color_map[last_label], alpha=0.1, label=f"{last_label.capitalize()} Background") - + plt.axvspan( + last_idx, + len(labels) - 1, + color=color_map[last_label], + alpha=0.1, + label=f"{last_label.capitalize()} Background", + ) + # Bar plot with clear separation for each token part for i in range(len(entropy)): - plt.bar(i, clipped_entropy[i], color=token_colors[i], alpha=alpha_values[i]) + plt.bar( + i, + clipped_entropy[i], + color=token_colors[i], + alpha=alpha_values[i]) plt.title(title) plt.xlabel("Token Index") @@ -51,20 +72,26 @@ def plot_entropy_bar(entropy, labels, title="Token Entropy", save_path="entropy_ plt.tight_layout() # Adding a legend to make distinction clear - plt.legend(handles=[plt.Line2D([0], [0], color=color_map["prompt"], lw=4), - plt.Line2D([0], [0], color=color_map["action"], lw=4), - plt.Line2D([0], [0], color=color_map["obs"], lw=4)], - labels=["Prompt", "Action", "Obs"], title="Token Type") - + plt.legend( + handles=[ + plt.Line2D([0], [0], color=color_map["prompt"], lw=4), + plt.Line2D([0], [0], color=color_map["action"], lw=4), + plt.Line2D([0], [0], color=color_map["obs"], lw=4), + ], + labels=["Prompt", "Action", "Obs"], + title="Token Type", + ) + # Grid lines for better readability - plt.grid(True, axis='y', linestyle='--', alpha=0.5) - + plt.grid(True, axis="y", linestyle="--", alpha=0.5) + plt.savefig(save_path, dpi=300) return save_path + def main( - file_path:str, - model_name:str = "Qwen/Qwen2.5-Math-1.5B", + file_path: str, + model_name: str = "Qwen/Qwen2.5-Math-1.5B", batch_size=4, vis_dir: str = "entropy_vis", ): @@ -74,47 +101,75 @@ def main( pad_token_id = tokenizer.pad_token_id # Read the JSON file - with open(file_path, 'r') as f: + with open(file_path, "r") as f: data = json.load(f) data = datasets.Dataset.from_list(data) - data = data.filter(lambda x: x['num_turn'] > 0, num_proc=8, desc="Filtering dataset with num_turn > 0") + data = data.filter( + lambda x: x["num_turn"] > 0, + num_proc=8, + desc="Filtering dataset with num_turn > 0", + ) print(data) - full_inputs = [x['prompt'] + x['response'] for x in data] - full_inputs_with_mask = [x['prompt'] + x['response_with_loss_mask'] for x in data] + full_inputs = [x["prompt"] + x["response"] for x in data] + full_inputs_with_mask = [x["prompt"] + + x["response_with_loss_mask"] for x in data] # Tokenize the inputs vis_dir = Path(vis_dir) vis_dir.mkdir(parents=True, exist_ok=True) vis_paths = [] - entropy_avgs = [] # list of sum entropy values, [0] for prompt, [1] for action 1, [2] for obs 1, [3] for action 2, [4] for obs 2, ... - for i in tqdm(range(0, len(full_inputs), batch_size), desc="Processing batches", total=len(full_inputs) // batch_size): - prompts = data['prompt'][i:i + batch_size] - batch = full_inputs[i:i + batch_size] - batch_with_mask = full_inputs_with_mask[i:i + batch_size] - inputs = tokenizer(batch, return_tensors='pt', padding="longest").to(model.device) - inputs_with_mask = tokenizer(batch_with_mask, return_tensors='pt', padding="longest").to(model.device) - attention_mask = inputs['attention_mask'] + # list of sum entropy values, [0] for prompt, [1] for action 1, [2] for + # obs 1, [3] for action 2, [4] for obs 2, ... + entropy_avgs = ([]) + for i in tqdm( + range(0, len(full_inputs), batch_size), + desc="Processing batches", + total=len(full_inputs) // batch_size, + ): + prompts = data["prompt"][i: i + batch_size] + batch = full_inputs[i: i + batch_size] + batch_with_mask = full_inputs_with_mask[i: i + batch_size] + inputs = tokenizer(batch, return_tensors="pt", padding="longest").to( + model.device + ) + inputs_with_mask = tokenizer( + batch_with_mask, return_tensors="pt", padding="longest" + ).to(model.device) + attention_mask = inputs["attention_mask"] # Get the model outputs with torch.no_grad(): outputs = model(**inputs) - logits = outputs.logits # [batch_size, seq_len, vocab_size] - probs = torch.softmax(logits, dim=-1) # [batch_size, seq_len, vocab_size] - log_probs = torch.log(probs + 1e-9) # [batch_size, seq_len, vocab_size] - batch_entropy = -(probs * log_probs * attention_mask.unsqueeze(-1)).sum(dim=-1) # [batch_size, seq_len] + logits = outputs.logits # [batch_size, seq_len, vocab_size] + # [batch_size, seq_len, vocab_size] + probs = torch.softmax(logits, dim=-1) + # [batch_size, seq_len, vocab_size] + log_probs = torch.log(probs + 1e-9) + batch_entropy = -(probs * log_probs * attention_mask.unsqueeze(-1)).sum( + dim=-1 + ) # [batch_size, seq_len] entrypy_list = [] - for j in tqdm(range(len(batch_entropy)), desc=f"Processing batch {i//batch_size}", leave=False, total=len(batch_entropy)): - effective_entry = batch_entropy[j][attention_mask[j] == 1].cpu().numpy() - labels = ["prompt"] * len(tokenizer.encode(prompts[j], add_special_tokens=False)) + for j in tqdm( + range(len(batch_entropy)), + desc=f"Processing batch {i // batch_size}", + leave=False, + total=len(batch_entropy), + ): + effective_entry = batch_entropy[j][attention_mask[j] == 1].cpu( + ).numpy() + labels = ["prompt"] * len( + tokenizer.encode(prompts[j], add_special_tokens=False) + ) labels += ["action"] * (len(effective_entry) - len(labels)) - masks = inputs_with_mask['input_ids'][j][attention_mask[j] == 1] + masks = inputs_with_mask["input_ids"][j][attention_mask[j] == 1] masks = (masks != pad_token_id).cpu().numpy() for k in range(len(labels)): if masks[k] == 0: labels[k] = "obs" - save_path = vis_dir / f"entropy_plot_sample_{i* batch_size + j}.png" + save_path = vis_dir / \ + f"entropy_plot_sample_{i * batch_size + j}.png" # plot_entropy_bar(effective_entry.cpu().numpy(), labels, title=f"Token Entropy for Batch {i//batch_size}, Sample {j}", save_path=save_path) # print(f"Saved plot to {save_path}") # Calculate average entropy for each type @@ -123,14 +178,15 @@ def main( avg_entropy = [] for k in range(len(labels)): if labels[k] != last_label: - avg_entropy.append(effective_entry[last_idx:k].mean().item()) + avg_entropy.append( + effective_entry[last_idx:k].mean().item()) last_idx = k last_label = labels[k] for k in range(len(avg_entropy)): if len(entropy_avgs) <= k: entropy_avgs.append([]) entropy_avgs[k].append(avg_entropy[k]) - + entrypy_list.append(effective_entry) vis_paths.append(save_path) @@ -139,11 +195,11 @@ def main( if i == 0: print(f"Average prompt entropy: {avg:.4f}") elif i % 2 == 1: - print(f"Average action {i//2 + 1} entropy: {avg:.4f}") + print(f"Average action {i // 2 + 1} entropy: {avg:.4f}") else: - print(f"Average obs {i//2} entropy: {avg:.4f}") + print(f"Average obs {i // 2} entropy: {avg:.4f}") + - if __name__ == "__main__": fire.Fire(main) @@ -157,4 +213,4 @@ def main( python scripts/visualize_entropy.py --file_path path/to/data.json --model_name Qwen/Qwen2.5-Math-1.5B --batch_size 1 python scripts/visualize_entropy.py --file_path /home/dongfu/WorkSpace/verl-tool/verl_step_records/torl-fsdp-agent-qwen_qwen2.5-math-1.5b-grpo-n16-b128-t1.0-lr1e-6debug/torl-step-1.json --model_name Qwen/Qwen2.5-Math-1.5B --batch_size 2 ``` -""" \ No newline at end of file +""" diff --git a/Agent0/executor_train/verl/Notice.txt b/Agent0/executor_train/verl/Notice.txt index ade439d..f86cb04 100644 --- a/Agent0/executor_train/verl/Notice.txt +++ b/Agent0/executor_train/verl/Notice.txt @@ -1 +1 @@ -Copyright 2023-2024 Bytedance Ltd. and/or its affiliates \ No newline at end of file +Copyright 2023-2026 Bytedance Ltd. and/or its affiliates \ No newline at end of file diff --git a/Agent0/executor_train/verl/docs/conf.py b/Agent0/executor_train/verl/docs/conf.py index d405288..e736435 100644 --- a/Agent0/executor_train/verl/docs/conf.py +++ b/Agent0/executor_train/verl/docs/conf.py @@ -1,100 +1,100 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -# Configuration file for the Sphinx documentation builder. -# -# This file only contains a selection of the most common options. For a full -# list see the documentation: -# https://www.sphinx-doc.org/en/master/usage/configuration.html - -# -- Path setup -------------------------------------------------------------- - -# If extensions (or modules to document with autodoc) are in another directory, -# add these directories to sys.path here. If the directory is relative to the -# documentation root, use os.path.abspath to make it absolute, like shown here. -# -# import os -# import sys -# sys.path.insert(0, os.path.abspath('.')) - - -# -- Project information ----------------------------------------------------- - -project = "verl" -copyright = "2024 ByteDance Seed Foundation MLSys Team" -author = "Guangming Sheng, Chi Zhang, Yanghua Peng, Haibin Lin" - - -# -- General configuration --------------------------------------------------- -# The master toctree document. -master_doc = "index" - -# Add any Sphinx extension module names here, as strings. They can be -# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom -# ones. -extensions = [ - "myst_parser", - "sphinx.ext.autodoc", - "sphinx.ext.autosummary", - "sphinx.ext.autosectionlabel", - "sphinx.ext.napoleon", - "sphinx.ext.viewcode", -] -# Use Google style docstrings instead of NumPy docstrings. -napoleon_google_docstring = True -napoleon_numpy_docstring = False - -# The suffix(es) of source filenames. -# You can specify multiple suffix as a list of string: -source_suffix = { - ".rst": "restructuredtext", - ".md": "markdown", -} - -# Add any paths that contain templates here, relative to this directory. -templates_path = ["_templates"] - -# The language for content autogenerated by Sphinx. Refer to documentation -# for a list of supported languages. -# -# This is also used if you do content translation via gettext catalogs. -# Usually you set "language" from the command line for these cases. -language = "en" - -# List of patterns, relative to source directory, that match files and -# directories to ignore when looking for source files. -# This pattern also affects html_static_path and html_extra_path. -exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] - - -# -- Options for HTML output ------------------------------------------------- - -# The theme to use for HTML and HTML Help pages. See the documentation for -# a list of builtin themes. -# -html_theme = "sphinx_rtd_theme" - -# Add any paths that contain custom static files (such as style sheets) here, -# relative to this directory. They are copied after the builtin static files, -# so a file named "default.css" will overwrite the builtin "default.css". -html_static_path = ["_static"] - -# Add the JavaScript file -html_js_files = [ - "js/runllm-widget.js", -] - -exclude_patterns += ["README.md", "README_vllm0.7.md"] - -suppress_warnings = ["ref.duplicate", "ref.myst"] +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +# Configuration file for the Sphinx documentation builder. +# +# This file only contains a selection of the most common options. For a full +# list see the documentation: +# https://www.sphinx-doc.org/en/master/usage/configuration.html + +# -- Path setup -------------------------------------------------------------- + +# If extensions (or modules to document with autodoc) are in another directory, +# add these directories to sys.path here. If the directory is relative to the +# documentation root, use os.path.abspath to make it absolute, like shown here. +# +# import os +# import sys +# sys.path.insert(0, os.path.abspath('.')) + + +# -- Project information ----------------------------------------------------- + +project = "verl" +copyright = "2024 ByteDance Seed Foundation MLSys Team" +author = "Guangming Sheng, Chi Zhang, Yanghua Peng, Haibin Lin" + + +# -- General configuration --------------------------------------------------- +# The master toctree document. +master_doc = "index" + +# Add any Sphinx extension module names here, as strings. They can be +# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom +# ones. +extensions = [ + "myst_parser", + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.autosectionlabel", + "sphinx.ext.napoleon", + "sphinx.ext.viewcode", +] +# Use Google style docstrings instead of NumPy docstrings. +napoleon_google_docstring = True +napoleon_numpy_docstring = False + +# The suffix(es) of source filenames. +# You can specify multiple suffix as a list of string: +source_suffix = { + ".rst": "restructuredtext", + ".md": "markdown", +} + +# Add any paths that contain templates here, relative to this directory. +templates_path = ["_templates"] + +# The language for content autogenerated by Sphinx. Refer to documentation +# for a list of supported languages. +# +# This is also used if you do content translation via gettext catalogs. +# Usually you set "language" from the command line for these cases. +language = "en" + +# List of patterns, relative to source directory, that match files and +# directories to ignore when looking for source files. +# This pattern also affects html_static_path and html_extra_path. +exclude_patterns = ["_build", "Thumbs.db", ".DS_Store"] + + +# -- Options for HTML output ------------------------------------------------- + +# The theme to use for HTML and HTML Help pages. See the documentation for +# a list of builtin themes. +# +html_theme = "sphinx_rtd_theme" + +# Add any paths that contain custom static files (such as style sheets) here, +# relative to this directory. They are copied after the builtin static files, +# so a file named "default.css" will overwrite the builtin "default.css". +html_static_path = ["_static"] + +# Add the JavaScript file +html_js_files = [ + "js/runllm-widget.js", +] + +exclude_patterns += ["README.md", "README_vllm0.7.md"] + +suppress_warnings = ["ref.duplicate", "ref.myst"] diff --git a/Agent0/executor_train/verl/examples/data_preprocess/aime2024_multiturn_w_tool.py b/Agent0/executor_train/verl/examples/data_preprocess/aime2024_multiturn_w_tool.py index e7d2835..b245592 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/aime2024_multiturn_w_tool.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/aime2024_multiturn_w_tool.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -54,7 +54,9 @@ def process_fn(example, idx): return process_fn - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True) local_dir = args.local_dir hdfs_dir = args.hdfs_dir diff --git a/Agent0/executor_train/verl/examples/data_preprocess/dapo_multiturn_w_tool.py b/Agent0/executor_train/verl/examples/data_preprocess/dapo_multiturn_w_tool.py index 12b0d09..06b9502 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/dapo_multiturn_w_tool.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/dapo_multiturn_w_tool.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -54,7 +54,9 @@ def process_fn(example, idx): return process_fn - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True) local_dir = args.local_dir hdfs_dir = args.hdfs_dir diff --git a/Agent0/executor_train/verl/examples/data_preprocess/full_hh_rlhf.py b/Agent0/executor_train/verl/examples/data_preprocess/full_hh_rlhf.py index 4625f28..629d409 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/full_hh_rlhf.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/full_hh_rlhf.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,7 +27,9 @@ from verl.utils.fs import copy, makedirs -def generate_sft_dataset(target_hdfs_path_dir, local_dir="~/data/full_hh_rlh/sft"): +def generate_sft_dataset( + target_hdfs_path_dir, + local_dir="~/data/full_hh_rlh/sft"): dataset = load_dataset("Dahoas/full-hh-rlhf") output = {"prompt": [], "response": []} for data in tqdm(dataset["train"]): @@ -55,14 +57,18 @@ def generate_sft_dataset(target_hdfs_path_dir, local_dir="~/data/full_hh_rlh/sft copy(local_path, hdfs_dir) -def generate_rm_dataset(target_hdfs_path_dir, local_dir="~/data/full_hh_rlh/rm"): +def generate_rm_dataset( + target_hdfs_path_dir, + local_dir="~/data/full_hh_rlh/rm"): train_dataset = load_dataset("Dahoas/full-hh-rlhf", split="train[:75%]") test_dataset = load_dataset("Dahoas/full-hh-rlhf", split="train[-25%:]") local_dir = os.path.expanduser(local_dir) os.makedirs(local_dir, exist_ok=True) - for dataset, name in zip([train_dataset, test_dataset], ["train", "test"], strict=True): + for dataset, name in zip( + [train_dataset, test_dataset], ["train", "test"], strict=True + ): output = {"prompt": [], "chosen": [], "rejected": []} for data in tqdm(dataset): # add chosen @@ -83,7 +89,9 @@ def generate_rm_dataset(target_hdfs_path_dir, local_dir="~/data/full_hh_rlh/rm") copy(local_path, hdfs_dir) -def generate_rl_dataset(target_hdfs_path_dir, local_dir="~/data/full_hh_rlhf/rl"): +def generate_rl_dataset( + target_hdfs_path_dir, + local_dir="~/data/full_hh_rlhf/rl"): dataset = load_dataset("Dahoas/full-hh-rlhf") train_dataset = dataset["train"] @@ -109,7 +117,9 @@ def process_fn(example, idx): return process_fn - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True) local_dir = os.path.expanduser(local_dir) local_path = os.path.join(local_dir, "train.parquet") train_dataset.to_parquet(local_path) @@ -123,17 +133,30 @@ def process_fn(example, idx): if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument("--split", type=str, choices=["sft", "rm", "rl"], required=True) + parser.add_argument( + "--split", + type=str, + choices=[ + "sft", + "rm", + "rl"], + required=True) parser.add_argument("--local_dir", type=str, default="~/data/full_hh_rlhf") parser.add_argument("--hdfs_dir", type=str, required=False, default=None) args = parser.parse_args() if args.split == "sft": - generate_sft_dataset(args.hdfs_dir, os.path.join(args.local_dir, args.split)) + generate_sft_dataset( + args.hdfs_dir, os.path.join( + args.local_dir, args.split)) elif args.split == "rm": - generate_rm_dataset(args.hdfs_dir, os.path.join(args.local_dir, args.split)) + generate_rm_dataset( + args.hdfs_dir, os.path.join( + args.local_dir, args.split)) elif args.split == "rl": - generate_rl_dataset(args.hdfs_dir, os.path.join(args.local_dir, args.split)) + generate_rl_dataset( + args.hdfs_dir, os.path.join( + args.local_dir, args.split)) else: raise NotImplementedError diff --git a/Agent0/executor_train/verl/examples/data_preprocess/geo3k.py b/Agent0/executor_train/verl/examples/data_preprocess/geo3k.py index 2df225d..aa036fd 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/geo3k.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/geo3k.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -39,8 +39,7 @@ instruction_following = ( r"You FIRST think about the reasoning process as an internal monologue and then provide the final answer. " r"The reasoning process MUST BE enclosed within tags. " - r"The final answer MUST BE put in \boxed{}." - ) + r"The final answer MUST BE put in \boxed{}.") # add a row to each data item that represents a unique id def make_map_fn(split): @@ -72,8 +71,12 @@ def process_fn(example, idx): return process_fn - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True, num_proc=8) - test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True, num_proc=8) + train_dataset = train_dataset.map( + function=make_map_fn("train"), with_indices=True, num_proc=8 + ) + test_dataset = test_dataset.map( + function=make_map_fn("test"), with_indices=True, num_proc=8 + ) local_dir = args.local_dir hdfs_dir = args.hdfs_dir diff --git a/Agent0/executor_train/verl/examples/data_preprocess/geo3k_multiturn_w_tool.py b/Agent0/executor_train/verl/examples/data_preprocess/geo3k_multiturn_w_tool.py index 6e00691..f6a3763 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/geo3k_multiturn_w_tool.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/geo3k_multiturn_w_tool.py @@ -1,7 +1,7 @@ -# Copyright 2023-2025 SGLang Team +# Copyright 2023-2026 SGLang Team # Copyright Amazon.com, Inc. or its affiliates. # Copyright 2025 Reallm Labs Ltd. or its affiliates -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -38,8 +38,7 @@ instruction_following = ( r"You FIRST think about the reasoning process as an internal monologue and then provide the final answer. " r"The reasoning process MUST BE enclosed within tags. " - r"The final answer MUST BE put in \boxed{}." - ) + r"The final answer MUST BE put in \boxed{}.") # add a row to each data item that represents a unique id def make_map_fn(split): @@ -88,8 +87,12 @@ def process_fn(example, idx): return process_fn - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True, num_proc=8) - test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True, num_proc=8) + train_dataset = train_dataset.map( + function=make_map_fn("train"), with_indices=True, num_proc=8 + ) + test_dataset = test_dataset.map( + function=make_map_fn("test"), with_indices=True, num_proc=8 + ) local_dir = args.local_dir hdfs_dir = args.hdfs_dir train_dataset.to_parquet(os.path.join(local_dir, "train.parquet")) diff --git a/Agent0/executor_train/verl/examples/data_preprocess/gsm8k.py b/Agent0/executor_train/verl/examples/data_preprocess/gsm8k.py index f39c4f0..35d73da 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/gsm8k.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/gsm8k.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -46,7 +46,9 @@ def extract_solution(solution_str): train_dataset = dataset["train"] test_dataset = dataset["test"] - instruction_following = 'Let\'s think step by step and output the final answer after "####".' + instruction_following = ( + 'Let\'s think step by step and output the final answer after "####".' + ) # add a row to each data item that represents a unique id def make_map_fn(split): @@ -78,8 +80,12 @@ def process_fn(example, idx): return process_fn - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) - test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True) + test_dataset = test_dataset.map( + function=make_map_fn("test"), + with_indices=True) local_dir = args.local_dir hdfs_dir = args.hdfs_dir diff --git a/Agent0/executor_train/verl/examples/data_preprocess/gsm8k_multiturn_w_interaction.py b/Agent0/executor_train/verl/examples/data_preprocess/gsm8k_multiturn_w_interaction.py index 718a874..82c5386 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/gsm8k_multiturn_w_interaction.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/gsm8k_multiturn_w_interaction.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -47,7 +47,9 @@ def extract_solution(solution_str): train_dataset = dataset["train"] test_dataset = dataset["test"] - instruction_following = "Let's think step by step and output the final answer after `####`." + instruction_following = ( + "Let's think step by step and output the final answer after `####`." + ) # add a row to each data item that represents a unique id def make_map_fn(split): @@ -92,8 +94,12 @@ def process_fn(example, idx): return process_fn - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) - test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True) + test_dataset = test_dataset.map( + function=make_map_fn("test"), + with_indices=True) local_dir = args.local_dir hdfs_dir = args.hdfs_dir diff --git a/Agent0/executor_train/verl/examples/data_preprocess/gsm8k_multiturn_w_tool.py b/Agent0/executor_train/verl/examples/data_preprocess/gsm8k_multiturn_w_tool.py index 400d885..d88e54b 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/gsm8k_multiturn_w_tool.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/gsm8k_multiturn_w_tool.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -47,7 +47,9 @@ def extract_solution(solution_str): train_dataset = dataset["train"] test_dataset = dataset["test"] - instruction_following = "Let's think step by step and output the final answer after `####`." + instruction_following = ( + "Let's think step by step and output the final answer after `####`." + ) # add a row to each data item that represents a unique id def make_map_fn(split): @@ -102,8 +104,12 @@ def process_fn(example, idx): return process_fn - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) - test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True) + test_dataset = test_dataset.map( + function=make_map_fn("test"), + with_indices=True) local_dir = args.local_dir hdfs_dir = args.hdfs_dir diff --git a/Agent0/executor_train/verl/examples/data_preprocess/hellaswag.py b/Agent0/executor_train/verl/examples/data_preprocess/hellaswag.py index 1b3f200..4a896d9 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/hellaswag.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/hellaswag.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -79,9 +79,15 @@ def process_fn(doc, idx): val_dataset = val_dataset.filter(lambda x: len(x["label"]) > 0) test_dataset = test_dataset.filter(lambda x: len(x["label"]) > 0) - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) - val_dataset = val_dataset.map(function=make_map_fn("validation"), with_indices=True) - test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True) + val_dataset = val_dataset.map( + function=make_map_fn("validation"), + with_indices=True) + test_dataset = test_dataset.map( + function=make_map_fn("test"), + with_indices=True) local_dir = args.local_dir hdfs_dir = args.hdfs_dir diff --git a/Agent0/executor_train/verl/examples/data_preprocess/math_dataset.py b/Agent0/executor_train/verl/examples/data_preprocess/math_dataset.py index e2e5d35..fe0bd12 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/math_dataset.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/math_dataset.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -44,7 +44,9 @@ def extract_solution(solution_str): train_dataset = dataset["train"] test_dataset = dataset["test"] - instruction_following = "Let's think step by step and output the final answer within \\boxed{}." + instruction_following = ( + "Let's think step by step and output the final answer within \\boxed{}." + ) # add a row to each data item that represents a unique id def make_map_fn(split): @@ -66,8 +68,12 @@ def process_fn(example, idx): return process_fn - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) - test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True) + test_dataset = test_dataset.map( + function=make_map_fn("test"), + with_indices=True) local_dir = args.local_dir hdfs_dir = args.hdfs_dir diff --git a/Agent0/executor_train/verl/examples/data_preprocess/multiturn.py b/Agent0/executor_train/verl/examples/data_preprocess/multiturn.py index 4bf0192..c4b256e 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/multiturn.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/multiturn.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -54,7 +54,10 @@ def main(): "content": "Quantum computing is a type of computing that uses quantum-mechanical phenomena, " "such as superposition and entanglement, to perform operations on data.", }, - {"role": "user", "content": "How is it different from classical computing?"}, + { + "role": "user", + "content": "How is it different from classical computing?", + }, { "role": "assistant", "content": "Classical computing uses bits that are either 0 or 1, while quantum computing uses " @@ -69,7 +72,10 @@ def main(): { "messages": [ {"role": "system", "content": "You are a helpful assistant."}, - {"role": "user", "content": "Write a simple Python function to calculate factorial."}, + { + "role": "user", + "content": "Write a simple Python function to calculate factorial.", + }, { "role": "assistant", "content": ( diff --git a/Agent0/executor_train/verl/examples/data_preprocess/preprocess_search_r1_dataset.py b/Agent0/executor_train/verl/examples/data_preprocess/preprocess_search_r1_dataset.py index a0c10d5..bdd2108 100644 --- a/Agent0/executor_train/verl/examples/data_preprocess/preprocess_search_r1_dataset.py +++ b/Agent0/executor_train/verl/examples/data_preprocess/preprocess_search_r1_dataset.py @@ -1,178 +1,208 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import argparse -import logging -import os -import tempfile - -import pandas as pd -from huggingface_hub import hf_hub_download -from huggingface_hub.utils import EntryNotFoundError - -from verl.utils.hdfs_io import copy, makedirs - -# Setup logging -logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s") -logger = logging.getLogger(__name__) - -# Configuration constants -DEFAULT_SYSTEM_CONTENT = "You are a helpful and harmless assistant." -DEFAULT_USER_CONTENT_PREFIX = ( - "Answer the given question. You must conduct reasoning inside and " - "first every time you get new information. After reasoning, if you find you lack " - "some knowledge, you can call a search engine by query " - "and it will return the top searched results between and " - ". You can search as many times as your want. If you find no " - "further external knowledge needed, you can directly provide the answer inside " - " and , without detailed illustrations. For example, " - " Beijing . Question: " -) - - -def process_single_row(row, current_split_name, row_index): - """ - Process a single row of data for SearchR1-like format. - - Args: - row: DataFrame row containing the original data - current_split_name: Name of the current split (train/test) - row_index: Index of the row in the DataFrame - - Returns: - pd.Series: Processed row data in the required format - """ - question = row.get("question", "") - - # Build prompt structure - user_content = user_content_prefix.rstrip("\n") + question - prompt = [{"role": "system", "content": system_content}, {"role": "user", "content": user_content}] - - # Extract ground truth from reward_model or fallback to golden_answers - reward_model_data = row.get("reward_model") - if isinstance(reward_model_data, dict) and "ground_truth" in reward_model_data: - ground_truth = reward_model_data.get("ground_truth") - else: - ground_truth = row.get("golden_answers", []) - - # Process data source - data_source_tagged = "searchR1_" + str(row.get("data_source", "")) - - # Build tools kwargs structure - tools_kwargs = { - "search": { - "create_kwargs": {"ground_truth": ground_truth, "question": question, "data_source": data_source_tagged} - } - } - - # Build complete extra_info structure - extra_info = { - "index": row_index, - "need_tools_kwargs": True, - "question": question, - "split": current_split_name, - "tools_kwargs": tools_kwargs, - } - - return pd.Series( - { - "data_source": data_source_tagged, - "prompt": prompt, - "ability": row.get("ability"), - "reward_model": reward_model_data, - "extra_info": extra_info, - "metadata": row.get("metadata"), - } - ) - - -def main(): - local_save_dir = os.path.expanduser(args.local_dir) - os.makedirs(local_save_dir, exist_ok=True) - - processed_files = [] - - # Download and process files using temporary directory - with tempfile.TemporaryDirectory() as tmp_download_dir: - for split in ["train", "test"]: - parquet_filename = f"{split}.parquet" - logger.info(f"Processing {split} split...") - - try: - # Download Parquet file from HuggingFace - logger.info(f"Downloading {parquet_filename} from {args.hf_repo_id}") - local_parquet_filepath = hf_hub_download( - repo_id=args.hf_repo_id, - filename=parquet_filename, - repo_type="dataset", - local_dir=tmp_download_dir, - local_dir_use_symlinks=False, - ) - - # Load and process Parquet file - df_raw = pd.read_parquet(local_parquet_filepath) - logger.info(f"Loaded {len(df_raw)} rows from {parquet_filename}") - - def apply_process_row(row, split_name=split): - return process_single_row(row, current_split_name=split_name, row_index=row.name) - - df_processed = df_raw.apply(apply_process_row, axis=1) - - # Save processed DataFrame - output_file_path = os.path.join(local_save_dir, f"{split}.parquet") - df_processed.to_parquet(output_file_path, index=False) - logger.info(f"Saved {len(df_processed)} processed rows to {output_file_path}") - processed_files.append(output_file_path) - - except EntryNotFoundError: - logger.warning(f"{parquet_filename} not found in repository {args.hf_repo_id}") - except Exception as e: - logger.error(f"Error processing {split} split: {e}") - - if not processed_files: - logger.warning("No data was processed or saved") - return - - logger.info(f"Successfully processed {len(processed_files)} files to {local_save_dir}") - - # Copy to HDFS if specified - if args.hdfs_dir: - try: - makedirs(args.hdfs_dir) - copy(src=local_save_dir, dst=args.hdfs_dir) - logger.info(f"Successfully copied files to HDFS: {args.hdfs_dir}") - except Exception as e: - logger.error(f"Error copying files to HDFS: {e}") - - -if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Download Search-R1 from HuggingFace, process, and save to Parquet.") - parser.add_argument( - "--hf_repo_id", default="PeterJinGo/nq_hotpotqa_train", help="HuggingFace dataset repository ID." - ) - parser.add_argument( - "--local_dir", - default="~/data/searchR1_processed_direct", - help="Local directory to save the processed Parquet files.", - ) - parser.add_argument("--hdfs_dir", default=None, help="Optional HDFS directory to copy the Parquet files to.") - - args = parser.parse_args() - - # System and user content configuration - system_content = DEFAULT_SYSTEM_CONTENT - user_content_prefix = DEFAULT_USER_CONTENT_PREFIX - - main() +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import argparse +import logging +import os +import tempfile + +import pandas as pd +from huggingface_hub import hf_hub_download +from huggingface_hub.utils import EntryNotFoundError + +from verl.utils.hdfs_io import copy, makedirs + +# Setup logging +logging.basicConfig( + level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s" +) +logger = logging.getLogger(__name__) + +# Configuration constants +DEFAULT_SYSTEM_CONTENT = "You are a helpful and harmless assistant." +DEFAULT_USER_CONTENT_PREFIX = ( + "Answer the given question. You must conduct reasoning inside and " + "first every time you get new information. After reasoning, if you find you lack " + "some knowledge, you can call a search engine by query " + "and it will return the top searched results between and " + ". You can search as many times as your want. If you find no " + "further external knowledge needed, you can directly provide the answer inside " + " and , without detailed illustrations. For example, " + " Beijing . Question: ") + + +def process_single_row(row, current_split_name, row_index): + """ + Process a single row of data for SearchR1-like format. + + Args: + row: DataFrame row containing the original data + current_split_name: Name of the current split (train/test) + row_index: Index of the row in the DataFrame + + Returns: + pd.Series: Processed row data in the required format + """ + question = row.get("question", "") + + # Build prompt structure + user_content = user_content_prefix.rstrip("\n") + question + prompt = [ + {"role": "system", "content": system_content}, + {"role": "user", "content": user_content}, + ] + + # Extract ground truth from reward_model or fallback to golden_answers + reward_model_data = row.get("reward_model") + if isinstance( + reward_model_data, + dict) and "ground_truth" in reward_model_data: + ground_truth = reward_model_data.get("ground_truth") + else: + ground_truth = row.get("golden_answers", []) + + # Process data source + data_source_tagged = "searchR1_" + str(row.get("data_source", "")) + + # Build tools kwargs structure + tools_kwargs = { + "search": { + "create_kwargs": { + "ground_truth": ground_truth, + "question": question, + "data_source": data_source_tagged, + } + } + } + + # Build complete extra_info structure + extra_info = { + "index": row_index, + "need_tools_kwargs": True, + "question": question, + "split": current_split_name, + "tools_kwargs": tools_kwargs, + } + + return pd.Series( + { + "data_source": data_source_tagged, + "prompt": prompt, + "ability": row.get("ability"), + "reward_model": reward_model_data, + "extra_info": extra_info, + "metadata": row.get("metadata"), + } + ) + + +def main(): + local_save_dir = os.path.expanduser(args.local_dir) + os.makedirs(local_save_dir, exist_ok=True) + + processed_files = [] + + # Download and process files using temporary directory + with tempfile.TemporaryDirectory() as tmp_download_dir: + for split in ["train", "test"]: + parquet_filename = f"{split}.parquet" + logger.info(f"Processing {split} split...") + + try: + # Download Parquet file from HuggingFace + logger.info( + f"Downloading {parquet_filename} from { + args.hf_repo_id}") + local_parquet_filepath = hf_hub_download( + repo_id=args.hf_repo_id, + filename=parquet_filename, + repo_type="dataset", + local_dir=tmp_download_dir, + local_dir_use_symlinks=False, + ) + + # Load and process Parquet file + df_raw = pd.read_parquet(local_parquet_filepath) + logger.info( + f"Loaded { + len(df_raw)} rows from {parquet_filename}") + + def apply_process_row(row, split_name=split): + return process_single_row( + row, current_split_name=split_name, row_index=row.name + ) + + df_processed = df_raw.apply(apply_process_row, axis=1) + + # Save processed DataFrame + output_file_path = os.path.join( + local_save_dir, f"{split}.parquet") + df_processed.to_parquet(output_file_path, index=False) + logger.info( + f"Saved { + len(df_processed)} processed rows to {output_file_path}") + processed_files.append(output_file_path) + + except EntryNotFoundError: + logger.warning( + f"{parquet_filename} not found in repository { + args.hf_repo_id}") + except Exception as e: + logger.error(f"Error processing {split} split: {e}") + + if not processed_files: + logger.warning("No data was processed or saved") + return + + logger.info( + f"Successfully processed { + len(processed_files)} files to {local_save_dir}") + + # Copy to HDFS if specified + if args.hdfs_dir: + try: + makedirs(args.hdfs_dir) + copy(src=local_save_dir, dst=args.hdfs_dir) + logger.info(f"Successfully copied files to HDFS: {args.hdfs_dir}") + except Exception as e: + logger.error(f"Error copying files to HDFS: {e}") + + +if __name__ == "__main__": + parser = argparse.ArgumentParser( + description="Download Search-R1 from HuggingFace, process, and save to Parquet.") + parser.add_argument( + "--hf_repo_id", + default="PeterJinGo/nq_hotpotqa_train", + help="HuggingFace dataset repository ID.", + ) + parser.add_argument( + "--local_dir", + default="~/data/searchR1_processed_direct", + help="Local directory to save the processed Parquet files.", + ) + parser.add_argument( + "--hdfs_dir", + default=None, + help="Optional HDFS directory to copy the Parquet files to.", + ) + + args = parser.parse_args() + + # System and user content configuration + system_content = DEFAULT_SYSTEM_CONTENT + user_content_prefix = DEFAULT_USER_CONTENT_PREFIX + + main() diff --git a/Agent0/executor_train/verl/examples/sglang_multiturn/search_r1_like/local_dense_retriever/download.py b/Agent0/executor_train/verl/examples/sglang_multiturn/search_r1_like/local_dense_retriever/download.py index 6fe5549..3da95bc 100644 --- a/Agent0/executor_train/verl/examples/sglang_multiturn/search_r1_like/local_dense_retriever/download.py +++ b/Agent0/executor_train/verl/examples/sglang_multiturn/search_r1_like/local_dense_retriever/download.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # Copyright 2025 Search-R1 Contributors # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,16 +13,28 @@ # WITHOUT 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/PeterGriffinJin/Search-R1/blob/main/scripts/download.py +# Adapted from +# https://github.com/PeterGriffinJin/Search-R1/blob/main/scripts/download.py import argparse from huggingface_hub import hf_hub_download -parser = argparse.ArgumentParser(description="Download files from a Hugging Face dataset repository.") -parser.add_argument("--repo_id", type=str, default="PeterJinGo/wiki-18-e5-index", help="Hugging Face repository ID") -parser.add_argument("--save_path", type=str, required=True, help="Local directory to save files") +parser = argparse.ArgumentParser( + description="Download files from a Hugging Face dataset repository." +) +parser.add_argument( + "--repo_id", + type=str, + default="PeterJinGo/wiki-18-e5-index", + help="Hugging Face repository ID", +) +parser.add_argument( + "--save_path", + type=str, + required=True, + help="Local directory to save files") args = parser.parse_args() diff --git a/Agent0/executor_train/verl/examples/sglang_multiturn/search_r1_like/local_dense_retriever/retrieval_server.py b/Agent0/executor_train/verl/examples/sglang_multiturn/search_r1_like/local_dense_retriever/retrieval_server.py index 2f67c14..f2251b9 100644 --- a/Agent0/executor_train/verl/examples/sglang_multiturn/search_r1_like/local_dense_retriever/retrieval_server.py +++ b/Agent0/executor_train/verl/examples/sglang_multiturn/search_r1_like/local_dense_retriever/retrieval_server.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # Copyright 2025 Search-R1 Contributors # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -13,7 +13,8 @@ # WITHOUT 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/PeterGriffinJin/Search-R1/blob/main/search_r1/search/retrieval_server.py +# Adapted from +# https://github.com/PeterGriffinJin/Search-R1/blob/main/search_r1/search/retrieval_server.py import argparse import json @@ -32,7 +33,9 @@ def load_corpus(corpus_path: str): - corpus = datasets.load_dataset("json", data_files=corpus_path, split="train", num_proc=4) + corpus = datasets.load_dataset( + "json", data_files=corpus_path, split="train", num_proc=4 + ) return corpus @@ -47,13 +50,21 @@ def load_model(model_path: str, use_fp16: bool = False): model.cuda() if use_fp16: model = model.half() - tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained( + model_path, use_fast=True, trust_remote_code=True + ) return model, tokenizer -def pooling(pooler_output, last_hidden_state, attention_mask=None, pooling_method="mean"): +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) + 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] @@ -64,14 +75,22 @@ def pooling(pooler_output, last_hidden_state, attention_mask=None, pooling_metho class Encoder: - def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16): + 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, self.tokenizer = load_model( + model_path=model_path, use_fp16=use_fp16 + ) self.model.eval() @torch.no_grad() @@ -89,25 +108,35 @@ def encode(self, query_list: list[str], is_query=True) -> np.ndarray: 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 + 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" + 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 + 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 ) - 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 + 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) @@ -133,13 +162,21 @@ def __init__(self, config): 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): + 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): + def batch_search( + self, + query_list: list[str], + num: int = None, + return_score: bool = False): return self._batch_search(query_list, num, return_score) @@ -173,7 +210,10 @@ def _search(self, query: str, num: int = None, return_score: bool = False): hits = hits[:num] if self.contain_doc: - all_contents = [json.loads(self.searcher.doc(hit.docid).raw())["contents"] for hit in hits] + all_contents = [ + json.loads(self.searcher.doc(hit.docid).raw())["contents"] + for hit in hits + ] results = [ { "title": content.split("\n")[0].strip('"'), @@ -190,7 +230,11 @@ def _search(self, query: str, num: int = None, return_score: bool = False): else: return results - def _batch_search(self, query_list: list[str], num: int = None, return_score: bool = False): + def _batch_search( + self, + query_list: list[str], + num: int = None, + return_score: bool = False): results = [] scores = [] for query in query_list: @@ -237,7 +281,11 @@ def _search(self, query: str, num: int = None, return_score: bool = False): else: return results - def _batch_search(self, query_list: list[str], num: int = None, return_score: bool = False): + 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: @@ -245,8 +293,13 @@ def _batch_search(self, query_list: list[str], num: int = None, return_score: bo 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] + 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() @@ -256,12 +309,21 @@ def _batch_search(self, query_list: list[str], num: int = None, return_score: bo 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))] + batch_results = [ + batch_results[i * num: (i + 1) * num] for i in range(len(batch_idxs)) + ] results.extend(batch_results) scores.extend(batch_scores) - del batch_emb, batch_scores, batch_idxs, query_batch, flat_idxs, batch_results + del ( + batch_emb, + batch_scores, + batch_idxs, + query_batch, + flat_idxs, + batch_results, + ) torch.cuda.empty_cache() if return_score: @@ -356,8 +418,7 @@ def retrieve_endpoint(request: QueryRequest): # Perform batch retrieval results, scores = retriever.batch_search( - query_list=request.queries, num=request.topk, return_score=request.return_scores - ) + query_list=request.queries, num=request.topk, return_score=request.return_scores) # Format response resp = [] @@ -374,9 +435,13 @@ def retrieve_endpoint(request: QueryRequest): if __name__ == "__main__": - parser = argparse.ArgumentParser(description="Launch the local faiss retriever.") + 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." + "--index_path", + type=str, + default="/home/peterjin/mnt/index/wiki-18/e5_Flat.index", + help="Corpus indexing file.", ) parser.add_argument( "--corpus_path", @@ -384,17 +449,31 @@ def retrieve_endpoint(request: QueryRequest): 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_name", type=str, default="e5", help="Name of the retriever model.") parser.add_argument( - "--retriever_model", type=str, default="intfloat/e5-base-v2", help="Path of the retriever model." + "--topk", + type=int, + default=3, + help="Number of retrieved passages for one query.", + ) + parser.add_argument( + "--retriever_name", + type=str, + default="e5", + help="Name of the retriever model.") + parser.add_argument( + "--retriever_model", + type=str, + default="intfloat/e5-base-v2", + help="Path of the retriever model.", + ) + parser.add_argument( + "--faiss_gpu", action="store_true", help="Use GPU for computation" ) - parser.add_argument("--faiss_gpu", action="store_true", help="Use GPU for computation") args = parser.parse_args() # 1) Build a config (could also parse from arguments). - # In real usage, you'd parse your CLI arguments or environment variables. + # In real usage, you'd parse your CLI arguments or environment variables. config = Config( retrieval_method=args.retriever_name, # or "dense" index_path=args.index_path, diff --git a/Agent0/executor_train/verl/examples/split_placement/main_ppo_split.py b/Agent0/executor_train/verl/examples/split_placement/main_ppo_split.py index c438e7a..e17f30a 100644 --- a/Agent0/executor_train/verl/examples/split_placement/main_ppo_split.py +++ b/Agent0/executor_train/verl/examples/split_placement/main_ppo_split.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -37,16 +37,19 @@ def _select_rm_score_fn(data_source): class RewardManager: def __init__(self, tokenizer, num_examine) -> None: self.tokenizer = tokenizer - self.num_examine = num_examine # the number of batches of decoded responses to print to the console + # the number of batches of decoded responses to print to the console + self.num_examine = num_examine def __call__(self, data: DataProto, return_dict: bool = False): """We will expand this function gradually based on the available datasets""" - # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn + # 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) + reward_tensor = torch.zeros_like( + data.batch["responses"], dtype=torch.float32) already_print_data_sources = {} @@ -57,11 +60,15 @@ def __call__(self, data: DataProto, return_dict: bool = False): prompt_length = prompt_ids.shape[-1] - valid_prompt_length = data_item.batch["attention_mask"][:prompt_length].sum() + 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_length = data_item.batch["attention_mask"][ + prompt_length: + ].sum() valid_response_ids = response_ids[:valid_response_length] # decode @@ -74,7 +81,9 @@ def __call__(self, data: DataProto, return_dict: bool = False): data_source = data_item.non_tensor_batch["data_source"] compute_score_fn = _select_rm_score_fn(data_source) - score = compute_score_fn(solution_str=sequences_str, ground_truth=ground_truth) + score = compute_score_fn( + solution_str=sequences_str, ground_truth=ground_truth + ) reward_tensor[i, valid_response_length - 1] = score if data_source not in already_print_data_sources: @@ -90,12 +99,17 @@ def __call__(self, data: DataProto, return_dict: bool = False): return reward_tensor -@hydra.main(config_path="config", config_name="ppo_trainer_split", version_base=None) +@hydra.main(config_path="config", + config_name="ppo_trainer_split", + version_base=None) 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"}}, + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN"}}, num_cpus=config.ray_init.num_cpus, ) @@ -111,7 +125,9 @@ def main_task(config): from verl.utils.fs import copy_to_local - pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values + pprint( + OmegaConf.to_container(config, resolve=True) + ) # resolve=True will eval symbol values OmegaConf.resolve(config) # download the checkpoint from hdfs @@ -152,13 +168,17 @@ def main_task(config): critic_pool_id = "critic_pool" if config.trainer.nnodes // 2 == 0 and config.trainer.n_gpus_per_node // 2 > 0: resource_pool_spec = { - actor_rollout_ref_pool_id: [config.trainer.n_gpus_per_node // 2] * config.trainer.nnodes, - critic_pool_id: [config.trainer.n_gpus_per_node // 2] * config.trainer.nnodes, + actor_rollout_ref_pool_id: [config.trainer.n_gpus_per_node // 2] + * config.trainer.nnodes, + critic_pool_id: [config.trainer.n_gpus_per_node // 2] + * config.trainer.nnodes, } else: resource_pool_spec = { - actor_rollout_ref_pool_id: [config.trainer.n_gpus_per_node] * (config.trainer.nnodes // 2), - critic_pool_id: [config.trainer.n_gpus_per_node] * (config.trainer.nnodes // 2), + actor_rollout_ref_pool_id: [config.trainer.n_gpus_per_node] + * (config.trainer.nnodes // 2), + critic_pool_id: [config.trainer.n_gpus_per_node] + * (config.trainer.nnodes // 2), } print(f"resource_pool_spec: {resource_pool_spec}") mapping = { @@ -192,7 +212,9 @@ def main_task(config): # Note that we always use function-based RM for validation val_reward_fn = RewardManager(tokenizer=tokenizer, num_examine=1) - resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, mapping=mapping + ) RayPPOTrainer.fit = fit trainer = RayPPOTrainer( diff --git a/Agent0/executor_train/verl/examples/split_placement/split_monkey_patch.py b/Agent0/executor_train/verl/examples/split_placement/split_monkey_patch.py index ef58509..150af48 100644 --- a/Agent0/executor_train/verl/examples/split_placement/split_monkey_patch.py +++ b/Agent0/executor_train/verl/examples/split_placement/split_monkey_patch.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -59,7 +59,9 @@ def fit(self): # 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): + 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) @@ -78,13 +80,16 @@ def fit(self): batch: DataProto = DataProto.from_single_dict(batch_dict) # pop those keys for generation - gen_batch = batch.pop(batch_keys=["input_ids", "attention_mask", "position_ids"]) + gen_batch = batch.pop( + batch_keys=["input_ids", "attention_mask", "position_ids"] + ) is_last_step = self.global_steps >= self.total_training_steps with marked_timer("step", timing_raw): # generate a batch with marked_timer("gen", timing_raw): - gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) + gen_batch_output = self.actor_rollout_wg.generate_sequences( + gen_batch) timing_raw.update(gen_batch_output.meta_info["timing"]) gen_batch_output.meta_info.pop("timing", None) @@ -92,13 +97,17 @@ def fit(self): with marked_timer("gen_max", timing_raw): gen_baseline_batch = deepcopy(gen_batch) gen_baseline_batch.meta_info["do_sample"] = False - gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch) + gen_baseline_output = self.actor_rollout_wg.generate_sequences( + gen_baseline_batch) batch = batch.union(gen_baseline_output) reward_baseline_tensor = self.reward_fn(batch) - reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1) + reward_baseline_tensor = reward_baseline_tensor.sum( + dim=-1) - batch.pop(batch_keys=list(gen_baseline_output.batch.keys())) + batch.pop( + batch_keys=list( + gen_baseline_output.batch.keys())) batch.batch["reward_baselines"] = reward_baseline_tensor @@ -108,7 +117,10 @@ def fit(self): [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.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) batch = batch.union(gen_batch_output) # Balance the number of valid tokens across DP ranks. @@ -119,17 +131,21 @@ def fit(self): 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_token_num"] = torch.sum( + batch.batch["attention_mask"], dim=-1 + ).tolist() # recompute old_log_probs with marked_timer("old_log_prob", timing_raw): - old_log_prob = self.actor_rollout_wg.compute_log_prob(batch) + old_log_prob = self.actor_rollout_wg.compute_log_prob( + batch) batch = batch.union(old_log_prob) if self.use_reference_policy: # compute reference log_prob with marked_timer("ref", timing_raw): - ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch) + ref_log_prob = self.ref_policy_wg.compute_ref_log_prob( + batch) batch = batch.union(ref_log_prob) # compute values @@ -154,14 +170,20 @@ def fit(self): # compute rewards. apply_kl_penalty if available if self.config.algorithm.use_kl_in_reward: batch, kl_metrics = apply_kl_penalty( - batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty + batch, + kl_ctrl=self.kl_ctrl_in_reward, + kl_penalty=self.config.algorithm.kl_penalty, ) metrics.update(kl_metrics) else: - batch.batch["token_level_rewards"] = batch.batch["token_level_scores"] + batch.batch["token_level_rewards"] = batch.batch[ + "token_level_scores" + ] # compute advantages, executed on the driver process - norm_adv_by_std_in_grpo = self.config.algorithm.get("norm_adv_by_std_in_grpo", True) + norm_adv_by_std_in_grpo = self.config.algorithm.get( + "norm_adv_by_std_in_grpo", True + ) batch = compute_advantage( batch, adv_estimator=self.config.algorithm.adv_estimator, @@ -175,7 +197,8 @@ def fit(self): if self.config.trainer.critic_warmup <= self.global_steps: # update actor with marked_timer("update_actor_call", timing_raw): - actor_output = self.actor_rollout_wg.update_actor(batch) + actor_output = self.actor_rollout_wg.update_actor( + batch) else: actor_output = None @@ -184,22 +207,30 @@ def fit(self): with marked_timer("update_critic_call", timing_raw): critic_output = self.critic_wg.update_critic(batch) - # NOTE: make sure you set blocking=False in update_actor and update_crtic in the worker class + # NOTE: make sure you set blocking=False in update_actor + # and update_crtic in the worker class with marked_timer("update_actor_critic", timing_raw): critic_output = critic_output.get() - critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"]) + critic_output_metrics = reduce_metrics( + critic_output.meta_info["metrics"] + ) metrics.update(critic_output_metrics) if actor_output is not None: actor_output = actor_output.get() - actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) + actor_output_metrics = reduce_metrics( + actor_output.meta_info["metrics"] + ) metrics.update(actor_output_metrics) # validate if ( self.val_reward_fn is not None and self.config.trainer.test_freq > 0 - and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0) + and ( + is_last_step + or self.global_steps % self.config.trainer.test_freq == 0 + ) ): with marked_timer("testing", timing_raw): val_metrics: dict = self._validate() @@ -208,14 +239,20 @@ def fit(self): metrics.update(val_metrics) if self.config.trainer.save_freq > 0 and ( - is_last_step or self.global_steps % self.config.trainer.save_freq == 0 + is_last_step + or self.global_steps % self.config.trainer.save_freq == 0 ): with marked_timer("save_checkpoint", timing_raw): self._save_checkpoint() # collect metrics - metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic)) - metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) + metrics.update( + compute_data_metrics(batch=batch, use_critic=self.use_critic) + ) + metrics.update( + compute_timing_metrics( + batch=batch, + timing_raw=timing_raw)) # TODO: make a canonical logger that supports various backend logger.log(data=metrics, step=self.global_steps) diff --git a/Agent0/executor_train/verl/recipe/char_count/create_dataset.py b/Agent0/executor_train/verl/recipe/char_count/create_dataset.py index 47571e0..d3e491a 100644 --- a/Agent0/executor_train/verl/recipe/char_count/create_dataset.py +++ b/Agent0/executor_train/verl/recipe/char_count/create_dataset.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -102,7 +102,8 @@ def create_prompt_response(min_length=3, max_length=5): full_output = [] for _ in range(total_number): - output = create_prompt_response(min_length=min_length, max_length=max_length) + output = create_prompt_response( + min_length=min_length, max_length=max_length) full_output.append(output) # random reorder @@ -138,9 +139,21 @@ def create_prompt_response(min_length=3, max_length=5): sft_test_dataset.to_parquet(os.path.join(folder, "test.parquet")) # build RL dataset - rl_train_dataset = {"prompt": [], "data_source": [], "ability": [], "reward_model": [], "extra_info": []} - - rl_test_dataset = {"prompt": [], "data_source": [], "ability": [], "reward_model": [], "extra_info": []} + rl_train_dataset = { + "prompt": [], + "data_source": [], + "ability": [], + "reward_model": [], + "extra_info": [], + } + + rl_test_dataset = { + "prompt": [], + "data_source": [], + "ability": [], + "reward_model": [], + "extra_info": [], + } from verl.utils.reward_score.math import last_boxed_only_string, remove_boxed @@ -158,7 +171,10 @@ def create_prompt_response(min_length=3, max_length=5): rl_train_dataset["data_source"].append("char_count") rl_train_dataset["ability"].append("other") rl_train_dataset["reward_model"].append( - {"style": "rule", "ground_truth": remove_boxed(last_boxed_only_string(response))} + { + "style": "rule", + "ground_truth": remove_boxed(last_boxed_only_string(response)), + } ) rl_train_dataset["extra_info"].append({"response": response}) @@ -176,7 +192,10 @@ def create_prompt_response(min_length=3, max_length=5): rl_test_dataset["data_source"].append("char_count") rl_test_dataset["ability"].append("other") rl_test_dataset["reward_model"].append( - {"style": "rule", "ground_truth": remove_boxed(last_boxed_only_string(response))} + { + "style": "rule", + "ground_truth": remove_boxed(last_boxed_only_string(response)), + } ) rl_test_dataset["extra_info"].append({"response": response}) diff --git a/Agent0/executor_train/verl/recipe/char_count/reward_function.py b/Agent0/executor_train/verl/recipe/char_count/reward_function.py index 9bdffe2..adbdd12 100644 --- a/Agent0/executor_train/verl/recipe/char_count/reward_function.py +++ b/Agent0/executor_train/verl/recipe/char_count/reward_function.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,7 +19,9 @@ from verl.utils.reward_score import math -def char_count_reward_function(data_source, solution_str, ground_truth, extra_info=None): +def char_count_reward_function( + data_source, solution_str, ground_truth, extra_info=None +): try: last_boxed_string = math.last_boxed_only_string(solution_str) if last_boxed_string is None: diff --git a/Agent0/executor_train/verl/recipe/dapo/dapo_ray_trainer.py b/Agent0/executor_train/verl/recipe/dapo/dapo_ray_trainer.py index d3d6dbc..d1a79be 100644 --- a/Agent0/executor_train/verl/recipe/dapo/dapo_ray_trainer.py +++ b/Agent0/executor_train/verl/recipe/dapo/dapo_ray_trainer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -73,7 +73,9 @@ def fit(self): # 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): + if self.val_reward_fn is not None and self.config.trainer.get( + "val_before_train", True + ): val_metrics = self._validate() assert val_metrics, f"{val_metrics=}" pprint(f"Initial validation metrics: {val_metrics}") @@ -82,7 +84,11 @@ def fit(self): return # add tqdm - progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress") + progress_bar = tqdm( + total=self.total_training_steps, + initial=self.global_steps, + desc="Training Progress", + ) # we start from step 1 self.global_steps += 1 @@ -116,22 +122,34 @@ def fit(self): # pop those keys for generation if "multi_modal_data" in new_batch.non_tensor_batch.keys(): gen_batch = new_batch.pop( - batch_keys=["input_ids", "attention_mask", "position_ids"], - non_tensor_batch_keys=["raw_prompt_ids", "multi_modal_data"], + batch_keys=[ + "input_ids", + "attention_mask", + "position_ids"], + non_tensor_batch_keys=[ + "raw_prompt_ids", + "multi_modal_data"], ) else: gen_batch = new_batch.pop( - batch_keys=["input_ids", "attention_mask", "position_ids"], + batch_keys=[ + "input_ids", + "attention_mask", + "position_ids"], non_tensor_batch_keys=["raw_prompt_ids"], ) - gen_batch = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) + gen_batch = gen_batch.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) is_last_step = self.global_steps >= self.total_training_steps with marked_timer("step", timing_raw): # generate a batch with marked_timer("gen", timing_raw, "red"): - gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) + gen_batch_output = self.actor_rollout_wg.generate_sequences( + gen_batch) timing_raw.update(gen_batch_output.meta_info["timing"]) gen_batch_output.meta_info.pop("timing", None) @@ -139,23 +157,34 @@ def fit(self): with marked_timer("gen_max", timing_raw, "red"): gen_baseline_batch = deepcopy(gen_batch) gen_baseline_batch.meta_info["do_sample"] = False - gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch) + gen_baseline_output = ( + self.actor_rollout_wg.generate_sequences( + gen_baseline_batch + ) + ) new_batch = new_batch.union(gen_baseline_output) reward_baseline_tensor = self.reward_fn(new_batch) - reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1) + reward_baseline_tensor = reward_baseline_tensor.sum( + dim=-1) - new_batch.pop(batch_keys=list(gen_baseline_output.batch.keys())) + new_batch.pop( + batch_keys=list( + gen_baseline_output.batch.keys())) new_batch.batch["reward_baselines"] = reward_baseline_tensor del gen_baseline_batch, gen_baseline_output new_batch.non_tensor_batch["uid"] = np.array( - [str(uuid.uuid4()) for _ in range(len(new_batch.batch))], dtype=object + [str(uuid.uuid4()) for _ in range(len(new_batch.batch))], + dtype=object, ) # repeat to align with repeated responses in rollout - new_batch = new_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) + new_batch = new_batch.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) new_batch = new_batch.union(gen_batch_output) with marked_timer("reward", timing_raw, "yellow"): @@ -164,15 +193,19 @@ def fit(self): # the results from reward model and rule-based results. if self.use_rm: # we first compute reward model score - reward_tensor = self.rm_wg.compute_rm_score(new_batch) + reward_tensor = self.rm_wg.compute_rm_score( + new_batch) new_batch = new_batch.union(reward_tensor) # we combine with rule-based rm reward_extra_infos_dict: dict[str, list] try: - reward_result = self.reward_fn(new_batch, return_dict=True) + reward_result = self.reward_fn( + new_batch, return_dict=True) reward_tensor = reward_result["reward_tensor"] - reward_extra_infos_dict = reward_result.get("reward_extra_info", {}) + reward_extra_infos_dict = reward_result.get( + "reward_extra_info", {} + ) except Exception as e: print(f"Error in reward_fn: {e}") reward_tensor = self.reward_fn(new_batch) @@ -182,19 +215,26 @@ def fit(self): if reward_extra_infos_dict: new_batch.non_tensor_batch.update( - {k: np.array(v) for k, v in reward_extra_infos_dict.items()} + { + k: np.array(v) + for k, v in reward_extra_infos_dict.items() + } ) # compute rewards. apply_kl_penalty if available if self.config.algorithm.use_kl_in_reward: new_batch, kl_metrics = apply_kl_penalty( - new_batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty + new_batch, + kl_ctrl=self.kl_ctrl_in_reward, + kl_penalty=self.config.algorithm.kl_penalty, ) - metrics.update( - kl_metrics - ) # TODO: This will be cleared if we use multiple genenration batches + # TODO: This will be cleared if we use multiple + # genenration batches + metrics.update(kl_metrics) else: - new_batch.batch["token_level_rewards"] = new_batch.batch["token_level_scores"] + new_batch.batch["token_level_rewards"] = new_batch.batch[ + "token_level_scores" + ] if not self.config.algorithm.filter_groups.enable: batch = new_batch @@ -204,23 +244,30 @@ def fit(self): if metric_name == "seq_final_reward": # Turn to numpy for easier filtering new_batch.non_tensor_batch["seq_final_reward"] = ( - new_batch.batch["token_level_rewards"].sum(dim=-1).numpy() + new_batch.batch["token_level_rewards"] + .sum(dim=-1) + .numpy() ) elif metric_name == "seq_reward": new_batch.non_tensor_batch["seq_reward"] = ( - new_batch.batch["token_level_scores"].sum(dim=-1).numpy() + new_batch.batch["token_level_scores"] + .sum(dim=-1) + .numpy() ) # Collect the sequence reward for each trajectory prompt_uid2metric_vals = defaultdict(list) for uid, metric_val in zip( - new_batch.non_tensor_batch["uid"], new_batch.non_tensor_batch[metric_name], strict=True + new_batch.non_tensor_batch["uid"], + new_batch.non_tensor_batch[metric_name], + strict=True, ): prompt_uid2metric_vals[uid].append(metric_val) prompt_uid2metric_std = {} for prompt_uid, metric_vals in prompt_uid2metric_vals.items(): - prompt_uid2metric_std[prompt_uid] = np.std(metric_vals) + prompt_uid2metric_std[prompt_uid] = np.std( + metric_vals) kept_prompt_uids = [ uid @@ -230,30 +277,44 @@ def fit(self): num_prompt_in_batch += len(kept_prompt_uids) kept_traj_idxs = [] - for idx, traj_from_prompt_uid in enumerate(new_batch.non_tensor_batch["uid"]): + for idx, traj_from_prompt_uid in enumerate( + new_batch.non_tensor_batch["uid"] + ): if traj_from_prompt_uid in kept_prompt_uids: kept_traj_idxs.append(idx) new_batch = new_batch[kept_traj_idxs] - batch = new_batch if batch is None else DataProto.concat([batch, new_batch]) + batch = ( + new_batch + if batch is None + else DataProto.concat([batch, new_batch]) + ) prompt_bsz = self.config.data.train_batch_size if num_prompt_in_batch < prompt_bsz: print(f"{num_prompt_in_batch=} < {prompt_bsz=}") - max_num_gen_batches = self.config.algorithm.filter_groups.max_num_gen_batches - if max_num_gen_batches <= 0 or num_gen_batches < max_num_gen_batches: + max_num_gen_batches = ( + self.config.algorithm.filter_groups.max_num_gen_batches) + if ( + max_num_gen_batches <= 0 + or num_gen_batches < max_num_gen_batches + ): print(f"{num_gen_batches=}. Keep generating...") progress_bar.update(1) continue else: raise ValueError( - f"{num_gen_batches=} >= {max_num_gen_batches=}." - + " Generated too many. Please check if your data are too difficult." - + " You could also try set max_num_gen_batches=0 to enable endless trials." - ) + f"{ + num_gen_batches=} >= { + max_num_gen_batches=}." + + " Generated too many. Please check if your data are too difficult." + + " You could also try set max_num_gen_batches=0 to enable endless trials.") else: # Align the batch - traj_bsz = self.config.data.train_batch_size * self.config.actor_rollout_ref.rollout.n + traj_bsz = ( + self.config.data.train_batch_size + * self.config.actor_rollout_ref.rollout.n + ) batch = batch[:traj_bsz] # === Updating === @@ -269,16 +330,27 @@ def fit(self): 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_token_num"] = torch.sum( + batch.batch["attention_mask"], dim=-1 + ).tolist() # recompute old_log_probs with marked_timer("old_log_prob", timing_raw, "blue"): - old_log_prob = self.actor_rollout_wg.compute_log_prob(batch) + old_log_prob = self.actor_rollout_wg.compute_log_prob( + batch) entropys = old_log_prob.batch["entropys"] response_masks = batch.batch["response_mask"] - loss_agg_mode = self.config.actor_rollout_ref.actor.loss_agg_mode - entropy_agg = agg_loss(loss_mat=entropys, loss_mask=response_masks, loss_agg_mode=loss_agg_mode) - old_log_prob_metrics = {"actor/entropy": entropy_agg.detach().item()} + loss_agg_mode = ( + self.config.actor_rollout_ref.actor.loss_agg_mode + ) + entropy_agg = agg_loss( + loss_mat=entropys, + loss_mask=response_masks, + loss_agg_mode=loss_agg_mode, + ) + old_log_prob_metrics = { + "actor/entropy": entropy_agg.detach().item() + } metrics.update(old_log_prob_metrics) old_log_prob.batch.pop("entropys") batch = batch.union(old_log_prob) @@ -286,7 +358,8 @@ def fit(self): if self.use_reference_policy: # compute reference log_prob with marked_timer("ref", timing_raw, "olive"): - ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch) + ref_log_prob = self.ref_policy_wg.compute_ref_log_prob( + batch) batch = batch.union(ref_log_prob) # compute values @@ -297,7 +370,9 @@ def fit(self): with marked_timer("adv", timing_raw, "brown"): # compute advantages, executed on the driver process - norm_adv_by_std_in_grpo = self.config.algorithm.get("norm_adv_by_std_in_grpo", True) + norm_adv_by_std_in_grpo = self.config.algorithm.get( + "norm_adv_by_std_in_grpo", True + ) batch = compute_advantage( batch, adv_estimator=self.config.algorithm.adv_estimator, @@ -311,22 +386,30 @@ def fit(self): if self.use_critic: with marked_timer("update_critic", timing_raw, "pink"): critic_output = self.critic_wg.update_critic(batch) - critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"]) + critic_output_metrics = reduce_metrics( + critic_output.meta_info["metrics"] + ) metrics.update(critic_output_metrics) # implement critic warmup if self.config.trainer.critic_warmup <= self.global_steps: # update actor with marked_timer("update_actor", timing_raw, "red"): - actor_output = self.actor_rollout_wg.update_actor(batch) - actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) + actor_output = self.actor_rollout_wg.update_actor( + batch) + actor_output_metrics = reduce_metrics( + actor_output.meta_info["metrics"] + ) metrics.update(actor_output_metrics) # validate if ( self.val_reward_fn is not None and self.config.trainer.test_freq > 0 - and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0) + and ( + is_last_step + or self.global_steps % self.config.trainer.test_freq == 0 + ) ): with marked_timer("testing", timing_raw, "green"): val_metrics: dict = self._validate() @@ -335,8 +418,8 @@ def fit(self): metrics.update(val_metrics) if self.config.trainer.save_freq > 0 and ( - is_last_step or self.global_steps % self.config.trainer.save_freq == 0 - ): + is_last_step or self.global_steps % + self.config.trainer.save_freq == 0): with marked_timer("save_checkpoint", timing_raw, "green"): self._save_checkpoint() @@ -351,11 +434,20 @@ def fit(self): self.rm_wg.stop_profile() # collect metrics - metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic)) - metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) + metrics.update( + compute_data_metrics( + batch=batch, + use_critic=self.use_critic)) + metrics.update( + compute_timing_metrics(batch=batch, timing_raw=timing_raw) + ) # TODO: implement actual tflpo and theoretical tflpo n_gpus = self.resource_pool_manager.get_n_gpus() - metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus)) + metrics.update( + compute_throughout_metrics( + batch=batch, timing_raw=timing_raw, n_gpus=n_gpus + ) + ) timing_raw = defaultdict(float) # clear timing metrics["train/num_gen_batches"] = num_gen_batches diff --git a/Agent0/executor_train/verl/recipe/dapo/main_dapo.py b/Agent0/executor_train/verl/recipe/dapo/main_dapo.py index 1ee7359..a569402 100644 --- a/Agent0/executor_train/verl/recipe/dapo/main_dapo.py +++ b/Agent0/executor_train/verl/recipe/dapo/main_dapo.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -28,7 +28,9 @@ from .dapo_ray_trainer import RayDAPOTrainer -@hydra.main(config_path="config", config_name="dapo_trainer", version_base=None) +@hydra.main(config_path="config", + config_name="dapo_trainer", + version_base=None) def main(config): run_ppo(config) @@ -38,7 +40,11 @@ def run_ppo(config) -> None: # this is for local ray cluster ray.init( runtime_env={ - "env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN", "VLLM_LOGGING_LEVEL": "WARN"} + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "WARN", + } }, num_cpus=config.ray_init.num_cpus, ) @@ -48,8 +54,12 @@ def run_ppo(config) -> None: and OmegaConf.select(config.trainer, "profile_steps") is not None and len(OmegaConf.select(config.trainer, "profile_steps")) > 0 ): - nsight_options = OmegaConf.to_container(config.trainer.controller_nsight_options) - runner = TaskRunner.options(runtime_env={"nsight": nsight_options}).remote() + nsight_options = OmegaConf.to_container( + config.trainer.controller_nsight_options + ) + runner = TaskRunner.options( + runtime_env={ + "nsight": nsight_options}).remote() else: runner = TaskRunner.remote() ray.get(runner.run.remote(config)) @@ -65,9 +75,14 @@ def run(self, config): from verl.utils.fs import copy_to_local - print(f"TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}") + print( + f"TaskRunner hostname: { + socket.gethostname()}, PID: { + os.getpid()}") - pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values + pprint( + OmegaConf.to_container(config, resolve=True) + ) # resolve=True will eval symbol values OmegaConf.resolve(config) # download the checkpoint from hdfs @@ -77,7 +92,9 @@ def run(self, config): from verl.utils import hf_processor, hf_tokenizer tokenizer = hf_tokenizer(local_path) - processor = hf_processor(local_path, use_fast=True) # used for multimodal LLM, could be none + processor = hf_processor( + local_path, use_fast=True + ) # used for multimodal LLM, could be none # define worker classes if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"}: @@ -90,7 +107,10 @@ def run(self, config): elif config.actor_rollout_ref.actor.strategy == "megatron": assert config.actor_rollout_ref.actor.strategy == config.critic.strategy from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup - from verl.workers.megatron_workers import ActorRolloutRefWorker, CriticWorker + from verl.workers.megatron_workers import ( + ActorRolloutRefWorker, + CriticWorker, + ) ray_worker_group_cls = NVMegatronRayWorkerGroup @@ -106,7 +126,9 @@ def run(self, config): global_pool_id = "global_pool" resource_pool_spec = { - global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + global_pool_id: [ + config.trainer.n_gpus_per_node] * + config.trainer.nnodes, } mapping = { Role.ActorRollout: global_pool_id, @@ -126,19 +148,25 @@ def run(self, config): from verl.workers.megatron_workers import RewardModelWorker else: raise NotImplementedError - role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker) + role_worker_mapping[Role.RewardModel] = ray.remote( + RewardModelWorker) mapping[Role.RewardModel] = global_pool_id # reference model - if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss: - role_worker_mapping[Role.RefPolicy] = ray.remote(ActorRolloutRefWorker) + if ( + config.algorithm.use_kl_in_reward + or config.actor_rollout_ref.actor.use_kl_loss + ): + role_worker_mapping[Role.RefPolicy] = ray.remote( + ActorRolloutRefWorker) mapping[Role.RefPolicy] = global_pool_id from verl.workers.reward_manager import get_reward_manager_cls # Note(haibin.lin): please make sure custom reward managers are imported and # registered via `verl.workers.reward_manager.register` - reward_manager_name = config.reward_model.get("reward_manager", "naive") + reward_manager_name = config.reward_model.get( + "reward_manager", "naive") reward_manager_cls = get_reward_manager_cls(reward_manager_name) compute_score = get_custom_reward_fn(config) @@ -160,7 +188,9 @@ def run(self, config): max_resp_len=config.data.max_response_length, overlong_buffer_cfg=config.reward_model.overlong_buffer, ) - resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, mapping=mapping + ) trainer = RayDAPOTrainer( config=config, diff --git a/Agent0/executor_train/verl/recipe/entropy/entropy_ray_trainer.py b/Agent0/executor_train/verl/recipe/entropy/entropy_ray_trainer.py index 0b0b043..8523f90 100644 --- a/Agent0/executor_train/verl/recipe/entropy/entropy_ray_trainer.py +++ b/Agent0/executor_train/verl/recipe/entropy/entropy_ray_trainer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -72,7 +72,9 @@ def fit(self): # 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): + if self.val_reward_fn is not None and self.config.trainer.get( + "val_before_train", True + ): val_metrics = self._validate() assert val_metrics, f"{val_metrics=}" pprint(f"Initial validation metrics: {val_metrics}") @@ -81,7 +83,11 @@ def fit(self): return # add tqdm - progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress") + progress_bar = tqdm( + total=self.total_training_steps, + initial=self.global_steps, + desc="Training Progress", + ) # we start from step 1 self.global_steps += 1 @@ -100,15 +106,28 @@ def fit(self): # pop those keys for generation if "multi_modal_inputs" in new_batch.non_tensor_batch.keys(): gen_batch = new_batch.pop( - batch_keys=["input_ids", "attention_mask", "position_ids"], - non_tensor_batch_keys=["raw_prompt_ids", "multi_modal_data", "multi_modal_inputs"], + batch_keys=[ + "input_ids", + "attention_mask", + "position_ids"], + non_tensor_batch_keys=[ + "raw_prompt_ids", + "multi_modal_data", + "multi_modal_inputs", + ], ) else: gen_batch = new_batch.pop( - batch_keys=["input_ids", "attention_mask", "position_ids"], + batch_keys=[ + "input_ids", + "attention_mask", + "position_ids"], non_tensor_batch_keys=["raw_prompt_ids"], ) - gen_batch = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) + gen_batch = gen_batch.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) is_last_step = self.global_steps >= self.total_training_steps @@ -118,31 +137,44 @@ def fit(self): # gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) with simple_timer("gen", timing_raw): if not self.async_rollout_mode: - gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) + gen_batch_output = self.actor_rollout_wg.generate_sequences( + gen_batch) else: - gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch) + gen_batch_output = ( + self.async_rollout_manager.generate_sequences(gen_batch)) if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX: with simple_timer("gen_max", timing_raw): gen_baseline_batch = deepcopy(gen_batch) gen_baseline_batch.meta_info["do_sample"] = False - gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch) + gen_baseline_output = ( + self.actor_rollout_wg.generate_sequences( + gen_baseline_batch + ) + ) new_batch = new_batch.union(gen_baseline_output) reward_baseline_tensor = self.reward_fn(new_batch) - reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1) + reward_baseline_tensor = reward_baseline_tensor.sum( + dim=-1) - new_batch.pop(batch_keys=list(gen_baseline_output.batch.keys())) + new_batch.pop( + batch_keys=list( + gen_baseline_output.batch.keys())) new_batch.batch["reward_baselines"] = reward_baseline_tensor del gen_baseline_batch, gen_baseline_output new_batch.non_tensor_batch["uid"] = np.array( - [str(uuid.uuid4()) for _ in range(len(new_batch.batch))], dtype=object + [str(uuid.uuid4()) for _ in range(len(new_batch.batch))], + dtype=object, ) # repeat to align with repeated responses in rollout - new_batch = new_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) + new_batch = new_batch.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) new_batch = new_batch.union(gen_batch_output) with simple_timer("reward", timing_raw): @@ -151,13 +183,15 @@ def fit(self): # the results from reward model and rule-based results. if self.use_rm: # we first compute reward model score - reward_tensor = self.rm_wg.compute_rm_score(new_batch) + reward_tensor = self.rm_wg.compute_rm_score( + new_batch) new_batch = new_batch.union(reward_tensor) # we combine with rule-based rm reward_extra_infos_dict: dict[str, list] try: - reward_result = self.reward_fn(new_batch, return_dict=True) + reward_result = self.reward_fn( + new_batch, return_dict=True) reward_tensor = reward_result["reward_tensor"] reward_extra_infos_dict = reward_result["reward_extra_info"] except Exception as e: @@ -170,19 +204,26 @@ def fit(self): print(f"{list(reward_extra_infos_dict.keys())=}") if reward_extra_infos_dict: new_batch.non_tensor_batch.update( - {k: np.array(v) for k, v in reward_extra_infos_dict.items()} + { + k: np.array(v) + for k, v in reward_extra_infos_dict.items() + } ) # compute rewards. apply_kl_penalty if available if self.config.algorithm.use_kl_in_reward: new_batch, kl_metrics = apply_kl_penalty( - new_batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty + new_batch, + kl_ctrl=self.kl_ctrl_in_reward, + kl_penalty=self.config.algorithm.kl_penalty, ) - metrics.update( - kl_metrics - ) # TODO: This will be cleared if we use multiple genenration batches + # TODO: This will be cleared if we use multiple + # genenration batches + metrics.update(kl_metrics) else: - new_batch.batch["token_level_rewards"] = new_batch.batch["token_level_scores"] + new_batch.batch["token_level_rewards"] = new_batch.batch[ + "token_level_scores" + ] if not self.config.algorithm.filter_groups.enable: batch = new_batch @@ -192,23 +233,30 @@ def fit(self): if metric_name == "seq_final_reward": # Turn to numpy for easier filtering new_batch.non_tensor_batch["seq_final_reward"] = ( - new_batch.batch["token_level_rewards"].sum(dim=-1).numpy() + new_batch.batch["token_level_rewards"] + .sum(dim=-1) + .numpy() ) elif metric_name == "seq_reward": new_batch.non_tensor_batch["seq_reward"] = ( - new_batch.batch["token_level_scores"].sum(dim=-1).numpy() + new_batch.batch["token_level_scores"] + .sum(dim=-1) + .numpy() ) # Collect the sequence reward for each trajectory prompt_uid2metric_vals = defaultdict(list) for uid, metric_val in zip( - new_batch.non_tensor_batch["uid"], new_batch.non_tensor_batch[metric_name], strict=True + new_batch.non_tensor_batch["uid"], + new_batch.non_tensor_batch[metric_name], + strict=True, ): prompt_uid2metric_vals[uid].append(metric_val) prompt_uid2metric_std = {} for prompt_uid, metric_vals in prompt_uid2metric_vals.items(): - prompt_uid2metric_std[prompt_uid] = np.std(metric_vals) + prompt_uid2metric_std[prompt_uid] = np.std( + metric_vals) kept_prompt_uids = [ uid @@ -218,29 +266,43 @@ def fit(self): num_prompt_in_batch += len(kept_prompt_uids) kept_traj_idxs = [] - for idx, traj_from_prompt_uid in enumerate(new_batch.non_tensor_batch["uid"]): + for idx, traj_from_prompt_uid in enumerate( + new_batch.non_tensor_batch["uid"] + ): if traj_from_prompt_uid in kept_prompt_uids: kept_traj_idxs.append(idx) new_batch = new_batch[kept_traj_idxs] - batch = new_batch if batch is None else DataProto.concat([batch, new_batch]) + batch = ( + new_batch + if batch is None + else DataProto.concat([batch, new_batch]) + ) prompt_bsz = self.config.data.train_batch_size if num_prompt_in_batch < prompt_bsz: print(f"{num_prompt_in_batch=} < {prompt_bsz=}") - max_num_gen_batches = self.config.algorithm.filter_groups.max_num_gen_batches - if max_num_gen_batches <= 0 or num_gen_batches < max_num_gen_batches: + max_num_gen_batches = ( + self.config.algorithm.filter_groups.max_num_gen_batches) + if ( + max_num_gen_batches <= 0 + or num_gen_batches < max_num_gen_batches + ): print(f"{num_gen_batches=}. Keep generating...") continue else: raise ValueError( - f"{num_gen_batches=} >= {max_num_gen_batches=}." - + " Generated too many. Please check if your data are too difficult." - + " You could also try set max_num_gen_batches=0 to enable endless trials." - ) + f"{ + num_gen_batches=} >= { + max_num_gen_batches=}." + + " Generated too many. Please check if your data are too difficult." + + " You could also try set max_num_gen_batches=0 to enable endless trials.") else: # Align the batch - traj_bsz = self.config.data.train_batch_size * self.config.actor_rollout_ref.rollout.n + traj_bsz = ( + self.config.data.train_batch_size + * self.config.actor_rollout_ref.rollout.n + ) print( f"Collected {num_prompt_in_batch} / {self.config.data.train_batch_size} prompt. " f"Collecting finished." @@ -253,22 +315,27 @@ def fit(self): # 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 + # Please take care when you implement group based adv + # computation such as GRPO and rloo if self.config.trainer.balance_batch: 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_token_num"] = torch.sum( + batch.batch["attention_mask"], dim=-1 + ).tolist() # recompute old_log_probs with simple_timer("old_log_prob", timing_raw): - old_log_prob = self.actor_rollout_wg.compute_log_prob(batch) + old_log_prob = self.actor_rollout_wg.compute_log_prob( + batch) batch = batch.union(old_log_prob) if self.use_reference_policy: # compute reference log_prob with simple_timer("ref", timing_raw): - ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch) + ref_log_prob = self.ref_policy_wg.compute_ref_log_prob( + batch) batch = batch.union(ref_log_prob) # compute values @@ -279,7 +346,9 @@ def fit(self): with simple_timer("adv", timing_raw): # compute advantages, executed on the driver process - norm_adv_by_std_in_grpo = self.config.algorithm.get("norm_adv_by_std_in_grpo", True) + norm_adv_by_std_in_grpo = self.config.algorithm.get( + "norm_adv_by_std_in_grpo", True + ) batch = compute_advantage( batch, adv_estimator=self.config.algorithm.adv_estimator, @@ -293,22 +362,30 @@ def fit(self): if self.use_critic: with simple_timer("update_critic", timing_raw): critic_output = self.critic_wg.update_critic(batch) - critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"]) + critic_output_metrics = reduce_metrics( + critic_output.meta_info["metrics"] + ) metrics.update(critic_output_metrics) # implement critic warmup if self.config.trainer.critic_warmup <= self.global_steps: # update actor with simple_timer("update_actor", timing_raw): - actor_output = self.actor_rollout_wg.update_actor(batch) - actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) + actor_output = self.actor_rollout_wg.update_actor( + batch) + actor_output_metrics = reduce_metrics( + actor_output.meta_info["metrics"] + ) metrics.update(actor_output_metrics) # validate if ( self.val_reward_fn is not None and self.config.trainer.test_freq > 0 - and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0) + and ( + is_last_step + or self.global_steps % self.config.trainer.test_freq == 0 + ) ): with simple_timer("testing", timing_raw): val_metrics: dict = self._validate() @@ -317,17 +394,26 @@ def fit(self): metrics.update(val_metrics) if self.config.trainer.save_freq > 0 and ( - is_last_step or self.global_steps % self.config.trainer.save_freq == 0 - ): + is_last_step or self.global_steps % + self.config.trainer.save_freq == 0): with simple_timer("save_checkpoint", timing_raw): self._save_checkpoint() # collect metrics - metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic)) - metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) + metrics.update( + compute_data_metrics( + batch=batch, + use_critic=self.use_critic)) + metrics.update( + compute_timing_metrics(batch=batch, timing_raw=timing_raw) + ) # TODO: implement actual tflpo and theoretical tflpo n_gpus = self.resource_pool_manager.get_n_gpus() - metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus)) + metrics.update( + compute_throughout_metrics( + batch=batch, timing_raw=timing_raw, n_gpus=n_gpus + ) + ) timing_raw = defaultdict(float) # clear timing metrics["train/num_gen_batches"] = num_gen_batches diff --git a/Agent0/executor_train/verl/recipe/entropy/main_entropy.py b/Agent0/executor_train/verl/recipe/entropy/main_entropy.py index a8bb0cb..28912a7 100644 --- a/Agent0/executor_train/verl/recipe/entropy/main_entropy.py +++ b/Agent0/executor_train/verl/recipe/entropy/main_entropy.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -22,7 +22,9 @@ from .reward import load_reward_manager -@hydra.main(config_path="config", config_name="entropy_trainer", version_base=None) +@hydra.main(config_path="config", + config_name="entropy_trainer", + version_base=None) def main(config): run_ppo(config) @@ -71,7 +73,9 @@ def run(self, config): from verl.utils.fs import copy_to_local - pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values + pprint( + OmegaConf.to_container(config, resolve=True) + ) # resolve=True will eval symbol values OmegaConf.resolve(config) # download the checkpoint from hdfs @@ -81,14 +85,21 @@ def run(self, config): from verl.utils import hf_processor, hf_tokenizer trust_remote_code = config.data.get("trust_remote_code", False) - tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code) - processor = hf_processor(local_path, use_fast=True) # used for multimodal LLM, could be none + tokenizer = hf_tokenizer( + local_path, trust_remote_code=trust_remote_code) + processor = hf_processor( + local_path, use_fast=True + ) # used for multimodal LLM, could be none # define worker classes if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"}: assert config.critic.strategy in {"fsdp", "fsdp2"} from verl.single_controller.ray import RayWorkerGroup - from verl.workers.fsdp_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker, CriticWorker + from verl.workers.fsdp_workers import ( + ActorRolloutRefWorker, + AsyncActorRolloutRefWorker, + CriticWorker, + ) actor_rollout_cls = ( AsyncActorRolloutRefWorker @@ -100,7 +111,10 @@ def run(self, config): elif config.actor_rollout_ref.actor.strategy == "megatron": assert config.actor_rollout_ref.actor.strategy == config.critic.strategy from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup - from verl.workers.megatron_workers import ActorRolloutRefWorker, CriticWorker + from verl.workers.megatron_workers import ( + ActorRolloutRefWorker, + CriticWorker, + ) actor_rollout_cls = ActorRolloutRefWorker ray_worker_group_cls = NVMegatronRayWorkerGroup @@ -117,7 +131,9 @@ def run(self, config): global_pool_id = "global_pool" resource_pool_spec = { - global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + global_pool_id: [ + config.trainer.n_gpus_per_node] * + config.trainer.nnodes, } mapping = { Role.ActorRollout: global_pool_id, @@ -137,12 +153,17 @@ def run(self, config): from verl.workers.megatron_workers import RewardModelWorker else: raise NotImplementedError - role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker) + role_worker_mapping[Role.RewardModel] = ray.remote( + RewardModelWorker) mapping[Role.RewardModel] = global_pool_id # use reference model - if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss: - role_worker_mapping[Role.RefPolicy] = ray.remote(ActorRolloutRefWorker) + if ( + config.algorithm.use_kl_in_reward + or config.actor_rollout_ref.actor.use_kl_loss + ): + role_worker_mapping[Role.RefPolicy] = ray.remote( + ActorRolloutRefWorker) mapping[Role.RefPolicy] = global_pool_id reward_kwargs = { @@ -151,15 +172,26 @@ def run(self, config): } cfg_reward_kwargs = config.reward_model.get("reward_kwargs", {}) reward_fn = load_reward_manager( - config, tokenizer, num_examine=0, **OmegaConf.merge(OmegaConf.create(reward_kwargs), cfg_reward_kwargs) + config, + tokenizer, + num_examine=0, + **OmegaConf.merge(OmegaConf.create(reward_kwargs), cfg_reward_kwargs), + ) + val_reward_fn = load_reward_manager( + config, tokenizer, num_examine=1, **reward_kwargs + ) + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, mapping=mapping ) - val_reward_fn = load_reward_manager(config, tokenizer, num_examine=1, **reward_kwargs) - resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) from verl.utils.dataset.rl_dataset import collate_fn - train_dataset = create_rl_dataset(config.data.train_files, config.data, tokenizer, processor) - val_dataset = create_rl_dataset(config.data.val_files, config.data, tokenizer, processor) + train_dataset = create_rl_dataset( + config.data.train_files, config.data, tokenizer, processor + ) + val_dataset = create_rl_dataset( + config.data.val_files, config.data, tokenizer, processor + ) train_sampler = create_rl_sampler(config.data, train_dataset) trainer = RayEntropyTrainer( config=config, @@ -194,15 +226,20 @@ def create_rl_dataset(data_paths, data_config, tokenizer, processor): from verl.utils.dataset.rl_dataset import RLHFDataset - if "custom_cls" in data_config and data_config.custom_cls.get("path", None) is not None: + if ( + "custom_cls" in data_config + and data_config.custom_cls.get("path", None) is not None + ): from verl.utils.import_utils import load_extern_type - dataset_cls = load_extern_type(data_config.custom_cls.path, data_config.custom_cls.name) + dataset_cls = load_extern_type( + data_config.custom_cls.path, data_config.custom_cls.name + ) if not issubclass(dataset_cls, Dataset): raise TypeError( - f"The custom dataset class '{data_config.custom_cls.name}' from '{data_config.custom_cls.path}' " - f"must inherit from torch.utils.data.Dataset" - ) + f"The custom dataset class '{ + data_config.custom_cls.name}' from '{ + data_config.custom_cls.path}' " f"must inherit from torch.utils.data.Dataset") else: dataset_cls = RLHFDataset print(f"Using dataset class: {dataset_cls.__name__}") @@ -234,7 +271,9 @@ def create_rl_sampler(data_config, dataset): if data_config.shuffle: train_dataloader_generator = torch.Generator() train_dataloader_generator.manual_seed(data_config.get("seed", 1)) - sampler = RandomSampler(data_source=dataset, generator=train_dataloader_generator) + sampler = RandomSampler( + data_source=dataset, generator=train_dataloader_generator + ) else: sampler = SequentialSampler(data_source=dataset) diff --git a/Agent0/executor_train/verl/recipe/entropy/reward.py b/Agent0/executor_train/verl/recipe/entropy/reward.py index 36b8b65..2d97d86 100644 --- a/Agent0/executor_train/verl/recipe/entropy/reward.py +++ b/Agent0/executor_train/verl/recipe/entropy/reward.py @@ -1,4 +1,4 @@ -# Copyright 2025 Individual Contributor: Thibaut Barroyer +# Copyright 2025-2026 Individual Contributor: Thibaut Barroyer # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -59,9 +59,13 @@ def load_reward_manager(config, tokenizer, num_examine, **reward_kwargs): if sandbox_url: sandbox_manager = multiprocessing.Manager() # Create a semaphore to control concurrent access to the sandbox - _concurrent_semaphore = sandbox_manager.Semaphore(sandbox_config.get("max_concurrent", 64)) + _concurrent_semaphore = sandbox_manager.Semaphore( + sandbox_config.get("max_concurrent", 64) + ) final_compute_score = partial( - _default_compute_score, sandbox_fusion_url=sandbox_url, concurrent_semaphore=_concurrent_semaphore + _default_compute_score, + sandbox_fusion_url=sandbox_url, + concurrent_semaphore=_concurrent_semaphore, ) else: final_compute_score = _default_compute_score @@ -82,5 +86,7 @@ def compute_reward_async(data: DataProto, config, tokenizer): Load the reward manager and compute the reward for a batch of data. This is meant to be run in a separate Ray worker. """ - reward_fn = load_reward_manager(config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {})) + reward_fn = load_reward_manager( + config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {}) + ) return compute_reward(data, reward_fn) diff --git a/Agent0/executor_train/verl/recipe/entropy/reward_score/__init__.py b/Agent0/executor_train/verl/recipe/entropy/reward_score/__init__.py index 7224bf3..8cd9c32 100644 --- a/Agent0/executor_train/verl/recipe/entropy/reward_score/__init__.py +++ b/Agent0/executor_train/verl/recipe/entropy/reward_score/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,7 +19,12 @@ def _default_compute_score( - data_source, solution_str, ground_truth, extra_info=None, sandbox_fusion_url=None, concurrent_semaphore=None + data_source, + solution_str, + ground_truth, + extra_info=None, + sandbox_fusion_url=None, + concurrent_semaphore=None, ): try: res = entropy_math.compute_score(solution_str, str(ground_truth)) diff --git a/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/__init__.py b/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/__init__.py index 1b2ba64..059b8c3 100644 --- a/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/__init__.py +++ b/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/__init__.py @@ -70,7 +70,7 @@ def mathd_normalize_answer(answer: Optional[str]) -> Optional[str]: answer = answer.strip() try: # Remove enclosing `\text{}`. - m = re.search("^\\\\text\{(?P.+?)\}$", answer) + m = re.search("^\\\\text\\{(?P.+?)\\}$", answer) if m is not None: answer = m.group("text").strip() return _strip_string(answer) @@ -306,7 +306,11 @@ def _fix_sqrt(string): # replace tfrac and dfrac with frac string = string.replace("tfrac", "frac") string = string.replace("dfrac", "frac") - string = string.replace("\\neq", "\\ne").replace("\\leq", "\\le").replace("\\geq", "\\ge") + string = ( + string.replace("\\neq", "\\ne") + .replace("\\leq", "\\le") + .replace("\\geq", "\\ge") + ) # print(string) # remove \left and \right @@ -324,8 +328,14 @@ def _fix_sqrt(string): for _ in range(2): for unit_text in unit_texts: # use regex, the prefix should be either the start of the string or a non-alphanumeric character - # the suffix should be either the end of the string or a non-alphanumeric character - _string = re.sub(r"(^|\W)" + unit_text + r"($|\W)", r"\1\2", string) + # the suffix should be either the end of the string or a + # non-alphanumeric character + _string = re.sub( + r"(^|\W)" + + unit_text + + r"($|\W)", + r"\1\2", + string) if _string != "": string = _string @@ -341,7 +351,7 @@ def _fix_sqrt(string): # remove percentage string = string.replace("\\%", "") - string = string.replace("\%", "") + string = string.replace("\\%", "") # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string string = string.replace(" .", " 0.") @@ -371,7 +381,8 @@ def _fix_sqrt(string): 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 + # 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 @@ -483,7 +494,16 @@ def suffixArray(s): line = ranks(s) n, k, ans, sa = len(s), 1, line, [0] * len(s) while k < n - 1: - line = ranks(list(zip_longest(line, islice(line, k, None), fillvalue=-1))) + line = ranks( + list( + zip_longest( + line, + islice( + line, + k, + None), + fillvalue=- + 1))) ans, k = line, k << 1 for i, k in enumerate(ans): sa[k] = i @@ -616,7 +636,8 @@ def _is_latex_equal(str1, str2): raise ValueError except Exception: # noqa try: - norm1, norm2 = normalize_final_answer(str1), normalize_final_answer(str2) + norm1, norm2 = normalize_final_answer( + str1), normalize_final_answer(str2) sym1, val1 = latex_eval(norm1) sym2, val2 = latex_eval(norm2) if sym1 == sym2 or val1 == val2: @@ -686,7 +707,9 @@ def is_value_equal(given_answer: str, ground_truth: str) -> bool: str_equal = ground_truth_normalized_mathd == given_answer_normalized_mathd try: - number_equal = float(ground_truth_normalized_mathd) == float(given_answer_normalized_mathd) + number_equal = float(ground_truth_normalized_mathd) == float( + given_answer_normalized_mathd + ) return str_equal or number_equal except Exception: return str_equal @@ -694,7 +717,7 @@ def is_value_equal(given_answer: str, ground_truth: str) -> bool: # sympy might hang -- we don't care about trying to be lenient in these cases BAD_SUBSTRINGS = ["^{", "^("] -BAD_REGEXES = ["\^[0-9]+\^", "\^[0-9][0-9]+"] +BAD_REGEXES = ["\\^[0-9]+\\^", "\\^[0-9][0-9]+"] TUPLE_CHARS = "()[]" @@ -703,7 +726,10 @@ def _sympy_parse(expr: str): py_expr = expr.replace("^", "**") return sympy_parser.parse_expr( py_expr, - transformations=(sympy_parser.standard_transformations + (sympy_parser.implicit_multiplication_application,)), + transformations=( + sympy_parser.standard_transformations + + (sympy_parser.implicit_multiplication_application,) + ), ) @@ -765,13 +791,13 @@ def _inject_implicit_mixed_number(step: str): e.g. 7 3/4 => 7+3/4 """ p1 = re.compile("([0-9]) +([0-9])") - step = p1.sub("\\1+\\2", step) ## implicit mults + step = p1.sub("\\1+\\2", step) # implicit mults return step def _strip_properly_formatted_commas(expr: str): # We want to be careful because we don't want to strip tuple commas - p1 = re.compile("(\d)(,)(\d\d\d)($|\D)") + p1 = re.compile("(\\d)(,)(\\d\\d\\d)($|\\D)") while True: next_expr = p1.sub("\\1\\3\\4", expr) if next_expr == expr: @@ -786,7 +812,7 @@ def _normalize(expr: str) -> str: return None # Remove enclosing `\text{}`. - m = re.search("^\\\\text\{(?P.+?)\}$", expr) + m = re.search("^\\\\text\\{(?P.+?)\\}$", expr) if m is not None: expr = m.group("text") @@ -819,8 +845,8 @@ def _normalize(expr: str) -> str: "inch", "yard", ]: - expr = re.sub(f"{unit}(es)?(s)? *(\^[0-9]+)?", "", expr) - expr = re.sub("\^ *\\\\circ", "", expr) + expr = re.sub(f"{unit}(es)?(s)? *(\\^[0-9]+)?", "", expr) + expr = re.sub("\\^ *\\\\circ", "", expr) if len(expr) > 0 and expr[0] == "{" and expr[-1] == "}": expr = expr[1:-1] @@ -861,7 +887,8 @@ def count_unknown_letters_in_expr(expr: str): def should_allow_eval(expr: str): - # we don't want to try parsing unknown text or functions of more than two variables + # we don't want to try parsing unknown text or functions of more than two + # variables if count_unknown_letters_in_expr(expr) > 2: return False @@ -932,7 +959,7 @@ def last_boxed_only_string(string): if right_brace_idx is None: retval = None else: - retval = string[idx : right_brace_idx + 1] + retval = string[idx: right_brace_idx + 1] return retval @@ -942,7 +969,7 @@ def remove_boxed(s): try: assert s[: len(left)] == left assert s[-1] == "}" - return s[len(left) : -1] + return s[len(left): -1] except Exception: return None @@ -971,13 +998,16 @@ def grade_answer_sympy(given_answer: str, ground_truth: str) -> bool: given_elems = split_tuple(given_normalized) if len(ground_truth_elems) > 1 and ( - ground_truth_normalized[0] != given_normalized[0] or ground_truth_normalized[-1] != given_normalized[-1] + ground_truth_normalized[0] != given_normalized[0] + or ground_truth_normalized[-1] != given_normalized[-1] ): is_correct = False elif len(ground_truth_elems) != len(given_elems): is_correct = False else: - for ground_truth_elem, given_elem in zip(ground_truth_elems, given_elems, strict=True): + for ground_truth_elem, given_elem in zip( + ground_truth_elems, given_elems, strict=True + ): if _is_frac(ground_truth_elem) and _is_frac(given_elem): # if fractions aren't reduced, then shouldn't be marked as correct # so, we don't want to allow sympy.simplify in this case @@ -987,7 +1017,8 @@ def grade_answer_sympy(given_answer: str, ground_truth: str) -> bool: # (no sympy.simplify) is_correct = False else: - is_correct = are_equal_under_sympy(ground_truth_elem, given_elem) + is_correct = are_equal_under_sympy( + ground_truth_elem, given_elem) if not is_correct: break @@ -1013,7 +1044,9 @@ def extract_answer(passage: str) -> str: def grade(model_answer: str, gt_answer: str, fast: bool = True): if "\\boxed" in gt_answer: gt_answer = extract_answer(gt_answer) - correct = grade_answer_mathd(model_answer, gt_answer) or grade_answer_sympy(model_answer, gt_answer) + correct = grade_answer_mathd( + model_answer, gt_answer) or grade_answer_sympy( + model_answer, gt_answer) if not fast: # This mode further uses math_verify to recall originally false positives. # Will be a bit slower, and sensitive to bad inputs. diff --git a/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/grader.py b/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/grader.py index 02507e3..de2da5c 100644 --- a/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/grader.py +++ b/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/grader.py @@ -125,7 +125,8 @@ def normalize(answer, pi) -> str: # checking if answer is % or \\% and removing % if isinstance(answer, str) and ( - bool(re.match(r"^\d+(\.\d+)?%$", answer)) or bool(re.match(r"^\d+(\.\d+)?\\%$", answer)) + bool(re.match(r"^\d+(\.\d+)?%$", answer)) + or bool(re.match(r"^\d+(\.\d+)?\\%$", answer)) ): return answer.replace("\\%", "").replace("%", "") @@ -148,21 +149,24 @@ def handle_base(x) -> str: def handle_pi(string, pi): - if isinstance(string, str) and "\pi" in string: + if isinstance(string, str) and "\\pi" in string: # Find the first occurrence of "\pi" - idx = string.find("\pi") + idx = string.find("\\pi") - # Iterate over the string and find all occurrences of "\pi" with a valid previous character + # Iterate over the string and find all occurrences of "\pi" with a + # valid previous character while idx != -1: if idx > 0 and string[idx - 1].isdigit(): - # Replace "\pi" with "*math.pi" if the previous character is a digit - string = string[:idx] + f"*{pi}" + string[idx + 3 :] + # Replace "\pi" with "*math.pi" if the previous character is a + # digit + string = string[:idx] + f"*{pi}" + string[idx + 3:] else: - # Replace "\pi" with "1*math.pi" if the previous character is not a digit - string = string[:idx] + f"1*{pi}" + string[idx + 3 :] + # Replace "\pi" with "1*math.pi" if the previous character is + # not a digit + string = string[:idx] + f"1*{pi}" + string[idx + 3:] # Find the next occurrence of "\pi" - idx = string.find("\pi", idx + 1) + idx = string.find("\\pi", idx + 1) # Evaluate the expression using eval() function with contextlib.suppress(Exception): @@ -188,7 +192,9 @@ def math_equal( prediction = normalize(prediction, pi) reference = normalize(reference, pi) - if isinstance(prediction, str) and len(prediction) > 1000: # handling weird corner-cases + if ( + isinstance(prediction, str) and len(prediction) > 1000 + ): # handling weird corner-cases prediction = prediction[:1000] # 0. string comparison @@ -203,7 +209,11 @@ def math_equal( prediction = is_digit(prediction)[1] reference = is_digit(reference)[1] # number questions - gt_result = [reference / 100, reference, reference * 100] if include_percentage else [reference] + gt_result = ( + [reference / 100, reference, reference * 100] + if include_percentage + else [reference] + ) for item in gt_result: try: if isclose(item, prediction, rel_tol=tolerance): @@ -221,12 +231,18 @@ def math_equal( reference = str(reference).strip() prediction = str(prediction).strip() - ## deal with [], (), {} + # deal with [], (), {} prediction = format_intervals(prediction) pred_str, ref_str = prediction, reference - if (prediction.startswith("[") and prediction.endswith("]") and not reference.startswith("(")) or ( - prediction.startswith("(") and prediction.endswith(")") and not reference.startswith("[") + if ( + prediction.startswith("[") + and prediction.endswith("]") + and not reference.startswith("(") + ) or ( + prediction.startswith("(") + and prediction.endswith(")") + and not reference.startswith("[") ): pred_str = pred_str.strip("[]()") ref_str = ref_str.strip("[]()") @@ -236,7 +252,7 @@ def math_equal( if pred_str == ref_str: return True - ## [a, b] vs. [c, d], return a==c and b==d + # [a, b] vs. [c, d], return a==c and b==d if ( prediction and reference @@ -260,18 +276,15 @@ def math_equal( ref_parts = [item.strip() for item in reference.split(",")] if len(pred_parts) == len(ref_parts): - return bool( - all( - [ - math_equal(pred_parts[i], ref_parts[i], include_percentage, tolerance) - for i in range(len(pred_parts)) - ] - ) - ) + return bool(all([math_equal(pred_parts[i], + ref_parts[i], + include_percentage, + tolerance) for i in range(len(pred_parts))])) # if we have point == tuple of values - if prediction.startswith("Point") and reference[0] == "(" and reference[-1] == ")": - pred_parts = prediction[prediction.find("(") + 1 : -1].split(",") + if prediction.startswith( + "Point") and reference[0] == "(" and reference[-1] == ")": + pred_parts = prediction[prediction.find("(") + 1: -1].split(",") ref_parts = reference[1:-1].split(",") if len(pred_parts) == len(ref_parts) and all( [ @@ -295,7 +308,11 @@ def math_equal( return True except Exception: pass - elif "\begin{pmatrix}" in reference and prediction.startswith("[") and prediction.endswith("]"): + elif ( + "\begin{pmatrix}" in reference + and prediction.startswith("[") + and prediction.endswith("]") + ): if isinstance(eval(prediction), list): try: pred_matrix = eval(prediction) @@ -307,7 +324,8 @@ def math_equal( .rstrip("\end{pmatrix}") ) # noqa: B005 ref_matrix_items = ref_matrix_items.split("\\") - ref_matrix_items = [row.split("&") if "&" in row else row for row in ref_matrix_items] + ref_matrix_items = [ + row.split("&") if "&" in row else row for row in ref_matrix_items] if len(pred_matrix) == len(ref_matrix_items) and all( [ math_equal(pred, ref, include_percentage, tolerance) diff --git a/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/math_normalize.py b/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/math_normalize.py index 74d94cc..52a5ec7 100644 --- a/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/math_normalize.py +++ b/Agent0/executor_train/verl/recipe/entropy/reward_score/entropy_math/math_normalize.py @@ -47,7 +47,7 @@ def normalize_answer(answer: Optional[str]) -> Optional[str]: answer = answer.strip() try: # Remove enclosing `\text{}`. - m = re.search("^\\\\text\{(?P.+?)\}$", answer) + m = re.search("^\\\\text\\{(?P.+?)\\}$", answer) if m is not None: answer = m.group("text").strip() return _strip_string(answer) @@ -157,7 +157,7 @@ def _strip_string(string): # remove percentage string = string.replace("\\%", "") - string = string.replace("\%", "") + string = string.replace("\\%", "") # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string string = string.replace(" .", " 0.") @@ -186,7 +186,8 @@ def _strip_string(string): 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 + # 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 diff --git a/Agent0/executor_train/verl/recipe/genrm_remote/reward_function.py b/Agent0/executor_train/verl/recipe/genrm_remote/reward_function.py index b2d3fbc..8cbe81f 100644 --- a/Agent0/executor_train/verl/recipe/genrm_remote/reward_function.py +++ b/Agent0/executor_train/verl/recipe/genrm_remote/reward_function.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -43,14 +43,16 @@ def get_response(problem, solution_str, ground_truth): - prompt = GENRM_PROMPT_TEMPLATE.format(problem=problem, solution=solution_str) + prompt = GENRM_PROMPT_TEMPLATE.format( + problem=problem, solution=solution_str) messages = [{"role": "user", "content": prompt}] for attempt in range(MAX_RETRIES): try: headers = {"Content-Type": "application/json"} chat_url = f"{BASE_URL}/v1/chat/completions" data = {"model": MODEL_NAME, "messages": messages} - output = requests.post(chat_url, headers=headers, json=data, timeout=30) + output = requests.post( + chat_url, headers=headers, json=data, timeout=30) response = output.json()["choices"][0]["message"]["content"] return response except Exception as e: @@ -81,7 +83,9 @@ def compute_score(data_source, solution_str, ground_truth, extra_info): split = extra_info["split"] from verl.utils.reward_score import default_compute_score - func_rm_score = default_compute_score(data_source, solution_str, ground_truth, extra_info) + func_rm_score = default_compute_score( + data_source, solution_str, ground_truth, extra_info + ) if split == "test": return func_rm_score @@ -96,13 +100,21 @@ def compute_score(data_source, solution_str, ground_truth, extra_info): return reward_score -def compute_score_batch(data_sources, solution_strs, ground_truths, extra_infos): +def compute_score_batch( + data_sources, + solution_strs, + ground_truths, + extra_infos): with ThreadPoolExecutor(max_workers=MAX_WORKERS) as executor: futures = [] for data_source, solution_str, ground_truth, extra_info in zip( - data_sources, solution_strs, ground_truths, extra_infos, strict=True - ): - future = executor.submit(compute_score, data_source, solution_str, ground_truth, extra_info) + data_sources, solution_strs, ground_truths, extra_infos, strict=True): + future = executor.submit( + compute_score, + data_source, + solution_str, + ground_truth, + extra_info) futures.append(future) results = [future.result() for future in futures] diff --git a/Agent0/executor_train/verl/recipe/minicpmo/rl_dataset.py b/Agent0/executor_train/verl/recipe/minicpmo/rl_dataset.py index 5ce15fb..db535be 100644 --- a/Agent0/executor_train/verl/recipe/minicpmo/rl_dataset.py +++ b/Agent0/executor_train/verl/recipe/minicpmo/rl_dataset.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -37,20 +37,28 @@ def build_transform(): - IMAGENET_INCEPTION_MEAN = (0.5, 0.5, 0.5) # timm.data.IMAGENET_INCEPTION_MEAN - IMAGENET_INCEPTION_STD = (0.5, 0.5, 0.5) # timm.data.IMAGENET_INCEPTION_STD + # timm.data.IMAGENET_INCEPTION_MEAN + IMAGENET_INCEPTION_MEAN = (0.5, 0.5, 0.5) + # timm.data.IMAGENET_INCEPTION_STD + IMAGENET_INCEPTION_STD = (0.5, 0.5, 0.5) return transforms.Compose( [ transforms.ToTensor(), - transforms.Normalize(mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD), + transforms.Normalize( + mean=IMAGENET_INCEPTION_MEAN, std=IMAGENET_INCEPTION_STD + ), ] ) def build_image_bound(input_ids, tokenizer, new_schema=True, logger=None): if new_schema: - start_cond = (input_ids == tokenizer.im_start_id) | (input_ids == tokenizer.slice_start_id) - end_cond = (input_ids == tokenizer.im_end_id) | (input_ids == tokenizer.slice_end_id) + start_cond = (input_ids == tokenizer.im_start_id) | ( + input_ids == tokenizer.slice_start_id + ) + end_cond = (input_ids == tokenizer.im_end_id) | ( + input_ids == tokenizer.slice_end_id + ) else: start_cond = input_ids == tokenizer.im_start_id end_cond = input_ids == tokenizer.im_end_id @@ -61,7 +69,9 @@ def build_image_bound(input_ids, tokenizer, new_schema=True, logger=None): logger.error("image start token != image end tokens") raise Exception("image start token != image end tokens") if len(image_start_tokens) > 0: - image_bound = torch.hstack([image_start_tokens.unsqueeze(-1), image_end_tokens.unsqueeze(-1)]) + image_bound = torch.hstack( + [image_start_tokens.unsqueeze(-1), image_end_tokens.unsqueeze(-1)] + ) else: image_bound = [] return image_bound @@ -92,7 +102,11 @@ def preprocess( assert "patch_size" in slice_config assert "max_slice_nums" in slice_config assert "scale_resolution" in slice_config - default_image_placeholder = tokenizer.im_start + tokenizer.unk_token * query_nums + tokenizer.im_end + default_image_placeholder = ( + tokenizer.im_start + + tokenizer.unk_token * + query_nums + + tokenizer.im_end) new_schema = False use_image_id = False if llm_type == "qwen": @@ -117,15 +131,21 @@ def preprocess( images.append(patches[i][j]) if use_image_id: image_placeholder = ( - f"{tokenizer.im_id_start}{image_id_cnt}{tokenizer.im_id_end}" + image_placeholder + f"{tokenizer.im_id_start}{image_id_cnt}{tokenizer.im_id_end}" + + image_placeholder ) image_id_cnt += 1 - image_placeholder += get_grid_placeholder(tokenizer, best_grid, query_nums, new_schema=new_schema) + image_placeholder += get_grid_placeholder( + tokenizer, best_grid, query_nums, new_schema=new_schema + ) image_placeholder_dict[img_name] = image_placeholder else: images.append(image) if use_image_id: - image_placeholder = f"{tokenizer.im_id_start}{image_id_cnt}{tokenizer.im_id_end}" + image_placeholder + image_placeholder = ( + f"{tokenizer.im_id_start}{image_id_cnt}{tokenizer.im_id_end}" + + image_placeholder + ) image_id_cnt += 1 else: image_placeholder = default_image_placeholder @@ -135,9 +155,13 @@ def preprocess( if len(images_dict) == 1 and "" in images_dict: if "" in conversations[0]["content"]: - conversations[0]["content"] = conversations[0]["content"].replace("", image_placeholder) + conversations[0]["content"] = conversations[0]["content"].replace( + "", image_placeholder + ) else: - conversations[0]["content"] = image_placeholder + "\n" + conversations[0]["content"] + conversations[0]["content"] = ( + image_placeholder + "\n" + conversations[0]["content"] + ) else: pattern = r"" new_conversations = [] @@ -157,7 +181,9 @@ def preprocess( conversations = new_conversations # TODO change role in conversation for different llm - prompt_with_chat_template = tokenizer.apply_chat_template(conversations, add_generation_prompt=True, tokenize=False) + prompt_with_chat_template = tokenizer.apply_chat_template( + conversations, add_generation_prompt=True, tokenize=False + ) input_ids, attention_mask = verl_F.tokenize_and_postprocess_data( prompt=prompt_with_chat_template, @@ -168,7 +194,8 @@ def preprocess( truncation=truncation, ) position_ids = compute_position_id_with_mask(attention_mask) - image_bound = build_image_bound(input_ids[0], tokenizer, new_schema, logger) + image_bound = build_image_bound( + input_ids[0], tokenizer, new_schema, logger) input_dict = { "input_ids": input_ids[0], @@ -198,11 +225,17 @@ def preprocess( return input_dict -def slice_image(image, max_slice_nums=9, scale_resolution=448, patch_size=14, never_split=False): +def slice_image( + image, + max_slice_nums=9, + scale_resolution=448, + patch_size=14, + never_split=False): original_size = image.size original_width, original_height = original_size log_ratio = math.log(original_width / original_height) - ratio = original_width * original_height / (scale_resolution * scale_resolution) + ratio = original_width * original_height / \ + (scale_resolution * scale_resolution) multiple = min(math.ceil(ratio), max_slice_nums) source_image = None @@ -211,7 +244,9 @@ def slice_image(image, max_slice_nums=9, scale_resolution=448, patch_size=14, ne if multiple <= 1 or never_split: # dont need to slice, upsample - best_size = find_best_resize(original_size, scale_resolution, patch_size, allow_upscale=True) + best_size = find_best_resize( + original_size, scale_resolution, patch_size, allow_upscale=True + ) source_image = image.resize(best_size, Image.Resampling.BICUBIC) else: candidate_split_grids_nums = [] @@ -221,7 +256,8 @@ def slice_image(image, max_slice_nums=9, scale_resolution=448, patch_size=14, ne candidate_split_grids_nums.append(i) # source image, down-sampling and ensure divided by patch_size - best_resize = find_best_resize(original_size, scale_resolution, patch_size) + best_resize = find_best_resize( + original_size, scale_resolution, patch_size) source_image = image.copy().resize(best_resize, Image.Resampling.BICUBIC) candidate_grids = [] @@ -241,7 +277,12 @@ def slice_image(image, max_slice_nums=9, scale_resolution=448, patch_size=14, ne best_grid = grid min_error = error - refine_size = get_refine_size(original_size, best_grid, scale_resolution, patch_size, allow_upscale=True) + refine_size = get_refine_size( + original_size, + best_grid, + scale_resolution, + patch_size, + allow_upscale=True) refine_image = image.resize(refine_size, Image.Resampling.BICUBIC) patches = split_to_patches(refine_image, best_grid) @@ -253,7 +294,11 @@ def ensure_divide(length, patch_size): return max(round(length / patch_size) * patch_size, patch_size) -def find_best_resize(original_size, scale_resolution, patch_size, allow_upscale=False): +def find_best_resize( + original_size, + scale_resolution, + patch_size, + allow_upscale=False): width, height = original_size if (width * height > scale_resolution * scale_resolution) or allow_upscale: r = width / height @@ -264,7 +309,9 @@ def find_best_resize(original_size, scale_resolution, patch_size, allow_upscale= return (best_width, best_height) -def get_refine_size(original_size, grid, scale_resolution, patch_size, allow_upscale=False): +def get_refine_size( + original_size, grid, scale_resolution, patch_size, allow_upscale=False +): width, height = original_size grid_x, grid_y = grid @@ -305,9 +352,17 @@ def split_to_patches(image, grid): def get_grid_placeholder(tokenizer, grid, query_num, new_schema=False): if new_schema: - image_placeholder = tokenizer.slice_start + tokenizer.unk_token * query_num + tokenizer.slice_end + image_placeholder = ( + tokenizer.slice_start + + tokenizer.unk_token * query_num + + tokenizer.slice_end + ) else: - image_placeholder = tokenizer.im_start + tokenizer.unk_token * query_num + tokenizer.im_end + image_placeholder = ( + tokenizer.im_start + + tokenizer.unk_token * + query_num + + tokenizer.im_end) cols = grid[0] rows = grid[1] @@ -320,7 +375,9 @@ def get_grid_placeholder(tokenizer, grid, query_num, new_schema=False): if new_schema: slice_placeholder = "\n".join(slices) else: - slice_placeholder = tokenizer.slice_start + "\n".join(slices) + tokenizer.slice_end + slice_placeholder = ( + tokenizer.slice_start + "\n".join(slices) + tokenizer.slice_end + ) return slice_placeholder @@ -330,10 +387,14 @@ def reshape_by_patch(image_tensor, patch_size): :param patch_size: :return: [3, patch_size, HW/patch_size] """ - patches = torch.nn.functional.unfold(image_tensor, (patch_size, patch_size), stride=(patch_size, patch_size)) + patches = torch.nn.functional.unfold( + image_tensor, (patch_size, patch_size), stride=(patch_size, patch_size) + ) patches = patches.reshape(image_tensor.size(0), patch_size, patch_size, -1) - patches = patches.permute(0, 1, 3, 2).reshape(image_tensor.size(0), patch_size, -1) + patches = patches.permute( + 0, 1, 3, 2).reshape( + image_tensor.size(0), patch_size, -1) return patches @@ -344,7 +405,12 @@ def init_minicpmo_config(processor, config): "patch_size": config.get("patch_size", 14), "query_nums": config.get("query_nums", 64), "slice_config": config.get( - "slice_config", {"max_slice_nums": 9, "patch_size": config.get("patch_size", 14), "scale_resolution": 448} + "slice_config", + { + "max_slice_nums": 9, + "patch_size": config.get("patch_size", 14), + "scale_resolution": 448, + }, ), "llm_type": config.get("llm_type", "qwen"), "batch_vision": config.get("batch_vision", True), @@ -353,7 +419,14 @@ def init_minicpmo_config(processor, config): def process_minicpmo_data( - row_dict, messages, tokenizer, minicpmo_config, image_key, max_prompt_length, truncation, logger + row_dict, + messages, + tokenizer, + minicpmo_config, + image_key, + max_prompt_length, + truncation, + logger, ): """Process data for MiniCPM-o model""" if len(row_dict[image_key]) == 1: @@ -379,7 +452,9 @@ def process_minicpmo_data( logger=logger, ) - raw_prompt = tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + raw_prompt = tokenizer.apply_chat_template( + messages, add_generation_prompt=True, tokenize=False + ) raw_prompt = raw_prompt.replace("", "(./)") return model_inputs, multi_modal_data, raw_prompt @@ -418,7 +493,9 @@ def __init__( self.processor = processor self.config = config - self.cache_dir = os.path.expanduser(config.get("cache_dir", "~/.cache/verl/rlhf")) + self.cache_dir = os.path.expanduser( + config.get("cache_dir", "~/.cache/verl/rlhf") + ) self.prompt_key = config.get("prompt_key", "prompt") self.image_key = config.get("image_key", "images") self.video_key = config.get("video_key", "videos") @@ -426,9 +503,12 @@ def __init__( self.return_raw_chat = config.get("return_raw_chat", False) self.return_full_prompt = config.get("return_full_prompt", False) self.truncation = config.get("truncation", "error") - self.filter_overlong_prompts = config.get("filter_overlong_prompts", True) + self.filter_overlong_prompts = config.get( + "filter_overlong_prompts", True) - self.num_workers = config.get("filter_overlong_prompts_workers", max(1, os.cpu_count() // 4)) + self.num_workers = config.get( + "filter_overlong_prompts_workers", max(1, os.cpu_count() // 4) + ) self.num_workers = min(self.num_workers, os.cpu_count()) self.use_shm = config.get("use_shm", False) self.chat_template_func = config.get("chat_template_func", None) @@ -442,17 +522,23 @@ def __init__( def _download(self, use_origin_parquet=False): from verl.utils.fs import copy_to_local - data_files = self.data_files if not use_origin_parquet else self.original_data_files + data_files = ( + self.data_files if not use_origin_parquet else self.original_data_files) for i, parquet_file in enumerate(data_files): - self.data_files[i] = copy_to_local(src=parquet_file, cache_dir=self.cache_dir, use_shm=self.use_shm) + self.data_files[i] = copy_to_local( + src=parquet_file, + cache_dir=self.cache_dir, + use_shm=self.use_shm) def _read_files_and_tokenize(self): dataframes = [] for parquet_file in self.data_files: # read parquet files and cache - dataframe = datasets.load_dataset("parquet", data_files=parquet_file)["train"] + dataframe = datasets.load_dataset( + "parquet", data_files=parquet_file)["train"] dataframes.append(dataframe) - self.dataframe: datasets.Dataset = datasets.concatenate_datasets(dataframes) + self.dataframe: datasets.Dataset = datasets.concatenate_datasets( + dataframes) print(f"dataset len: {len(self.dataframe)}") @@ -460,10 +546,14 @@ def resume_dataset_state(self): self.serialize_dataset = not hasattr(self, "original_data_files") # resume dataframe if not it's serialized in data.pt if not self.serialize_dataset: - self._download(use_origin_parquet=True) # download and resume from original parquet files + self._download( + use_origin_parquet=True + ) # download and resume from original parquet files self._read_files_and_tokenize() else: - print(r"old dataloader ckpt file is used, please train from scratch for better ckpt performance") + print( + r"old dataloader ckpt file is used, please train from scratch for better ckpt performance" + ) def __len__(self): return len(self.dataframe) @@ -494,12 +584,17 @@ def __getitem__(self, item): attention_mask = model_inputs.pop("attention_mask") position_ids = model_inputs.pop("position_ids") - # There's a trap here, multi_modal_inputs has to be a dict, not BatchFeature + # There's a trap here, multi_modal_inputs has to be a dict, not + # BatchFeature row_dict["multi_modal_data"] = multi_modal_data row_dict["multi_modal_inputs"] = dict(model_inputs) else: - raw_prompt = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) - model_inputs = self.tokenizer(raw_prompt, return_tensors="pt", add_special_tokens=False) + raw_prompt = self.tokenizer.apply_chat_template( + messages, add_generation_prompt=True, tokenize=False + ) + model_inputs = self.tokenizer( + raw_prompt, return_tensors="pt", add_special_tokens=False + ) input_ids = model_inputs.pop("input_ids") attention_mask = model_inputs.pop("attention_mask") position_ids = compute_position_id_with_mask(attention_mask) @@ -508,18 +603,24 @@ def __getitem__(self, item): row_dict["attention_mask"] = attention_mask row_dict["position_ids"] = position_ids - raw_prompt_ids = self.tokenizer.encode(raw_prompt, add_special_tokens=False) + raw_prompt_ids = self.tokenizer.encode( + raw_prompt, add_special_tokens=False) if len(raw_prompt_ids) > self.max_prompt_length: if self.truncation == "left": - raw_prompt_ids = raw_prompt_ids[-self.max_prompt_length :] + raw_prompt_ids = raw_prompt_ids[-self.max_prompt_length:] elif self.truncation == "right": raw_prompt_ids = raw_prompt_ids[: self.max_prompt_length] elif self.truncation == "middle": left_half = self.max_prompt_length // 2 right_half = self.max_prompt_length - left_half - raw_prompt_ids = raw_prompt_ids[:left_half] + raw_prompt_ids[-right_half:] + raw_prompt_ids = ( + raw_prompt_ids[:left_half] + raw_prompt_ids[-right_half:] + ) elif self.truncation == "error": - raise RuntimeError(f"Prompt length {len(raw_prompt_ids)} is longer than {self.max_prompt_length}.") + raise RuntimeError( + f"Prompt length { + len(raw_prompt_ids)} is longer than { + self.max_prompt_length}.") row_dict["raw_prompt_ids"] = raw_prompt_ids # encode prompts without chat template @@ -533,10 +634,18 @@ def __getitem__(self, item): # add index for each prompt index = row_dict.get("extra_info", {}).get("index", 0) tools_kwargs = row_dict.get("extra_info", {}).get("tools_kwargs", {}) - interaction_kwargs = row_dict.get("extra_info", {}).get("interaction_kwargs", {}) - need_tools_kwargs = row_dict.get("extra_info", {}).get("need_tools_kwargs", self.need_tools_kwargs) + interaction_kwargs = row_dict.get("extra_info", {}).get( + "interaction_kwargs", {} + ) + need_tools_kwargs = row_dict.get("extra_info", {}).get( + "need_tools_kwargs", self.need_tools_kwargs + ) if need_tools_kwargs and not tools_kwargs: - logger.warning("tools_kwargs is empty for index {}, data source: {}", index, row_dict["data_source"]) + logger.warning( + "tools_kwargs is empty for index {}, data source: {}", + index, + row_dict["data_source"], + ) row_dict["index"] = index row_dict["tools_kwargs"] = tools_kwargs row_dict["interaction_kwargs"] = interaction_kwargs diff --git a/Agent0/executor_train/verl/recipe/prime/main_prime.py b/Agent0/executor_train/verl/recipe/prime/main_prime.py index 6bf7f5e..22d7a5c 100644 --- a/Agent0/executor_train/verl/recipe/prime/main_prime.py +++ b/Agent0/executor_train/verl/recipe/prime/main_prime.py @@ -12,7 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -35,7 +35,9 @@ from .prime_ray_trainer import RayPRIMETrainer -@hydra.main(config_path="config", config_name="prime_trainer", version_base=None) +@hydra.main(config_path="config", + config_name="prime_trainer", + version_base=None) def main(config): run_prime(config) @@ -44,7 +46,10 @@ def run_prime(config, compute_score=None): if not ray.is_initialized(): # this is for local ray cluster ray.init( - runtime_env={"env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN"}}, + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN"}}, num_cpus=config.ray_init.num_cpus, ) @@ -60,7 +65,9 @@ def main_task(config, compute_score=None): from verl.utils.fs import copy_local_path_from_hdfs - pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values + pprint( + OmegaConf.to_container(config, resolve=True) + ) # resolve=True will eval symbol values OmegaConf.resolve(config) # download the checkpoint from hdfs @@ -97,7 +104,9 @@ def main_task(config, compute_score=None): global_pool_id = "global_pool" resource_pool_spec = { - global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + global_pool_id: [ + config.trainer.n_gpus_per_node] * + config.trainer.nnodes, } mapping = { Role.ActorRollout: global_pool_id, @@ -111,7 +120,8 @@ def main_task(config, compute_score=None): if config.reward_model.enable: from .prime_fsdp_workers import PRIMERewardModelWorker - role_worker_mapping[Role.RewardModel] = ray.remote(PRIMERewardModelWorker) + role_worker_mapping[Role.RewardModel] = ray.remote( + PRIMERewardModelWorker) mapping[Role.RewardModel] = global_pool_id reward_manager_name = config.reward_model.get("reward_manager", "naive") @@ -125,12 +135,18 @@ def main_task(config, compute_score=None): reward_manager_cls = PrimeRewardManager else: raise NotImplementedError - reward_fn = reward_manager_cls(tokenizer=tokenizer, num_examine=0, compute_score=compute_score) + reward_fn = reward_manager_cls( + tokenizer=tokenizer, num_examine=0, compute_score=compute_score + ) # Note that we always use function-based RM for validation - val_reward_fn = reward_manager_cls(tokenizer=tokenizer, num_examine=1, compute_score=compute_score) + val_reward_fn = reward_manager_cls( + tokenizer=tokenizer, num_examine=1, compute_score=compute_score + ) - resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, mapping=mapping + ) trainer = RayPRIMETrainer( config=config, diff --git a/Agent0/executor_train/verl/recipe/prime/prime_core_algos.py b/Agent0/executor_train/verl/recipe/prime/prime_core_algos.py index 8256712..a602788 100644 --- a/Agent0/executor_train/verl/recipe/prime/prime_core_algos.py +++ b/Agent0/executor_train/verl/recipe/prime/prime_core_algos.py @@ -18,7 +18,9 @@ import verl.utils.torch_functional as verl_F -def compute_rloo_advantage_return(data: verl.DataProto, response_mask: torch.Tensor, n_samples, config): +def compute_rloo_advantage_return( + data: verl.DataProto, response_mask: torch.Tensor, n_samples, config +): # calculate rloo reward on different reward sources, and sum again def masked_rloo(reward_tensor_original, mask_tensor): reward_tensor = reward_tensor_original.clone() @@ -26,15 +28,21 @@ def masked_rloo(reward_tensor_original, mask_tensor): for start_pos in range(0, reward_tensor.shape[0], n_samples): cur_rewards_mean = torch.cat( [ - reward_tensor[pos : pos + 1][mask_tensor[pos : pos + 1]].mean(dim=0, keepdim=True) + reward_tensor[pos: pos + 1][mask_tensor[pos: pos + 1]].mean( + dim=0, keepdim=True + ) for pos in range(start_pos, start_pos + n_samples) ], dim=0, ) cur_rewards_sum = cur_rewards_mean.sum() cur_reward_baseline = cur_rewards_sum / (n_samples - 1) - reward_tensor[start_pos : start_pos + n_samples][mask_tensor[start_pos : start_pos + n_samples]] = ( - reward_tensor[start_pos : start_pos + n_samples][mask_tensor[start_pos : start_pos + n_samples]] + reward_tensor[start_pos: start_pos + n_samples][ + mask_tensor[start_pos: start_pos + n_samples] + ] = ( + reward_tensor[start_pos: start_pos + n_samples][ + mask_tensor[start_pos: start_pos + n_samples] + ] * (n_samples / (n_samples - 1)) - cur_reward_baseline ) @@ -48,30 +56,53 @@ def masked_rloo(reward_tensor_original, mask_tensor): reward_tensor = data.batch["rm_scores"] reward_mask = response_mask.bool() - reward_tensors.append(masked_rloo(reward_tensor, reward_mask) * config.algorithm.reward_dpo_coef) + reward_tensors.append( + masked_rloo(reward_tensor, reward_mask) + * config.algorithm.reward_dpo_coef + ) if "acc" in data.batch.keys() and config.algorithm.reward_gt_coef != 0.0: - reward_tensor = torch.zeros_like(response_mask, dtype=torch.float32) + reward_tensor = torch.zeros_like( + response_mask, dtype=torch.float32) reward_mask = torch.zeros_like(response_mask, dtype=torch.bool) prompt_ids = data.batch["prompts"] prompt_length = prompt_ids.shape[-1] - valid_response_length = data.batch["attention_mask"][:, prompt_length:].sum(-1) + valid_response_length = data.batch["attention_mask"][:, + prompt_length:].sum(-1) reward_mask[ - torch.arange(0, valid_response_length.shape[0], dtype=torch.long, device=valid_response_length.device), + torch.arange( + 0, + valid_response_length.shape[0], + dtype=torch.long, + device=valid_response_length.device, + ), valid_response_length - 1, ] = True reward_tensor[ - torch.arange(0, valid_response_length.shape[0], dtype=torch.long, device=valid_response_length.device), + torch.arange( + 0, + valid_response_length.shape[0], + dtype=torch.long, + device=valid_response_length.device, + ), valid_response_length - 1, ] = data.batch["acc"] - reward_tensors.append(masked_rloo(reward_tensor, reward_mask) * config.algorithm.reward_gt_coef) + reward_tensors.append( + masked_rloo(reward_tensor, reward_mask) + * config.algorithm.reward_gt_coef + ) final_reward_tensor = sum(reward_tensors) - returns = (final_reward_tensor * response_mask).flip(dims=[-1]).cumsum(dim=-1).flip(dims=[-1]) + returns = ( + (final_reward_tensor * response_mask) + .flip(dims=[-1]) + .cumsum(dim=-1) + .flip(dims=[-1]) + ) advantages = returns.clone() advantages = verl_F.masked_whiten(advantages, response_mask) @@ -80,24 +111,33 @@ def masked_rloo(reward_tensor_original, mask_tensor): def compute_ce_dpo_loss_rm(token_level_scores, acc, response_mask, beta): - cur_scores = ((token_level_scores * response_mask).sum(dim=1) * beta).sigmoid() + cur_scores = ( + (token_level_scores * + response_mask).sum( + dim=1) * + beta).sigmoid() cur_dpo_loss = torch.nn.functional.binary_cross_entropy(cur_scores, acc) return cur_dpo_loss -def compute_detach_dpo_loss_rm(token_level_scores, acc, Q_bc, acc_bc, response_mask, beta, bon_mode="none"): +def compute_detach_dpo_loss_rm( + token_level_scores, acc, Q_bc, acc_bc, response_mask, beta, bon_mode="none" +): # we always assume that the BoN size equals n_samples # mode1: use acc as rm # mode2: use Q as rm cur_Q = (token_level_scores * response_mask).sum(dim=1) * beta other_Q = torch.zeros_like(cur_Q) for i in range(token_level_scores.shape[0]): - Q_chosen = Q_bc[i][acc_bc[i] < acc[i]] if acc[i] > 0 else Q_bc[i][acc_bc[i] > acc[i]] + Q_chosen = (Q_bc[i][acc_bc[i] < acc[i]] if acc[i] + > 0 else Q_bc[i][acc_bc[i] > acc[i]]) if len(Q_chosen) > 0: other_Q[i] = Q_chosen.mean() * beta else: other_Q[i] = 0 - dpo_loss = -torch.log(torch.sigmoid((cur_Q - other_Q) * ((acc > 0).float() * 2 - 1))) + dpo_loss = -torch.log( + torch.sigmoid((cur_Q - other_Q) * ((acc > 0).float() * 2 - 1)) + ) if bon_mode == "none": dpo_loss = dpo_loss.mean() else: @@ -105,10 +145,14 @@ def compute_detach_dpo_loss_rm(token_level_scores, acc, Q_bc, acc_bc, response_m n_samples = acc_bc.shape[1] if bon_mode == "bon_rm": for i in range(token_level_scores.shape[0]): - weight[i] = n_samples * torch.pow((Q_bc[i] * beta <= cur_Q[i]).float().mean(), n_samples - 1) + weight[i] = n_samples * torch.pow( + (Q_bc[i] * beta <= cur_Q[i]).float().mean(), n_samples - 1 + ) elif bon_mode == "bon_acc": for i in range(token_level_scores.shape[0]): - weight[i] = n_samples * torch.pow((acc_bc[i] <= acc[i]).float().mean(), n_samples - 1) + weight[i] = n_samples * torch.pow( + (acc_bc[i] <= acc[i]).float().mean(), n_samples - 1 + ) else: raise NotImplementedError dpo_loss = (dpo_loss * weight).sum() @@ -120,22 +164,28 @@ def compute_dpo_accuracy(token_level_scores, acc, response_mask, n_samples): dpo_acc = [] for start_id in range(0, token_level_scores.shape[0], n_samples): cur_scores = ( - token_level_scores[start_id : start_id + n_samples] * response_mask[start_id : start_id + n_samples] + token_level_scores[start_id: start_id + n_samples] + * response_mask[start_id: start_id + n_samples] ).sum(dim=1) def get_upper_triangle(tensor_x): diff_matrix = tensor_x.unsqueeze(1) - tensor_x.unsqueeze(0) - upper_tri_indices = torch.triu(torch.ones_like(diff_matrix).bool(), diagonal=1) + upper_tri_indices = torch.triu( + torch.ones_like(diff_matrix).bool(), diagonal=1 + ) return diff_matrix[upper_tri_indices] - cur_acc_diff = get_upper_triangle(acc[start_id : start_id + n_samples]) # in range [-1,1] + cur_acc_diff = get_upper_triangle( + acc[start_id: start_id + n_samples] + ) # in range [-1,1] cur_score_diff = get_upper_triangle(cur_scores) # in R cur_score_prediction = (cur_score_diff > 0).float() # in [0,1] if cur_acc_diff.abs().sum() == 0: cur_acc = torch.zeros_like(cur_score_prediction[0]) + 0.5 else: cur_acc = ( - ((cur_score_diff > 0) == (cur_acc_diff > 0)).float() * cur_acc_diff.abs() + ((cur_score_diff > 0) == (cur_acc_diff > 0)).float() + * cur_acc_diff.abs() ).sum() / cur_acc_diff.abs().sum() dpo_acc.append(cur_acc.unsqueeze(0)) @@ -143,5 +193,16 @@ def get_upper_triangle(tensor_x): return torch.cat(dpo_acc, dim=0).mean() -def compute_dpo_abs_accuracy(token_level_scores, acc, response_mask, n_samples): - return (torch.sign((token_level_scores * response_mask).sum(dim=-1)) == torch.sign(acc * 2 - 1)).float().mean() +def compute_dpo_abs_accuracy( + token_level_scores, + acc, + response_mask, + n_samples): + return ( + ( + torch.sign((token_level_scores * response_mask).sum(dim=-1)) + == torch.sign(acc * 2 - 1) + ) + .float() + .mean() + ) diff --git a/Agent0/executor_train/verl/recipe/prime/prime_dp_rm.py b/Agent0/executor_train/verl/recipe/prime/prime_dp_rm.py index c9cc060..c42983a 100644 --- a/Agent0/executor_train/verl/recipe/prime/prime_dp_rm.py +++ b/Agent0/executor_train/verl/recipe/prime/prime_dp_rm.py @@ -36,17 +36,27 @@ class DataParallelPRIMERewardModel: - def __init__(self, config, reward_module: nn.Module, ref_module: nn.Module, reward_optimizer: optim.Optimizer): + def __init__( + self, + config, + reward_module: nn.Module, + ref_module: nn.Module, + reward_optimizer: optim.Optimizer, + ): self.config = config self.reward_module = reward_module self.ref_module = ref_module self.reward_optimizer = reward_optimizer - self.use_remove_padding = self.config.model.get("use_remove_padding", False) + self.use_remove_padding = self.config.model.get( + "use_remove_padding", False) print(f"Reward model use_remove_padding={self.use_remove_padding}") - self.use_fused_kernels = self.config.model.get("use_fused_kernels", False) + self.use_fused_kernels = self.config.model.get( + "use_fused_kernels", False) print(f"Reward model use_fused_kernels={self.use_fused_kernels}") - self.ulysses_sequence_parallel_size = self.config.get("ulysses_sequence_parallel_size", 1) + self.ulysses_sequence_parallel_size = self.config.get( + "ulysses_sequence_parallel_size", 1 + ) def _forward_micro_batch(self, micro_batch, prompt_length): input_ids = micro_batch["input_ids"] @@ -55,7 +65,8 @@ def _forward_micro_batch(self, micro_batch, prompt_length): position_ids = micro_batch["position_ids"] num_actions = micro_batch["input_ids"].shape[-1] - prompt_length - max_positions = micro_batch["attention_mask"][:, prompt_length:].sum(-1) + max_positions = micro_batch["attention_mask"][:, + prompt_length:].sum(-1) if self.use_remove_padding: input_ids_rmpad, indices, *_ = unpad_input( @@ -69,16 +80,21 @@ def _forward_micro_batch(self, micro_batch, prompt_length): ).transpose(0, 1) # for compute the log_prob - input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=1) # (1, total_nnz) + 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.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 + 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, None, self.ulysses_sequence_parallel_size) input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0) output = self.reward_module( @@ -101,10 +117,14 @@ def _forward_micro_batch(self, micro_batch, prompt_length): ) if self.ulysses_sequence_parallel_size > 1: - rm_log_labels = gather_outpus_and_unpad(rm_log_labels, gather_dim=0, unpad_dim=0, padding_size=pad_size) + rm_log_labels = gather_outpus_and_unpad( + rm_log_labels, gather_dim=0, unpad_dim=0, padding_size=pad_size) rm_log_labels = pad_input( - hidden_states=rm_log_labels.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen - ).squeeze(-1)[:, -num_actions - 1 : -1] + hidden_states=rm_log_labels.unsqueeze(-1), + indices=indices, + batch=batch_size, + seqlen=seqlen, + ).squeeze(-1)[:, -num_actions - 1: -1] else: output = self.reward_module( @@ -124,13 +144,17 @@ def _forward_micro_batch(self, micro_batch, prompt_length): rm_log_prob = torch.nn.functional.log_softmax( rm_output_logits[:, :-1, :], dim=-1 ) # (batch_size, seq_length, vocab_size) - rm_log_labels = rm_log_prob.gather(dim=-1, index=micro_batch["input_ids"][:, 1:].unsqueeze(-1)).squeeze( + rm_log_labels = rm_log_prob.gather( + dim=-1, index=micro_batch["input_ids"][:, 1:].unsqueeze(-1) + ).squeeze( -1 ) # (batch, seq_length) if self.ref_module is not None: # do not have to pad again - with torch.no_grad(), torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16): + with torch.no_grad(), torch.autocast( + device_type=get_device_name(), dtype=torch.bfloat16 + ): if self.ulysses_sequence_parallel_size > 1 and self.use_remove_padding: ref_output = self.ref_module( input_ids=input_ids_rmpad, @@ -140,21 +164,23 @@ def _forward_micro_batch(self, micro_batch, prompt_length): ) if self.use_fused_kernels: - ref_log_labels = ref_output.log_probs.squeeze(0) # (total_nnz,) + ref_log_labels = ref_output.log_probs.squeeze( + 0) # (total_nnz,) ref_log_labels = ref_log_labels.to(torch.float32) else: ref_output_logits = ref_output.logits.squeeze(0) ref_log_labels = verl_F.logprobs_from_logits( - logits=ref_output_logits, labels=input_ids_rmpad_rolled - ) + logits=ref_output_logits, labels=input_ids_rmpad_rolled) ref_log_labels = gather_outpus_and_unpad( - ref_log_labels, gather_dim=0, unpad_dim=0, padding_size=pad_size - ) + ref_log_labels, gather_dim=0, unpad_dim=0, padding_size=pad_size) ref_log_labels = pad_input( - hidden_states=ref_log_labels.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen - ).squeeze(-1)[:, -num_actions - 1 : -1] + hidden_states=ref_log_labels.unsqueeze(-1), + indices=indices, + batch=batch_size, + seqlen=seqlen, + ).squeeze(-1)[:, -num_actions - 1: -1] else: ref_output = self.ref_module( input_ids=micro_batch["input_ids"], @@ -164,7 +190,9 @@ def _forward_micro_batch(self, micro_batch, prompt_length): ) if self.use_fused_kernels: - ref_log_labels = ref_output.log_probs[:, :-1] # (batch_size, seq_length) + ref_log_labels = ref_output.log_probs[ + :, :-1 + ] # (batch_size, seq_length) ref_log_labels = ref_log_labels.to(torch.float32) else: @@ -174,17 +202,21 @@ def _forward_micro_batch(self, micro_batch, prompt_length): ) # (batch_size, seq_length, vocab_size) ref_log_labels = ref_log_prob.gather( dim=-1, index=micro_batch["input_ids"][:, 1:].unsqueeze(-1) - ).squeeze(-1) # (batch, seq_length) + ).squeeze( + -1 + ) # (batch, seq_length) else: ref_log_labels = micro_batch["old_log_probs"] ref_log_labels.to(rm_log_labels.dtype) - q = rm_log_labels[:, -num_actions:] - ref_log_labels[:, -num_actions:] # this is actually diff of q + q = ( + rm_log_labels[:, -num_actions:] - ref_log_labels[:, -num_actions:] + ) # this is actually diff of q # trim unnecessary logprobs here for i in range(micro_batch["input_ids"].shape[0]): - q[i, max_positions[i] :] = 0 + q[i, max_positions[i]:] = 0 # reward computation does not need gradient. only q needs with torch.no_grad(): @@ -204,8 +236,10 @@ def _forward_micro_batch(self, micro_batch, prompt_length): # outcome reward to calculate V for i in range(q.shape[0]): if self.config.prime_use_gt: - q_[i, max_positions[i] - 1] = acc[i] - q_[i, : max_positions[i] - 1].sum() - q_[i, max_positions[i] :] = 0 + q_[i, max_positions[i] - 1] = ( + acc[i] - q_[i, : max_positions[i] - 1].sum() + ) + q_[i, max_positions[i]:] = 0 for t in reversed(range(num_actions)): delta = q_[:, t] @@ -216,10 +250,14 @@ def _forward_micro_batch(self, micro_batch, prompt_length): if self.config.prime_granularity == "token": for i in range(micro_batch["input_ids"].shape[0]): - token_level_score[i, : max_positions[i] - 1] = r[i, : max_positions[i] - 1] + token_level_score[i, : max_positions[i] - 1] = r[ + i, : max_positions[i] - 1 + ] elif self.config.prime_granularity == "whole": for i in range(micro_batch["input_ids"].shape[0]): - token_level_score[i, max_positions[i] - 1] = r[i, : max_positions[i]] + token_level_score[i, max_positions[i] - 1] = r[ + i, : max_positions[i] + ] else: raise NotImplementedError @@ -229,33 +267,52 @@ def _optimizer_step(self): assert self.config.model.optim.grad_clip is not None if isinstance(self.reward_module, FSDP): - grad_norm = self.reward_module.clip_grad_norm_(self.config.model.optim.grad_clip) + grad_norm = self.reward_module.clip_grad_norm_( + self.config.model.optim.grad_clip + ) else: grad_norm = torch.nn.utils.clip_grad_norm_( - self.reward_module.parameters(), max_norm=self.config.model.optim.grad_clip + self.reward_module.parameters(), + max_norm=self.config.model.optim.grad_clip, ) self.reward_optimizer.step() return grad_norm def prime_norm(self, token_level_scores): if self.config.prime_norm == "batch_norm": - reverse_cumsum = torch.cumsum(token_level_scores.flip(dims=[1]), dim=-1).flip(dims=[1]) - token_level_scores = token_level_scores / (reverse_cumsum.abs().max() + 1e-6) + reverse_cumsum = torch.cumsum( + token_level_scores.flip(dims=[1]), dim=-1 + ).flip(dims=[1]) + token_level_scores = token_level_scores / ( + reverse_cumsum.abs().max() + 1e-6 + ) return token_level_scores def compute_rm_score(self, data: DataProto): self.reward_module.eval() self.ref_module.eval() micro_batch_size = data.meta_info["micro_batch_size"] - select_keys = ["responses", "input_ids", "attention_mask", "position_ids", "acc"] + select_keys = [ + "responses", + "input_ids", + "attention_mask", + "position_ids", + "acc", + ] batch = data.select(batch_keys=select_keys).batch use_dynamic_bsz = data.meta_info["use_dynamic_bsz"] - prompt_length = data.batch["input_ids"].shape[-1] - data.batch["responses"].shape[-1] + prompt_length = ( + data.batch["input_ids"].shape[-1] - data.batch["responses"].shape[-1] + ) 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) + 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) @@ -263,7 +320,8 @@ def compute_rm_score(self, data: DataProto): q_lst = [] for micro_batch in micro_batches: with torch.no_grad(): - rm_score, q = self._forward_micro_batch(micro_batch, prompt_length) + rm_score, q = self._forward_micro_batch( + micro_batch, prompt_length) rm_scores_lst.append(rm_score) q_lst.append(q) rm_scores = torch.concat(rm_scores_lst, dim=0) @@ -273,8 +331,11 @@ def compute_rm_score(self, data: DataProto): if use_dynamic_bsz: indices = list(itertools.chain.from_iterable(indices)) - assert len(indices) == rm_scores.size(0), f"{len(indices)} vs. {rm_scores.size()}" - revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long) + assert len(indices) == rm_scores.size( + 0 + ), f"{len(indices)} vs. {rm_scores.size()}" + revert_indices = torch.tensor( + get_reverse_idx(indices), dtype=torch.long) rm_scores = rm_scores[revert_indices] return ( @@ -293,7 +354,14 @@ def update_rm(self, data: DataProto): beta = self.config.model.get("beta_train", 0.05) - select_keys = ["input_ids", "responses", "attention_mask", "position_ids", "acc", "prompts"] + select_keys = [ + "input_ids", + "responses", + "attention_mask", + "position_ids", + "acc", + "prompts", + ] for key in ["Q_bc", "acc_bc"]: if key in data.batch.keys(): @@ -311,11 +379,19 @@ def update_rm(self, data: DataProto): # 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) + 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.micro_batch_size_per_gpu) - self.gradient_accumulation = self.config.mini_batch_size // self.config.micro_batch_size_per_gpu + micro_batches = mini_batch.split( + self.config.micro_batch_size_per_gpu) + self.gradient_accumulation = ( + self.config.mini_batch_size // + self.config.micro_batch_size_per_gpu) self.reward_optimizer.zero_grad() @@ -335,15 +411,23 @@ def update_rm(self, data: DataProto): q_lst.append(q.detach()) if self.config.model.loss_type == "ce": - dpo_loss = compute_ce_dpo_loss_rm(q, acc, response_mask=response_mask, beta=beta) + dpo_loss = compute_ce_dpo_loss_rm( + q, acc, response_mask=response_mask, beta=beta + ) elif self.config.model.loss_type == "dpo": # the implementation of dpo is actually detached, which means we have to know the average # value of w/l reward before the update. dpo_loss = compute_detach_dpo_loss_rm( - q, acc, Q_bc=data["Q_bc"], acc_bc=data["acc_bc"], response_mask=response_mask, beta=beta + q, + acc, + Q_bc=data["Q_bc"], + acc_bc=data["acc_bc"], + response_mask=response_mask, + beta=beta, ) elif self.config.model.loss_type == "bon_acc": - # change the original distribution of each sample to BoN distribution, then update reward model + # change the original distribution of each sample to BoN + # distribution, then update reward model dpo_loss = compute_detach_dpo_loss_rm( q, acc, @@ -370,7 +454,8 @@ def update_rm(self, data: DataProto): if self.config.use_dynamic_bsz: # relative to the dynamic bsz - loss = dpo_loss * (len(data) / self.config.ppo_mini_batch_size) + loss = dpo_loss * \ + (len(data) / self.config.ppo_mini_batch_size) else: loss = dpo_loss / self.gradient_accumulation diff --git a/Agent0/executor_train/verl/recipe/prime/prime_fsdp_workers.py b/Agent0/executor_train/verl/recipe/prime/prime_fsdp_workers.py index e353404..f97261b 100644 --- a/Agent0/executor_train/verl/recipe/prime/prime_fsdp_workers.py +++ b/Agent0/executor_train/verl/recipe/prime/prime_fsdp_workers.py @@ -61,28 +61,42 @@ def __init__(self, config): world_size = torch.distributed.get_world_size() fsdp_size = self.config.model.fsdp_config.fsdp_size - self.device_mesh = create_device_mesh(world_size=world_size, fsdp_size=fsdp_size) + self.device_mesh = create_device_mesh( + world_size=world_size, fsdp_size=fsdp_size + ) self.ulysses_device_mesh = None - self.ulysses_sequence_parallel_size = self.config.get("ulysses_sequence_parallel_size", 1) + 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( - get_device_name(), mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=["dp", "sp"] + get_device_name(), + mesh_shape=(dp, self.ulysses_sequence_parallel_size), + mesh_dim_names=["dp", "sp"], ) - self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) + 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_optimizer = self.config.model.fsdp_config.optimizer_offload # normalize config - self.config.mini_batch_size //= torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size + self.config.mini_batch_size //= ( + torch.distributed.get_world_size() // + self.ulysses_sequence_parallel_size) if self.config.micro_batch_size is not None: - self.config.micro_batch_size //= torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size + self.config.micro_batch_size //= ( + torch.distributed.get_world_size() + // self.ulysses_sequence_parallel_size + ) self.config.micro_batch_size_per_gpu = self.config.micro_batch_size - assert self.config.mini_batch_size % self.config.micro_batch_size_per_gpu == 0 + assert (self.config.mini_batch_size % + self.config.micro_batch_size_per_gpu == 0) def _build_reward_ref_model_optimizer(self, config): # the following line is necessary @@ -96,11 +110,16 @@ def _build_reward_ref_model_optimizer(self, config): local_path = copy_local_path_from_hdfs(config.model.path) 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)) + 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 = 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, @@ -116,10 +135,14 @@ def _build_reward_ref_model_optimizer(self, config): from transformers import AutoConfig, AutoModelForCausalLM trust_remote_code = False - reward_model_config = AutoConfig.from_pretrained(local_path, trust_remote_code=trust_remote_code) + reward_model_config = AutoConfig.from_pretrained( + local_path, trust_remote_code=trust_remote_code + ) reward_model_config.num_labels = 1 - init_context = get_init_weight_context_manager(use_meta_tensor=not reward_model_config.tie_word_embeddings) + init_context = get_init_weight_context_manager( + use_meta_tensor=not reward_model_config.tie_word_embeddings + ) with init_context(), warnings.catch_warnings(): warnings.simplefilter("ignore") reward_model_config.classifier_dropout = 0.0 @@ -132,16 +155,23 @@ def _build_reward_ref_model_optimizer(self, config): trust_remote_code=trust_remote_code, ) - fused_kernel_options = config.model.get("fused_kernel_options", None) + fused_kernel_options = config.model.get( + "fused_kernel_options", None) fused_kernels_backend = ( - fused_kernel_options.get("impl_backend", None) if fused_kernel_options is not None else None + fused_kernel_options.get("impl_backend", None) + if fused_kernel_options is not None + else None ) apply_monkey_patch( model=reward_module, ulysses_sp_size=self.ulysses_sequence_parallel_size, - use_remove_padding=config.model.get("use_remove_padding", False), - use_fused_kernels=config.model.get("use_fused_kernels", False), + use_remove_padding=config.model.get( + "use_remove_padding", + False), + use_fused_kernels=config.model.get( + "use_fused_kernels", + False), fused_kernels_backend=fused_kernels_backend, ) @@ -149,7 +179,9 @@ def _build_reward_ref_model_optimizer(self, config): reward_module.to(torch_dtype) if config.model.get("enable_gradient_checkpointing", False): - reward_module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + reward_module.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) if self.rank == 0: print_model_size(reward_module) @@ -158,17 +190,29 @@ def _build_reward_ref_model_optimizer(self, 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")) + 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) + mixed_precision = MixedPrecision( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + buffer_dtype=buffer_dtype, + ) - auto_wrap_policy = get_fsdp_wrap_policy(module=reward_module, config=self.config.model.fsdp_config.wrap_policy) + auto_wrap_policy = get_fsdp_wrap_policy( + module=reward_module, + config=self.config.model.fsdp_config.wrap_policy) log_gpu_memory_usage("Before reward model FSDP", logger=None) @@ -180,7 +224,9 @@ def _build_reward_ref_model_optimizer(self, config): reward_model_config.classifier_dropout = 0.0 reward_model_config.hidden_dropout = "0" ref_module = AutoModelForCausalLM.from_pretrained( - pretrained_model_name_or_path=copy_local_path_from_hdfs(config.model.ref_path), + pretrained_model_name_or_path=copy_local_path_from_hdfs( + config.model.ref_path + ), torch_dtype=torch_dtype, config=reward_model_config, attn_implementation="flash_attention_2", @@ -230,10 +276,13 @@ def _build_reward_ref_model_optimizer(self, config): total_steps = config.model.optim.get("total_training_steps", 0) num_warmup_steps = int(config.model.optim.get("lr_warmup_steps", -1)) if num_warmup_steps < 0: - num_warmup_steps_ratio = config.model.optim.get("lr_warmup_steps_ratio", 0.0) + num_warmup_steps_ratio = config.model.optim.get( + "lr_warmup_steps_ratio", 0.0 + ) num_warmup_steps = int(num_warmup_steps_ratio * total_steps) - print(f"Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}") + print( + f"Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}") from verl.utils.torch_functional import get_constant_schedule_with_warmup @@ -250,9 +299,12 @@ def init_model(self): from .prime_dp_rm import DataParallelPRIMERewardModel - self.reward_module, self.ref_module, self.reward_optimizer, self.reward_lr_scheduler = ( - self._build_reward_ref_model_optimizer(config=self.config) - ) + ( + self.reward_module, + self.ref_module, + self.reward_optimizer, + self.reward_lr_scheduler, + ) = self._build_reward_ref_model_optimizer(config=self.config) if self._is_offload_param: offload_fsdp_model_to_cpu(self.reward_module) @@ -295,14 +347,25 @@ def compute_rm_score(self, data: DataProto): response_mask = data.batch["attention_mask"][:, prompt_length:] acc = data.batch["acc"] - dpo_acc = compute_dpo_accuracy(rm_scores, acc, response_mask=response_mask, n_samples=data.meta_info["n"]) - dpo_acc_abs = compute_dpo_abs_accuracy(rm_scores, acc, response_mask, n_samples=data.meta_info["n"]) + dpo_acc = compute_dpo_accuracy( + rm_scores, + acc, + response_mask=response_mask, + n_samples=data.meta_info["n"], + ) + dpo_acc_abs = compute_dpo_abs_accuracy( + rm_scores, acc, response_mask, n_samples=data.meta_info["n"] + ) metrics["reward_model/dpo_acc"] = dpo_acc.detach().item() metrics["reward_model/dpo_acc_abs"] = dpo_acc_abs.detach().item() - output = DataProto.from_dict(tensors={"rm_scores": rm_scores, "q": q}, meta_info={"metrics": metrics}) - output = self.ulysses_sharding_manager.postprocess_data(data=output) + output = DataProto.from_dict( + tensors={ + "rm_scores": rm_scores, "q": q}, meta_info={ + "metrics": metrics}) + output = self.ulysses_sharding_manager.postprocess_data( + data=output) output = output.to("cpu") if self._is_offload_param: @@ -317,7 +380,9 @@ def update_rm(self, data: DataProto): load_fsdp_model_to_gpu(self.ref_module) load_fsdp_model_to_gpu(self.reward_module) if self._is_offload_optimizer: - load_fsdp_optimizer(optimizer=self.reward_optimizer, device_id=get_device_id()) + load_fsdp_optimizer( + optimizer=self.reward_optimizer, device_id=get_device_id() + ) # perform forward computation with self.ulysses_sharding_manager: @@ -334,15 +399,24 @@ def update_rm(self, data: DataProto): acc = data.batch["acc"] dpo_acc_before = compute_dpo_accuracy( - rm_scores, acc, response_mask=response_mask, n_samples=data.meta_info["n"] + rm_scores, + acc, + response_mask=response_mask, + n_samples=data.meta_info["n"], + ) + dpo_acc_abs = compute_dpo_abs_accuracy( + rm_scores, acc, response_mask, n_samples=data.meta_info["n"] ) - dpo_acc_abs = compute_dpo_abs_accuracy(rm_scores, acc, response_mask, n_samples=data.meta_info["n"]) metrics["reward_model/dpo_acc_before"] = dpo_acc_before.detach().item() metrics["reward_model/dpo_acc_abs_before"] = dpo_acc_abs.detach().item() - output = DataProto.from_dict(tensors={"rm_scores": rm_scores}, meta_info={"metrics": metrics}) - output = self.ulysses_sharding_manager.postprocess_data(data=output) + output = DataProto.from_dict( + tensors={ + "rm_scores": rm_scores}, meta_info={ + "metrics": metrics}) + output = self.ulysses_sharding_manager.postprocess_data( + data=output) if self._is_offload_param: offload_fsdp_model_to_cpu(self.reward_module) @@ -353,14 +427,19 @@ def update_rm(self, data: DataProto): return output @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None): + def save_checkpoint( + self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None + ): import torch if self._is_offload_param: load_fsdp_model_to_gpu(self.reward_module) self.checkpoint_manager.save_checkpoint( - local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep + local_path=local_path, + hdfs_path=hdfs_path, + global_step=global_step, + max_ckpt_to_keep=max_ckpt_to_keep, ) torch.distributed.barrier() @@ -374,7 +453,9 @@ def load_checkpoint(self, local_path, del_local_after_load=True): if self._is_offload_param: load_fsdp_model_to_gpu(self.reward_module) - self.checkpoint_manager.load_checkpoint(local_path=local_path, del_local_after_load=del_local_after_load) + self.checkpoint_manager.load_checkpoint( + local_path=local_path, del_local_after_load=del_local_after_load + ) torch.distributed.barrier() if self._is_offload_param: diff --git a/Agent0/executor_train/verl/recipe/prime/prime_ray_trainer.py b/Agent0/executor_train/verl/recipe/prime/prime_ray_trainer.py index a5ad964..6a1581f 100644 --- a/Agent0/executor_train/verl/recipe/prime/prime_ray_trainer.py +++ b/Agent0/executor_train/verl/recipe/prime/prime_ray_trainer.py @@ -30,7 +30,12 @@ from verl.single_controller.ray import RayWorkerGroup from verl.trainer.ppo.core_algos import agg_loss from verl.trainer.ppo.metric_utils import _compute_response_info -from verl.trainer.ppo.ray_trainer import RayPPOTrainer, ResourcePoolManager, Role, WorkerType +from verl.trainer.ppo.ray_trainer import ( + RayPPOTrainer, + ResourcePoolManager, + Role, + WorkerType, +) from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path from verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn from verl.utils.metric import reduce_metrics @@ -61,8 +66,10 @@ def compute_data_metrics(batch, use_critic=True): 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() + 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) @@ -95,7 +102,9 @@ def compute_data_metrics(batch, use_critic=True): "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(), + "critic/vf_explained_var": (1.0 - return_diff_var / (return_var + 1e-5)) + .detach() + .item(), } if use_critic else {} @@ -104,14 +113,20 @@ def compute_data_metrics(batch, use_critic=True): "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()) + "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(), + "prompt_length/clip_ratio": torch.mean( + torch.eq(prompt_length, max_prompt_length).float() + ) + .detach() + .item(), } return metrics @@ -129,15 +144,20 @@ def compute_timing_metrics(batch, timing_raw): 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"]}, - } + num_tokens_of_section = {"gen": num_response_tokens, + **{name: num_overall_tokens for name in ["ref", + "values", + "adv", + "update_critic", + "update_actor"]}, + } 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] + 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()) }, } @@ -185,26 +205,35 @@ def _create_dataloader(self, *args, **kwargs): # TODO: we have to make sure the batch size is divisible by the dp size self.train_dataset = RLHFDataset( - data_files=self.config.data.train_files, tokenizer=self.tokenizer, config=self.config.data + data_files=self.config.data.train_files, + tokenizer=self.tokenizer, + config=self.config.data, ) # use sampler for better ckpt resume if self.config.data.shuffle: train_dataloader_generator = torch.Generator() - train_dataloader_generator.manual_seed(self.config.data.get("seed", 1)) - sampler = RandomSampler(data_source=self.train_dataset, generator=train_dataloader_generator) + train_dataloader_generator.manual_seed( + self.config.data.get("seed", 1)) + sampler = RandomSampler( + data_source=self.train_dataset, + generator=train_dataloader_generator) else: sampler = SequentialSampler(data_source=self.train_dataset) self.train_dataloader = DataLoader( dataset=self.train_dataset, - batch_size=int(self.config.data.train_batch_size * self.config.data.oversample_factor), + batch_size=int( + self.config.data.train_batch_size * + self.config.data.oversample_factor), drop_last=True, collate_fn=collate_fn, sampler=sampler, ) self.val_dataset = RLHFDataset( - data_files=self.config.data.val_files, tokenizer=self.tokenizer, config=self.config.data + data_files=self.config.data.val_files, + tokenizer=self.tokenizer, + config=self.config.data, ) self.val_dataloader = DataLoader( dataset=self.val_dataset, @@ -220,8 +249,11 @@ def _create_dataloader(self, *args, **kwargs): print(f"Size of train dataloader: {len(self.train_dataloader)}") print(f"Size of val dataloader: {len(self.val_dataloader)}") - # inject total_training_steps to actor/critic optim_config. This is hacky. - total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs + # 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 @@ -231,21 +263,28 @@ def _create_dataloader(self, *args, **kwargs): 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.actor_rollout_ref.actor.optim.total_training_steps = ( + total_training_steps + ) self.config.critic.optim.total_training_steps = total_training_steps def _save_checkpoint(self): # path: given_path + `/global_step_{global_steps}` + `/actor` local_global_step_folder = os.path.join( - self.config.trainer.default_local_dir, f"global_step_{self.global_steps}" - ) + self.config.trainer.default_local_dir, + f"global_step_{ + self.global_steps}") print(f"local_global_step_folder: {local_global_step_folder}") actor_local_path = os.path.join(local_global_step_folder, "actor") actor_remote_path = ( None if self.config.trainer.default_hdfs_dir is None - else os.path.join(self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", "actor") + else os.path.join( + self.config.trainer.default_hdfs_dir, + f"global_step_{self.global_steps}", + "actor", + ) ) self.actor_rollout_wg.save_checkpoint( actor_local_path, @@ -254,11 +293,16 @@ def _save_checkpoint(self): ) if self.use_rm: - reward_local_path = os.path.join(local_global_step_folder, "reward") + reward_local_path = os.path.join( + local_global_step_folder, "reward") reward_remote_path = ( None if self.config.trainer.default_hdfs_dir is None - else os.path.join(self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", "reward") + else os.path.join( + self.config.trainer.default_hdfs_dir, + f"global_step_{self.global_steps}", + "reward", + ) ) self.rm_wg.save_checkpoint( reward_local_path, @@ -267,15 +311,19 @@ def _save_checkpoint(self): ) # save dataloader - dataloader_local_path = os.path.join(local_global_step_folder, "data.pt") + dataloader_local_path = os.path.join( + local_global_step_folder, "data.pt") import dill - torch.save(self.train_dataloader, dataloader_local_path, pickle_module=dill) + torch.save( + self.train_dataloader, + dataloader_local_path, + pickle_module=dill) # latest checkpointed iteration tracker (for atomic usage) local_latest_checkpointed_iteration = os.path.join( - self.config.trainer.default_local_dir, "latest_checkpointed_iteration.txt" - ) + self.config.trainer.default_local_dir, + "latest_checkpointed_iteration.txt") with open(local_latest_checkpointed_iteration, "w") as f: f.write(str(self.global_steps)) @@ -287,11 +335,16 @@ def _load_checkpoint(self): if self.config.trainer.default_hdfs_dir is not None: NotImplementedError("load from hdfs is not implemented yet") else: - checkpoint_folder = self.config.trainer.default_local_dir # TODO: check path + checkpoint_folder = ( + self.config.trainer.default_local_dir + ) # TODO: check path if not os.path.isabs(checkpoint_folder): working_dir = os.getcwd() - checkpoint_folder = os.path.join(working_dir, checkpoint_folder) - global_step_folder = find_latest_ckpt_path(checkpoint_folder) # None if no latest + checkpoint_folder = os.path.join( + working_dir, checkpoint_folder) + global_step_folder = find_latest_ckpt_path( + checkpoint_folder + ) # None if no latest # find global_step_folder if self.config.trainer.resume_mode == "auto": @@ -300,14 +353,17 @@ def _load_checkpoint(self): return 0 else: if self.config.trainer.resume_mode == "resume_path": - assert isinstance(self.config.trainer.resume_from_path, str), "resume ckpt must be str type" - assert "global_step_" in self.config.trainer.resume_from_path, ( - "resume ckpt must specify the global_steps" - ) + assert isinstance( + self.config.trainer.resume_from_path, str + ), "resume ckpt must be str type" + assert ( + "global_step_" in self.config.trainer.resume_from_path + ), "resume ckpt must specify the global_steps" global_step_folder = self.config.trainer.resume_from_path if not os.path.isabs(global_step_folder): working_dir = os.getcwd() - global_step_folder = os.path.join(working_dir, global_step_folder) + global_step_folder = os.path.join( + working_dir, global_step_folder) print(f"Load from checkpoint folder: {global_step_folder}") # set global step self.global_steps = int(global_step_folder.split("global_step_")[-1]) @@ -319,11 +375,15 @@ def _load_checkpoint(self): reward_path = os.path.join(global_step_folder, "reward") # load actor self.actor_rollout_wg.load_checkpoint( - actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load + actor_path, + del_local_after_load=self.config.trainer.del_local_ckpt_after_load, ) # load rm if self.use_rm: - self.rm_wg.load_checkpoint(reward_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load) + self.rm_wg.load_checkpoint( + reward_path, + del_local_after_load=self.config.trainer.del_local_ckpt_after_load, + ) # load dataloader, # TODO: from remote not implemented yet @@ -356,7 +416,9 @@ def fit(self): # 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): + if self.val_reward_fn is not None and self.config.trainer.get( + "val_before_train", True + ): val_metrics = self._validate() assert val_metrics, f"{val_metrics=}" pprint(f"Initial validation metrics: {val_metrics}") @@ -375,13 +437,19 @@ def fit(self): batch: DataProto = DataProto.from_single_dict(batch_dict) # pop those keys for generation - gen_batch = batch.pop(batch_keys=["input_ids", "attention_mask", "position_ids"]) - gen_batch = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) + gen_batch = batch.pop( + batch_keys=["input_ids", "attention_mask", "position_ids"] + ) + gen_batch = gen_batch.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) with simple_timer("step", timing_raw): # generate a batch with simple_timer("gen", timing_raw): - gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) + gen_batch_output = self.actor_rollout_wg.generate_sequences( + gen_batch) timing_raw.update(gen_batch_output.meta_info["timing"]) gen_batch_output.meta_info.pop("timing", None) @@ -389,23 +457,34 @@ def fit(self): with simple_timer("gen_max", timing_raw): gen_baseline_batch = deepcopy(gen_batch) gen_baseline_batch.meta_info["do_sample"] = False - gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch) + gen_baseline_output = ( + self.actor_rollout_wg.generate_sequences( + gen_baseline_batch + ) + ) batch = batch.union(gen_baseline_output) reward_baseline_tensor = self.reward_fn(batch) - reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1) + reward_baseline_tensor = reward_baseline_tensor.sum( + dim=-1) - batch.pop(batch_keys=list(gen_baseline_output.batch.keys())) + batch.pop( + batch_keys=list( + gen_baseline_output.batch.keys())) batch.batch["reward_baselines"] = reward_baseline_tensor del gen_baseline_batch, gen_baseline_output batch.non_tensor_batch["uid"] = np.array( - [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object + [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.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) batch = batch.union(gen_batch_output) # Balance the number of valid tokens across DP ranks. @@ -417,7 +496,9 @@ def fit(self): 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_token_num"] = torch.sum( + batch.batch["attention_mask"], dim=-1 + ).tolist() # verify with simple_timer("verify", timing_raw): @@ -425,7 +506,8 @@ def fit(self): metrics["acc"] = statistics.mean(scores) # filter the batch. 1/oversample_factor samples will be kept. - # If there is a filter, prompts passing it will be prioritized. + # If there is a filter, prompts passing it will be + # prioritized. batch = self.filter_and_downsample(scores, batch) batch.meta_info["n"] = self.config.actor_rollout_ref.rollout.n @@ -433,12 +515,21 @@ def fit(self): # recompute old_log_probs with simple_timer("old_log_prob", timing_raw): - old_log_prob = self.actor_rollout_wg.compute_log_prob(batch) + old_log_prob = self.actor_rollout_wg.compute_log_prob( + batch) entropys = old_log_prob.batch["entropys"] response_masks = compute_response_mask(batch) - loss_agg_mode = self.config.actor_rollout_ref.actor.loss_agg_mode - entropy_agg = agg_loss(loss_mat=entropys, loss_mask=response_masks, loss_agg_mode=loss_agg_mode) - old_log_prob_metrics = {"actor/entropy": entropy_agg.detach().item()} + loss_agg_mode = ( + self.config.actor_rollout_ref.actor.loss_agg_mode + ) + entropy_agg = agg_loss( + loss_mat=entropys, + loss_mask=response_masks, + loss_agg_mode=loss_agg_mode, + ) + old_log_prob_metrics = { + "actor/entropy": entropy_agg.detach().item() + } metrics.update(old_log_prob_metrics) old_log_prob.batch.pop("entropys") batch = batch.union(old_log_prob) @@ -446,27 +537,39 @@ def fit(self): if self.use_reference_policy: # compute reference log_prob with simple_timer("ref", timing_raw): - ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch) + ref_log_prob = self.ref_policy_wg.compute_ref_log_prob( + batch) batch = batch.union(ref_log_prob) with simple_timer("adv", timing_raw): if self.use_rm: - update_style = self.config.reward_model.model.get("update", "none") + update_style = self.config.reward_model.model.get( + "update", "none" + ) if update_style == "none": # only run forward - reward_output = self.rm_wg.compute_rm_score(batch) - elif update_style == "after": # update and directly return the reward + reward_output = self.rm_wg.compute_rm_score( + batch) + elif ( + update_style == "after" + ): # update and directly return the reward reward_output = self.rm_wg.update_rm(batch) - elif update_style == "before": # update reward model, and then run forward + elif ( + update_style == "before" + ): # update reward model, and then run forward reward_output = self.rm_wg.update_rm(batch) if "metrics" in reward_output.meta_info.keys(): - reward_output_metrics = reduce_metrics(reward_output.meta_info["metrics"]) + reward_output_metrics = reduce_metrics( + reward_output.meta_info["metrics"] + ) metrics.update(reward_output_metrics) - reward_output = self.rm_wg.compute_rm_score(batch) + reward_output = self.rm_wg.compute_rm_score( + batch) elif ( update_style == "reverse" ): # run forward to calculate statistics, then update reward model - reward_output = self.rm_wg.compute_rm_score(batch) + reward_output = self.rm_wg.compute_rm_score( + batch) # broadcast q and acc tensor to each result bc_td = DataProto.from_dict( tensors={ @@ -489,18 +592,25 @@ def fit(self): raise NotImplementedError batch = batch.union(reward_output) if "metrics" in reward_output.meta_info.keys(): - reward_output_metrics = reduce_metrics(reward_output.meta_info["metrics"]) + reward_output_metrics = reduce_metrics( + reward_output.meta_info["metrics"] + ) metrics.update(reward_output_metrics) # compute advantages, executed on the driver process batch = compute_advantage( - batch, adv_estimator=self.config.algorithm.adv_estimator, config=self.config + batch, + adv_estimator=self.config.algorithm.adv_estimator, + config=self.config, ) # update actor with simple_timer("update_actor", timing_raw): - actor_output = self.actor_rollout_wg.update_actor(batch) - actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) + actor_output = self.actor_rollout_wg.update_actor( + batch) + actor_output_metrics = reduce_metrics( + actor_output.meta_info["metrics"] + ) metrics.update(actor_output_metrics) # validate @@ -513,13 +623,19 @@ def fit(self): val_metrics: dict = self._validate() metrics.update(val_metrics) - if self.config.trainer.save_freq > 0 and self.global_steps % self.config.trainer.save_freq == 0: + if (self.config.trainer.save_freq > 0 and self.global_steps % + self.config.trainer.save_freq == 0): with simple_timer("save_checkpoint", timing_raw): self._save_checkpoint() # collect metrics - metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic)) - metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) + metrics.update( + compute_data_metrics( + batch=batch, + use_critic=self.use_critic)) + metrics.update( + compute_timing_metrics(batch=batch, timing_raw=timing_raw) + ) # TODO: make a canonical logger that supports various backend logger.log(data=metrics, step=self.global_steps) @@ -532,10 +648,8 @@ def fit(self): val_metrics = self._validate() pprint(f"Final validation metrics: {val_metrics}") logger.log(data=val_metrics, step=self.global_steps) - if ( - self.config.trainer.save_freq > 0 - and (self.global_steps - 1) % self.config.trainer.save_freq != 0 - ): + if (self.config.trainer.save_freq > 0 and ( + self.global_steps - 1) % self.config.trainer.save_freq != 0): with simple_timer("save_checkpoint", timing_raw): self._save_checkpoint() return @@ -559,15 +673,19 @@ def filter_and_downsample(self, scores, batch: DataProto): if self.config.data.filter_truncate: length_matrix = ( - batch.batch["attention_mask"][:, -batch.batch["responses"].shape[-1] :] + batch.batch["attention_mask"][:, -batch.batch["responses"].shape[-1]:] .sum(dim=-1) .reshape(-1, n_samples) ) length_tensor = torch.max(length_matrix, dim=-1)[0] - filter_mask[length_tensor >= self.config.data.max_response_length - 1] = False + filter_mask[length_tensor >= + self.config.data.max_response_length - 1] = (False) reorder_index = torch.argsort(filter_mask, descending=True) - reorder_index = (reorder_index.unsqueeze(-1) * n_samples + torch.arange(0, n_samples).unsqueeze(0)).view(-1) + reorder_index = ( + reorder_index.unsqueeze(-1) * n_samples + + torch.arange(0, n_samples).unsqueeze(0) + ).view(-1) batch.reorder( reorder_index[: int(len(batch) // self.config.data.oversample_factor)] ) # this operation is inplace diff --git a/Agent0/executor_train/verl/recipe/r1/__init__.py b/Agent0/executor_train/verl/recipe/r1/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/recipe/r1/__init__.py +++ b/Agent0/executor_train/verl/recipe/r1/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/recipe/r1/data_process.py b/Agent0/executor_train/verl/recipe/r1/data_process.py index fb41c81..daacc80 100644 --- a/Agent0/executor_train/verl/recipe/r1/data_process.py +++ b/Agent0/executor_train/verl/recipe/r1/data_process.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -44,9 +44,15 @@ def process_aime2024(example): print(f"Loading the {data_source} dataset from huggingface...", flush=True) dataset = load_dataset(data_source, split="train") map_fn = partial( - example_map_fn, process_fn=process_aime2024, data_source=data_source, ability="English", split="test" + example_map_fn, + process_fn=process_aime2024, + data_source=data_source, + ability="English", + split="test", + ) + dataset = dataset.map( + map_fn, with_indices=True, remove_columns=dataset.column_names ) - dataset = dataset.map(map_fn, with_indices=True, remove_columns=dataset.column_names) return dataset @@ -56,16 +62,23 @@ def build_gpqa_dimond_dataset(): GPQA_QUERY_TEMPLATE = ( "Answer the following multiple choice question. The last line of your response should be of the following " "format: 'Answer: $LETTER' (without quotes) where LETTER is one of ABCD. Think step by step before " - "answering.\n\n{Question}\n\nA) {A}\nB) {B}\nC) {C}\nD) {D}" - ) + "answering.\n\n{Question}\n\nA) {A}\nB) {B}\nC) {C}\nD) {D}") def process_gpqa_diamond(example): - choices = [example["Incorrect Answer 1"], example["Incorrect Answer 2"], example["Incorrect Answer 3"]] + choices = [ + example["Incorrect Answer 1"], + example["Incorrect Answer 2"], + example["Incorrect Answer 3"], + ] random.shuffle(choices) gold_index = random.randint(0, 3) choices.insert(gold_index, example["Correct Answer"]) query_prompt = GPQA_QUERY_TEMPLATE.format( - A=choices[0], B=choices[1], C=choices[2], D=choices[3], Question=example["Question"] + A=choices[0], + B=choices[1], + C=choices[2], + D=choices[3], + Question=example["Question"], ) gold_choice = "ABCD"[gold_index] return query_prompt, gold_choice @@ -75,9 +88,15 @@ def process_gpqa_diamond(example): dataset = load_dataset(data_source, "gpqa_diamond", split="train") map_fn = partial( - example_map_fn, process_fn=process_gpqa_diamond, data_source=data_source, ability="Math", split="test" + example_map_fn, + process_fn=process_gpqa_diamond, + data_source=data_source, + ability="Math", + split="test", + ) + dataset = dataset.map( + map_fn, with_indices=True, remove_columns=dataset.column_names ) - dataset = dataset.map(map_fn, with_indices=True, remove_columns=dataset.column_names) return dataset @@ -90,15 +109,27 @@ def process_cnmo2024(example): dataset_en = load_dataset(data_source, "v202412_CNMO_en", split="test") map_fn_en = partial( - example_map_fn, process_fn=process_cnmo2024, data_source="opencompass/cnmo2024_en", ability="Math", split="test" + example_map_fn, + process_fn=process_cnmo2024, + data_source="opencompass/cnmo2024_en", + ability="Math", + split="test", + ) + dataset_en = dataset_en.map( + map_fn_en, with_indices=True, remove_columns=dataset_en.column_names ) - dataset_en = dataset_en.map(map_fn_en, with_indices=True, remove_columns=dataset_en.column_names) dataset_zh = load_dataset(data_source, "v202412_CNMO_cn", split="test") map_fn_zh = partial( - example_map_fn, process_fn=process_cnmo2024, data_source="opencompass/cnmo2024_zh", ability="Math", split="test" + example_map_fn, + process_fn=process_cnmo2024, + data_source="opencompass/cnmo2024_zh", + ability="Math", + split="test", + ) + dataset_zh = dataset_zh.map( + map_fn_zh, with_indices=True, remove_columns=dataset_zh.column_names ) - dataset_zh = dataset_zh.map(map_fn_zh, with_indices=True, remove_columns=dataset_zh.column_names) dataset = concatenate_datasets([dataset_en, dataset_zh]) return dataset @@ -112,7 +143,8 @@ def build_livecodebench_dataset(): def process_livecodebench(example): # Construct Query Prompt - # From https://github.com/LiveCodeBench/LiveCodeBench/blob/998c52d394b836f15fff3b9a29866191108ff81b/lcb_runner/prompts/code_generation.py#L140 + # From + # https://github.com/LiveCodeBench/LiveCodeBench/blob/998c52d394b836f15fff3b9a29866191108ff81b/lcb_runner/prompts/code_generation.py#L140 query_prompt = ( f"You will be given a question (problem specification) and will generate a correct Python program " f"that matches the specification and passes all tests.\n\nQuestion: {example['question_content']}\n\n" @@ -127,8 +159,7 @@ def process_livecodebench(example): "Read the inputs from stdin solve the problem and write the answer to stdout (do not directly test " "on the sample inputs). Enclose your code within delimiters as follows. Ensure that when the python " "program runs, it reads the inputs, runs the algorithm and writes output to STDOUT." - "```python\n# YOUR CODE HERE\n```" - ) + "```python\n# YOUR CODE HERE\n```") # Construct test cases public_test_cases = json.loads(example["public_test_cases"]) @@ -137,8 +168,10 @@ def process_livecodebench(example): except Exception as e: print(f"Error loading private test cases: {e}") private_test_cases = json.loads( - pickle.loads(zlib.decompress(base64.b64decode(example["private_test_cases"].encode("utf-8")))) - ) + pickle.loads( + zlib.decompress( + base64.b64decode( + example["private_test_cases"].encode("utf-8"))))) full_test_cases = public_test_cases + private_test_cases metadata = json.loads(example["metadata"]) @@ -147,19 +180,33 @@ def process_livecodebench(example): "outputs": [t["output"] for t in full_test_cases], "fn_name": metadata.get("func_name", None), } - text_cases_compressed = base64.b64encode(zlib.compress(pickle.dumps(json.dumps(test_cases)))).decode("utf-8") + text_cases_compressed = base64.b64encode( + zlib.compress(pickle.dumps(json.dumps(test_cases))) + ).decode("utf-8") return query_prompt, text_cases_compressed data_source = "livecodebench/code_generation_lite" print(f"Loading the {data_source} dataset from huggingface...", flush=True) dataset = load_dataset(data_source, split="test") # R1 Evaluation use LiveCodeBench 24.08-25.01 - dataset = dataset.filter(lambda line: "2024-08-00T00:00:00" <= line["contest_date"] < "2025-01-00T00:00:00") + dataset = dataset.filter( + lambda line: "2024-08-00T00:00:00" + <= line["contest_date"] + < "2025-01-00T00:00:00" + ) map_fn = partial( - example_map_fn, process_fn=process_livecodebench, data_source=data_source, ability="Code", split="test" + example_map_fn, + process_fn=process_livecodebench, + data_source=data_source, + ability="Code", + split="test", ) - dataset = dataset.map(map_fn, with_indices=True, remove_columns=dataset.column_names, num_proc=8) + dataset = dataset.map( + map_fn, + with_indices=True, + remove_columns=dataset.column_names, + num_proc=8) return dataset @@ -182,7 +229,8 @@ def process_livecodebench(example): if args.tasks.lower() == "all": args.tasks = SUPPORTED_TASKS else: - args.tasks = [task.strip() for task in args.tasks.split(",") if task.strip()] + args.tasks = [task.strip() + for task in args.tasks.split(",") if task.strip()] for task in args.tasks: if task not in SUPPORTED_TASKS: raise NotImplementedError(f"{task} has not been supported.") diff --git a/Agent0/executor_train/verl/recipe/r1/main_eval.py b/Agent0/executor_train/verl/recipe/r1/main_eval.py index b9c0379..ebcff27 100644 --- a/Agent0/executor_train/verl/recipe/r1/main_eval.py +++ b/Agent0/executor_train/verl/recipe/r1/main_eval.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -56,7 +56,8 @@ def main(config): # Create remote tasks remote_tasks = [ - process_item.remote(config, data_sources[i], responses[i], reward_model_data[i]) for i in range(total) + process_item.remote(config, data_sources[i], responses[i], reward_model_data[i]) + for i in range(total) ] # Process results as they come in diff --git a/Agent0/executor_train/verl/recipe/r1/reward_score.py b/Agent0/executor_train/verl/recipe/r1/reward_score.py index 2010665..a35e0a3 100644 --- a/Agent0/executor_train/verl/recipe/r1/reward_score.py +++ b/Agent0/executor_train/verl/recipe/r1/reward_score.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,7 +14,11 @@ def reward_func(data_source, solution_str, ground_truth, extra_info=None): - if data_source in ["Maxwell-Jia/AIME_2024", "opencompass/cnmo2024_en", "opencompass/cnmo2024_zh"]: + if data_source in [ + "Maxwell-Jia/AIME_2024", + "opencompass/cnmo2024_en", + "opencompass/cnmo2024_zh", + ]: from recipe.r1.tasks import math return math.compute_score(solution_str, ground_truth) @@ -22,7 +26,10 @@ def reward_func(data_source, solution_str, ground_truth, extra_info=None): from recipe.r1.tasks import gpqa return gpqa.compute_score(solution_str, ground_truth) - elif data_source in ["livecodebench/code_generation_lite", "livecodebench/code_generation"]: + elif data_source in [ + "livecodebench/code_generation_lite", + "livecodebench/code_generation", + ]: from recipe.r1.tasks import livecodebench return livecodebench.compute_score(solution_str, ground_truth) diff --git a/Agent0/executor_train/verl/recipe/r1/tasks/__init__.py b/Agent0/executor_train/verl/recipe/r1/tasks/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/recipe/r1/tasks/__init__.py +++ b/Agent0/executor_train/verl/recipe/r1/tasks/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/recipe/r1/tasks/gpqa.py b/Agent0/executor_train/verl/recipe/r1/tasks/gpqa.py index 65b37e9..f6dc206 100644 --- a/Agent0/executor_train/verl/recipe/r1/tasks/gpqa.py +++ b/Agent0/executor_train/verl/recipe/r1/tasks/gpqa.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,7 +14,8 @@ import re -# Extraction Template from https://github.com/openai/simple-evals/blob/90e3e821cabba2aeb6be651dcb662b253df04225/common.py#L25 +# Extraction Template from +# https://github.com/openai/simple-evals/blob/90e3e821cabba2aeb6be651dcb662b253df04225/common.py#L25 ANSWER_PATTERN_MULTICHOICE = r"(?i)Answer[ \t]*:[ \t]*\$?([A-D])\$?" diff --git a/Agent0/executor_train/verl/recipe/r1/tasks/livecodebench.py b/Agent0/executor_train/verl/recipe/r1/tasks/livecodebench.py index f0cbab6..e40c611 100644 --- a/Agent0/executor_train/verl/recipe/r1/tasks/livecodebench.py +++ b/Agent0/executor_train/verl/recipe/r1/tasks/livecodebench.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,7 +23,8 @@ def _temp_run(in_outs, generation, debug, result, metadata_list, timeout): - res, metadata = run_test(in_outs, test=generation, debug=debug, timeout=timeout) + res, metadata = run_test(in_outs, test=generation, + debug=debug, timeout=timeout) result.append(res) metadata_list.append(metadata) @@ -60,11 +61,17 @@ def compute_score(completion, test_cases): in_outs = json.loads(test_cases) except Exception as e: print(f"Error loading test cases: {e}") - in_outs = json.loads(pickle.loads(zlib.decompress(base64.b64decode(test_cases.encode("utf-8"))))) + in_outs = json.loads( + pickle.loads( + zlib.decompress( + base64.b64decode( + test_cases.encode("utf-8"))))) success = False try: - res, metadata = check_correctness(in_outs=in_outs, generation=solution, timeout=6, debug=False) + res, metadata = check_correctness( + in_outs=in_outs, generation=solution, timeout=6, debug=False + ) success = all(map(lambda x: x is True, res)) except Exception: pass diff --git a/Agent0/executor_train/verl/recipe/r1/tasks/math.py b/Agent0/executor_train/verl/recipe/r1/tasks/math.py index 5ecde54..e27ede9 100644 --- a/Agent0/executor_train/verl/recipe/r1/tasks/math.py +++ b/Agent0/executor_train/verl/recipe/r1/tasks/math.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,14 +17,16 @@ from math_verify.metric import math_metric from math_verify.parser import ExprExtractionConfig, LatexExtractionConfig except ImportError: - print("To use Math-Verify, please install it first by running `pip install math-verify`.") + print( + "To use Math-Verify, please install it first by running `pip install math-verify`." + ) def compute_score(model_output: str, ground_truth: str) -> bool: verify_func = math_metric( - gold_extraction_target=(LatexExtractionConfig(),), - pred_extraction_target=(ExprExtractionConfig(), LatexExtractionConfig()), - ) + gold_extraction_target=( + LatexExtractionConfig(),), pred_extraction_target=( + ExprExtractionConfig(), LatexExtractionConfig()), ) ret_score = 0.0 # Wrap the ground truth in \boxed{} format for verification diff --git a/Agent0/executor_train/verl/recipe/retool/retool.py b/Agent0/executor_train/verl/recipe/retool/retool.py index b4d6028..ed60605 100644 --- a/Agent0/executor_train/verl/recipe/retool/retool.py +++ b/Agent0/executor_train/verl/recipe/retool/retool.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,13 +32,16 @@ def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): self.code_pattern = re.compile(r"```python(.*?)```", re.DOTALL) @rollout_trace_op - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: code = parameters["code"] matches = self.code_pattern.findall(code) if matches: code = matches[0].strip() - # NOTE: some script may not explicitly print result, we need to add a print statement to the end of the script + # NOTE: some script may not explicitly print result, we need to add a + # print statement to the end of the script lines = code.split("\n") for i, line in reversed(list(enumerate(lines))): if line == "": @@ -53,12 +56,16 @@ async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) if not isinstance(code, str): code = str(code) - result = await self.execution_pool.execute.remote(self.execute_code, instance_id, code, timeout, language) + result = await self.execution_pool.execute.remote( + self.execute_code, instance_id, code, timeout, language + ) # sandbox has no score or metrics, use Nones return result, None, None -answer_format = """\nThe answer format must be: \\boxed{'The final answer goes here.'}""" +answer_format = ( + """\nThe answer format must be: \\boxed{'The final answer goes here.'}""" +) class CustomRLHFDataset(RLHFDataset): @@ -70,14 +77,19 @@ def _read_files_and_tokenize(self): # read parquet files and cache dataframe = datasets.load_dataset(parquet_file)["train"] data_source = "/".join(parquet_file.split("/")[-2:]) - if data_source in ["Maxwell-Jia/AIME_2024", "yentinglin/aime_2025"]: + if data_source in [ + "Maxwell-Jia/AIME_2024", + "yentinglin/aime_2025"]: dataframe = dataframe.map( - self.map_fn, fn_kwargs={"data_source": data_source}, remove_columns=dataframe.column_names + self.map_fn, + fn_kwargs={"data_source": data_source}, + remove_columns=dataframe.column_names, ) else: dataframe = dataframe.map(self.map_fn2, num_proc=16) dataframes.append(dataframe) - self.dataframe: datasets.Dataset = datasets.concatenate_datasets(dataframes) + self.dataframe: datasets.Dataset = datasets.concatenate_datasets( + dataframes) print(f"dataset len: {len(self.dataframe)}") @@ -89,7 +101,8 @@ def map_fn(self, row: dict, *, data_source: str = None): prompt = problem + answer_format data = { - "data_source": data_source.split("/")[1].lower(), # aime_2024, aime_2025 + # aime_2024, aime_2025 + "data_source": data_source.split("/")[1].lower(), "prompt": [{"role": "user", "content": prompt}], "ability": "MATH", "reward_model": {"ground_truth": str(answer)}, @@ -106,7 +119,8 @@ def map_fn2(self, row: dict): def compute_score(data_source, solution_str, ground_truth, extra_info): # use \\boxed{...} answer - result = math_dapo.compute_score(solution_str, ground_truth, strict_box_verify=True) + result = math_dapo.compute_score( + solution_str, ground_truth, strict_box_verify=True) # encourage model to call tools num_turns = extra_info["num_turns"] diff --git a/Agent0/executor_train/verl/recipe/retool/retool_multi_turn_sft_preprocess.py b/Agent0/executor_train/verl/recipe/retool/retool_multi_turn_sft_preprocess.py index 201ee68..1e63f6d 100644 --- a/Agent0/executor_train/verl/recipe/retool/retool_multi_turn_sft_preprocess.py +++ b/Agent0/executor_train/verl/recipe/retool/retool_multi_turn_sft_preprocess.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -37,7 +37,9 @@ def main(): shuffled_train_dataset = train_dataset.shuffle(seed=args.seed) split_idx = int(len(shuffled_train_dataset) * args.train_ratio) train_dataset = shuffled_train_dataset.select(range(split_idx)) - test_dataset = shuffled_train_dataset.select(range(split_idx, len(shuffled_train_dataset))) + test_dataset = shuffled_train_dataset.select( + range(split_idx, len(shuffled_train_dataset)) + ) # add a row to each data item that represents a unique id def make_map_fn(split): @@ -58,8 +60,12 @@ def process_fn(example, idx): return process_fn - train_dataset = train_dataset.map(function=make_map_fn("train"), with_indices=True) - test_dataset = test_dataset.map(function=make_map_fn("test"), with_indices=True) + train_dataset = train_dataset.map( + function=make_map_fn("train"), + with_indices=True) + test_dataset = test_dataset.map( + function=make_map_fn("test"), + with_indices=True) # Create output directory local_dir = os.path.expanduser(args.local_dir) diff --git a/Agent0/executor_train/verl/recipe/retool/retool_sft_preprocess.py b/Agent0/executor_train/verl/recipe/retool/retool_sft_preprocess.py index 0a46c15..0677f1a 100644 --- a/Agent0/executor_train/verl/recipe/retool/retool_sft_preprocess.py +++ b/Agent0/executor_train/verl/recipe/retool/retool_sft_preprocess.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,7 +33,7 @@ def extract_code_message(content: str) -> tuple[dict[str, Any], str]: j = content.find(stop) assert j > i - code = content[i + len(start) : j] + code = content[i + len(start): j] matches = code_pattern.findall(code) if matches: code = matches[0].strip() @@ -51,7 +51,7 @@ def extract_code_message(content: str) -> tuple[dict[str, Any], str]: }, ], } - return message, content[j + len(stop) :] + return message, content[j + len(stop):] def extract_answer_message(content: str) -> tuple[dict[str, Any], str]: @@ -62,12 +62,12 @@ def extract_answer_message(content: str) -> tuple[dict[str, Any], str]: j = content.find(stop) assert j > i - answer = content[:i] + content[i + len(start) : j] + answer = content[:i] + content[i + len(start): j] message = { "role": "assistant", "content": answer.strip(), } - return message, content[j + len(stop) :] + return message, content[j + len(stop):] def extract_interpreter_message(content: str) -> tuple[dict[str, Any], str]: @@ -78,12 +78,12 @@ def extract_interpreter_message(content: str) -> tuple[dict[str, Any], str]: j = content.find(stop) assert j > i - interpreter = content[i + len(start) : j] + interpreter = content[i + len(start): j] message = { "role": "tool", "content": interpreter.strip(), } - return message, content[j + len(stop) :] + return message, content[j + len(stop):] def process(row: dict, *, tools: str): @@ -94,7 +94,12 @@ def process(row: dict, *, tools: str): start = "*user question:*" i = content.find(start) assert i != -1 - prompt = content[i + len(start) :].replace("", "").replace("", "").strip() + prompt = ( + content[i + len(start):] + .replace("", "") + .replace("", "") + .strip() + ) messages.append( { "role": "user", @@ -125,7 +130,8 @@ def process(row: dict, *, tools: str): if __name__ == "__main__": tools_config_file = "recipe/retool/sandbox_fusion_tool_config.yaml" tools_config = OmegaConf.load(tools_config_file) - tool_schema = OmegaConf.to_container(tools_config["tools"][0]["tool_schema"]) + tool_schema = OmegaConf.to_container( + tools_config["tools"][0]["tool_schema"]) tools = json.dumps([tool_schema]) data = datasets.load_dataset("JoeYing/ReTool-SFT")["train"] diff --git a/Agent0/executor_train/verl/recipe/spin/core_algos.py b/Agent0/executor_train/verl/recipe/spin/core_algos.py index c48027e..97593e9 100644 --- a/Agent0/executor_train/verl/recipe/spin/core_algos.py +++ b/Agent0/executor_train/verl/recipe/spin/core_algos.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -50,8 +50,14 @@ def get_kl_controller(kl_ctrl): if kl_ctrl.type == "fixed": return FixedKLController(kl_coef=kl_ctrl.kl_coef) elif kl_ctrl.type == "adaptive": - assert kl_ctrl.horizon > 0, f"horizon must be larger than 0. Got {kl_ctrl.horizon}" - return AdaptiveKLController(init_kl_coef=kl_ctrl.kl_coef, target_kl=kl_ctrl.target_kl, horizon=kl_ctrl.horizon) + assert ( + kl_ctrl.horizon > 0 + ), f"horizon must be larger than 0. Got {kl_ctrl.horizon}" + return AdaptiveKLController( + init_kl_coef=kl_ctrl.kl_coef, + target_kl=kl_ctrl.target_kl, + horizon=kl_ctrl.horizon, + ) else: raise NotImplementedError @@ -79,15 +85,19 @@ def compute_onlinedpo_pref( # print(f"---- [DEBUG] Inside compute_onlinedpo_pref ----") if token_level_rewards.shape[0] % 2 != 0 or response_mask.shape[0] % 2 != 0: raise ValueError( - f"Input tensor batch dimension must be even for pair comparison, got shapes: " - f"{token_level_rewards.shape}, {response_mask.shape}" - ) + f"Input tensor batch dimension must be even for pair comparison, got shapes: " f"{ + token_level_rewards.shape}, { + response_mask.shape}") if token_level_rewards.shape != response_mask.shape: - raise ValueError(f"Shape mismatch between rewards {token_level_rewards.shape} and mask {response_mask.shape}") + raise ValueError( + f"Shape mismatch between rewards { + token_level_rewards.shape} and mask { + response_mask.shape}") # 1. Calculate Sequence Scores scores = (token_level_rewards * response_mask).sum(dim=-1) - # print(f" Calculated sequence scores shape: {scores.shape}") # [batch_size * 2] + # print(f" Calculated sequence scores shape: {scores.shape}") # + # [batch_size * 2] # 2. Reshape scores to group pairs: [batch_size, 2] try: @@ -95,15 +105,20 @@ def compute_onlinedpo_pref( except RuntimeError as e: print(f"ERROR reshaping scores (shape {scores.shape}) into pairs: {e}") raise e - print(f" Reshaped score pairs shape: {score_pairs.shape}") # [batch_size, 2] + print( + f" Reshaped score pairs shape: { + score_pairs.shape}") # [batch_size, 2] # 3. Compare scores to find which index (0 or 1) is the winner within each pair # winner_indices[i] = 0 if score_pairs[i, 0] >= score_pairs[i, 1] else 1 - winner_indices = torch.argmax(score_pairs, dim=1) # 0 if first is max, 1 if second is max + winner_indices = torch.argmax( + score_pairs, dim=1 + ) # 0 if first is max, 1 if second is max # Handle ties explicitly if argmax behavior isn't guaranteed (usually picks first max) # Alternatively: winner_mask_original = score_pairs[:, 0] >= score_pairs[:, 1] # print(f" Winner indices shape: {winner_indices.shape}") # [batch_size] - # print(f" Number where Response 2 (index 1) is preferred: {winner_indices.sum().item()}") # Counts number of 1s + # print(f" Number where Response 2 (index 1) is preferred: + # {winner_indices.sum().item()}") # Counts number of 1s # 4. Create the final [batch_size * 2] mask num_pairs = score_pairs.shape[0] @@ -112,12 +127,15 @@ def compute_onlinedpo_pref( # full_indices = torch.arange(full_batch_size, device=scores.device) # Create indices corresponding to the winner within each pair's original index # E.g., if winner_indices is [0, 1, 0], pair_indices is [0, 1, 2] - # winner_global_indices = (pair_indices * 2) + winner_indices -> [ (0*2)+0, (1*2)+1, (2*2)+0 ] -> [0, 3, 4] + # winner_global_indices = (pair_indices * 2) + winner_indices -> [ + # (0*2)+0, (1*2)+1, (2*2)+0 ] -> [0, 3, 4] pair_indices = torch.arange(num_pairs, device=scores.device) winner_global_indices = (pair_indices * 2) + winner_indices # Create boolean mask - True at the winner's position - output_preference_mask = torch.zeros(full_batch_size, dtype=torch.bool, device=scores.device) + output_preference_mask = torch.zeros( + full_batch_size, dtype=torch.bool, device=scores.device + ) output_preference_mask[winner_global_indices] = True # print(f" Output preference mask shape: {output_preference_mask.shape}") # Should be [batch_size * 2] @@ -149,18 +167,23 @@ def compute_online_dpo_loss( logits = pi_logratios - ref_logratios if loss_type == "sigmoid": - losses = -F.logsigmoid(beta * logits) * (1 - label_smoothing) - F.logsigmoid(-beta * logits) * label_smoothing + losses = ( + -F.logsigmoid(beta * logits) * (1 - label_smoothing) + - F.logsigmoid(-beta * logits) * label_smoothing + ) elif loss_type == "ipo": losses = (logits - 1 / (2 * beta)) ** 2 else: - raise ValueError(f"Unsupported loss_type: {loss_type}. Choose 'sigmoid', 'ipo', or 'hinge'.") + raise ValueError( + f"Unsupported loss_type: {loss_type}. Choose 'sigmoid', 'ipo', or 'hinge'.") return losses.mean() def get_batch_logps( - logits: torch.FloatTensor, labels: torch.LongTensor, average_log_prob: bool = False -) -> torch.FloatTensor: + logits: torch.FloatTensor, + labels: torch.LongTensor, + average_log_prob: bool = False) -> torch.FloatTensor: """ Compute the log probabilities of the given labels under the given logits. @@ -184,7 +207,9 @@ def get_batch_logps( # Calculate per token log probability loss_fct = torch.nn.CrossEntropyLoss(ignore_index=-100, reduction="none") - per_token_logps = -loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) + per_token_logps = -loss_fct( + shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1) + ) per_token_logps = per_token_logps.view( shift_logits.size(0), shift_logits.size(1) ) # Reshape back to (batch_size, seq_len-1) diff --git a/Agent0/executor_train/verl/recipe/spin/dp_actor.py b/Agent0/executor_train/verl/recipe/spin/dp_actor.py index 35caa29..72d3b6e 100644 --- a/Agent0/executor_train/verl/recipe/spin/dp_actor.py +++ b/Agent0/executor_train/verl/recipe/spin/dp_actor.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -53,38 +53,56 @@ def compute_log_prob(self, data: DataProto) -> torch.Tensor: 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 silent error + temperature = data.meta_info[ + "temperature" + ] # temperature must be in the data.meta_info to avoid silent error use_dynamic_bsz = data.meta_info["use_dynamic_bsz"] - select_keys = ["responses", "input_ids", "attention_mask", "position_ids"] + select_keys = [ + "responses", + "input_ids", + "attention_mask", + "position_ids"] batch = data.select(batch_keys=select_keys).batch has_multi_modal_inputs = "multi_modal_inputs" in data.non_tensor_batch.keys() if has_multi_modal_inputs: num_micro_batches = data.batch.batch_size[0] // micro_batch_size non_tensor_select_keys = ["multi_modal_inputs"] - micro_batches = data.select(select_keys, non_tensor_select_keys).chunk(num_micro_batches) + micro_batches = data.select( + select_keys, non_tensor_select_keys).chunk(num_micro_batches) elif 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) + 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: if isinstance(micro_batch, DataProto): - micro_batch = {**micro_batch.batch, **micro_batch.non_tensor_batch} + micro_batch = { + **micro_batch.batch, + **micro_batch.non_tensor_batch} with torch.no_grad(): - _, log_probs = self._forward_micro_batch(micro_batch, temperature=temperature) + _, 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) + 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 @@ -98,15 +116,20 @@ def update_policy_dpo_with_ref(self, data: DataProto): # --- Retrieve necessary data --- try: - # Expects batch prepared by fit_dpo loop, including reference log probs + # Expects batch prepared by fit_dpo loop, including reference log + # probs batch_td = data.batch chosen_labels = batch_td["chosen_labels"] rejected_labels = batch_td["rejected_labels"] # ... other needed tensors like chosen/rejected input_ids, attention_mask, position_ids ... # === Get PRE-CALCULATED reference log probs from input data === - reference_chosen_logps = batch_td["reference_chosen_logps"] # Should be sequence-level logps - reference_rejected_logps = batch_td["reference_rejected_logps"] # Should be sequence-level logps + reference_chosen_logps = batch_td[ + "reference_chosen_logps" + ] # Should be sequence-level logps + reference_rejected_logps = batch_td[ + "reference_rejected_logps" + ] # Should be sequence-level logps # ============================================================ # Get DPO params from meta_info @@ -115,14 +138,20 @@ def update_policy_dpo_with_ref(self, data: DataProto): loss_type = data.meta_info.get("dpo_loss_type", "sigmoid") label_smoothing = data.meta_info.get("dpo_label_smoothing", 0.0) # reference_free should now be False as we provide ref logps - reference_free = data.meta_info.get("reference_free", False) # Default False + reference_free = data.meta_info.get( + "reference_free", False + ) # Default False except KeyError as e: - print(f"ERROR: Missing required key for DPO update (in update_policy_dpo): {e}") - print(f"Available keys in data.batch: {list(batch_td.keys())}") # Debug print + print( + f"ERROR: Missing required key for DPO update (in update_policy_dpo): {e}") + print( + f"Available keys in data.batch: {list(batch_td.keys())}" + ) # Debug print return {} # Return empty metrics on error except Exception as e_data: - print(f"ERROR accessing data for DPO update (in update_policy_dpo): {e_data}") + print( + f"ERROR accessing data for DPO update (in update_policy_dpo): {e_data}") return {} # --- Micro-batching Setup --- @@ -130,7 +159,9 @@ def update_policy_dpo_with_ref(self, data: DataProto): if micro_batch_size is None: # Fallback or default if not set, or raise error micro_batch_size = 1 # Example fallback, adjust as needed - print(f"Warning: 'ppo_micro_batch_size_per_gpu' not set, defaulting to {micro_batch_size}") + print( + f"Warning: 'ppo_micro_batch_size_per_gpu' not set, defaulting to {micro_batch_size}" + ) # raise ValueError("Config 'ppo_micro_batch_size_per_gpu' must be set.") # Ensure chosen_input_ids exists before getting shape @@ -141,7 +172,10 @@ def update_policy_dpo_with_ref(self, data: DataProto): if bsz == 0: print("Warning: DPO batch size is 0 in update_policy_dpo. Skipping update.") - return {"actor/dpo_loss": 0.0, "actor/grad_norm": 0.0} # Return zero metrics if batch is empty + return { + "actor/dpo_loss": 0.0, + "actor/grad_norm": 0.0, + } # Return zero metrics if batch is empty num_micro_batches = math.ceil(bsz / micro_batch_size) gradient_accumulation_steps = num_micro_batches @@ -162,7 +196,8 @@ def update_policy_dpo_with_ref(self, data: DataProto): continue # Slice the full DPO batch into micro-batches - # Important: Slice ALL required tensors, including labels and inputs + # Important: Slice ALL required tensors, including labels and + # inputs micro_batch_chosen_labels = chosen_labels[start_idx:end_idx] micro_batch_rejected_labels = rejected_labels[start_idx:end_idx] micro_batch_chosen_inputs = { @@ -170,47 +205,68 @@ def update_policy_dpo_with_ref(self, data: DataProto): "attention_mask": batch_td["chosen_attention_mask"][start_idx:end_idx], } if "chosen_position_ids" in batch_td: - micro_batch_chosen_inputs["position_ids"] = batch_td["chosen_position_ids"][start_idx:end_idx] + micro_batch_chosen_inputs["position_ids"] = batch_td[ + "chosen_position_ids" + ][start_idx:end_idx] micro_batch_rejected_inputs = { "input_ids": batch_td["rejected_input_ids"][start_idx:end_idx], - "attention_mask": batch_td["rejected_attention_mask"][start_idx:end_idx], + "attention_mask": batch_td["rejected_attention_mask"][ + start_idx:end_idx + ], } if "rejected_position_ids" in batch_td: - micro_batch_rejected_inputs["position_ids"] = batch_td["rejected_position_ids"][start_idx:end_idx] + micro_batch_rejected_inputs["position_ids"] = batch_td[ + "rejected_position_ids" + ][start_idx:end_idx] # Determine autocast dtype - autocast_dtype = torch.bfloat16 # Or get dynamically from config/FSDP settings + autocast_dtype = ( + torch.bfloat16 + ) # Or get dynamically from config/FSDP settings # --- Autocast Forward Pass --- with torch.autocast(device_type=get_device_name(), dtype=autocast_dtype): # --- Step 1: Forward pass for CURRENT policy log probs (with grad) --- - policy_chosen_outputs = self.actor_module(**micro_batch_chosen_inputs, use_cache=False) - policy_rejected_outputs = self.actor_module(**micro_batch_rejected_inputs, use_cache=False) + policy_chosen_outputs = self.actor_module( + **micro_batch_chosen_inputs, use_cache=False + ) + policy_rejected_outputs = self.actor_module( + **micro_batch_rejected_inputs, use_cache=False + ) # --- Step 2: Calculate CURRENT policy log probs using get_batch_logps --- policy_chosen_logps = get_batch_logps( - policy_chosen_outputs.logits, micro_batch_chosen_labels, average_log_prob=False + policy_chosen_outputs.logits, + micro_batch_chosen_labels, + average_log_prob=False, ) policy_rejected_logps = get_batch_logps( - policy_rejected_outputs.logits, micro_batch_rejected_labels, average_log_prob=False + policy_rejected_outputs.logits, + micro_batch_rejected_labels, + average_log_prob=False, ) # --- Step 3: Retrieve PRE-CALCULATED reference log probs (NO grad needed) --- - # Slice the full batch reference logps for the current micro-batch + # Slice the full batch reference logps for the current + # micro-batch micro_ref_chosen_logps = reference_chosen_logps[start_idx:end_idx] micro_ref_rejected_logps = reference_rejected_logps[start_idx:end_idx] # --- The ActorAsRef calculation block is REMOVED --- # --- Step 4: Calculate DPO Logits and Loss --- pi_logratios = policy_chosen_logps - policy_rejected_logps - ref_logratios = micro_ref_chosen_logps - micro_ref_rejected_logps # Uses pre-calculated values + ref_logratios = ( + micro_ref_chosen_logps - micro_ref_rejected_logps + ) # Uses pre-calculated values logits = pi_logratios - ref_logratios # DPO logits loss = compute_online_dpo_loss( policy_chosen_logps=policy_chosen_logps, # Has grad policy_rejected_logps=policy_rejected_logps, # Has grad - reference_chosen_logps=micro_ref_chosen_logps, # No grad (from input) - reference_rejected_logps=micro_ref_rejected_logps, # No grad (from input) + # No grad (from input) + reference_chosen_logps=micro_ref_chosen_logps, + # No grad (from input) + reference_rejected_logps=micro_ref_rejected_logps, beta=beta, label_smoothing=label_smoothing, loss_type=loss_type, @@ -223,21 +279,28 @@ def update_policy_dpo_with_ref(self, data: DataProto): # --- Accumulate Metrics --- total_loss += loss.item() # Unscaled loss accumulated_metrics["actor/dpo_loss_batch"].append(loss.item()) - accumulated_metrics["actor/dpo_logits_batch"].append(logits.mean().item()) - # Accumulate policy and reference log probs/ratios if needed for debugging - accumulated_metrics["actor/policy_chosen_logps_batch"].append(policy_chosen_logps.mean().item()) - accumulated_metrics["actor/policy_rejected_logps_batch"].append(policy_rejected_logps.mean().item()) - accumulated_metrics["actor/reference_chosen_logps_batch"].append(micro_ref_chosen_logps.mean().item()) - accumulated_metrics["actor/reference_rejected_logps_batch"].append( - micro_ref_rejected_logps.mean().item() + accumulated_metrics["actor/dpo_logits_batch"].append( + logits.mean().item() + ) + # Accumulate policy and reference log probs/ratios if needed + # for debugging + accumulated_metrics["actor/policy_chosen_logps_batch"].append( + policy_chosen_logps.mean().item() ) + accumulated_metrics["actor/policy_rejected_logps_batch"].append( + policy_rejected_logps.mean().item()) + accumulated_metrics["actor/reference_chosen_logps_batch"].append( + micro_ref_chosen_logps.mean().item()) + accumulated_metrics["actor/reference_rejected_logps_batch"].append( + micro_ref_rejected_logps.mean().item()) # --- Backward Pass (outside autocast) --- # Check if loss requires grad before backward if scaled_loss.requires_grad: scaled_loss.backward() else: - print(f"Warning: Scaled loss at micro-batch {i} does not require grad. Skipping backward.") + print( + f"Warning: Scaled loss at micro-batch {i} does not require grad. Skipping backward.") # --- End Micro-batch Loop --- @@ -248,31 +311,46 @@ def update_policy_dpo_with_ref(self, data: DataProto): if num_micro_batches > 0 and bsz > 0: # Check if any processing happened metrics["actor/dpo_loss"] = total_loss / num_micro_batches metrics["actor/grad_norm"] = ( - grad_norm.item() if torch.is_tensor(grad_norm) and torch.isfinite(grad_norm) else float("inf") + grad_norm.item() + if torch.is_tensor(grad_norm) and torch.isfinite(grad_norm) + else float("inf") ) # Average other accumulated metrics for key, val_list in accumulated_metrics.items(): if val_list: metrics[key.replace("_batch", "")] = np.mean(val_list) - # Calculate accuracy / rewards / margins based on averaged logprobs if desired + # Calculate accuracy / rewards / margins based on averaged logprobs + # if desired if ( "actor/policy_chosen_logps" in metrics and "actor/policy_rejected_logps" in metrics and "actor/reference_chosen_logps" in metrics and "actor/reference_rejected_logps" in metrics ): - policy_ratio_mean = metrics["actor/policy_chosen_logps"] - metrics["actor/policy_rejected_logps"] - ref_ratio_mean = metrics["actor/reference_chosen_logps"] - metrics["actor/reference_rejected_logps"] + policy_ratio_mean = ( + metrics["actor/policy_chosen_logps"] + - metrics["actor/policy_rejected_logps"] + ) + ref_ratio_mean = ( + metrics["actor/reference_chosen_logps"] + - metrics["actor/reference_rejected_logps"] + ) logits_mean = policy_ratio_mean - ref_ratio_mean metrics["actor/rewards_chosen"] = beta * ( - metrics["actor/policy_chosen_logps"] - metrics["actor/reference_chosen_logps"] + metrics["actor/policy_chosen_logps"] + - metrics["actor/reference_chosen_logps"] ) metrics["actor/rewards_rejected"] = beta * ( - metrics["actor/policy_rejected_logps"] - metrics["actor/reference_rejected_logps"] + metrics["actor/policy_rejected_logps"] + - metrics["actor/reference_rejected_logps"] ) - metrics["actor/rewards_accuracies"] = float(logits_mean > 0) # Mean accuracy proxy - metrics["actor/rewards_margins"] = metrics["actor/rewards_chosen"] - metrics["actor/rewards_rejected"] + metrics["actor/rewards_accuracies"] = float( + logits_mean > 0 + ) # Mean accuracy proxy + metrics["actor/rewards_margins"] = ( + metrics["actor/rewards_chosen"] - + metrics["actor/rewards_rejected"]) else: # Handle case where no micro-batches were run (e.g., bsz=0) metrics["actor/dpo_loss"] = 0.0 diff --git a/Agent0/executor_train/verl/recipe/spin/fsdp_workers.py b/Agent0/executor_train/verl/recipe/spin/fsdp_workers.py index e8a43e0..a6ca3d8 100644 --- a/Agent0/executor_train/verl/recipe/spin/fsdp_workers.py +++ b/Agent0/executor_train/verl/recipe/spin/fsdp_workers.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,7 +31,12 @@ from verl.single_controller.base.decorator import Dispatch, register from verl.utils import hf_tokenizer from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager -from verl.utils.device import get_device_id, get_device_name, get_nccl_backend, get_torch_device +from verl.utils.device import ( + get_device_id, + get_device_name, + get_nccl_backend, + get_torch_device, +) from verl.utils.flops_counter import FlopsCounter from verl.utils.fs import copy_to_local from verl.utils.fsdp_utils import ( @@ -55,10 +60,14 @@ def create_device_mesh(world_size, fsdp_size): if fsdp_size < 0 or fsdp_size >= world_size: - device_mesh = init_device_mesh(get_device_name(), mesh_shape=(world_size,), mesh_dim_names=["fsdp"]) + device_mesh = init_device_mesh( + get_device_name(), mesh_shape=( + world_size,), mesh_dim_names=["fsdp"]) else: device_mesh = init_device_mesh( - get_device_name(), mesh_shape=(world_size // fsdp_size, fsdp_size), mesh_dim_names=["ddp", "fsdp"] + get_device_name(), + mesh_shape=(world_size // fsdp_size, fsdp_size), + mesh_dim_names=["ddp", "fsdp"], ) return device_mesh @@ -71,21 +80,27 @@ def get_sharding_strategy(device_mesh): elif device_mesh.ndim == 2: sharding_strategy = ShardingStrategy.HYBRID_SHARD else: - raise NotImplementedError(f"Get device mesh ndim={device_mesh.ndim}, but only support 1 or 2") + raise NotImplementedError( + f"Get device mesh ndim={device_mesh.ndim}, but only support 1 or 2" + ) return sharding_strategy class SPINRolloutRefWorker(ActorRolloutRefWorker): @register(dispatch_mode=Dispatch.ONE_TO_ALL) def init_model(self): - from recipe.spin.dp_actor import SPINDataParallelPPOActor as DataParallelPPOActor + from recipe.spin.dp_actor import ( + SPINDataParallelPPOActor as 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())) + 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) use_fused_kernels = self.config.model.get("use_fused_kernels", False) @@ -98,19 +113,28 @@ def init_model(self): 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, - use_fused_kernels=use_fused_kernels, - enable_gradient_checkpointing=self.config.model.get("enable_gradient_checkpointing", False), - trust_remote_code=self.config.model.get("trust_remote_code", False), - use_liger=self.config.model.get("use_liger", False), - role="actor", - ) + ( + 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, + use_fused_kernels=use_fused_kernels, + enable_gradient_checkpointing=self.config.model.get( + "enable_gradient_checkpointing", + False), + trust_remote_code=self.config.model.get( + "trust_remote_code", + False), + use_liger=self.config.model.get( + "use_liger", + False), + role="actor", ) # get the original unwrapped module @@ -118,7 +142,9 @@ def init_model(self): if self._is_offload_optimizer: offload_fsdp_optimizer(optimizer=self.actor_optimizer) - log_gpu_memory_usage("After offload actor optimizer during init", logger=logger) + log_gpu_memory_usage( + "After offload actor optimizer during init", logger=logger + ) # load from checkpoint if self._is_actor or self._is_ref: OmegaConf.set_struct(self.config.actor, True) @@ -126,7 +152,9 @@ def init_model(self): self.config.actor.use_remove_padding = use_remove_padding self.config.actor.use_fused_kernels = use_fused_kernels self.actor = DataParallelPPOActor( - config=self.config.actor, actor_module=self.actor_module_fsdp, actor_optimizer=self.actor_optimizer + config=self.config.actor, + actor_module=self.actor_module_fsdp, + actor_optimizer=self.actor_optimizer, ) if self._is_rollout: @@ -142,20 +170,27 @@ def init_model(self): override_model_config=override_model_config, use_remove_padding=use_remove_padding, use_fused_kernels=use_fused_kernels, - trust_remote_code=self.config.model.get("trust_remote_code", False), - use_liger=self.config.model.get("use_liger", False), + trust_remote_code=self.config.model.get( + "trust_remote_code", + False), + use_liger=self.config.model.get( + "use_liger", + False), role="ref", )[0] OmegaConf.set_struct(self.config.ref, True) with open_dict(self.config.ref): self.config.ref.use_remove_padding = use_remove_padding self.config.ref.use_fused_kernels = use_fused_kernels - self.ref_policy = DataParallelPPOActor(config=self.config.ref, actor_module=self.ref_module_fsdp) + self.ref_policy = DataParallelPPOActor( + config=self.config.ref, actor_module=self.ref_module_fsdp + ) self.checkpoint_manager = FSDPCheckpointManager( model=self.actor_module_fsdp, optimizer=self.actor.actor_optimizer, lr_scheduler=self.actor_lr_scheduler, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), checkpoint_config=self.config.actor.checkpoint, ) @@ -165,7 +200,8 @@ def init_model(self): model=self.actor_module_fsdp, optimizer=self.actor.actor_optimizer, lr_scheduler=self.actor_lr_scheduler, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), checkpoint_config=self.config.actor.checkpoint, ) @@ -205,8 +241,12 @@ def compute_log_prob(self, data: DataProto): # Support all hardwares data = data.to(get_device_id()) # we should always recompute old_log_probs when it is HybridEngine - data.meta_info["micro_batch_size"] = self.config.rollout.log_prob_micro_batch_size_per_gpu - data.meta_info["max_token_len"] = self.config.rollout.log_prob_max_token_len_per_gpu + data.meta_info["micro_batch_size"] = ( + self.config.rollout.log_prob_micro_batch_size_per_gpu + ) + data.meta_info["max_token_len"] = ( + self.config.rollout.log_prob_max_token_len_per_gpu + ) data.meta_info["use_dynamic_bsz"] = self.config.rollout.log_prob_use_dynamic_bsz data.meta_info["temperature"] = self.config.rollout.temperature # perform recompute log_prob @@ -214,7 +254,8 @@ def compute_log_prob(self, data: DataProto): data = self.ulysses_sharding_manager.preprocess_data(data) output = self.actor.compute_log_prob(data=data) output = DataProto.from_dict( - tensors={"old_log_probs": output}, meta_info={"temperature": self.config.rollout.temperature} + tensors={"old_log_probs": output}, + meta_info={"temperature": self.config.rollout.temperature}, ) output = self.ulysses_sharding_manager.postprocess_data(output) @@ -243,24 +284,33 @@ def update_actor_dpo(self, data: DataProto): assert self._is_actor # Make sure this worker has the actor role if self.actor is None: - raise RuntimeError("Actor instance (self.actor) not initialized in worker.") + raise RuntimeError( + "Actor instance (self.actor) not initialized in worker.") # --- FSDP State Management --- if self._is_offload_param: load_fsdp_model_to_gpu(self.actor_module_fsdp) if self._is_offload_optimizer: - load_fsdp_optimizer(optimizer=self.actor_optimizer, device_id=get_device_id()) + load_fsdp_optimizer( + optimizer=self.actor_optimizer, device_id=get_device_id() + ) - log_gpu_memory_usage("Before update policy (DPO via PPO path)", logger=logger) + log_gpu_memory_usage( + "Before update policy (DPO via PPO path)", + logger=logger) # --- Ulysses Sharding (if used) --- with self.ulysses_sharding_manager: data = self.ulysses_sharding_manager.preprocess_data(data=data) # --- Call the core update method (now containing DPO logic) --- - with Timer(name="update_policy_dpo_via_ppo", logger=None) as timer: # Use a distinct timer name + with Timer( + name="update_policy_dpo_via_ppo", logger=None + ) as timer: # Use a distinct timer name # Calls the modified update_policy method - metrics = self.actor.update_policy_dpo_with_ref(data=data) # <-- THIS CALLS THE MODIFIED FUNCTION + metrics = self.actor.update_policy_dpo_with_ref( + data=data + ) # <-- THIS CALLS THE MODIFIED FUNCTION delta_time = timer.last # --- Add Performance Metrics --- @@ -268,23 +318,38 @@ def update_actor_dpo(self, data: DataProto): metrics["perf/approx_tokens_processed"] = torch.sum( data.batch.get("attention_mask", torch.tensor(0)) ).item() # Approx tokens - metrics["perf/max_memory_allocated_gb"] = get_torch_device().max_memory_allocated() / (1024**3) - metrics["perf/max_memory_reserved_gb"] = get_torch_device().max_memory_reserved() / (1024**3) - metrics["perf/cpu_memory_used_gb"] = psutil.virtual_memory().used / (1024**3) + metrics["perf/max_memory_allocated_gb"] = ( + get_torch_device().max_memory_allocated() / (1024**3) + ) + metrics["perf/max_memory_reserved_gb"] = ( + get_torch_device().max_memory_reserved() / (1024**3) + ) + metrics["perf/cpu_memory_used_gb"] = psutil.virtual_memory().used / ( + 1024**3 + ) global_num_tokens = data.meta_info["global_token_num"] - estimated_flops, promised_flops = self.flops_counter.estimate_flops(global_num_tokens, delta_time) - metrics["perf/mfu/actor"] = estimated_flops * self.config.ppo_epochs / promised_flops / self.world_size + estimated_flops, promised_flops = self.flops_counter.estimate_flops( + global_num_tokens, delta_time) + metrics["perf/mfu/actor"] = ( + estimated_flops + * self.config.ppo_epochs + / promised_flops + / self.world_size + ) # --- LR Scheduler Step --- lr = self.actor_lr_scheduler.get_last_lr()[0] metrics["actor/lr"] = lr self.actor_lr_scheduler.step() - log_gpu_memory_usage("After update policy (DPO via PPO path)", logger=logger) + log_gpu_memory_usage( + "After update policy (DPO via PPO path)", logger=logger + ) # --- Prepare Output --- output = DataProto(meta_info={"metrics": metrics}) - output = self.ulysses_sharding_manager.postprocess_data(data=output) + output = self.ulysses_sharding_manager.postprocess_data( + data=output) output = output.to("cpu") # --- FSDP State Management (Offload) --- @@ -315,19 +380,28 @@ def __init__(self, config): from torch.distributed.device_mesh import init_device_mesh fsdp_size = self.config.model.fsdp_config.fsdp_size - self.device_mesh = create_device_mesh(world_size=world_size, fsdp_size=fsdp_size) + self.device_mesh = create_device_mesh( + world_size=world_size, fsdp_size=fsdp_size + ) self.ulysses_device_mesh = None - self.ulysses_sequence_parallel_size = self.config.get("ulysses_sequence_parallel_size", 1) + 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( - get_device_name(), mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=["dp", "sp"] + get_device_name(), + mesh_shape=(dp, self.ulysses_sequence_parallel_size), + mesh_dim_names=["dp", "sp"], ) - self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) + self.ulysses_sharding_manager = FSDPUlyssesShardingManager( + self.ulysses_device_mesh + ) - self.use_remove_padding = self.config.model.get("use_remove_padding", False) + self.use_remove_padding = self.config.model.get( + "use_remove_padding", False) # normalize config if self.config.micro_batch_size is not None: @@ -347,20 +421,27 @@ def _build_model(self, config): self._do_switch_chat_template = False else: self._do_switch_chat_template = True - input_tokenizer_local_path = copy_to_local(config.model.input_tokenizer) + input_tokenizer_local_path = copy_to_local( + config.model.input_tokenizer) self.input_tokenizer = hf_tokenizer( - input_tokenizer_local_path, trust_remote_code=config.model.get("trust_remote_code", False) + 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), ) - 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 = AutoConfig.from_pretrained( + local_path, trust_remote_code=trust_remote_code + ) model_config.num_labels = 1 - # note that we have to create model in fp32. Otherwise, the optimizer is in bf16, which is incorrect + # 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, mesh=self.device_mesh - ) + use_meta_tensor=not model_config.tie_word_embeddings, mesh=self.device_mesh) with init_context(), warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -373,14 +454,22 @@ def _build_model(self, config): trust_remote_code=trust_remote_code, ) - if config.model.get("use_remove_padding", False) or self.ulysses_sequence_parallel_size > 1: + if ( + config.model.get("use_remove_padding", False) + or self.ulysses_sequence_parallel_size > 1 + ): from verl.models.transformers.monkey_patch import apply_monkey_patch - apply_monkey_patch(model=reward_module, ulysses_sp_size=self.ulysses_sequence_parallel_size) + apply_monkey_patch( + model=reward_module, + ulysses_sp_size=self.ulysses_sequence_parallel_size, + ) reward_module.to(torch.bfloat16) - auto_wrap_policy = get_fsdp_wrap_policy(module=reward_module, config=self.config.model.fsdp_config) + auto_wrap_policy = get_fsdp_wrap_policy( + module=reward_module, config=self.config.model.fsdp_config + ) fsdp_mesh = self.device_mesh sharding_strategy = get_sharding_strategy(fsdp_mesh) @@ -407,11 +496,21 @@ def init_model(self): self.reward_module = self._build_model(config=self.config) def _forward_micro_batch(self, micro_batch): - from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input + from flash_attn.bert_padding import ( + index_first_axis, + pad_input, + rearrange, + unpad_input, + ) - from verl.utils.ulysses import gather_outpus_and_unpad, ulysses_pad_and_slice_inputs + from verl.utils.ulysses import ( + gather_outpus_and_unpad, + ulysses_pad_and_slice_inputs, + ) - with torch.no_grad(), torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16): + with torch.no_grad(), torch.autocast( + device_type=get_device_name(), dtype=torch.bfloat16 + ): input_ids = micro_batch["input_ids"] batch_size, seqlen = input_ids.shape attention_mask = micro_batch["attention_mask"] @@ -421,22 +520,32 @@ def _forward_micro_batch(self, micro_batch): 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) + 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 + 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 + 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 + # 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 + 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) @@ -444,20 +553,26 @@ def _forward_micro_batch(self, micro_batch): # 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 - ) + 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) + 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, use_cache=False + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + use_cache=False, ) 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,) + 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 @@ -467,8 +582,12 @@ def _expand_to_token_level(self, data: DataProto, scores: torch.Tensor): 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) + 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 @@ -495,7 +614,9 @@ def _switch_chat_template(self, data: DataProto): # 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_length = data.batch["attention_mask"][i][ + -response_length: + ].sum() valid_response_ids = response_ids[:valid_response_length] # decode @@ -512,12 +633,16 @@ def _switch_chat_template(self, data: DataProto): # for debugging purpose print(f"Switch template. chat: {prompt_with_chat_template}") - # the maximum length is actually determined by the reward model itself + # 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 - model_inputs = target_tokenizer(prompt_with_chat_template, return_tensors="pt", add_special_tokens=False) + model_inputs = target_tokenizer( + prompt_with_chat_template, + return_tensors="pt", + add_special_tokens=False) input_ids, attention_mask = verl_F.postprocess_data( input_ids=model_inputs["input_ids"], attention_mask=model_inputs["attention_mask"], @@ -535,7 +660,11 @@ def _switch_chat_template(self, data: DataProto): 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} + rm_inputs = { + "input_ids": rm_input_ids, + "attention_mask": rm_attention_mask, + "position_ids": rm_position_ids, + } return DataProto.from_dict(rm_inputs) @@ -565,15 +694,23 @@ def compute_rm_score(self, data: DataProto): # perform forward computation with self.ulysses_sharding_manager: - rm_data = self.ulysses_sharding_manager.preprocess_data(data=rm_data) + 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) + 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_per_gpu) + micro_batches = rm_data.batch.split( + self.config.micro_batch_size_per_gpu + ) output = [] for micro_batch in micro_batches: rm_score = self._forward_micro_batch(micro_batch) @@ -582,14 +719,21 @@ def compute_rm_score(self, data: DataProto): 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) + 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) + # 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) # https://pytorch.org/docs/stable/notes/fsdp.html#fsdp-notes # unshard the root FSDP module diff --git a/Agent0/executor_train/verl/recipe/spin/main_spin.py b/Agent0/executor_train/verl/recipe/spin/main_spin.py index 9a879ee..fbbbbaa 100644 --- a/Agent0/executor_train/verl/recipe/spin/main_spin.py +++ b/Agent0/executor_train/verl/recipe/spin/main_spin.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -22,7 +22,9 @@ from verl.trainer.ppo.reward import get_custom_reward_fn -@hydra.main(config_path="config", config_name="spin_trainer", version_base=None) +@hydra.main(config_path="config", + config_name="spin_trainer", + version_base=None) def main(config): run_ppo(config) @@ -30,12 +32,18 @@ def main(config): def run_ppo(config) -> None: # TODO(linjunrong.ocss884): this ENV is left for resolving SGLang conflict with ray devices # isolation, will solve in the future - os.environ["ENSURE_CUDA_VISIBLE_DEVICES"] = os.environ.get("CUDA_VISIBLE_DEVICES", "") + os.environ["ENSURE_CUDA_VISIBLE_DEVICES"] = os.environ.get( + "CUDA_VISIBLE_DEVICES", "" + ) if not ray.is_initialized(): # this is for local ray cluster ray.init( runtime_env={ - "env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN", "VLLM_LOGGING_LEVEL": "WARN"} + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "WARN", + } } ) @@ -53,7 +61,9 @@ def run(self, config): from verl.utils.fs import copy_to_local - pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values + pprint( + OmegaConf.to_container(config, resolve=True) + ) # resolve=True will eval symbol values OmegaConf.resolve(config) # download the checkpoint from hdfs @@ -63,8 +73,11 @@ def run(self, config): from verl.utils import hf_processor, hf_tokenizer trust_remote_code = config.data.get("trust_remote_code", False) - tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code) - processor = hf_processor(local_path, use_fast=True) # used for multimodal LLM, could be none + tokenizer = hf_tokenizer( + local_path, trust_remote_code=trust_remote_code) + processor = hf_processor( + local_path, use_fast=True + ) # used for multimodal LLM, could be none # define worker classes if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"}: @@ -94,7 +107,9 @@ def run(self, config): global_pool_id = "global_pool" resource_pool_spec = { - global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + global_pool_id: [ + config.trainer.n_gpus_per_node] * + config.trainer.nnodes, } mapping = { Role.ActorRollout: global_pool_id, @@ -108,7 +123,8 @@ def run(self, config): from verl.workers.megatron_workers import RewardModelWorker else: raise NotImplementedError - role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker) + role_worker_mapping[Role.RewardModel] = ray.remote( + RewardModelWorker) mapping[Role.RewardModel] = global_pool_id # use reference model @@ -121,7 +137,8 @@ def run(self, config): # Note(haibin.lin): please make sure custom reward managers are imported and # registered via `verl.workers.reward_manager.register` - reward_manager_name = config.reward_model.get("reward_manager", "naive") + reward_manager_name = config.reward_model.get( + "reward_manager", "naive") reward_manager_cls = get_reward_manager_cls(reward_manager_name) compute_score = get_custom_reward_fn(config) @@ -136,9 +153,14 @@ def run(self, config): # Note that we always use function-based RM for validation val_reward_fn = reward_manager_cls( - tokenizer=tokenizer, num_examine=1, compute_score=compute_score, reward_fn_key=config.data.reward_fn_key + tokenizer=tokenizer, + num_examine=1, + compute_score=compute_score, + reward_fn_key=config.data.reward_fn_key, + ) + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, mapping=mapping ) - resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) trainer = RaySPINTrainer( config=config, diff --git a/Agent0/executor_train/verl/recipe/spin/spin_trainer.py b/Agent0/executor_train/verl/recipe/spin/spin_trainer.py index fa435db..ca56426 100644 --- a/Agent0/executor_train/verl/recipe/spin/spin_trainer.py +++ b/Agent0/executor_train/verl/recipe/spin/spin_trainer.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -36,7 +36,11 @@ 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 RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.single_controller.ray import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) from verl.single_controller.ray.base import create_colocated_worker_cls from verl.trainer.ppo.metric_utils import ( compute_throughout_metrics, @@ -46,7 +50,10 @@ ) from verl.trainer.ppo.ray_trainer import Role from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path -from verl.utils.seqlen_balancing import get_seqlen_balanced_partitions, log_seqlen_unbalance +from verl.utils.seqlen_balancing import ( + get_seqlen_balanced_partitions, + log_seqlen_unbalance, +) from verl.utils.torch_functional import masked_mean from verl.utils.tracking import ValidationGenerationsLogger @@ -75,7 +82,8 @@ class ResourcePoolManager: resource_pool_spec: dict[str, list[int]] mapping: dict[Role, str] - resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict) + 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(): @@ -84,7 +92,10 @@ def create_resource_pool(self): # For Megatron backend, we recommend using max_colocate_count>1 that can utilize different # WorkerGroup for different models resource_pool = RayResourcePool( - process_on_nodes=process_on_nodes, use_gpu=True, max_colocate_count=1, name_prefix=resource_pool_name + 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 @@ -96,24 +107,37 @@ def get_resource_pool(self, role: Role) -> RayResourcePool: def get_n_gpus(self) -> int: """Get the number of gpus in this cluster.""" - return sum([n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes]) + return sum( + [ + n_gpus + for process_on_nodes in self.resource_pool_spec.values() + for n_gpus in process_on_nodes + ] + ) def _check_resource_available(self): """Check if the resource pool can be satisfied in this ray cluster.""" node_available_resources = ray.state.available_resources_per_node() - node_available_gpus = {node: node_info.get("GPU", 0) for node, node_info in node_available_resources.items()} + node_available_gpus = { + node: node_info.get("GPU", 0) + for node, node_info in node_available_resources.items() + } # check total required gpus can be satisfied total_available_gpus = sum(node_available_gpus.values()) total_required_gpus = sum( - [n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes] + [ + n_gpus + for process_on_nodes in self.resource_pool_spec.values() + for n_gpus in process_on_nodes + ] ) if total_available_gpus < total_required_gpus: raise ValueError( - f"Total available GPUs {total_available_gpus} is less than total desired GPUs {total_required_gpus}" - ) + f"Total available GPUs {total_available_gpus} is less than total desired GPUs {total_required_gpus}") - # check each resource pool can be satisfied, O(#resource_pools * #nodes) + # check each resource pool can be satisfied, O(#resource_pools * + # #nodes) for resource_pool_name, process_on_nodes in self.resource_pool_spec.items(): num_gpus, num_nodes = process_on_nodes[0], len(process_on_nodes) for node, available_gpus in node_available_gpus.items(): @@ -138,12 +162,16 @@ def _compute_response_info(batch: DataProto) -> dict[str, Any]: # This is simplified - real implementation might use attention masks # to get actual lengths per sample. batch_size = batch.batch.batch_size[0] - prompt_lengths_tensor = torch.full((batch_size,), prompt_len, dtype=torch.float32, device=batch.batch.device) - response_lengths_tensor = torch.full((batch_size,), resp_len, dtype=torch.float32, device=batch.batch.device) + prompt_lengths_tensor = torch.full( + (batch_size,), prompt_len, dtype=torch.float32, device=batch.batch.device) + response_lengths_tensor = torch.full( + (batch_size,), resp_len, dtype=torch.float32, device=batch.batch.device) - # Try getting actual lengths from attention mask if possible (more accurate) + # Try getting actual lengths from attention mask if possible (more + # accurate) if "response_mask" in batch.batch: - response_lengths_tensor = batch.batch["response_mask"].sum(dim=1).float() + response_lengths_tensor = batch.batch["response_mask"].sum( + dim=1).float() # if "attention_mask" in batch.batch and "response_mask" in batch.batch: # full_mask = batch.batch["attention_mask"] # resp_mask = batch.batch["response_mask"] @@ -152,7 +180,9 @@ def _compute_response_info(batch: DataProto) -> dict[str, Any]: # Example: prompt_lengths_tensor = full_mask.sum(dim=1).float() - response_lengths_tensor # Fallback to using prompt shape if mask logic is complex: prompt_lengths_tensor = torch.tensor( - [batch.batch["prompts"].shape[1]] * batch_size, dtype=torch.float32, device=batch.batch.device + [batch.batch["prompts"].shape[1]] * batch_size, + dtype=torch.float32, + device=batch.batch.device, ) return { @@ -162,11 +192,20 @@ def _compute_response_info(batch: DataProto) -> dict[str, Any]: "max_prompt_length": prompt_len, # Or from config if fixed padding } except KeyError as e: - print(f"Warning: Missing key in _compute_response_info: {e}. Returning defaults.") + print( + f"Warning: Missing key in _compute_response_info: {e}. Returning defaults.") # Return default/dummy values if keys are missing b_size = batch.batch.batch_size[0] if batch.batch.batch_size else 1 - max_resp = batch.batch.get("responses").shape[1] if batch.batch.get("responses") is not None else 0 - max_prompt = batch.batch.get("prompts").shape[1] if batch.batch.get("prompts") is not None else 0 + max_resp = ( + batch.batch.get("responses").shape[1] + if batch.batch.get("responses") is not None + else 0 + ) + max_prompt = ( + batch.batch.get("prompts").shape[1] + if batch.batch.get("prompts") is not None + else 0 + ) return { "prompt_length": torch.zeros(b_size), "response_length": torch.zeros(b_size), @@ -187,7 +226,10 @@ def compute_dpo_data_metrics(batch: DataProto) -> dict[str, Any]: metrics = {} try: # --- Scores and Rewards (from reward_fn) --- - if "token_level_scores" in batch.batch and batch.batch["token_level_scores"] is not None: + if ( + "token_level_scores" in batch.batch + and batch.batch["token_level_scores"] is not None + ): sequence_score = batch.batch["token_level_scores"].sum(-1) metrics.update( { @@ -199,7 +241,10 @@ def compute_dpo_data_metrics(batch: DataProto) -> dict[str, Any]: else: print("DEBUG compute_dpo_data_metrics: 'token_level_scores' not found.") - if "token_level_rewards" in batch.batch and batch.batch["token_level_rewards"] is not None: + if ( + "token_level_rewards" in batch.batch + and batch.batch["token_level_rewards"] is not None + ): sequence_reward = batch.batch["token_level_rewards"].sum(-1) metrics.update( { @@ -222,8 +267,13 @@ def compute_dpo_data_metrics(batch: DataProto) -> dict[str, Any]: else: print("DEBUG compute_dpo_data_metrics: 'chosen_logps' not found.") - if "rejected_logps" in batch.batch and batch.batch["rejected_logps"] is not None: - metrics["actor/rejected_logps"] = batch.batch["rejected_logps"].mean().item() + if ( + "rejected_logps" in batch.batch + and batch.batch["rejected_logps"] is not None + ): + metrics["actor/rejected_logps"] = ( + batch.batch["rejected_logps"].mean().item() + ) else: print("DEBUG compute_dpo_data_metrics: 'rejected_logps' not found.") @@ -232,26 +282,34 @@ def compute_dpo_data_metrics(batch: DataProto) -> dict[str, Any]: # prefs_mask = batch.batch["preferences"] # Shape [batch_size * n] # Calculate accuracy based on RM scores (assuming higher score -> True in mask) # Requires chosen/rejected scores to be available or recalculated - # This is complex here, better calculated in the main loop or update function + # This is complex here, better calculated in the main loop or update + # function # --- Length Metrics --- response_info = _compute_response_info(batch) prompt_length = response_info["prompt_length"] response_length = response_info["response_length"] max_response_length = response_info["max_response_length"] - max_prompt_length = response_info["max_prompt_length"] # Use calculated or from config + max_prompt_length = response_info[ + "max_prompt_length" + ] # Use calculated or from config metrics.update( { "response_length/mean": torch.mean(response_length).item(), "response_length/max": torch.max(response_length).item(), "response_length/min": torch.min(response_length).item(), - "response_length/clip_ratio": torch.mean(torch.eq(response_length, max_response_length).float()).item(), + "response_length/clip_ratio": torch.mean( + torch.eq(response_length, max_response_length).float() + ).item(), "prompt_length/mean": torch.mean(prompt_length).item(), "prompt_length/max": torch.max(prompt_length).item(), "prompt_length/min": torch.min(prompt_length).item(), - # Prompt clip ratio might need adjustment based on how max_prompt_length is defined - "prompt_length/clip_ratio": torch.mean(torch.eq(prompt_length, max_prompt_length).float()).item(), + # Prompt clip ratio might need adjustment based on how + # max_prompt_length is defined + "prompt_length/clip_ratio": torch.mean( + torch.eq(prompt_length, max_prompt_length).float() + ).item(), } ) @@ -261,11 +319,14 @@ def compute_dpo_data_metrics(batch: DataProto) -> dict[str, Any]: print(f"ERROR in compute_dpo_data_metrics: {e}") traceback.print_exc() - print(f"---- [DEBUG] Calculated DPO Data Metrics: {list(metrics.keys())} ----") + print( + f"---- [DEBUG] Calculated DPO Data Metrics: {list(metrics.keys())} ----") return metrics -def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty="kl"): +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"] @@ -274,23 +335,33 @@ def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, response_mask = attention_mask[:, -response_length:] # compute kl between ref_policy and current policy - # When apply_kl_penalty, algorithm.use_kl_in_reward=True, so the reference model has been enabled. + # When apply_kl_penalty, algorithm.use_kl_in_reward=True, so the reference + # model has been enabled. kld = core_algos.kl_penalty( - data.batch["old_log_probs"], data.batch["ref_log_prob"], kl_penalty=kl_penalty - ) # (batch_size, response_length) + 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 token_level_rewards = token_level_scores - beta * kld - current_kl = masked_mean(kld, mask=response_mask, axis=-1) # average over sequence + 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 + # 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 = {"actor/reward_kl_penalty": current_kl, "actor/reward_kl_penalty_coeff": beta} + metrics = { + "actor/reward_kl_penalty": current_kl, + "actor/reward_kl_penalty_coeff": beta, + } return data, metrics @@ -315,22 +386,29 @@ def compute_onlineDPO_pref(data: DataProto): mask_tensor = data.batch.get("response_mask") if rewards_tensor is None or mask_tensor is None: - print(" ERROR: Missing 'token_level_rewards' or 'response_mask' in input data!") + print( + " ERROR: Missing 'token_level_rewards' or 'response_mask' in input data!" + ) # Handle error case - maybe return original data or raise? # Returning original data for now to potentially allow skipping return data try: - preferences = core_algos.compute_onlinedpo_pref(token_level_rewards=rewards_tensor, response_mask=mask_tensor) + preferences = core_algos.compute_onlinedpo_pref( + token_level_rewards=rewards_tensor, response_mask=mask_tensor + ) # Store the result data.batch["preferences"] = preferences except AttributeError: - print("ERROR: Function 'compute_online_dpo_preference' not found in core_algos.py!") + print( + "ERROR: Function 'compute_online_dpo_preference' not found in core_algos.py!" + ) # Assign dummy value or raise error data.batch["preferences"] = None # Indicate failure except Exception as e_pref: - print(f"ERROR during core_algos.compute_online_dpo_preference: {e_pref}") + print( + f"ERROR during core_algos.compute_online_dpo_preference: {e_pref}") import traceback traceback.print_exc() @@ -382,7 +460,9 @@ def __init__( 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()=}" + 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 @@ -396,11 +476,17 @@ def __init__( # define in-reward KL control # kl loss control currently not suppoorted if config.algorithm.use_kl_in_reward: - self.kl_ctrl_in_reward = core_algos.get_kl_controller(config.algorithm.kl_ctrl) + self.kl_ctrl_in_reward = core_algos.get_kl_controller( + config.algorithm.kl_ctrl + ) self.use_critic = False self._validate_config() - self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler) + self._create_dataloader( + train_dataset, + val_dataset, + collate_fn, + train_sampler) def _validate_config(self): config = self.config @@ -408,13 +494,16 @@ def _validate_config(self): n_gpus = config.trainer.n_gpus_per_node * config.trainer.nnodes # 1. Check total batch size for data correctness - real_train_batch_size = config.data.train_batch_size * config.actor_rollout_ref.rollout.n - assert real_train_batch_size % n_gpus == 0, ( - f"real_train_batch_size ({real_train_batch_size}) must be divisible by total n_gpus ({n_gpus})." + real_train_batch_size = ( + config.data.train_batch_size * config.actor_rollout_ref.rollout.n ) + assert ( + real_train_batch_size % n_gpus == 0 + ), f"real_train_batch_size ({real_train_batch_size}) must be divisible by total n_gpus ({n_gpus})." # A helper function to check "micro_batch_size" vs "micro_batch_size_per_gpu" - # We throw an error if the user sets both. The new convention is "..._micro_batch_size_per_gpu". + # We throw an error if the user sets both. The new convention is + # "..._micro_batch_size_per_gpu". def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): settings = { "actor_rollout_ref.actor": "micro_batch_size", @@ -449,14 +538,16 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): ) if self.use_reference_policy: - # reference: log_prob_micro_batch_size vs. log_prob_micro_batch_size_per_gpu + # reference: log_prob_micro_batch_size vs. + # log_prob_micro_batch_size_per_gpu check_mutually_exclusive( config.actor_rollout_ref.ref.log_prob_micro_batch_size, config.actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu, "actor_rollout_ref.ref", ) - # The rollout section also has log_prob_micro_batch_size vs. log_prob_micro_batch_size_per_gpu + # The rollout section also has log_prob_micro_batch_size vs. + # log_prob_micro_batch_size_per_gpu check_mutually_exclusive( config.actor_rollout_ref.rollout.log_prob_micro_batch_size, config.actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu, @@ -466,13 +557,17 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): if self.use_critic and not config.critic.use_dynamic_bsz: # Check for critic micro-batch size conflicts check_mutually_exclusive( - config.critic.ppo_micro_batch_size, config.critic.ppo_micro_batch_size_per_gpu, "critic" + config.critic.ppo_micro_batch_size, + config.critic.ppo_micro_batch_size_per_gpu, + "critic", ) # Check for reward model micro-batch size conflicts if config.reward_model.enable and not config.reward_model.use_dynamic_bsz: check_mutually_exclusive( - config.reward_model.micro_batch_size, config.reward_model.micro_batch_size_per_gpu, "reward_model" + config.reward_model.micro_batch_size, + config.reward_model.micro_batch_size_per_gpu, + "reward_model", ) # Actor @@ -481,15 +576,22 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): # ppo_mini_batch_size is divisible by ppo_micro_batch_size # ppo_micro_batch_size * sequence_parallel_size >= n_gpus if not config.actor_rollout_ref.actor.use_dynamic_bsz: - assert config.data.train_batch_size >= config.actor_rollout_ref.actor.ppo_mini_batch_size - sp_size = config.actor_rollout_ref.actor.get("ulysses_sequence_parallel_size", 1) + assert ( + config.data.train_batch_size + >= config.actor_rollout_ref.actor.ppo_mini_batch_size + ) + sp_size = config.actor_rollout_ref.actor.get( + "ulysses_sequence_parallel_size", 1 + ) if config.actor_rollout_ref.actor.ppo_micro_batch_size is not None: assert ( config.actor_rollout_ref.actor.ppo_mini_batch_size % config.actor_rollout_ref.actor.ppo_micro_batch_size == 0 ) - assert config.actor_rollout_ref.actor.ppo_micro_batch_size * sp_size >= n_gpus + assert ( + config.actor_rollout_ref.actor.ppo_micro_batch_size * + sp_size >= n_gpus) assert config.actor_rollout_ref.actor.loss_agg_mode in [ "token-mean", @@ -497,7 +599,10 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): "seq-mean-token-mean", ], f"Invalid loss_agg_mode: {config.actor_rollout_ref.actor.loss_agg_mode}" - if config.algorithm.use_kl_in_reward and config.actor_rollout_ref.actor.use_kl_loss: + if ( + config.algorithm.use_kl_in_reward + and config.actor_rollout_ref.actor.use_kl_loss + ): print("NOTICE: You have both enabled in-reward kl and kl loss.") # critic @@ -505,40 +610,51 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): assert config.data.train_batch_size >= config.critic.ppo_mini_batch_size sp_size = config.critic.get("ulysses_sequence_parallel_size", 1) if config.critic.ppo_micro_batch_size is not None: - assert config.critic.ppo_mini_batch_size % config.critic.ppo_micro_batch_size == 0 + assert ( + config.critic.ppo_mini_batch_size + % config.critic.ppo_micro_batch_size + == 0 + ) assert config.critic.ppo_micro_batch_size * sp_size >= n_gpus - # Check if use_remove_padding is enabled when using sequence parallelism for fsdp + # Check if use_remove_padding is enabled when using sequence + # parallelism for fsdp if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"}: if ( - config.actor_rollout_ref.actor.get("ulysses_sequence_parallel_size", 1) > 1 - or config.actor_rollout_ref.ref.get("ulysses_sequence_parallel_size", 1) > 1 - ): - assert config.actor_rollout_ref.model.use_remove_padding, ( - "When using sequence parallelism for actor/ref policy, you must enable `use_remove_padding`." - ) + config.actor_rollout_ref.actor.get( + "ulysses_sequence_parallel_size", + 1) > 1 or config.actor_rollout_ref.ref.get( + "ulysses_sequence_parallel_size", + 1) > 1): + assert ( + config.actor_rollout_ref.model.use_remove_padding + ), "When using sequence parallelism for actor/ref policy, you must enable `use_remove_padding`." if self.use_critic and config.critic.strategy in {"fsdp", "fsdp2"}: if config.critic.get("ulysses_sequence_parallel_size", 1) > 1: - assert config.critic.model.use_remove_padding, ( - "When using sequence parallelism for critic, you must enable `use_remove_padding`." - ) + assert ( + config.critic.model.use_remove_padding + ), "When using sequence parallelism for critic, you must enable `use_remove_padding`." if config.data.get("val_batch_size", None) is not None: print( "WARNING: val_batch_size is deprecated. Validation datasets are sent to inference engines " - "as a whole batch, which will schedule the memory themselves." - ) + "as a whole batch, which will schedule the memory themselves.") # check eval config if config.actor_rollout_ref.rollout.val_kwargs.do_sample: - assert config.actor_rollout_ref.rollout.temperature > 0, ( - "validation gen temperature should be greater than 0 when enabling do_sample" - ) + assert ( + config.actor_rollout_ref.rollout.temperature > 0 + ), "validation gen temperature should be greater than 0 when enabling do_sample" print("[validate_config] All configuration checks passed successfully!") - def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampler): + def _create_dataloader( + self, + train_dataset, + val_dataset, + collate_fn, + train_sampler): """ Creates the train and validation dataloaders. """ @@ -547,16 +663,23 @@ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampl if train_dataset is None: train_dataset = create_rl_dataset( - self.config.data.train_files, self.config.data, self.tokenizer, self.processor + self.config.data.train_files, + self.config.data, + self.tokenizer, + self.processor, ) if val_dataset is None: val_dataset = create_rl_dataset( - self.config.data.val_files, self.config.data, self.tokenizer, self.processor + self.config.data.val_files, + self.config.data, + self.tokenizer, + self.processor, ) self.train_dataset, self.val_dataset = train_dataset, val_dataset if train_sampler is None: - train_sampler = create_rl_sampler(self.config.data, self.train_dataset) + train_sampler = create_rl_sampler( + self.config.data, self.train_dataset) if collate_fn is None: from verl.utils.dataset.rl_dataset import collate_fn as default_collate_fn @@ -564,7 +687,9 @@ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampl self.train_dataloader = StatefulDataLoader( dataset=self.train_dataset, - batch_size=self.config.data.get("gen_batch_size", self.config.data.train_batch_size), + batch_size=self.config.data.get( + "gen_batch_size", self.config.data.train_batch_size + ), num_workers=self.config.data.get("dataloader_num_workers", 8), drop_last=True, collate_fn=collate_fn, @@ -592,7 +717,9 @@ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampl f"Size of val dataloader: {len(self.val_dataloader)}" ) - total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs + 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 @@ -603,12 +730,17 @@ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampl try: OmegaConf.set_struct(self.config, True) with open_dict(self.config): - if OmegaConf.select(self.config, "actor_rollout_ref.actor.optim"): - self.config.actor_rollout_ref.actor.optim.total_training_steps = total_training_steps + if OmegaConf.select( + self.config, + "actor_rollout_ref.actor.optim"): + self.config.actor_rollout_ref.actor.optim.total_training_steps = ( + total_training_steps) if OmegaConf.select(self.config, "critic.optim"): self.config.critic.optim.total_training_steps = total_training_steps except Exception as e: - print(f"Warning: Could not set total_training_steps in config. Structure missing? Error: {e}") + print( + f"Warning: Could not set total_training_steps in config. Structure missing? Error: {e}" + ) def _maybe_log_val_generations(self, inputs, outputs, scores): """Log a table of validation samples to the configured logger (wandb or swanlab)""" @@ -632,7 +764,9 @@ def _maybe_log_val_generations(self, inputs, outputs, scores): samples = samples[:generations_to_log] # Log to each configured logger - self.validation_generations_logger.log(self.config.trainer.logger, samples, self.global_steps) + self.validation_generations_logger.log( + self.config.trainer.logger, samples, self.global_steps + ) def _validate(self): data_source_lst = [] @@ -648,23 +782,32 @@ def _validate(self): # repeat test batch test_batch = test_batch.repeat( - repeat_times=self.config.actor_rollout_ref.rollout.val_kwargs.n, interleave=True + repeat_times=self.config.actor_rollout_ref.rollout.val_kwargs.n, + interleave=True, ) # 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": + if ( + self.config.reward_model.enable + and test_batch[0].non_tensor_batch["reward_model"]["style"] == "model" + ): return {} # Store original inputs input_ids = test_batch.batch["input_ids"] # TODO: Can we keep special tokens except for padding tokens? - input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids] + input_texts = [ + self.tokenizer.decode(ids, skip_special_tokens=True) + for ids in input_ids + ] sample_inputs.extend(input_texts) batch_keys_to_pop = ["input_ids", "attention_mask", "position_ids"] non_tensor_batch_keys_to_pop = ["raw_prompt_ids"] if "multi_modal_inputs" in test_batch.non_tensor_batch: - non_tensor_batch_keys_to_pop.extend(["multi_modal_data", "multi_modal_inputs"]) + non_tensor_batch_keys_to_pop.extend( + ["multi_modal_data", "multi_modal_inputs"] + ) if "raw_prompt" in test_batch.non_tensor_batch: non_tensor_batch_keys_to_pop.append("raw_prompt") if "tools_kwargs" in test_batch.non_tensor_batch: @@ -684,19 +827,28 @@ def _validate(self): print(f"test_gen_batch meta info: {test_gen_batch.meta_info}") # 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_gen_batch_padded, pad_size = pad_dataproto_to_divisor( + test_gen_batch, self.actor_rollout_wg.world_size + ) if not self.async_rollout_mode: - test_output_gen_batch_padded = self.actor_rollout_wg.generate_sequences(test_gen_batch_padded) + test_output_gen_batch_padded = self.actor_rollout_wg.generate_sequences( + test_gen_batch_padded) else: - test_output_gen_batch_padded = self.async_rollout_manager.generate_sequences(test_gen_batch_padded) + test_output_gen_batch_padded = ( + self.async_rollout_manager.generate_sequences(test_gen_batch_padded)) # unpad - test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size) + test_output_gen_batch = unpad_dataproto( + test_output_gen_batch_padded, pad_size=pad_size + ) print("validation generation end") # Store generated outputs output_ids = test_output_gen_batch.batch["responses"] - output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids] + output_texts = [ + self.tokenizer.decode(ids, skip_special_tokens=True) + for ids in output_ids + ] sample_outputs.extend(output_texts) test_batch = test_batch.union(test_output_gen_batch) @@ -712,9 +864,15 @@ def _validate(self): for key, lst in result["reward_extra_info"].items(): reward_extra_infos_dict[key].extend(lst) - data_source_lst.append(test_batch.non_tensor_batch.get("data_source", ["unknown"] * reward_tensor.shape[0])) + data_source_lst.append( + test_batch.non_tensor_batch.get( + "data_source", ["unknown"] * reward_tensor.shape[0] + ) + ) - self._maybe_log_val_generations(inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores) + self._maybe_log_val_generations( + inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores + ) # dump generations val_data_dir = self.config.trainer.get("validation_data_dir", None) @@ -728,13 +886,21 @@ def _validate(self): ) for key_info, lst in reward_extra_infos_dict.items(): - assert len(lst) == 0 or len(lst) == len(sample_scores), f"{key_info}: {len(lst)=}, {len(sample_scores)=}" + assert len(lst) == 0 or len(lst) == len( + sample_scores + ), f"{key_info}: {len(lst)=}, {len(sample_scores)=}" data_sources = np.concatenate(data_source_lst, axis=0) - print(f"DEBUG: Data sources shape: {data_sources.shape}") # Added Print - print(f"DEBUG: reward_extra_infos_dict keys before processing: {reward_extra_infos_dict.keys()}") # Added Print + print( + f"DEBUG: Data sources shape: { + data_sources.shape}") # Added Print + print( + f"DEBUG: reward_extra_infos_dict keys before processing: { + reward_extra_infos_dict.keys()}") # Added Print - data_src2var2metric2val = process_validation_metrics(data_sources, sample_inputs, reward_extra_infos_dict) + data_src2var2metric2val = process_validation_metrics( + data_sources, sample_inputs, reward_extra_infos_dict + ) print( f"DEBUG: Output of process_validation_metrics (data_src2var2metric2val): {data_src2var2metric2val}" ) # Added Print @@ -742,11 +908,19 @@ def _validate(self): for data_source, var2metric2val in data_src2var2metric2val.items(): core_var = "acc" if "acc" in var2metric2val else "reward" for var_name, metric2val in var2metric2val.items(): - n_max = max([int(name.split("@")[-1].split("/")[0]) for name in metric2val.keys()]) + n_max = max( + [ + int(name.split("@")[-1].split("/")[0]) + for name in metric2val.keys() + ] + ) for metric_name, metric_val in metric2val.items(): if ( (var_name == core_var) - and any(metric_name.startswith(pfx) for pfx in ["mean", "maj", "best"]) + and any( + metric_name.startswith(pfx) + for pfx in ["mean", "maj", "best"] + ) and (f"@{n_max}" in metric_name) ): metric_sec = "val-core" @@ -761,39 +935,55 @@ 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()} + 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) + 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 + self.resource_pool_to_cls[resource_pool][ + "actor_rollout" + ] = actor_rollout_cls else: raise NotImplementedError # create critic if self.use_critic: - 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) + 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 # create reference policy if needed if self.use_reference_policy: - resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy) + 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.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) + 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 @@ -801,15 +991,21 @@ def init_workers(self): # 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. + # See + # https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb + # for more information. all_wg = {} self.wg_dicts = [] wg_kwargs = {} # Setting up kwargs for RayWorkerGroup - if OmegaConf.select(self.config.trainer, "ray_wait_register_center_timeout") is not None: - wg_kwargs["ray_wait_register_center_timeout"] = self.config.trainer.ray_wait_register_center_timeout + if (OmegaConf.select(self.config.trainer, + "ray_wait_register_center_timeout") is not None): + wg_kwargs["ray_wait_register_center_timeout"] = ( + self.config.trainer.ray_wait_register_center_timeout + ) for resource_pool, class_dict in self.resource_pool_to_cls.items(): - worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict) + 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, @@ -818,7 +1014,8 @@ def init_workers(self): ) 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 + # 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: @@ -833,15 +1030,17 @@ def init_workers(self): 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 + # 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() def _save_checkpoint(self): # path: given_path + `/global_step_{global_steps}` + `/actor` local_global_step_folder = os.path.join( - self.config.trainer.default_local_dir, f"global_step_{self.global_steps}" - ) + self.config.trainer.default_local_dir, + f"global_step_{ + self.global_steps}") print(f"local_global_step_folder: {local_global_step_folder}") actor_local_path = os.path.join(local_global_step_folder, "actor") @@ -849,46 +1048,67 @@ def _save_checkpoint(self): actor_remote_path = ( None if self.config.trainer.default_hdfs_dir is None - else os.path.join(self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", "actor") + else os.path.join( + self.config.trainer.default_hdfs_dir, + f"global_step_{self.global_steps}", + "actor", + ) ) - remove_previous_ckpt_in_save = self.config.trainer.get("remove_previous_ckpt_in_save", False) + remove_previous_ckpt_in_save = self.config.trainer.get( + "remove_previous_ckpt_in_save", False + ) if remove_previous_ckpt_in_save: print( "Warning: remove_previous_ckpt_in_save is deprecated, set max_actor_ckpt_to_keep=1 and " - "max_critic_ckpt_to_keep=1 instead" - ) + "max_critic_ckpt_to_keep=1 instead") max_actor_ckpt_to_keep = ( - self.config.trainer.get("max_actor_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1 + self.config.trainer.get("max_actor_ckpt_to_keep", None) + if not remove_previous_ckpt_in_save + else 1 ) max_critic_ckpt_to_keep = ( - self.config.trainer.get("max_critic_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1 + self.config.trainer.get("max_critic_ckpt_to_keep", None) + if not remove_previous_ckpt_in_save + else 1 ) self.actor_rollout_wg.save_checkpoint( - actor_local_path, actor_remote_path, self.global_steps, max_ckpt_to_keep=max_actor_ckpt_to_keep + actor_local_path, + actor_remote_path, + self.global_steps, + max_ckpt_to_keep=max_actor_ckpt_to_keep, ) if self.use_critic: - critic_local_path = os.path.join(local_global_step_folder, "critic") + critic_local_path = os.path.join( + local_global_step_folder, "critic") critic_remote_path = ( None if self.config.trainer.default_hdfs_dir is None - else os.path.join(self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", "critic") + else os.path.join( + self.config.trainer.default_hdfs_dir, + f"global_step_{self.global_steps}", + "critic", + ) ) self.critic_wg.save_checkpoint( - critic_local_path, critic_remote_path, self.global_steps, max_ckpt_to_keep=max_critic_ckpt_to_keep + critic_local_path, + critic_remote_path, + self.global_steps, + max_ckpt_to_keep=max_critic_ckpt_to_keep, ) # save dataloader - dataloader_local_path = os.path.join(local_global_step_folder, "data.pt") + dataloader_local_path = os.path.join( + local_global_step_folder, "data.pt") dataloader_state_dict = self.train_dataloader.state_dict() torch.save(dataloader_state_dict, dataloader_local_path) # latest checkpointed iteration tracker (for atomic usage) local_latest_checkpointed_iteration = os.path.join( - self.config.trainer.default_local_dir, "latest_checkpointed_iteration.txt" - ) + self.config.trainer.default_local_dir, + "latest_checkpointed_iteration.txt") with open(local_latest_checkpointed_iteration, "w") as f: f.write(str(self.global_steps)) @@ -900,11 +1120,16 @@ def _load_checkpoint(self): if self.config.trainer.default_hdfs_dir is not None: raise NotImplementedError("load from hdfs is not implemented yet") else: - checkpoint_folder = self.config.trainer.default_local_dir # TODO: check path + checkpoint_folder = ( + self.config.trainer.default_local_dir + ) # TODO: check path if not os.path.isabs(checkpoint_folder): working_dir = os.getcwd() - checkpoint_folder = os.path.join(working_dir, checkpoint_folder) - global_step_folder = find_latest_ckpt_path(checkpoint_folder) # None if no latest + checkpoint_folder = os.path.join( + working_dir, checkpoint_folder) + global_step_folder = find_latest_ckpt_path( + checkpoint_folder + ) # None if no latest # find global_step_folder if self.config.trainer.resume_mode == "auto": @@ -913,14 +1138,17 @@ def _load_checkpoint(self): return 0 else: if self.config.trainer.resume_mode == "resume_path": - assert isinstance(self.config.trainer.resume_from_path, str), "resume ckpt must be str type" - assert "global_step_" in self.config.trainer.resume_from_path, ( - "resume ckpt must specify the global_steps" - ) + assert isinstance( + self.config.trainer.resume_from_path, str + ), "resume ckpt must be str type" + assert ( + "global_step_" in self.config.trainer.resume_from_path + ), "resume ckpt must specify the global_steps" global_step_folder = self.config.trainer.resume_from_path if not os.path.isabs(global_step_folder): working_dir = os.getcwd() - global_step_folder = os.path.join(working_dir, global_step_folder) + global_step_folder = os.path.join( + working_dir, global_step_folder) print(f"Load from checkpoint folder: {global_step_folder}") # set global step self.global_steps = int(global_step_folder.split("global_step_")[-1]) @@ -932,37 +1160,51 @@ def _load_checkpoint(self): critic_path = os.path.join(global_step_folder, "critic") # load actor self.actor_rollout_wg.load_checkpoint( - actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load + actor_path, + del_local_after_load=self.config.trainer.del_local_ckpt_after_load, ) # load critic if self.use_critic: self.critic_wg.load_checkpoint( - critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load + critic_path, + del_local_after_load=self.config.trainer.del_local_ckpt_after_load, ) # load dataloader, # TODO: from remote not implemented yet dataloader_local_path = os.path.join(global_step_folder, "data.pt") if os.path.exists(dataloader_local_path): - dataloader_state_dict = torch.load(dataloader_local_path, weights_only=False) + dataloader_state_dict = torch.load( + dataloader_local_path, weights_only=False + ) self.train_dataloader.load_state_dict(dataloader_state_dict) else: - print(f"Warning: No dataloader state found at {dataloader_local_path}, will start from scratch") + print( + f"Warning: No dataloader state found at {dataloader_local_path}, will start from scratch" + ) - def _balance_batch(self, batch: DataProto, metrics, logging_prefix="global_seqlen"): + 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 = batch.batch["attention_mask"].view(batch_size, -1).sum(-1).tolist() # (train_batch_size,) + global_seqlen_lst = ( + batch.batch["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]) + # 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 + seqlen_list=global_seqlen_lst, + partitions=global_partition_lst, + prefix=logging_prefix, ) metrics.update(global_balance_stats) @@ -985,7 +1227,9 @@ def fit_dpo(self): # Renamed for clarity as standard PPO loop 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, throw_on_missing=False), + config=OmegaConf.to_container( + self.config, resolve=True, throw_on_missing=False + ), ) except Exception as e: print(f"Warning: Failed to initialize logger: {e}") @@ -993,30 +1237,39 @@ def fit_dpo(self): # Renamed for clarity as standard PPO loop self.global_steps = 0 # Load checkpoint before doing anything loaded_step = self._load_checkpoint() - self.global_steps = loaded_step + 1 if loaded_step is not None and loaded_step > 0 else 1 + self.global_steps = ( + loaded_step + + 1 if loaded_step is not None and loaded_step > 0 else 1) print( - f"Starting Online DPO training from global step {self.global_steps}. " - f"Total steps: {self.total_training_steps}" - ) - print(f"Reference model update frequency: {self.config.trainer.get('ref_update_freq', 'Not Set')}") + f"Starting Online DPO training from global step { + self.global_steps}. " f"Total steps: { + self.total_training_steps}") + print( + f"Reference model update frequency: { + self.config.trainer.get( + 'ref_update_freq', + 'Not Set')}") # Check if reference policy is configured correctly for this mode if not self.use_reference_policy: print( "WARNING: 'use_reference_policy' is False. Periodic reference model update requires a " - "reference policy worker. DPO updates might fail or use incorrect logic." - ) + "reference policy worker. DPO updates might fail or use incorrect logic.") # Consider raising an error if strict adherence is required: # raise ValueError("Periodic reference model update requires 'use_reference_policy' to be True " # "and a configured reference worker.") # Perform validation before training - if self.val_reward_fn is not None and self.config.trainer.get("val_before_train", True): + if self.val_reward_fn is not None and self.config.trainer.get( + "val_before_train", True + ): print("Running validation before Online DPO training...") val_metrics = self._validate() pprint(f"Initial validation metrics: {val_metrics}") if logger and val_metrics: - logger.log(data=val_metrics, step=max(0, self.global_steps - 1)) + logger.log( + data=val_metrics, step=max( + 0, self.global_steps - 1)) if self.config.trainer.get("val_only", False): print("Validation only mode enabled. Exiting training.") if logger and hasattr(logger, "finish"): @@ -1053,12 +1306,15 @@ def fit_dpo(self): # Renamed for clarity as standard PPO loop metrics = {} timing_raw = {} step_timer = Timer(logger=None) - ref_log_prob_computed = False # Flag to track if ref log probs were computed + ref_log_prob_computed = ( + False # Flag to track if ref log probs were computed + ) try: # Outer try-except for the whole step step_timer.start() with _timer("step", timing_raw): - batch: DataProto = DataProto.from_single_dict(batch_dict) + batch: DataProto = DataProto.from_single_dict( + batch_dict) current_batch_size = batch.batch.batch_size[0] print( f"\n[Step {self.global_steps}, Batch {batch_idx}] Processing batch size: " @@ -1066,88 +1322,127 @@ def fit_dpo(self): # Renamed for clarity as standard PPO loop ) # --- Reference Model Update --- - ref_update_freq = self.config.trainer.get("ref_update_freq", -1) + ref_update_freq = self.config.trainer.get( + "ref_update_freq", -1) if ( self.use_reference_policy and ref_update_freq > 0 and self.global_steps % ref_update_freq == 0 ): - print(f"\n[Step {self.global_steps}] Updating Reference Model Weights from Actor...") + print( + f"\n[Step { + self.global_steps}] Updating Reference Model Weights from Actor...") try: # --- This requires careful implementation with FSDP --- # 1. Save actor state dict (potentially to CPU memory or disk) # This needs to be done collectively across actor worker ranks. # The checkpoint_manager might be adaptable, or use FSDP APIs directly. - # Example placeholder using a conceptual save/load mechanism: - actor_state_path = "/tmp/actor_state_mid" # Temporary path - self.actor_rollout_wg.save_checkpoint(actor_state_path) # Adapt save logic + # Example placeholder using a conceptual + # save/load mechanism: + actor_state_path = ( + "/tmp/actor_state_mid" # Temporary path + ) + self.actor_rollout_wg.save_checkpoint( + actor_state_path + ) # Adapt save logic # 2. Load the state dict onto the reference model worker group - # This also needs collective loading on the ref worker ranks. - self.ref_policy_wg.load_checkpoint(actor_state_path, None, True) # Adapt load logic - - print(f"[Step {self.global_steps}] Reference Model Weights Updated.") + # This also needs collective loading on the ref + # worker ranks. + self.ref_policy_wg.load_checkpoint( + actor_state_path, None, True + ) # Adapt load logic + + print( + f"[Step {self.global_steps}] Reference Model Weights Updated." + ) # Optionally remove the temporary state file - # os.remove(actor_state_path) # Needs rank-aware removal or shared storage + # os.remove(actor_state_path) # Needs + # rank-aware removal or shared storage except Exception as sync_e: - print(f"ERROR during reference model sync at step {self.global_steps}: {sync_e}") + print( + f"ERROR during reference model sync at step { + self.global_steps}: {sync_e}") traceback.print_exc() # Pop keys for generation pop_batch_keys = ["input_ids", "attention_mask"] if "position_ids" in batch.batch: pop_batch_keys.append("position_ids") - pop_non_tensor_keys = ["raw_prompt_ids"] if "raw_prompt_ids" in batch.non_tensor_batch else [] + pop_non_tensor_keys = ( + ["raw_prompt_ids"] + if "raw_prompt_ids" in batch.non_tensor_batch + else [] + ) if "multi_modal_inputs" in batch.non_tensor_batch.keys(): - pop_non_tensor_keys.extend(["multi_modal_data", "multi_modal_inputs"]) + pop_non_tensor_keys.extend( + ["multi_modal_data", "multi_modal_inputs"] + ) original_non_tensor_data = batch.non_tensor_batch gen_batch = batch.pop( batch_keys=pop_batch_keys, non_tensor_batch_keys=pop_non_tensor_keys, ) gen_batch = gen_batch.repeat( - repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True - ) + repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True, ) # (Add Debug prints for gen_batch if needed) # Generate sequences (chosen/rejected pairs) with _timer("gen", timing_raw): try: - gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) + gen_batch_output = ( + self.actor_rollout_wg.generate_sequences(gen_batch)) # (Add Debug prints for gen_batch_output if needed) except Exception as gen_e: - print(f"\n!!!!!!!! ERROR DURING GENERATION (Step {self.global_steps}) !!!!!!!!") + print( + f"\n!!!!!!!! ERROR DURING GENERATION (Step { + self.global_steps}) !!!!!!!!") print(gen_e) traceback.print_exc() - print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!") + print( + "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" + ) step_timer.stop() continue # Combine original prompts with generated sequences - batch.non_tensor_batch = original_non_tensor_data # Restore non-tensor data + batch.non_tensor_batch = ( + original_non_tensor_data # Restore non-tensor data + ) batch.non_tensor_batch["uid"] = np.array( - [str(uuid.uuid4()) for _ in range(current_batch_size)], dtype=object + [str(uuid.uuid4()) for _ in range(current_batch_size)], + dtype=object, ) - batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) + batch = batch.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True, ) batch = batch.union(gen_batch_output) # (Add Debug prints after union if needed) - # Compute response mask (needed for ref logprob calc and DPO prep) - batch.batch["response_mask"] = compute_response_mask(batch) + # Compute response mask (needed for ref logprob calc + # and DPO prep) + batch.batch["response_mask"] = compute_response_mask( + batch) if self.config.trainer.balance_batch: self._balance_batch(batch, metrics=metrics) - batch.meta_info["global_token_num"] = torch.sum(batch.batch["attention_mask"], dim=-1).tolist() + batch.meta_info["global_token_num"] = torch.sum( + batch.batch["attention_mask"], dim=-1 + ).tolist() # --- Compute Log Probs for the CURRENT policy (used for KL if enabled, or ActorAsRef # fallback) --- # Note: For pure DPO with external ref, this 'old_log_probs' might not be strictly needed - # unless used for other metrics or a fallback. Keep it for now. + # unless used for other metrics or a fallback. Keep it + # for now. with _timer("policy_log_prob", timing_raw): - policy_log_prob_output = self.actor_rollout_wg.compute_log_prob(batch) - batch = batch.union(policy_log_prob_output) # Adds 'old_log_probs' + policy_log_prob_output = ( + self.actor_rollout_wg.compute_log_prob(batch) + ) + batch = batch.union( + policy_log_prob_output + ) # Adds 'old_log_probs' # (Debug prints for old_log_probs) # --- Compute Log Probs using the EXTERNAL Reference Model --- @@ -1156,129 +1451,192 @@ def fit_dpo(self): # Renamed for clarity as standard PPO loop # print(f"---- [Step {self.global_steps}] DEBUG DPO: Calling compute_ref_log_prob ----") try: # 'batch' contains interleaved chosen/rejected sequences - ref_log_prob_output = self.ref_policy_wg.compute_ref_log_prob( - batch + ref_log_prob_output = ( + self.ref_policy_wg.compute_ref_log_prob(batch) ) # Returns DataProto with 'ref_log_prob' - batch = batch.union( - ref_log_prob_output - ) # Adds 'ref_log_prob' key [batch_size * n, seq_len] + # Adds 'ref_log_prob' key [batch_size * n, + # seq_len] + batch = batch.union(ref_log_prob_output) ref_log_prob_computed = True # Mark success # print(f"---- [Step {self.global_steps}] DEBUG DPO: ref_log_prob tensor shape: " - # f"{batch.batch['ref_log_prob'].shape} ----") + # f"{batch.batch['ref_log_prob'].shape} + # ----") except Exception as ref_e: - print(f"ERROR computing reference log probs at step {self.global_steps}: {ref_e}") + print( + f"ERROR computing reference log probs at step { + self.global_steps}: {ref_e}") traceback.print_exc() - batch.batch["ref_log_prob"] = None # Mark as failed + # Mark as failed + batch.batch["ref_log_prob"] = None ref_log_prob_computed = False else: print( "Warning: Skipping external reference log prob calculation as use_reference_policy " - "is False." - ) - # DPO update will likely fail unless ActorAsRef logic is re-enabled in dp_actor + "is False.") + # DPO update will likely fail unless ActorAsRef + # logic is re-enabled in dp_actor # --- Compute Rewards/Scores (used to determine preference) --- with _timer("reward_calc", timing_raw): # (Reward calculation logic using RM or reward_fn as before) # ... Ensure this calculates 'token_level_rewards' or similar ... if self.use_rm: - reward_tensor_rm = self.rm_wg.compute_rm_score(batch) - batch = batch.union(reward_tensor_rm) # Adds 'rm_scores' + reward_tensor_rm = self.rm_wg.compute_rm_score( + batch) + batch = batch.union( + reward_tensor_rm + ) # Adds 'rm_scores' reward_extra_infos_dict = {} try: if self.reward_fn is None: # print(f"---- [DEBUG Step {self.global_steps}] ERROR: self.reward_fn is None! " # f"Using dummy rewards. ----") - # Use rm_scores if available, otherwise zeros + # Use rm_scores if available, otherwise + # zeros reward_tensor = batch.batch.get( - "rm_scores", torch.zeros_like(batch.batch["response_mask"], dtype=torch.float32) + "rm_scores", + torch.zeros_like( + batch.batch["response_mask"], + dtype=torch.float32, + ), ) else: - reward_result = self.reward_fn(batch, return_dict=True) - reward_tensor = reward_result["reward_tensor"] # Final combined reward - reward_extra_infos_dict = reward_result.get("reward_extra_info", {}) + reward_result = self.reward_fn( + batch, return_dict=True + ) + reward_tensor = reward_result[ + "reward_tensor" + ] # Final combined reward + reward_extra_infos_dict = reward_result.get( + "reward_extra_info", {}) except Exception: # print(f'---- [DEBUG Step {self.global_steps}] Error in reward_fn call: {e}. ' # f'Using dummy rewards. ----') traceback.print_exc() - reward_tensor = torch.zeros_like(batch.batch["response_mask"], dtype=torch.float32) + reward_tensor = torch.zeros_like( + batch.batch["response_mask"], dtype=torch.float32) reward_extra_infos_dict = {} - # Use 'token_level_rewards' as the key for preference calculation + # Use 'token_level_rewards' as the key for + # preference calculation batch.batch["token_level_rewards"] = reward_tensor if reward_extra_infos_dict: batch.non_tensor_batch.update( - {k: np.array(v) for k, v in reward_extra_infos_dict.items()} + { + k: np.array(v) + for k, v in reward_extra_infos_dict.items() + } ) # --- Determine Preferences --- - # Uses 'token_level_rewards' to determine chosen/rejected based on score - batch = compute_onlineDPO_pref(batch) # Adds 'preferences' key + # Uses 'token_level_rewards' to determine + # chosen/rejected based on score + batch = compute_onlineDPO_pref( + batch) # Adds 'preferences' key # --- Prepare DPO Batch --- dpo_update_batch_proto = None # Initialize with _timer("prepare_dpo_batch", timing_raw): try: - if "preferences" not in batch.batch or batch.batch["preferences"] is None: - raise ValueError("'preferences' key missing or None after compute_onlineDPO_pref.") + if ( + "preferences" not in batch.batch + or batch.batch["preferences"] is None + ): + raise ValueError( + "'preferences' key missing or None after compute_onlineDPO_pref." + ) - # Check if reference log probs were computed successfully (if needed) - if self.use_reference_policy and not ref_log_prob_computed: - raise ValueError("Reference log probs required but failed to compute.") + # Check if reference log probs were computed + # successfully (if needed) + if ( + self.use_reference_policy + and not ref_log_prob_computed + ): + raise ValueError( + "Reference log probs required but failed to compute." + ) # Check required base keys - required_keys = ["input_ids", "attention_mask", "response_mask"] + required_keys = [ + "input_ids", + "attention_mask", + "response_mask", + ] for rk in required_keys: if rk not in batch.batch or batch.batch[rk] is None: - raise KeyError(f"Required key '{rk}' missing from batch for DPO prep.") + raise KeyError( + f"Required key '{rk}' missing from batch for DPO prep.") - preferences_mask = batch.batch["preferences"] # Shape [batch_size * n] + preferences_mask = batch.batch[ + "preferences" + ] # Shape [batch_size * n] not_preferences_mask = ~preferences_mask # Gather Chosen/Rejected Base Tensors - chosen_input_ids = batch.batch["input_ids"][preferences_mask] - chosen_attention_mask = batch.batch["attention_mask"][preferences_mask] - rejected_input_ids = batch.batch["input_ids"][not_preferences_mask] - rejected_attention_mask = batch.batch["attention_mask"][not_preferences_mask] + chosen_input_ids = batch.batch["input_ids"][ + preferences_mask + ] + chosen_attention_mask = batch.batch["attention_mask"][ + preferences_mask + ] + rejected_input_ids = batch.batch["input_ids"][ + not_preferences_mask + ] + rejected_attention_mask = batch.batch["attention_mask"][ + not_preferences_mask + ] chosen_position_ids = ( batch.batch.get("position_ids")[preferences_mask] if "position_ids" in batch.batch else None ) rejected_position_ids = ( - batch.batch.get("position_ids")[not_preferences_mask] + batch.batch.get("position_ids")[ + not_preferences_mask + ] if "position_ids" in batch.batch else None ) # Create Labels - print("WARNING: Creating DPO labels using configured max_prompt_length...") + print( + "WARNING: Creating DPO labels using configured max_prompt_length..." + ) prompt_len = self.config.data.max_prompt_length chosen_labels = chosen_input_ids.clone() chosen_labels[:, :prompt_len] = -100 rejected_labels = rejected_input_ids.clone() rejected_labels[:, :prompt_len] = -100 - # Calculate and Gather Reference Log Probs (Sequence Level) + # Calculate and Gather Reference Log Probs + # (Sequence Level) if self.use_reference_policy: - ref_log_prob_tensor = batch.batch["ref_log_prob"] # Token level [bsz * n, seq_len] + ref_log_prob_tensor = batch.batch[ + "ref_log_prob" + ] # Token level [bsz * n, seq_len] response_mask_full = batch.batch[ "response_mask" ] # Response mask [bsz * n, seq_len] - ref_sequence_logps = (ref_log_prob_tensor * response_mask_full).sum( + ref_sequence_logps = ( + ref_log_prob_tensor * response_mask_full + ).sum( dim=-1 ) # Sequence level [bsz * n] - reference_chosen_logps = ref_sequence_logps[preferences_mask] - reference_rejected_logps = ref_sequence_logps[not_preferences_mask] + reference_chosen_logps = ref_sequence_logps[ + preferences_mask + ] + reference_rejected_logps = ref_sequence_logps[ + not_preferences_mask + ] else: # If not using external ref, DPO needs ActorAsRef logic in dp_actor - # We won't add the keys here, dp_actor will handle it (or fail if not modified) + # We won't add the keys here, dp_actor will + # handle it (or fail if not modified) print( "Info: Not adding explicit reference logps to DPO batch " - "(use_reference_policy=False)." - ) + "(use_reference_policy=False).") reference_chosen_logps = None # Explicitly None reference_rejected_logps = None @@ -1293,110 +1651,160 @@ def fit_dpo(self): # Renamed for clarity as standard PPO loop } # Conditionally add reference logps if computed if reference_chosen_logps is not None: - dpo_tensors["reference_chosen_logps"] = reference_chosen_logps + dpo_tensors["reference_chosen_logps"] = ( + reference_chosen_logps + ) if reference_rejected_logps is not None: - dpo_tensors["reference_rejected_logps"] = reference_rejected_logps + dpo_tensors["reference_rejected_logps"] = ( + reference_rejected_logps + ) # Add position ids if they exist if chosen_position_ids is not None: - dpo_tensors["chosen_position_ids"] = chosen_position_ids + dpo_tensors["chosen_position_ids"] = ( + chosen_position_ids + ) if rejected_position_ids is not None: - dpo_tensors["rejected_position_ids"] = rejected_position_ids + dpo_tensors["rejected_position_ids"] = ( + rejected_position_ids + ) # Prepare Meta Info dpo_meta = { - "dpo_beta": OmegaConf.select(self.config.algorithm, "dpo_beta", default=0.1), + "dpo_beta": OmegaConf.select( + self.config.algorithm, "dpo_beta", default=0.1 + ), "dpo_loss_type": OmegaConf.select( - self.config.algorithm, "dpo_loss_type", default="sigmoid" + self.config.algorithm, + "dpo_loss_type", + default="sigmoid", ), "dpo_label_smoothing": OmegaConf.select( - self.config.algorithm, "dpo_label_smoothing", default=0.0 + self.config.algorithm, + "dpo_label_smoothing", + default=0.0, ), "use_reference_policy": self.use_reference_policy, "reference_free": not self.use_reference_policy, # False if using external ref "global_step": self.global_steps, } - dpo_update_batch_proto = DataProto.from_dict(tensors=dpo_tensors, meta_info=dpo_meta) + dpo_update_batch_proto = DataProto.from_dict( + tensors=dpo_tensors, meta_info=dpo_meta + ) # print(f"---- [Step {self.global_steps}] DEBUG DPO: Prepared DPO Update Batch ----") # print(f" Keys: {list(dpo_update_batch_proto.batch.keys())}") # print(f" Meta Info: {dpo_meta}") except Exception as e_prep: - print(f"ERROR preparing DPO batch at step {self.global_steps}: {e_prep}") + print( + f"ERROR preparing DPO batch at step { + self.global_steps}: {e_prep}") traceback.print_exc() dpo_update_batch_proto = None # Skip update on error # --- Actor Update Step --- actor_output = None - if self.config.trainer.critic_warmup <= self.global_steps and dpo_update_batch_proto: + if (self.config.trainer.critic_warmup <= + self.global_steps and dpo_update_batch_proto): with _timer("update_actor", timing_raw): # Pass the batch containing reference log probs (if computed) - # The modified update_actor_dpo expects them if reference_free=False - actor_output = self.actor_rollout_wg.update_actor_dpo(dpo_update_batch_proto) + # The modified update_actor_dpo expects them if + # reference_free=False + actor_output = self.actor_rollout_wg.update_actor_dpo( + dpo_update_batch_proto) if actor_output and "metrics" in actor_output.meta_info: - metrics.update(reduce_metrics(actor_output.meta_info["metrics"])) + metrics.update( + reduce_metrics( + actor_output.meta_info["metrics"])) elif dpo_update_batch_proto is None: print( - f"Skipping actor update at step {self.global_steps} due to DPO batch preparation error." - ) + f"Skipping actor update at step { + self.global_steps} due to DPO batch preparation error.") # --- Validation and Saving --- - test_freq = OmegaConf.select(self.config.trainer, "test_freq", default=-1) + test_freq = OmegaConf.select( + self.config.trainer, "test_freq", default=-1 + ) is_last_step = self.global_steps >= self.total_training_steps - if ( - self.val_reward_fn is not None - and test_freq > 0 - and (is_last_step or self.global_steps % test_freq == 0) - ): - print(f"\nRunning DPO validation at step {self.global_steps}...") + if (self.val_reward_fn is not None and test_freq > 0 and ( + is_last_step or self.global_steps % test_freq == 0)): + print( + f"\nRunning DPO validation at step { + self.global_steps}...") val_timing_raw = {} with _timer("testing", val_timing_raw): val_metrics: dict = self._validate() if is_last_step: last_val_metrics = val_metrics if val_metrics: - metrics["time/validation_run"] = val_timing_raw.get("testing", 0) + metrics["time/validation_run"] = val_timing_raw.get( + "testing", 0) metrics.update(val_metrics) else: print("Validation skipped or returned no metrics.") - save_freq = OmegaConf.select(self.config.trainer, "save_freq", default=-1) - if save_freq > 0 and (is_last_step or self.global_steps % save_freq == 0): - print(f"\nSaving DPO checkpoint at step {self.global_steps}...") + save_freq = OmegaConf.select( + self.config.trainer, "save_freq", default=-1 + ) + if save_freq > 0 and ( + is_last_step or self.global_steps % save_freq == 0 + ): + print( + f"\nSaving DPO checkpoint at step { + self.global_steps}...") with _timer("save_checkpoint", timing_raw): self._save_checkpoint() # Saves actor (and potentially critic if used elsewhere) - metrics["time/save_checkpoint"] = timing_raw.get("save_checkpoint", 0) + metrics["time/save_checkpoint"] = timing_raw.get( + "save_checkpoint", 0 + ) # --- End main step timer context --- # --- Metrics calculation AFTER the 'step' timer block --- - metrics.update(compute_dpo_data_metrics(batch=batch)) # Use DPO-specific metrics - metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) + metrics.update( + compute_dpo_data_metrics(batch=batch) + ) # Use DPO-specific metrics + metrics.update( + compute_timing_metrics( + batch=batch, + timing_raw=timing_raw)) n_gpus = self.resource_pool_manager.get_n_gpus() if "step" in timing_raw: - metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus)) + metrics.update( + compute_throughout_metrics( + batch=batch, + timing_raw=timing_raw, + n_gpus=n_gpus)) else: print( - f"Warning: 'step' key missing from timing_raw at step {self.global_steps}. " - f"Skipping throughput." - ) + f"Warning: 'step' key missing from timing_raw at step { + self.global_steps}. " f"Skipping throughput.") step_timer.stop() metrics["time/step"] = step_timer.last # Log metrics - log_freq = OmegaConf.select(self.config.trainer, "log_freq", default=1) + log_freq = OmegaConf.select( + self.config.trainer, "log_freq", default=1 + ) if logger and self.global_steps % log_freq == 0: log_payload = metrics.copy() # Add learning rate to log payload if actor_output and "actor/lr" in metrics: log_payload["actor/lr"] = metrics["actor/lr"] - print(f"[Step {self.global_steps} DPO] Logging Step Payload Keys: {list(log_payload.keys())}") + print( + f"[Step { + self.global_steps} DPO] Logging Step Payload Keys: { + list( + log_payload.keys())}") try: - logger.log(data=log_payload, step=self.global_steps) + logger.log( + data=log_payload, step=self.global_steps) except Exception as e: - print(f"Logging failed at step {self.global_steps}: {e}") + print( + f"Logging failed at step { + self.global_steps}: {e}") # Update progress bar postfix_metrics = { @@ -1407,16 +1815,22 @@ def fit_dpo(self): # Renamed for clarity as standard PPO loop progress_bar.set_postfix(postfix_metrics) except Exception as step_e: - print(f"\n!!!!!!!! ERROR DURING DPO Step {self.global_steps} !!!!!!!!") + print( + f"\n!!!!!!!! ERROR DURING DPO Step { + self.global_steps} !!!!!!!!") print(f"Caught Exception: {step_e}") traceback.print_exc() - print("!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!") + print( + "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!" + ) step_timer.stop() should_stop = True break if is_last_step or should_stop: - print(f"Stopping DPO training at step {self.global_steps}.") + print( + f"Stopping DPO training at step { + self.global_steps}.") break self.global_steps += 1 @@ -1427,7 +1841,8 @@ def fit_dpo(self): # Renamed for clarity as standard PPO loop try: self.train_dataloader.reset() except Exception as e: - print(f"Warning: Failed to reset train dataloader state: {e}") + print( + f"Warning: Failed to reset train dataloader state: {e}") if should_stop: break @@ -1436,13 +1851,20 @@ def fit_dpo(self): # Renamed for clarity as standard PPO loop final_step = max(0, self.global_steps - 1) print(f"Online DPO Training finished at step {final_step}.") # Save final checkpoint - save_freq = OmegaConf.select(self.config.trainer, "save_freq", default=-1) - if not self.config.trainer.get("val_only", False) and (save_freq <= 0 or final_step % save_freq != 0): + save_freq = OmegaConf.select( + self.config.trainer, "save_freq", default=-1) + if not self.config.trainer.get("val_only", False) and ( + save_freq <= 0 or final_step % save_freq != 0 + ): print(f"Saving final DPO checkpoint at step {final_step}...") self._save_checkpoint() # Final validation run - if self.val_reward_fn and last_val_metrics is None and not self.config.trainer.get("val_only", False): + if ( + self.val_reward_fn + and last_val_metrics is None + and not self.config.trainer.get("val_only", False) + ): print("Running final validation...") last_val_metrics = self._validate() if last_val_metrics and logger: diff --git a/Agent0/executor_train/verl/recipe/sppo/__init__.py b/Agent0/executor_train/verl/recipe/sppo/__init__.py index bc88468..2d57b11 100644 --- a/Agent0/executor_train/verl/recipe/sppo/__init__.py +++ b/Agent0/executor_train/verl/recipe/sppo/__init__.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/recipe/sppo/dp_actor.py b/Agent0/executor_train/verl/recipe/sppo/dp_actor.py index df14c0b..1b400bf 100644 --- a/Agent0/executor_train/verl/recipe/sppo/dp_actor.py +++ b/Agent0/executor_train/verl/recipe/sppo/dp_actor.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -63,10 +63,19 @@ def update_policy(self, data: DataProto): # make sure we are in training mode self.actor_module.train() - temperature = data.meta_info["temperature"] # temperature must be in the data.meta_info to avoid slient error + temperature = data.meta_info[ + "temperature" + ] # temperature must be in the data.meta_info to avoid slient error multi_turn = data.meta_info.get("multi_turn", False) - select_keys = ["responses", "input_ids", "attention_mask", "position_ids", "old_log_probs", "seq_level_rewards"] + select_keys = [ + "responses", + "input_ids", + "attention_mask", + "position_ids", + "old_log_probs", + "seq_level_rewards", + ] if multi_turn: select_keys.append("loss_mask") if self.config.use_kl_loss: @@ -77,9 +86,12 @@ def update_policy(self, data: DataProto): # Split to make minibatch iterator for updating the actor # See PPO paper for details. https://arxiv.org/abs/1707.06347 if has_multi_modal_inputs: - num_mini_batches = data.batch.batch_size[0] // self.config.ppo_mini_batch_size + num_mini_batches = ( + data.batch.batch_size[0] // self.config.ppo_mini_batch_size + ) non_tensor_select_keys = ["multi_modal_inputs"] - dataloader = data.select(select_keys, non_tensor_select_keys).chunk(num_mini_batches) + dataloader = data.select( + select_keys, non_tensor_select_keys).chunk(num_mini_batches) else: dataloader = batch.split(self.config.ppo_mini_batch_size) @@ -90,28 +102,47 @@ def update_policy(self, data: DataProto): mini_batch = data if has_multi_modal_inputs: self.gradient_accumulation = ( - self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu + self.config.ppo_mini_batch_size + // self.config.ppo_micro_batch_size_per_gpu ) - num_micro_batches = mini_batch.batch.batch_size[0] // self.config.ppo_micro_batch_size_per_gpu - micro_batches = data.select(select_keys, non_tensor_select_keys).chunk(num_micro_batches) + num_micro_batches = ( + mini_batch.batch.batch_size[0] + // self.config.ppo_micro_batch_size_per_gpu + ) + micro_batches = data.select( + select_keys, non_tensor_select_keys + ).chunk(num_micro_batches) elif 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) + 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: self.gradient_accumulation = ( - self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu + self.config.ppo_mini_batch_size + // self.config.ppo_micro_batch_size_per_gpu ) # split batch into micro_batches - micro_batches = mini_batch.split(self.config.ppo_micro_batch_size_per_gpu) + micro_batches = mini_batch.split( + self.config.ppo_micro_batch_size_per_gpu + ) self.actor_optimizer.zero_grad() for data in micro_batches: # Support all hardwares if isinstance(data, DataProto): - data = {**data.batch.to(get_device_id()), **data.non_tensor_batch} + data = { + **data.batch.to(get_device_id()), + **data.non_tensor_batch, + } else: - data = data.to(get_device_id()) # actor device is cpu when using offload + data = data.to( + get_device_id() + ) # actor device is cpu when using offload responses = data["responses"] response_length = responses.size(1) attention_mask = data["attention_mask"] @@ -132,7 +163,9 @@ def update_policy(self, data: DataProto): if entropy_coeff != 0: calculate_entropy = True entropy, log_prob = self._forward_micro_batch( - micro_batch=data, temperature=temperature, calculate_entropy=calculate_entropy + micro_batch=data, + temperature=temperature, + calculate_entropy=calculate_entropy, ) pg_loss, log_ratios, preference = compute_sppo_loss( @@ -145,7 +178,11 @@ def update_policy(self, data: DataProto): ) if entropy_coeff != 0: - entropy_loss = agg_loss(loss_mat=entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode) + entropy_loss = agg_loss( + loss_mat=entropy, + loss_mask=response_mask, + loss_agg_mode=loss_agg_mode, + ) # compute policy loss policy_loss = pg_loss - entropy_loss * entropy_coeff @@ -156,10 +193,14 @@ def update_policy(self, data: DataProto): ref_log_prob = data["ref_log_prob"] # compute kl loss kld = kl_penalty( - logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type + logprob=log_prob, + ref_logprob=ref_log_prob, + kl_penalty=self.config.kl_loss_type, ) kl_loss = agg_loss( - loss_mat=kld, loss_mask=response_mask, loss_agg_mode=self.config.loss_agg_mode + loss_mat=kld, + loss_mask=response_mask, + loss_agg_mode=self.config.loss_agg_mode, ) policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef @@ -168,7 +209,9 @@ def update_policy(self, data: DataProto): if self.config.use_dynamic_bsz: # relative to the dynamic bsz - loss = policy_loss * (len(data) / self.config.ppo_mini_batch_size) + loss = policy_loss * ( + len(data) / self.config.ppo_mini_batch_size + ) else: loss = policy_loss / self.gradient_accumulation loss.backward() diff --git a/Agent0/executor_train/verl/recipe/sppo/main_sppo.py b/Agent0/executor_train/verl/recipe/sppo/main_sppo.py index d99f4f2..d17e14c 100644 --- a/Agent0/executor_train/verl/recipe/sppo/main_sppo.py +++ b/Agent0/executor_train/verl/recipe/sppo/main_sppo.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,7 +27,9 @@ from .sppo_ray_trainer import RaySPPOTrainer -@hydra.main(config_path="config", config_name="sppo_trainer", version_base=None) +@hydra.main(config_path="config", + config_name="sppo_trainer", + version_base=None) def main(config): run_ppo(config) @@ -35,12 +37,18 @@ def main(config): def run_ppo(config) -> None: # TODO(linjunrong.ocss884): this ENV is left for resolving SGLang conflict with ray devices # isolation, will solve in the future - os.environ["ENSURE_CUDA_VISIBLE_DEVICES"] = os.environ.get("CUDA_VISIBLE_DEVICES", "") + os.environ["ENSURE_CUDA_VISIBLE_DEVICES"] = os.environ.get( + "CUDA_VISIBLE_DEVICES", "" + ) if not ray.is_initialized(): # this is for local ray cluster ray.init( runtime_env={ - "env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN", "VLLM_LOGGING_LEVEL": "WARN"} + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN", + "VLLM_LOGGING_LEVEL": "WARN", + } }, num_cpus=config.ray_init.num_cpus, ) @@ -59,7 +67,9 @@ def run(self, config): from verl.utils.fs import copy_to_local - pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values + pprint( + OmegaConf.to_container(config, resolve=True) + ) # resolve=True will eval symbol values OmegaConf.resolve(config) # download the checkpoint from hdfs @@ -69,8 +79,11 @@ def run(self, config): from verl.utils import hf_processor, hf_tokenizer trust_remote_code = config.data.get("trust_remote_code", False) - tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code) - processor = hf_processor(local_path, use_fast=True) # used for multimodal LLM, could be none + tokenizer = hf_tokenizer( + local_path, trust_remote_code=trust_remote_code) + processor = hf_processor( + local_path, use_fast=True + ) # used for multimodal LLM, could be none # define worker classes if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"}: @@ -102,7 +115,9 @@ def run(self, config): global_pool_id = "global_pool" resource_pool_spec = { - global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + global_pool_id: [ + config.trainer.n_gpus_per_node] * + config.trainer.nnodes, } mapping = { Role.ActorRollout: global_pool_id, @@ -121,19 +136,29 @@ def run(self, config): from verl.workers.megatron_workers import RewardModelWorker else: raise NotImplementedError - role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker) + role_worker_mapping[Role.RewardModel] = ray.remote( + RewardModelWorker) mapping[Role.RewardModel] = global_pool_id # use reference model - if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss: - role_worker_mapping[Role.RefPolicy] = ray.remote(SPPOActorRolloutRefWorker) + if ( + config.algorithm.use_kl_in_reward + or config.actor_rollout_ref.actor.use_kl_loss + ): + role_worker_mapping[Role.RefPolicy] = ray.remote( + SPPOActorRolloutRefWorker) mapping[Role.RefPolicy] = global_pool_id reward_fn = load_reward_manager( - config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {}) + config, + tokenizer, + num_examine=0, + **config.reward_model.get("reward_kwargs", {}) ) val_reward_fn = load_reward_manager(config, tokenizer, num_examine=1) - resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, mapping=mapping + ) trainer = RaySPPOTrainer( config=config, diff --git a/Agent0/executor_train/verl/recipe/sppo/sppo_ray_trainer.py b/Agent0/executor_train/verl/recipe/sppo/sppo_ray_trainer.py index 15e2f9c..fa9c443 100644 --- a/Agent0/executor_train/verl/recipe/sppo/sppo_ray_trainer.py +++ b/Agent0/executor_train/verl/recipe/sppo/sppo_ray_trainer.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -48,7 +48,9 @@ from verl.utils.tracking import ValidationGenerationsLogger -def softmean(x: torch.Tensor, beta: float, dim: int = -1, keepdim: bool = False) -> torch.Tensor: +def softmean( + x: torch.Tensor, beta: float, dim: int = -1, keepdim: bool = False +) -> torch.Tensor: """ Compute SoftMean_ฮฒ(x) = (1/ฮฒ) * log( (1/n) * ฮฃ exp(ฮฒ * x_i) ) Falls back to arithmetic mean when ฮฒ=0. @@ -107,7 +109,9 @@ def __init__( 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()=}" + 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 @@ -120,12 +124,18 @@ def __init__( # define in-reward KL control # kl loss control currently not supported if config.algorithm.use_kl_in_reward: - self.kl_ctrl_in_reward = core_algos.get_kl_controller(config.algorithm.kl_ctrl) + self.kl_ctrl_in_reward = core_algos.get_kl_controller( + config.algorithm.kl_ctrl + ) self.use_critic = False self._validate_config() - self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler) + self._create_dataloader( + train_dataset, + val_dataset, + collate_fn, + train_sampler) def fit(self): """ @@ -152,7 +162,9 @@ def fit(self): # 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): + 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) @@ -160,7 +172,11 @@ def fit(self): return # add tqdm - progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress") + progress_bar = tqdm( + total=self.total_training_steps, + initial=self.global_steps, + desc="Training Progress", + ) # we start from step 1 self.global_steps += 1 @@ -173,7 +189,8 @@ def fit(self): batch: DataProto = DataProto.from_single_dict(batch_dict) # pop those keys for generation - batch_keys_to_pop = ["input_ids", "attention_mask", "position_ids"] + batch_keys_to_pop = [ + "input_ids", "attention_mask", "position_ids"] non_tensor_batch_keys_to_pop = ["raw_prompt_ids"] if "multi_modal_data" in batch.non_tensor_batch: non_tensor_batch_keys_to_pop.append("multi_modal_data") @@ -185,7 +202,10 @@ def fit(self): batch_keys=batch_keys_to_pop, non_tensor_batch_keys=non_tensor_batch_keys_to_pop, ) - gen_batch = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) + gen_batch = gen_batch.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) is_last_step = self.global_steps >= self.total_training_steps @@ -193,9 +213,11 @@ def fit(self): # generate a batch with simple_timer("gen", timing_raw): if not self.async_rollout_mode: - gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) + gen_batch_output = self.actor_rollout_wg.generate_sequences( + gen_batch) else: - gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch) + gen_batch_output = ( + self.async_rollout_manager.generate_sequences(gen_batch)) timing_raw.update(gen_batch_output.meta_info["timing"]) gen_batch_output.meta_info.pop("timing", None) @@ -203,23 +225,34 @@ def fit(self): with simple_timer("gen_max", timing_raw): gen_baseline_batch = deepcopy(gen_batch) gen_baseline_batch.meta_info["do_sample"] = False - gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch) + gen_baseline_output = ( + self.actor_rollout_wg.generate_sequences( + gen_baseline_batch + ) + ) batch = batch.union(gen_baseline_output) reward_baseline_tensor = self.reward_fn(batch) - reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1) + reward_baseline_tensor = reward_baseline_tensor.sum( + dim=-1) - batch.pop(batch_keys=list(gen_baseline_output.batch.keys())) + batch.pop( + batch_keys=list( + gen_baseline_output.batch.keys())) batch.batch["reward_baselines"] = reward_baseline_tensor del gen_baseline_batch, gen_baseline_output batch.non_tensor_batch["uid"] = np.array( - [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object + [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.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) batch = batch.union(gen_batch_output) batch.batch["response_mask"] = compute_response_mask(batch) @@ -232,7 +265,9 @@ def fit(self): 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_token_num"] = torch.sum( + batch.batch["attention_mask"], dim=-1 + ).tolist() with simple_timer("reward", timing_raw): # compute reward model score @@ -241,18 +276,28 @@ def fit(self): batch = batch.union(reward_tensor) if self.config.reward_model.launch_reward_fn_async: - future_reward = compute_reward_async.remote(batch, self.config, self.tokenizer) + future_reward = compute_reward_async.remote( + batch, self.config, self.tokenizer + ) else: - reward_tensor, reward_extra_infos_dict = compute_reward(batch, self.reward_fn) + reward_tensor, reward_extra_infos_dict = compute_reward( + batch, self.reward_fn) # recompute old_log_probs with simple_timer("old_log_prob", timing_raw): - old_log_prob = self.actor_rollout_wg.compute_log_prob(batch) + old_log_prob = self.actor_rollout_wg.compute_log_prob( + batch) entropys = old_log_prob.batch["entropys"] response_masks = batch.batch["response_mask"] loss_agg_mode = self.config.actor_rollout_ref.actor.loss_agg_mode - entropy_agg = agg_loss(loss_mat=entropys, loss_mask=response_masks, loss_agg_mode=loss_agg_mode) - old_log_prob_metrics = {"actor/entropy": entropy_agg.detach().item()} + entropy_agg = agg_loss( + loss_mat=entropys, + loss_mask=response_masks, + loss_agg_mode=loss_agg_mode, + ) + old_log_prob_metrics = { + "actor/entropy": entropy_agg.detach().item() + } metrics.update(old_log_prob_metrics) old_log_prob.batch.pop("entropys") batch = batch.union(old_log_prob) @@ -260,7 +305,8 @@ def fit(self): if self.use_reference_policy: # compute reference log_prob with simple_timer("ref", timing_raw): - ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch) + ref_log_prob = self.ref_policy_wg.compute_ref_log_prob( + batch) batch = batch.union(ref_log_prob) # compute values @@ -273,21 +319,30 @@ def fit(self): # we combine with rule-based rm reward_extra_infos_dict: dict[str, list] if self.config.reward_model.launch_reward_fn_async: - reward_tensor, reward_extra_infos_dict = ray.get(future_reward) + reward_tensor, reward_extra_infos_dict = ray.get( + future_reward) batch.batch["token_level_scores"] = reward_tensor if reward_extra_infos_dict: - batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()}) + batch.non_tensor_batch.update( + {k: np.array(v) for k, v in reward_extra_infos_dict.items()} + ) # compute rewards. apply_kl_penalty if available if self.config.algorithm.use_kl_in_reward: batch, kl_metrics = apply_kl_penalty( - batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty + batch, + kl_ctrl=self.kl_ctrl_in_reward, + kl_penalty=self.config.algorithm.kl_penalty, ) metrics.update(kl_metrics) else: - batch.batch["token_level_rewards"] = batch.batch["token_level_scores"] - batch.batch["seq_level_rewards"] = batch.batch["token_level_scores"] + batch.batch["token_level_rewards"] = batch.batch[ + "token_level_scores" + ] + batch.batch["seq_level_rewards"] = batch.batch[ + "token_level_scores" + ] beta = self.config.algorithm.sppo_eta batch = compute_advantage(batch, beta=beta) @@ -296,26 +351,39 @@ def fit(self): if self.use_critic: with simple_timer("update_critic", timing_raw): critic_output = self.critic_wg.update_critic(batch) - critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"]) + critic_output_metrics = reduce_metrics( + critic_output.meta_info["metrics"] + ) metrics.update(critic_output_metrics) # implement critic warmup if self.config.trainer.critic_warmup <= self.global_steps: # update actor with simple_timer("update_actor", timing_raw): - batch.meta_info["multi_turn"] = self.config.actor_rollout_ref.rollout.multi_turn.enable - actor_output = self.actor_rollout_wg.update_actor(batch) - actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) + batch.meta_info["multi_turn"] = ( + self.config.actor_rollout_ref.rollout.multi_turn.enable) + actor_output = self.actor_rollout_wg.update_actor( + batch) + actor_output_metrics = reduce_metrics( + actor_output.meta_info["metrics"] + ) metrics.update(actor_output_metrics) # Log rollout generations if enabled - rollout_data_dir = self.config.trainer.get("rollout_data_dir", None) + rollout_data_dir = self.config.trainer.get( + "rollout_data_dir", None) if rollout_data_dir: with simple_timer("dump_rollout_generations", timing_raw): print(batch.batch.keys()) - inputs = self.tokenizer.batch_decode(batch.batch["prompts"], skip_special_tokens=True) - outputs = self.tokenizer.batch_decode(batch.batch["responses"], skip_special_tokens=True) - scores = batch.batch["token_level_scores"].sum(-1).cpu().tolist() + inputs = self.tokenizer.batch_decode( + batch.batch["prompts"], skip_special_tokens=True + ) + outputs = self.tokenizer.batch_decode( + batch.batch["responses"], skip_special_tokens=True + ) + scores = ( + batch.batch["token_level_scores"].sum(-1).cpu().tolist() + ) self._dump_generations( inputs=inputs, outputs=outputs, @@ -328,7 +396,10 @@ def fit(self): if ( self.val_reward_fn is not None and self.config.trainer.test_freq > 0 - and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0) + and ( + is_last_step + or self.global_steps % self.config.trainer.test_freq == 0 + ) ): with simple_timer("testing", timing_raw): val_metrics: dict = self._validate() @@ -337,7 +408,8 @@ def fit(self): metrics.update(val_metrics) if self.config.trainer.save_freq > 0 and ( - is_last_step or self.global_steps % self.config.trainer.save_freq == 0 + is_last_step + or self.global_steps % self.config.trainer.save_freq == 0 ): with simple_timer("save_checkpoint", timing_raw): self._save_checkpoint() diff --git a/Agent0/executor_train/verl/recipe/sppo/sppo_worker.py b/Agent0/executor_train/verl/recipe/sppo/sppo_worker.py index fbe3a6e..0f48314 100644 --- a/Agent0/executor_train/verl/recipe/sppo/sppo_worker.py +++ b/Agent0/executor_train/verl/recipe/sppo/sppo_worker.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -45,7 +45,9 @@ def init_model(self): from omegaconf import OmegaConf - override_model_config = OmegaConf.to_container(self.config.model.get("override_config", OmegaConf.create())) + 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) use_fused_kernels = self.config.model.get("use_fused_kernels", False) @@ -58,19 +60,28 @@ def init_model(self): 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, - use_fused_kernels=use_fused_kernels, - enable_gradient_checkpointing=self.config.model.get("enable_gradient_checkpointing", False), - trust_remote_code=self.config.model.get("trust_remote_code", False), - use_liger=self.config.model.get("use_liger", False), - role="actor", - ) + ( + 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, + use_fused_kernels=use_fused_kernels, + enable_gradient_checkpointing=self.config.model.get( + "enable_gradient_checkpointing", + False), + trust_remote_code=self.config.model.get( + "trust_remote_code", + False), + use_liger=self.config.model.get( + "use_liger", + False), + role="actor", ) # get the original unwrapped module @@ -78,11 +89,15 @@ def init_model(self): if self._is_offload_param: offload_fsdp_model_to_cpu(self.actor_module_fsdp) - log_gpu_memory_usage("After offload actor model during init", logger=logger) + log_gpu_memory_usage( + "After offload actor model 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) + 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) @@ -90,7 +105,9 @@ def init_model(self): self.config.actor.use_remove_padding = use_remove_padding self.config.actor.use_fused_kernels = use_fused_kernels self.actor = DataParallelSPPOActor( - config=self.config.actor, actor_module=self.actor_module_fsdp, actor_optimizer=self.actor_optimizer + config=self.config.actor, + actor_module=self.actor_module_fsdp, + actor_optimizer=self.actor_optimizer, ) if self._is_rollout: @@ -106,15 +123,21 @@ def init_model(self): override_model_config=override_model_config, use_remove_padding=use_remove_padding, use_fused_kernels=use_fused_kernels, - trust_remote_code=self.config.model.get("trust_remote_code", False), - use_liger=self.config.model.get("use_liger", False), + trust_remote_code=self.config.model.get( + "trust_remote_code", + False), + use_liger=self.config.model.get( + "use_liger", + False), role="ref", )[0] OmegaConf.set_struct(self.config.ref, True) with open_dict(self.config.ref): self.config.ref.use_remove_padding = use_remove_padding self.config.ref.use_fused_kernels = use_fused_kernels - self.ref_policy = DataParallelSPPOActor(config=self.config.ref, actor_module=self.ref_module_fsdp) + self.ref_policy = DataParallelSPPOActor( + config=self.config.ref, actor_module=self.ref_module_fsdp + ) if self._is_actor: self.flops_counter = FlopsCounter(self.actor_model_config) @@ -122,6 +145,7 @@ def init_model(self): model=self.actor_module_fsdp, optimizer=self.actor.actor_optimizer, lr_scheduler=self.actor_lr_scheduler, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), checkpoint_config=self.config.actor.checkpoint, ) diff --git a/Agent0/executor_train/verl/scripts/__init__.py b/Agent0/executor_train/verl/scripts/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/scripts/__init__.py +++ b/Agent0/executor_train/verl/scripts/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/scripts/converter_hf_to_mcore.py b/Agent0/executor_train/verl/scripts/converter_hf_to_mcore.py index b3101a6..a808d7d 100644 --- a/Agent0/executor_train/verl/scripts/converter_hf_to_mcore.py +++ b/Agent0/executor_train/verl/scripts/converter_hf_to_mcore.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -35,11 +35,30 @@ def _init_args(): parser = argparse.ArgumentParser() - parser.add_argument("--hf_model_path", type=str, required=True, help="The path for the huggingface model") - parser.add_argument("--output_path", type=str, required=True, help="The path for the output mcore model") - parser.add_argument("--use_cpu_initialization", action="store_true", help="Whether to use cpu initialization") - parser.add_argument("--test", action="store_true", help="Whether to test the conversion") - parser.add_argument("--trust_remote_code", action="store_true", help="Whether to trust remote code") + parser.add_argument( + "--hf_model_path", + type=str, + required=True, + help="The path for the huggingface model", + ) + parser.add_argument( + "--output_path", + type=str, + required=True, + help="The path for the output mcore model", + ) + parser.add_argument( + "--use_cpu_initialization", + action="store_true", + help="Whether to use cpu initialization", + ) + parser.add_argument( + "--test", action="store_true", help="Whether to test the conversion" + ) + parser.add_argument( + "--trust_remote_code", + action="store_true", + help="Whether to trust remote code") args = parser.parse_args() return args @@ -54,7 +73,9 @@ def test_conversion(megatron_model_provider, tfconfig, output_path, model): transformer_config=tfconfig, ) ref_state_dict = model_test[0].module.sharded_state_dict() - dist_checkpointing.load(ref_state_dict, output_path, strict=StrictHandling.ASSUME_OK_UNEXPECTED) + dist_checkpointing.load( + ref_state_dict, output_path, strict=StrictHandling.ASSUME_OK_UNEXPECTED + ) dut_state_dict = model[0].module.state_dict() for name in dut_state_dict.keys(): @@ -68,7 +89,9 @@ def test_conversion(megatron_model_provider, tfconfig, output_path, model): ref_data = ref_data.data.view(ref_data.local_shape) else: ref_data = ref_data.data - assert dut_data.shape == ref_data.shape, f"{name=} {dut_data.shape=} {ref_data.shape=}" + assert ( + dut_data.shape == ref_data.shape + ), f"{name=} {dut_data.shape=} {ref_data.shape=}" assert (dut_data == ref_data).all(), f"{name} is not equal" print(f"{name} is equal") else: @@ -84,7 +107,9 @@ def test_conversion(megatron_model_provider, tfconfig, output_path, model): ref_data = ref_data.data if name in dut_state_dict: dut_data = dut_state_dict[name].data - assert dut_data.shape == ref_data.shape, f"{name=} {dut_data.shape=} {ref_data.shape=}" + assert ( + dut_data.shape == ref_data.shape + ), f"{name=} {dut_data.shape=} {ref_data.shape=}" assert (dut_data == ref_data).all(), f"{name} is not equal" print(f"{name} is equal") else: @@ -92,56 +117,103 @@ def test_conversion(megatron_model_provider, tfconfig, output_path, model): print("Conversion test passed!") -def convert_checkpoint_from_transformers_to_megatron(hf_model, model, hf_config): +def convert_checkpoint_from_transformers_to_megatron( + hf_model, model, hf_config): num_attention_heads = hf_config.num_attention_heads num_key_value_heads = hf_config.num_key_value_heads hidden_dim = hf_config.hidden_size - head_dim = getattr(hf_config, "head_dim", hidden_dim // num_attention_heads) + head_dim = getattr( + hf_config, + "head_dim", + hidden_dim // + num_attention_heads) if num_attention_heads != num_key_value_heads: print("[WARNING] Converting GQA model") - has_qkv_bias = getattr(hf_config, "qkv_bias", False) or getattr(hf_config, "attention_bias", False) - has_share_expert = getattr(hf_config, "shared_expert_intermediate_size", None) + has_qkv_bias = getattr(hf_config, "qkv_bias", False) or getattr( + hf_config, "attention_bias", False + ) + has_share_expert = getattr( + hf_config, "shared_expert_intermediate_size", None) with torch.no_grad(): - model.embedding.word_embeddings.weight.copy_(hf_model.model.embed_tokens.weight) - for layer, hf_layer in zip(model.decoder.layers, hf_model.model.layers, strict=True): - layer.self_attention.linear_qkv.layer_norm_weight.copy_(hf_layer.input_layernorm.weight) + model.embedding.word_embeddings.weight.copy_( + hf_model.model.embed_tokens.weight) + for layer, hf_layer in zip( + model.decoder.layers, hf_model.model.layers, strict=True + ): + layer.self_attention.linear_qkv.layer_norm_weight.copy_( + hf_layer.input_layernorm.weight + ) q = hf_layer.self_attn.q_proj.weight.view( - [num_key_value_heads, head_dim * num_attention_heads // num_key_value_heads, -1] + [ + num_key_value_heads, + head_dim * num_attention_heads // num_key_value_heads, + -1, + ] + ) + k = hf_layer.self_attn.k_proj.weight.view( + [num_key_value_heads, head_dim, -1] + ) + v = hf_layer.self_attn.v_proj.weight.view( + [num_key_value_heads, head_dim, -1] ) - k = hf_layer.self_attn.k_proj.weight.view([num_key_value_heads, head_dim, -1]) - v = hf_layer.self_attn.v_proj.weight.view([num_key_value_heads, head_dim, -1]) qkv = torch.cat([q, k, v], dim=1).view(-1, hidden_dim).contiguous() layer.self_attention.linear_qkv.weight.copy_(qkv) if has_qkv_bias: - q_bias = hf_layer.self_attn.q_proj.bias.view([num_key_value_heads, -1]) - k_bias = hf_layer.self_attn.k_proj.bias.view([num_key_value_heads, -1]) - v_bias = hf_layer.self_attn.v_proj.bias.view([num_key_value_heads, -1]) - qkv_bias = torch.cat([q_bias, k_bias, v_bias], dim=1).view(-1).contiguous() + q_bias = hf_layer.self_attn.q_proj.bias.view( + [num_key_value_heads, -1]) + k_bias = hf_layer.self_attn.k_proj.bias.view( + [num_key_value_heads, -1]) + v_bias = hf_layer.self_attn.v_proj.bias.view( + [num_key_value_heads, -1]) + qkv_bias = (torch.cat( + [q_bias, k_bias, v_bias], dim=1).view(-1).contiguous()) layer.self_attention.linear_qkv.bias.copy_(qkv_bias) if hasattr(hf_layer.self_attn, "q_norm"): - layer.self_attention.q_layernorm.weight.copy_(hf_layer.self_attn.q_norm.weight.data) - layer.self_attention.k_layernorm.weight.copy_(hf_layer.self_attn.k_norm.weight.data) + layer.self_attention.q_layernorm.weight.copy_( + hf_layer.self_attn.q_norm.weight.data + ) + layer.self_attention.k_layernorm.weight.copy_( + hf_layer.self_attn.k_norm.weight.data + ) - layer.self_attention.linear_proj.weight.copy_(hf_layer.self_attn.o_proj.weight) - layer.pre_mlp_layernorm.weight.copy_(hf_layer.post_attention_layernorm.weight) + layer.self_attention.linear_proj.weight.copy_( + hf_layer.self_attn.o_proj.weight + ) + layer.pre_mlp_layernorm.weight.copy_( + hf_layer.post_attention_layernorm.weight + ) layer.mlp.router.weight.copy_(hf_layer.mlp.gate.weight) for idx, hf_expert in enumerate(hf_layer.mlp.experts): - fc1_weight = torch.cat([hf_expert.gate_proj.weight, hf_expert.up_proj.weight]) - layer.mlp.experts.linear_fc1._parameters[f"weight{idx}"].copy_(fc1_weight) - layer.mlp.experts.linear_fc2._parameters[f"weight{idx}"].copy_(hf_expert.down_proj.weight) + fc1_weight = torch.cat( + [hf_expert.gate_proj.weight, hf_expert.up_proj.weight] + ) + layer.mlp.experts.linear_fc1._parameters[f"weight{idx}"].copy_( + fc1_weight + ) + layer.mlp.experts.linear_fc2._parameters[f"weight{idx}"].copy_( + hf_expert.down_proj.weight + ) if has_share_expert: - layer.mlp.shared_experts.gate_weight.copy_(hf_layer.mlp.shared_expert_gate.weight) + layer.mlp.shared_experts.gate_weight.copy_( + hf_layer.mlp.shared_expert_gate.weight + ) shared_fc1_weight = torch.cat( - [hf_layer.mlp.shared_expert.gate_proj.weight, hf_layer.mlp.shared_expert.up_proj.weight] + [ + hf_layer.mlp.shared_expert.gate_proj.weight, + hf_layer.mlp.shared_expert.up_proj.weight, + ] + ) + layer.mlp.shared_experts.linear_fc1.weight.copy_( + shared_fc1_weight) + layer.mlp.shared_experts.linear_fc2.weight.copy_( + hf_layer.mlp.shared_expert.down_proj.weight ) - layer.mlp.shared_experts.linear_fc1.weight.copy_(shared_fc1_weight) - layer.mlp.shared_experts.linear_fc2.weight.copy_(hf_layer.mlp.shared_expert.down_proj.weight) model.decoder.final_layernorm.weight.copy_(hf_model.model.norm.weight) model.output_layer.weight.copy_(hf_model.lm_head.weight) @@ -154,14 +226,19 @@ def safe_copy( ): if not skip_dtype_assert: if src_tensor.dtype != dst_tensor.dtype: - raise ValueError(f"Get source dtype {src_tensor.dtype}, but target dtype {dst_tensor.dtype}") + raise ValueError( + f"Get source dtype { + src_tensor.dtype}, but target dtype { + dst_tensor.dtype}") assert src_tensor.shape == dst_tensor.shape dst_tensor.data.copy_(src_tensor.data) return src_tensor.numel() @torch.inference_mode() -def convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl(hfmodel, mgmodel, hf_config): +def convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl( + hfmodel, mgmodel, hf_config +): mgmodel = mgmodel.bfloat16() hfmodel = hfmodel.bfloat16() num_attention_heads = hf_config.num_attention_heads @@ -176,22 +253,37 @@ def convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl(hfmodel, mgmodel vision_num_query_groups = mgvision.config.num_query_groups vision_head_dim = vision_hidden_size // mgvision.config.num_attention_heads copied_numel = 0 - safe_copy(hfvision.rotary_pos_emb.inv_freq, mgvision.rotary_pos_emb.inv_freq) - copied_numel += safe_copy(hfvision.patch_embed.proj.weight, mgvision.patch_embed.proj.weight) - for hfblock, mgblock in zip(hfvision.blocks, mgvision.decoder.layers, strict=True): + safe_copy(hfvision.rotary_pos_emb.inv_freq, + mgvision.rotary_pos_emb.inv_freq) + copied_numel += safe_copy( + hfvision.patch_embed.proj.weight, mgvision.patch_embed.proj.weight + ) + for hfblock, mgblock in zip( + hfvision.blocks, mgvision.decoder.layers, strict=True): # norm1 --> linear_qkv.norm - copied_numel += safe_copy(hfblock.norm1.weight, mgblock.self_attention.linear_qkv.layer_norm_weight) + copied_numel += safe_copy(hfblock.norm1.weight, + mgblock.self_attention.linear_qkv.layer_norm_weight) # norm2 --> mlp.linear_fc1.norm - copied_numel += safe_copy(hfblock.norm2.weight, mgblock.mlp.linear_fc1.layer_norm_weight) + copied_numel += safe_copy( + hfblock.norm2.weight, mgblock.mlp.linear_fc1.layer_norm_weight + ) # qkv --> self_attention.linear_qkv converted_weight = ( - hfblock.attn.qkv.weight.view(3, vision_num_query_groups, -1, vision_head_dim, vision_hidden_size) - .transpose(0, 1) - .flatten(1, 2) - .reshape(-1, vision_hidden_size) - .contiguous() + hfblock.attn.qkv.weight.view( + 3, + vision_num_query_groups, + -1, + vision_head_dim, + vision_hidden_size) .transpose( + 0, + 1) .flatten( + 1, + 2) .reshape( + -1, + vision_hidden_size) .contiguous()) + copied_numel += safe_copy( + converted_weight, mgblock.self_attention.linear_qkv.weight ) - copied_numel += safe_copy(converted_weight, mgblock.self_attention.linear_qkv.weight) converted_bias = ( hfblock.attn.qkv.bias.view(3, vision_num_query_groups, -1) .transpose(0, 1) @@ -199,57 +291,111 @@ def convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl(hfmodel, mgmodel .view(-1) .contiguous() ) - copied_numel += safe_copy(converted_bias, mgblock.self_attention.linear_qkv.bias) + copied_numel += safe_copy( + converted_bias, mgblock.self_attention.linear_qkv.bias + ) # proj --> self_attention.linear_proj - copied_numel += safe_copy(hfblock.attn.proj.weight, mgblock.self_attention.linear_proj.weight) - copied_numel += safe_copy(hfblock.attn.proj.bias, mgblock.self_attention.linear_proj.bias) + copied_numel += safe_copy( + hfblock.attn.proj.weight, mgblock.self_attention.linear_proj.weight + ) + copied_numel += safe_copy( + hfblock.attn.proj.bias, mgblock.self_attention.linear_proj.bias + ) # mlp --> mlp: gate - fc1_weight = torch.cat([hfblock.mlp.gate_proj.weight, hfblock.mlp.up_proj.weight]) - fc1_bias = torch.cat([hfblock.mlp.gate_proj.bias, hfblock.mlp.up_proj.bias]) + fc1_weight = torch.cat( + [hfblock.mlp.gate_proj.weight, hfblock.mlp.up_proj.weight] + ) + fc1_bias = torch.cat( + [hfblock.mlp.gate_proj.bias, hfblock.mlp.up_proj.bias]) copied_numel += safe_copy(fc1_weight, mgblock.mlp.linear_fc1.weight) copied_numel += safe_copy(fc1_bias, mgblock.mlp.linear_fc1.bias) - copied_numel += safe_copy(hfblock.mlp.down_proj.weight, mgblock.mlp.linear_fc2.weight) - copied_numel += safe_copy(hfblock.mlp.down_proj.bias, mgblock.mlp.linear_fc2.bias) + copied_numel += safe_copy( + hfblock.mlp.down_proj.weight, mgblock.mlp.linear_fc2.weight + ) + copied_numel += safe_copy( + hfblock.mlp.down_proj.bias, mgblock.mlp.linear_fc2.bias + ) # 2. vision projector hfprojector = hfvision.merger mgprojector = mgvision.projection - copied_numel += safe_copy(hfprojector.ln_q.weight, mgvision.decoder.final_layernorm.weight) + copied_numel += safe_copy( + hfprojector.ln_q.weight, mgvision.decoder.final_layernorm.weight + ) - copied_numel += safe_copy(hfprojector.mlp[0].weight, mgprojector.encoder.linear_fc1.weight) - copied_numel += safe_copy(hfprojector.mlp[0].bias, mgprojector.encoder.linear_fc1.bias) - copied_numel += safe_copy(hfprojector.mlp[2].weight, mgprojector.encoder.linear_fc2.weight) - copied_numel += safe_copy(hfprojector.mlp[2].bias, mgprojector.encoder.linear_fc2.bias) + copied_numel += safe_copy( + hfprojector.mlp[0].weight, mgprojector.encoder.linear_fc1.weight + ) + copied_numel += safe_copy( + hfprojector.mlp[0].bias, mgprojector.encoder.linear_fc1.bias + ) + copied_numel += safe_copy( + hfprojector.mlp[2].weight, mgprojector.encoder.linear_fc2.weight + ) + copied_numel += safe_copy( + hfprojector.mlp[2].bias, mgprojector.encoder.linear_fc2.bias + ) n_params = sum([t.numel() for t in hfvision.state_dict().values()]) assert n_params == copied_numel # 3. llm [just Qwen2] hfllm = hfmodel.model mgllm = mgmodel.language_model copied_numel = 0 - copied_numel += safe_copy(hfllm.embed_tokens.weight, mgllm.embedding.word_embeddings.weight) - for mglayer, hflayer in zip(mgllm.decoder.layers, hfllm.layers, strict=True): - copied_numel += safe_copy(hflayer.input_layernorm.weight, mglayer.self_attention.linear_qkv.layer_norm_weight) + copied_numel += safe_copy( + hfllm.embed_tokens.weight, mgllm.embedding.word_embeddings.weight + ) + for mglayer, hflayer in zip( + mgllm.decoder.layers, hfllm.layers, strict=True): + copied_numel += safe_copy( + hflayer.input_layernorm.weight, + mglayer.self_attention.linear_qkv.layer_norm_weight, + ) - q_proj_weight = hflayer.self_attn.q_proj.weight.view(num_query_groups, -1, head_dim, hidden_size) - k_proj_weight = hflayer.self_attn.k_proj.weight.view(num_query_groups, -1, head_dim, hidden_size) - v_proj_weight = hflayer.self_attn.v_proj.weight.view(num_query_groups, -1, head_dim, hidden_size) - qkv_proj = torch.cat([q_proj_weight, k_proj_weight, v_proj_weight], dim=1).view(-1, hidden_size).contiguous() - copied_numel += safe_copy(qkv_proj, mglayer.self_attention.linear_qkv.weight) + q_proj_weight = hflayer.self_attn.q_proj.weight.view( + num_query_groups, -1, head_dim, hidden_size + ) + k_proj_weight = hflayer.self_attn.k_proj.weight.view( + num_query_groups, -1, head_dim, hidden_size + ) + v_proj_weight = hflayer.self_attn.v_proj.weight.view( + num_query_groups, -1, head_dim, hidden_size + ) + qkv_proj = ( + torch.cat([q_proj_weight, k_proj_weight, v_proj_weight], dim=1) + .view(-1, hidden_size) + .contiguous() + ) + copied_numel += safe_copy(qkv_proj, + mglayer.self_attention.linear_qkv.weight) q_proj_bias = hflayer.self_attn.q_proj.bias.view(num_query_groups, -1) k_proj_bias = hflayer.self_attn.k_proj.bias.view(num_query_groups, -1) v_proj_bias = hflayer.self_attn.v_proj.bias.view(num_query_groups, -1) - qkv_bias = torch.cat([q_proj_bias, k_proj_bias, v_proj_bias], dim=1).view(-1).contiguous() - copied_numel += safe_copy(qkv_bias, mglayer.self_attention.linear_qkv.bias) - copied_numel += safe_copy(hflayer.self_attn.o_proj.weight, mglayer.self_attention.linear_proj.weight) + qkv_bias = ( + torch.cat([q_proj_bias, k_proj_bias, v_proj_bias], dim=1) + .view(-1) + .contiguous() + ) + copied_numel += safe_copy(qkv_bias, + mglayer.self_attention.linear_qkv.bias) + copied_numel += safe_copy(hflayer.self_attn.o_proj.weight, + mglayer.self_attention.linear_proj.weight) - fc1_weight = torch.cat([hflayer.mlp.gate_proj.weight, hflayer.mlp.up_proj.weight]) + fc1_weight = torch.cat( + [hflayer.mlp.gate_proj.weight, hflayer.mlp.up_proj.weight] + ) copied_numel += safe_copy(fc1_weight, mglayer.mlp.linear_fc1.weight) - copied_numel += safe_copy(hflayer.mlp.down_proj.weight, mglayer.mlp.linear_fc2.weight) - copied_numel += safe_copy(hflayer.post_attention_layernorm.weight, mglayer.mlp.linear_fc1.layer_norm_weight) + copied_numel += safe_copy( + hflayer.mlp.down_proj.weight, mglayer.mlp.linear_fc2.weight + ) + copied_numel += safe_copy( + hflayer.post_attention_layernorm.weight, + mglayer.mlp.linear_fc1.layer_norm_weight, + ) - copied_numel += safe_copy(hfllm.norm.weight, mgllm.decoder.final_layernorm.weight) + copied_numel += safe_copy(hfllm.norm.weight, + mgllm.decoder.final_layernorm.weight) if not hf_config.tie_word_embeddings: safe_copy(hfmodel.lm_head.weight, mgllm.output_layer.weight) @@ -259,68 +405,122 @@ def convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl(hfmodel, mgmodel @torch.no_grad() -def convert_checkpoint_from_transformers_to_megatron_dpskv3(hf_model, model, hf_config, tfconfig): +def convert_checkpoint_from_transformers_to_megatron_dpskv3( + hf_model, model, hf_config, tfconfig +): warnings.warn("MTP model is not supported yet", stacklevel=2) numel: int = 0 - numel += safe_copy(hf_model.model.embed_tokens.weight, model.embedding.word_embeddings.weight) + numel += safe_copy(hf_model.model.embed_tokens.weight, + model.embedding.word_embeddings.weight) print(f"{numel=}") - for layer_idx, (layer, hf_layer) in enumerate(zip(model.decoder.layers, hf_model.model.layers, strict=True)): + for layer_idx, (layer, hf_layer) in enumerate( + zip(model.decoder.layers, hf_model.model.layers, strict=True) + ): numel_cur: int = numel - numel += safe_copy(hf_layer.input_layernorm.weight, layer.input_layernorm.weight) + numel += safe_copy( + hf_layer.input_layernorm.weight, layer.input_layernorm.weight + ) if hf_config.q_lora_rank is None: - numel += safe_copy(hf_layer.self_attn.q_proj.weight, layer.self_attention.linear_q_proj.weight) + numel += safe_copy( + hf_layer.self_attn.q_proj.weight, + layer.self_attention.linear_q_proj.weight, + ) else: - numel += safe_copy(hf_layer.self_attn.q_a_proj.weight, layer.self_attention.linear_q_down_proj.weight) - numel += safe_copy(hf_layer.self_attn.q_b_proj.weight, layer.self_attention.linear_q_up_proj.weight) numel += safe_copy( - hf_layer.self_attn.q_a_layernorm.weight, layer.self_attention.linear_q_up_proj.layer_norm_weight + hf_layer.self_attn.q_a_proj.weight, + layer.self_attention.linear_q_down_proj.weight, + ) + numel += safe_copy( + hf_layer.self_attn.q_b_proj.weight, + layer.self_attention.linear_q_up_proj.weight, + ) + numel += safe_copy( + hf_layer.self_attn.q_a_layernorm.weight, + layer.self_attention.linear_q_up_proj.layer_norm_weight, ) numel += safe_copy( - hf_layer.self_attn.kv_a_proj_with_mqa.weight, layer.self_attention.linear_kv_down_proj.weight + hf_layer.self_attn.kv_a_proj_with_mqa.weight, + layer.self_attention.linear_kv_down_proj.weight, + ) + numel += safe_copy( + hf_layer.self_attn.kv_b_proj.weight, + layer.self_attention.linear_kv_up_proj.weight, ) - numel += safe_copy(hf_layer.self_attn.kv_b_proj.weight, layer.self_attention.linear_kv_up_proj.weight) numel += safe_copy( - hf_layer.self_attn.kv_a_layernorm.weight, layer.self_attention.linear_kv_up_proj.layer_norm_weight + hf_layer.self_attn.kv_a_layernorm.weight, + layer.self_attention.linear_kv_up_proj.layer_norm_weight, ) - numel += safe_copy(hf_layer.self_attn.o_proj.weight, layer.self_attention.linear_proj.weight) + numel += safe_copy(hf_layer.self_attn.o_proj.weight, + layer.self_attention.linear_proj.weight) if not hasattr(layer.mlp, "router"): - numel += safe_copy(hf_layer.post_attention_layernorm.weight, layer.mlp.linear_fc1.layer_norm_weight) numel += safe_copy( - torch.cat([hf_layer.mlp.gate_proj.weight, hf_layer.mlp.up_proj.weight]), layer.mlp.linear_fc1.weight + hf_layer.post_attention_layernorm.weight, + layer.mlp.linear_fc1.layer_norm_weight, + ) + numel += safe_copy( + torch.cat([hf_layer.mlp.gate_proj.weight, hf_layer.mlp.up_proj.weight]), + layer.mlp.linear_fc1.weight, + ) + numel += safe_copy( + hf_layer.mlp.down_proj.weight, layer.mlp.linear_fc2.weight ) - numel += safe_copy(hf_layer.mlp.down_proj.weight, layer.mlp.linear_fc2.weight) else: - numel += safe_copy(hf_layer.mlp.gate.weight, layer.mlp.router.weight) + numel += safe_copy(hf_layer.mlp.gate.weight, + layer.mlp.router.weight) # NOTE: the e_score_correction_bias in mcore model will be initialized with bfloat16 and \ - # recover to fp32 in the first forward. There is always a diff in the bias between two models (~0.3%) + # recover to fp32 in the first forward. There is always a diff in + # the bias between two models (~0.3%) numel += safe_copy( - hf_layer.mlp.gate.e_score_correction_bias, layer.mlp.router.expert_bias, skip_dtype_assert=True + hf_layer.mlp.gate.e_score_correction_bias, + layer.mlp.router.expert_bias, + skip_dtype_assert=True, ) if tfconfig.moe_grouped_gemm: for i, hf_expert in enumerate(hf_layer.mlp.experts): - fc1_weight = torch.cat([hf_expert.gate_proj.weight, hf_expert.up_proj.weight]) - linear_fc1_weighti = getattr(layer.mlp.experts.linear_fc1, "weight" + str(i)) + fc1_weight = torch.cat( + [hf_expert.gate_proj.weight, hf_expert.up_proj.weight] + ) + linear_fc1_weighti = getattr( + layer.mlp.experts.linear_fc1, "weight" + str(i) + ) numel += safe_copy(fc1_weight, linear_fc1_weighti) - linear_fc2_weighti = getattr(layer.mlp.experts.linear_fc2, "weight" + str(i)) - numel += safe_copy(hf_expert.down_proj.weight, linear_fc2_weighti) + linear_fc2_weighti = getattr( + layer.mlp.experts.linear_fc2, "weight" + str(i) + ) + numel += safe_copy(hf_expert.down_proj.weight, + linear_fc2_weighti) else: for i, hf_expert in enumerate(hf_layer.mlp.experts): expert = layer.mlp.experts.local_experts[i] - fc1_weight = torch.cat([hf_expert.gate_proj.weight, hf_expert.up_proj.weight]) + fc1_weight = torch.cat( + [hf_expert.gate_proj.weight, hf_expert.up_proj.weight] + ) numel += safe_copy(fc1_weight, expert.linear_fc1.weight) - numel += safe_copy(hf_expert.down_proj.weight, expert.linear_fc2.weight) - numel += safe_copy(hf_layer.post_attention_layernorm.weight, layer.pre_mlp_layernorm.weight) + numel += safe_copy( + hf_expert.down_proj.weight, expert.linear_fc2.weight + ) + numel += safe_copy(hf_layer.post_attention_layernorm.weight, + layer.pre_mlp_layernorm.weight) shared_fc1_weight = torch.cat( - [hf_layer.mlp.shared_experts.gate_proj.weight, hf_layer.mlp.shared_experts.up_proj.weight] + [ + hf_layer.mlp.shared_experts.gate_proj.weight, + hf_layer.mlp.shared_experts.up_proj.weight, + ] + ) + numel += safe_copy( + shared_fc1_weight, layer.mlp.shared_experts.linear_fc1.weight + ) + numel += safe_copy( + hf_layer.mlp.shared_experts.down_proj.weight, + layer.mlp.shared_experts.linear_fc2.weight, ) - numel += safe_copy(shared_fc1_weight, layer.mlp.shared_experts.linear_fc1.weight) - numel += safe_copy(hf_layer.mlp.shared_experts.down_proj.weight, layer.mlp.shared_experts.linear_fc2.weight) print(f"{layer_idx=} {numel=} numel this layer={numel - numel_cur}") - numel += safe_copy(hf_model.model.norm.weight, model.decoder.final_layernorm.weight) + numel += safe_copy(hf_model.model.norm.weight, + model.decoder.final_layernorm.weight) if not hf_config.tie_word_embeddings: numel += safe_copy(hf_model.lm_head.weight, model.output_layer.weight) @@ -333,7 +533,13 @@ def noop_context() -> Any: yield -def convert_hf_to_mcore(hf_model_path, output_path, use_cpu_initialization=False, test=False, trust_remote_code=False): +def convert_hf_to_mcore( + hf_model_path, + output_path, + use_cpu_initialization=False, + test=False, + trust_remote_code=False, +): os.makedirs(output_path, exist_ok=True) if len(os.listdir(output_path)) > 0 and not test: print(f"Output path {output_path} is not empty, skipping conversion") @@ -375,7 +581,9 @@ def megatron_model_provider(pre_process, post_process): ) return parallel_model - context: Callable[..., ContextManager] = init_empty_weights if use_cpu_initialization else noop_context + context: Callable[..., ContextManager] = ( + init_empty_weights if use_cpu_initialization else noop_context + ) with context(): model = get_model( model_provider_func=megatron_model_provider, @@ -395,29 +603,44 @@ def megatron_model_provider(pre_process, post_process): # init hf model if "Qwen2_5_VLForConditionalGeneration" in hf_config.architectures: hf_model = AutoModelForImageTextToText.from_pretrained( - hf_model_path, torch_dtype=torch.bfloat16, trust_remote_code=trust_remote_code + hf_model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=trust_remote_code, ) else: hf_model = AutoModelForCausalLM.from_pretrained( - hf_model_path, torch_dtype=torch.bfloat16, trust_remote_code=trust_remote_code + hf_model_path, + torch_dtype=torch.bfloat16, + trust_remote_code=trust_remote_code, ) hf_state_dict = hf_model.state_dict() # load hf state dict to megatron model if "Qwen2MoeForCausalLM" in hf_config.architectures: - convert_checkpoint_from_transformers_to_megatron(hf_model, model[0].module, hf_config) + convert_checkpoint_from_transformers_to_megatron( + hf_model, model[0].module, hf_config + ) elif "Qwen2_5_VLForConditionalGeneration" in hf_config.architectures: - convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl(hf_model, model[0].module, hf_config) + convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl( + hf_model, model[0].module, hf_config + ) elif "DeepseekV3ForCausalLM" in hf_config.architectures: numel: int = convert_checkpoint_from_transformers_to_megatron_dpskv3( hf_model, model[0].module, hf_config, tfconfig=tfconfig ) if numel != hf_model.num_parameters(): - warnings.warn(f"numel mismatch: {numel=} != {hf_model.num_parameters()=}", stacklevel=1) + warnings.warn( + f"numel mismatch: {numel=} != {hf_model.num_parameters()=}", + stacklevel=1, + ) elif "Qwen3MoeForCausalLM" in hf_config.architectures: - convert_checkpoint_from_transformers_to_megatron(hf_model, model[0].module, hf_config) + convert_checkpoint_from_transformers_to_megatron( + hf_model, model[0].module, hf_config + ) else: - assert not use_cpu_initialization, "use_cpu_initialization is only supported for MoE model" + assert ( + not use_cpu_initialization + ), "use_cpu_initialization is only supported for MoE model" from verl.models.mcore.loader import load_state_dict_to_megatron_gptmodel load_state_dict_to_megatron_gptmodel( @@ -433,7 +656,12 @@ def megatron_model_provider(pre_process, post_process): # save megatron model if len(os.listdir(output_path)) == 0: - dist_checkpointing.save(megatron_state_dict, output_path, sharded_strategy=None, async_sharded_save=False) + dist_checkpointing.save( + megatron_state_dict, + output_path, + sharded_strategy=None, + async_sharded_save=False, + ) if test: test_conversion(megatron_model_provider, tfconfig, output_path, model) @@ -441,5 +669,9 @@ def megatron_model_provider(pre_process, post_process): if __name__ == "__main__": args = _init_args() convert_hf_to_mcore( - args.hf_model_path, args.output_path, args.use_cpu_initialization, args.test, args.trust_remote_code + args.hf_model_path, + args.output_path, + args.use_cpu_initialization, + args.test, + args.trust_remote_code, ) diff --git a/Agent0/executor_train/verl/scripts/diagnose.py b/Agent0/executor_train/verl/scripts/diagnose.py index 174b1f9..3cd40f0 100644 --- a/Agent0/executor_train/verl/scripts/diagnose.py +++ b/Agent0/executor_train/verl/scripts/diagnose.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -43,8 +43,7 @@ "cn": { "PYPI(douban)": "https://pypi.douban.com/", "Conda(tsinghua)": "https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/", - } -} + }} def test_connection(name, url, timeout=10): @@ -61,10 +60,18 @@ def test_connection(name, url, timeout=10): try: _ = urlopen(url, timeout=timeout) except Exception as e: - print("Error open {}: {}, {}, DNS finished in {} sec.".format(name, url, e, dns_elapsed)) + print( + "Error open {}: {}, {}, DNS finished in {} sec.".format( + name, url, e, dns_elapsed + ) + ) return load_elapsed = time.time() - start - print("Timing for {}: {}, DNS: {:.4f} sec, LOAD: {:.4f} sec.".format(name, url, dns_elapsed, load_elapsed)) + print( + "Timing for {}: {}, DNS: {:.4f} sec, LOAD: {:.4f} sec.".format( + name, url, dns_elapsed, load_elapsed + ) + ) def check_python(): @@ -88,7 +95,9 @@ def check_pip(): def _get_current_git_commit(): try: - result = subprocess.run(["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True) + result = subprocess.run( + ["git", "rev-parse", "HEAD"], capture_output=True, text=True, check=True + ) return result.stdout.strip() except subprocess.CalledProcessError as e: print(f"Error running git command: {e.stderr.strip()}") @@ -162,7 +171,12 @@ def check_network(args): else: import warnings - warnings.warn("Region {} do not need specific test, please refer to global sites.".format(r), stacklevel=2) + warnings.warn( + "Region {} do not need specific test, please refer to global sites.".format( + r + ), + stacklevel=2, + ) for name, url in URLS.items(): test_connection(name, url, args.timeout) @@ -170,7 +184,13 @@ def check_network(args): def check_environment(): print("----------Environment----------") for k, v in os.environ.items(): - if k.startswith("VERL_") or k.startswith("OMP_") or k.startswith("KMP_") or k == "CC" or k == "CXX": + if ( + k.startswith("VERL_") + or k.startswith("OMP_") + or k.startswith("KMP_") + or k == "CC" + or k == "CXX" + ): print('{}="{}"'.format(k, v)) @@ -191,8 +211,10 @@ def check_cuda_versions(): print(f"CUDA Runtime : {cuda_runtime_version}") import subprocess - nvcc_output = subprocess.check_output(["nvcc", "--version"]).decode("utf-8") - cuda_compiler_version = next((line for line in nvcc_output.splitlines() if "release" in line), None) + nvcc_output = subprocess.check_output( + ["nvcc", "--version"]).decode("utf-8") + cuda_compiler_version = next( + (line for line in nvcc_output.splitlines() if "release" in line), None) if cuda_compiler_version: print(f"CUDA Compiler : {cuda_compiler_version.strip()}") else: @@ -219,7 +241,11 @@ def _get_gpu_info(): """ try: result = subprocess.run( - ["nvidia-smi", "--query-gpu=gpu_name,memory.total", "--format=csv,noheader,nounits"], + [ + "nvidia-smi", + "--query-gpu=gpu_name,memory.total", + "--format=csv,noheader,nounits", + ], capture_output=True, text=True, check=True, @@ -247,7 +273,10 @@ def _get_system_info(): """ cpu_memory = _get_cpu_memory() gpu_count, gpu_info = _get_gpu_info() - return {"cpu_memory": cpu_memory, "gpu_count": gpu_count, "gpu_info": gpu_info} + return { + "cpu_memory": cpu_memory, + "gpu_count": gpu_count, + "gpu_info": gpu_info} def check_system_info(): @@ -268,9 +297,21 @@ def parse_args(): ) choices = ["python", "pip", "verl", "system", "os", "environment"] for choice in choices: - parser.add_argument("--" + choice, default=1, type=int, help="Diagnose {}.".format(choice)) - parser.add_argument("--network", default=0, type=int, help="Diagnose network.") - parser.add_argument("--hardware", default=0, type=int, help="Diagnose hardware.") + parser.add_argument( + "--" + choice, + default=1, + type=int, + help="Diagnose {}.".format(choice)) + parser.add_argument( + "--network", + default=0, + type=int, + help="Diagnose network.") + parser.add_argument( + "--hardware", + default=0, + type=int, + help="Diagnose hardware.") parser.add_argument( "--region", default="", @@ -278,7 +319,12 @@ def parse_args(): help="Additional sites in which region(s) to test. \ Specify 'cn' for example to test mirror sites in China.", ) - parser.add_argument("--timeout", default=10, type=int, help="Connection test timeout threshold, 0 to disable.") + parser.add_argument( + "--timeout", + default=10, + type=int, + help="Connection test timeout threshold, 0 to disable.", + ) args = parser.parse_args() return args diff --git a/Agent0/executor_train/verl/scripts/init_random_model.py b/Agent0/executor_train/verl/scripts/init_random_model.py index 2804bc2..509afee 100644 --- a/Agent0/executor_train/verl/scripts/init_random_model.py +++ b/Agent0/executor_train/verl/scripts/init_random_model.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,7 +14,7 @@ # limitations under the License. """ -This script override a model with custom config and random weights, mainly for create small models for +This script override a model with custom config and random weights, mainly for create small models for debugging purposes. Usage: @@ -31,41 +31,66 @@ import warnings from typing import Any -from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, PretrainedConfig +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoTokenizer, + PretrainedConfig, +) def _init_args(): parser = argparse.ArgumentParser() - parser.add_argument("--hf_model_path", type=str, required=True, help="The path for the huggingface model") - parser.add_argument("--new_config_path", type=str, required=True, help="The path for the new config file") - parser.add_argument("--output_path", type=str, required=True, help="The path for the output random model") + parser.add_argument( + "--hf_model_path", + type=str, + required=True, + help="The path for the huggingface model", + ) + parser.add_argument( + "--new_config_path", + type=str, + required=True, + help="The path for the new config file", + ) + parser.add_argument( + "--output_path", + type=str, + required=True, + help="The path for the output random model", + ) args = parser.parse_args() return args def check_output_path(output_path: str): if os.path.exists(output_path): - warnings.warn(f"Output path '{output_path}' already exists. Will do nothing.", stacklevel=2) + warnings.warn( + f"Output path '{output_path}' already exists. Will do nothing.", + stacklevel=2, + ) exit() else: os.makedirs(output_path, exist_ok=True) print(f"Output path '{output_path}' created.") -def check_configs(original_config: dict[str, Any], new_config: dict[str, Any]) -> bool: +def check_configs( + original_config: dict[str, Any], new_config: dict[str, Any]) -> bool: """ Check if the original config and new config are compatible. This is a placeholder function; actual implementation may vary based on requirements. """ # Example check: ensure 'model_type' is the same - if new_config.get("model_type", None) is not None and original_config.get("model_type") != new_config.get( + if new_config.get("model_type", None) is not None and original_config.get( "model_type" - ): + ) != new_config.get("model_type"): raise RuntimeError("Model types do not match.") for key in new_config: if key not in original_config: warnings.warn( - f"Key '{key}' in new config does not exist in original config, may not take effect.", stacklevel=2 + f"Key '{key}' in new config does not exist in original config, may not take effect.", + stacklevel=2, ) @@ -91,5 +116,7 @@ def init_random_model(hf_model_path, new_config_path, output_path): args = _init_args() check_output_path(args.output_path) init_random_model( - hf_model_path=args.hf_model_path, new_config_path=args.new_config_path, output_path=args.output_path + hf_model_path=args.hf_model_path, + new_config_path=args.new_config_path, + output_path=args.output_path, ) diff --git a/Agent0/executor_train/verl/scripts/legacy_model_merger.py b/Agent0/executor_train/verl/scripts/legacy_model_merger.py index 8a5224a..56d7a83 100644 --- a/Agent0/executor_train/verl/scripts/legacy_model_merger.py +++ b/Agent0/executor_train/verl/scripts/legacy_model_merger.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -87,7 +87,8 @@ class ModelMergerConfig: hf_upload: bool = field(init=False) def __post_init__(self): - self.hf_upload = self.operation == "merge" and bool(self.hf_upload_path) + self.hf_upload = self.operation == "merge" and bool( + self.hf_upload_path) if self.operation == "test": self.target_dir = None self.hf_upload_path = None @@ -105,7 +106,8 @@ def __init__(self, config: ModelMergerConfig): ) self.hf_model_config_path = config.hf_model_path - self.model_config = AutoConfig.from_pretrained(self.hf_model_config_path) + self.model_config = AutoConfig.from_pretrained( + self.hf_model_config_path) def get_transformers_auto_model_class(self): if "ForTokenClassification" in self.model_config.architectures[0]: @@ -115,7 +117,9 @@ def get_transformers_auto_model_class(self): elif "ForConditionalGeneration" in self.model_config.architectures[0]: return AutoModelForVision2Seq - raise NotImplementedError(f"Unknown architecture {self.model_config.architectures}") + raise NotImplementedError( + f"Unknown architecture {self.model_config.architectures}" + ) def patch_model_generation_config(self, model): """ @@ -126,11 +130,13 @@ def patch_model_generation_config(self, model): """ if model.can_generate(): try: - model.generation_config = GenerationConfig.from_pretrained(self.hf_model_config_path) + model.generation_config = GenerationConfig.from_pretrained( + self.hf_model_config_path + ) except OSError: print( - f"Warning: Generation config file not found in {self.hf_model_config_path}, using a generation config created from the model config." - ) + f"Warning: Generation config file not found in { + self.hf_model_config_path}, using a generation config created from the model config.") return model def save_lora_adapter(self, state_dict: dict[str, torch.Tensor]): @@ -143,7 +149,8 @@ def save_lora_adapter(self, state_dict: dict[str, torch.Tensor]): Note: This function change the 'state_dict' in place. """ - lora_params_names = [name for name in state_dict.keys() if "lora_" in name] + lora_params_names = [ + name for name in state_dict.keys() if "lora_" in name] if len(lora_params_names) == 0: return None @@ -163,22 +170,34 @@ def save_lora_adapter(self, state_dict: dict[str, torch.Tensor]): target_modules.add(lora_key.split(".")[-3]) lora_params[lora_key] = state_dict.pop(name) - lora_rank = min(lora_params[lora_key].shape[0], lora_params[lora_key].shape[1]) + lora_rank = min( + lora_params[lora_key].shape[0], + lora_params[lora_key].shape[1]) peft_dict = { "r": lora_rank, - "lora_alpha": 0, # lora_alpha is not set. An error should be raised to inform the user to set it manually. + # lora_alpha is not set. An error should be raised to inform the + # user to set it manually. + "lora_alpha": 0, "target_modules": list(target_modules), } peft_config = peft.LoraConfig(**peft_dict).to_dict() - peft_config["task_type"] = peft_config["task_type"].value if peft_config["task_type"] else None - peft_config["peft_type"] = peft_config["peft_type"].value if peft_config["peft_type"] else None + peft_config["task_type"] = ( + peft_config["task_type"].value if peft_config["task_type"] else None) + peft_config["peft_type"] = ( + peft_config["peft_type"].value if peft_config["peft_type"] else None) peft_config["target_modules"] = list(peft_config["target_modules"]) lora_path = os.path.join(self.config.target_dir, "lora_adapter") os.makedirs(lora_path, exist_ok=True) - with open(os.path.join(lora_path, "adapter_config.json"), "w", encoding="utf-8") as f: + with open( + os.path.join(lora_path, "adapter_config.json"), "w", encoding="utf-8" + ) as f: json.dump(peft_config, f, ensure_ascii=False, indent=4) - save_file(lora_params, os.path.join(lora_path, "adapter_model.safetensors")) + save_file( + lora_params, + os.path.join( + lora_path, + "adapter_model.safetensors")) for name in list(state_dict.keys()): key = ( @@ -193,7 +212,9 @@ def save_lora_adapter(self, state_dict: dict[str, torch.Tensor]): def save_hf_model_and_tokenizer(self, state_dict: dict[str, torch.Tensor]): auto_model_class = self.get_transformers_auto_model_class() with init_empty_weights(): - model = auto_model_class.from_config(self.model_config, torch_dtype=torch.bfloat16) + model = auto_model_class.from_config( + self.model_config, torch_dtype=torch.bfloat16 + ) model.to_empty(device="cpu") model = self.patch_model_generation_config(model) @@ -219,8 +240,16 @@ def upload_to_huggingface(self): from huggingface_hub import HfApi api = HfApi() - api.create_repo(repo_id=self.config.hf_upload_path, private=self.config.private, exist_ok=True) - api.upload_folder(folder_path=self.config.target_dir, repo_id=self.config.hf_upload_path, repo_type="model") + api.create_repo( + repo_id=self.config.hf_upload_path, + private=self.config.private, + exist_ok=True, + ) + api.upload_folder( + folder_path=self.config.target_dir, + repo_id=self.config.hf_upload_path, + repo_type="model", + ) @abstractmethod def merge_and_save(self): @@ -235,17 +264,21 @@ def _get_world_size(self) -> int: if match: return int(match.group(1)) raise FileNotFoundError( - f"Could not determine world size. No file matching 'model_world_size_(\d+)_rank_0.pt' found in {self.config.local_dir}" - ) + f"Could not determine world size. No file matching 'model_world_size_(\\d+)_rank_0.pt' found in { + self.config.local_dir}") def _load_rank_zero_state_dict(self, world_size: int) -> dict: return torch.load( - Path(self.config.local_dir) / f"model_world_size_{world_size}_rank_0.pt", + Path( + self.config.local_dir) / + f"model_world_size_{world_size}_rank_0.pt", map_location="cpu", weights_only=False, ) - def _extract_device_mesh_info(self, state_dict: dict, world_size: int) -> tuple[np.ndarray, tuple[str, ...]]: + def _extract_device_mesh_info( + self, state_dict: dict, world_size: int + ) -> tuple[np.ndarray, tuple[str, ...]]: """ Retrieves sharding information (device_mesh, mesh_dim_names) from a DTensor in the state_dict. If no DTensor is found, infers a simple FSDP mesh based on world_size. @@ -269,7 +302,10 @@ def _calculate_shard_configuration( self, mesh: np.ndarray, mesh_dim_names: tuple[str, ...] ) -> tuple[int, tuple[int, ...]]: """Calculates the total number of shards and the shape of the device mesh.""" - assert mesh_dim_names in (("fsdp",), ("ddp", "fsdp")), f"Unsupported mesh_dim_names {mesh_dim_names}" + assert mesh_dim_names in ( + ("fsdp",), + ("ddp", "fsdp"), + ), f"Unsupported mesh_dim_names {mesh_dim_names}" if "tp" in mesh_dim_names: # TODO: "tp" is not supported yet due to the above assert @@ -281,7 +317,9 @@ def _calculate_shard_configuration( return total_shards, mesh_shape - def _merge_by_placement(self, tensors: list[torch.Tensor], placement: Placement) -> torch.Tensor: + def _merge_by_placement( + self, tensors: list[torch.Tensor], placement: Placement + ) -> torch.Tensor: """Merges a list of tensors based on their DTensor placement""" if placement.is_replicate(): return tensors[0] @@ -293,19 +331,35 @@ def _merge_by_placement(self, tensors: list[torch.Tensor], placement: Placement) raise NotImplementedError(f"Unsupported placement: {placement}") def _load_and_merge_state_dicts( - self, world_size: int, total_shards: int, mesh_shape: tuple[int, ...], mesh_dim_names: tuple[str, ...] + self, + world_size: int, + total_shards: int, + mesh_shape: tuple[int, ...], + mesh_dim_names: tuple[str, ...], ) -> dict[str, torch.Tensor]: model_state_dict_lst = [None] * total_shards def process_one_shard(rank: int, model_state_dict_lst: list): - model_path = Path(self.config.local_dir) / f"model_world_size_{world_size}_rank_{rank}.pt" - state_dict = torch.load(model_path, map_location="cpu", weights_only=False) + model_path = ( + Path(self.config.local_dir) + / f"model_world_size_{world_size}_rank_{rank}.pt" + ) + state_dict = torch.load( + model_path, + map_location="cpu", + weights_only=False) model_state_dict_lst[rank] = state_dict return state_dict with ThreadPoolExecutor(max_workers=min(32, os.cpu_count())) as executor: - futures = [executor.submit(process_one_shard, rank, model_state_dict_lst) for rank in range(total_shards)] - for future in tqdm(futures, desc=f"Loading {total_shards} FSDP shards", total=total_shards): + futures = [ + executor.submit(process_one_shard, rank, model_state_dict_lst) + for rank in range(total_shards) + ] + for future in tqdm( + futures, + desc=f"Loading {total_shards} FSDP shards", + total=total_shards): future.result() # Merge state dicts from all shards @@ -346,7 +400,8 @@ def process_one_shard(rank: int, model_state_dict_lst: list): # 1-D list, FSDP without TP assert len(placements) == 1 shards = state_dict[key] - state_dict[key] = self._merge_by_placement(shards, placements[0]) + state_dict[key] = self._merge_by_placement( + shards, placements[0]) else: # 2-D list, FSDP + TP raise NotImplementedError("FSDP + TP is not supported yet") @@ -359,17 +414,25 @@ def merge_and_save(self): world_size = self._get_world_size() rank_zero_state_dict = self._load_rank_zero_state_dict(world_size) - mesh, mesh_dim_names = self._extract_device_mesh_info(rank_zero_state_dict, world_size) + mesh, mesh_dim_names = self._extract_device_mesh_info( + rank_zero_state_dict, world_size + ) print(f"Got device mesh {mesh}, mesh_dim_names {mesh_dim_names}") - total_shards, mesh_shape = self._calculate_shard_configuration(mesh, mesh_dim_names) - print(f"Processing model shards with {total_shards} {mesh_shape} in total") + total_shards, mesh_shape = self._calculate_shard_configuration( + mesh, mesh_dim_names + ) + print( + f"Processing model shards with {total_shards} {mesh_shape} in total") - merged_state_dict = self._load_and_merge_state_dicts(world_size, total_shards, mesh_shape, mesh_dim_names) + merged_state_dict = self._load_and_merge_state_dicts( + world_size, total_shards, mesh_shape, mesh_dim_names + ) if self.config.operation == "test": if not self.config.test_hf_dir: - raise ValueError("test_hf_dir must be provided for test operation") + raise ValueError( + "test_hf_dir must be provided for test operation") self._test_state_dict(merged_state_dict) elif self.config.operation == "merge": self.save_hf_model_and_tokenizer(merged_state_dict) @@ -381,7 +444,9 @@ def merge_and_save(self): def _test_state_dict(self, state_dict: dict[str, torch.Tensor]): auto_model_class = self.get_transformers_auto_model_class() - hf_model = auto_model_class.from_pretrained(self.config.test_hf_dir, torch_dtype=torch.bfloat16) + hf_model = auto_model_class.from_pretrained( + self.config.test_hf_dir, torch_dtype=torch.bfloat16 + ) hf_state_dict = hf_model.state_dict() del hf_model @@ -389,39 +454,53 @@ def _test_state_dict(self, state_dict: dict[str, torch.Tensor]): collected_keys = set(state_dict.keys()) missing_keys = hf_model_keys - collected_keys - assert len(missing_keys) == 0, f"Missing keys in collected state dict: {list(sorted(missing_keys))}" + assert ( + len(missing_keys) == 0), f"Missing keys in collected state dict: { + list( + sorted(missing_keys))}" extra_keys = collected_keys - hf_model_keys - assert len(extra_keys) == 0, f"Extra keys in collected state dict: {list(sorted(extra_keys))}" + assert ( + len(extra_keys) == 0 + ), f"Extra keys in collected state dict: {list(sorted(extra_keys))}" for key in hf_model_keys: hf_shape = hf_state_dict[key].shape collected_shape = state_dict[key].shape - assert hf_shape == collected_shape, ( - f"Shape mismatch for key '{key}': original {hf_shape} vs collected {collected_shape}" - ) + assert ( + hf_shape == collected_shape + ), f"Shape mismatch for key '{key}': original {hf_shape} vs collected {collected_shape}" hf_dtype = hf_state_dict[key].dtype collected_dtype = state_dict[key].dtype - assert hf_dtype == collected_dtype, ( - f"Dtype mismatch for key '{key}': original {hf_dtype} vs collected {collected_dtype}" - ) + assert ( + hf_dtype == collected_dtype + ), f"Dtype mismatch for key '{key}': original {hf_dtype} vs collected {collected_dtype}" - torch.testing.assert_close(hf_state_dict[key], state_dict[key], atol=1e-6, rtol=1e-6) + torch.testing.assert_close( + hf_state_dict[key], state_dict[key], atol=1e-6, rtol=1e-6 + ) - print("FSDP checks passed: The merged state_dict matches the hf model saved by FSDPCheckpointManager.") + print( + "FSDP checks passed: The merged state_dict matches the hf model saved by FSDPCheckpointManager." + ) class MegatronModelMerger(BaseModelMerger): def __init__(self, config: ModelMergerConfig): - from verl.utils.megatron_utils import get_hf_config_and_tokenizer_checkpoint_path + from verl.utils.megatron_utils import ( + get_hf_config_and_tokenizer_checkpoint_path, + ) - config.hf_model_config_path = get_hf_config_and_tokenizer_checkpoint_path(config.local_dir) + config.hf_model_config_path = get_hf_config_and_tokenizer_checkpoint_path( + config.local_dir) super().__init__(config) self.params_mapping = { # megatron core gpt model name, huggingface model name - # NOTICE: It's a little bit tricky, when 2 keys have the same prefix, we need to make sure the longer key within the containing relationship is processed first. + # NOTICE: It's a little bit tricky, when 2 keys have the same + # prefix, we need to make sure the longer key within the containing + # relationship is processed first. "embedding.word_embeddings": "model.embed_tokens", # attn "self_attention.linear_qkv.layer_norm_weight": "input_layernorm.weight", @@ -456,7 +535,8 @@ def __init__(self, config: ModelMergerConfig): "output_layer": "lm_head", } - def _get_tp_pp_rank_from_sharded_dir(self, sharded_dir: str) -> tuple[int, int]: + def _get_tp_pp_rank_from_sharded_dir( + self, sharded_dir: str) -> tuple[int, int]: tp_rank = pp_rank = None rank_list = sharded_dir.split("_")[2:] if re.match(r"mp_rank_(\d\d)_(\d\d\d)", sharded_dir): @@ -466,11 +546,15 @@ def _get_tp_pp_rank_from_sharded_dir(self, sharded_dir: str) -> tuple[int, int]: tp_rank = int(rank_list[0]) pp_rank = 0 - assert tp_rank is not None and pp_rank is not None, f"Invalid sharded dir {sharded_dir}" + assert ( + tp_rank is not None and pp_rank is not None + ), f"Invalid sharded dir {sharded_dir}" return tp_rank, pp_rank - def _check_megatron_checkpoint_path(self, model_path: str) -> tuple[list[str], int, int]: + def _check_megatron_checkpoint_path( + self, model_path: str + ) -> tuple[list[str], int, int]: """ Validates the Megatron checkpoint structure (presence of 'model.pt' in sharded directories). Determines TP and PP sizes from directory names. @@ -479,8 +563,11 @@ def _check_megatron_checkpoint_path(self, model_path: str) -> tuple[list[str], i pp_size = 0 sharded_dirs = sorted(os.listdir(model_path)) for sharded_dir in sharded_dirs: - assert "model.pt" in os.listdir(Path(model_path) / sharded_dir), f"model.pt not found in {sharded_dir}" - tp_rank, pp_rank = self._get_tp_pp_rank_from_sharded_dir(sharded_dir) + assert "model.pt" in os.listdir( + Path(model_path) / sharded_dir + ), f"model.pt not found in {sharded_dir}" + tp_rank, pp_rank = self._get_tp_pp_rank_from_sharded_dir( + sharded_dir) tp_size = max(tp_size, tp_rank + 1) pp_size = max(pp_size, pp_rank + 1) return sharded_dirs, tp_size, pp_size @@ -514,15 +601,23 @@ def _merge_across_tp( num_q_per_kv = config.num_attention_heads // config.num_key_value_heads assert tp_data[0].shape[0] % (num_q_per_kv + 2) == 0 kv_size_per_tp = tp_data[0].shape[0] // (num_q_per_kv + 2) - split_size = [kv_size_per_tp * num_q_per_kv, kv_size_per_tp, kv_size_per_tp] + split_size = [ + kv_size_per_tp * + num_q_per_kv, + kv_size_per_tp, + kv_size_per_tp] for infer_param in tp_data: num_query_groups_per_partition = config.num_key_value_heads // tp_size for chunk in infer_param.chunk(num_query_groups_per_partition): split_size = [ - kv_size_per_tp * num_q_per_kv // num_query_groups_per_partition, - kv_size_per_tp // num_query_groups_per_partition, - kv_size_per_tp // num_query_groups_per_partition, + kv_size_per_tp * + num_q_per_kv // + num_query_groups_per_partition, + kv_size_per_tp // + num_query_groups_per_partition, + kv_size_per_tp // + num_query_groups_per_partition, ] q, k, v = chunk.split(split_size) q_lst.append(q) @@ -533,7 +628,12 @@ def _merge_across_tp( k = torch.cat(k_lst, dim=0) v = torch.cat(v_lst, dim=0) return [q, k, v] - elif "layer_norm" in key or "layernorm" in key or "router" in key or ("output_layer" in key and is_value_model): + elif ( + "layer_norm" in key + or "layernorm" in key + or "router" in key + or ("output_layer" in key and is_value_model) + ): return tp_data[0] else: dim = 0 @@ -541,20 +641,33 @@ def _merge_across_tp( dim = 1 return torch.cat(tp_data, dim=dim) - def _load_state_dicts( - self, model_ckpt_path: str, sharded_dirs: list[str], tp_size: int, pp_size: int - ) -> list[list[dict]]: - model_state_dict_lst = [[None for _ in range(tp_size)] for _ in range(pp_size)] + def _load_state_dicts(self, + model_ckpt_path: str, + sharded_dirs: list[str], + tp_size: int, + pp_size: int) -> list[list[dict]]: + model_state_dict_lst = [ + [None for _ in range(tp_size)] for _ in range(pp_size)] def _process_one_megatron_shard(sharded_dir: str): model_file_path = Path(model_ckpt_path) / sharded_dir / "model.pt" - state_dict = torch.load(model_file_path, map_location="cpu", weights_only=False) - tp_rank, pp_rank = self._get_tp_pp_rank_from_sharded_dir(sharded_dir) + state_dict = torch.load( + model_file_path, map_location="cpu", weights_only=False + ) + tp_rank, pp_rank = self._get_tp_pp_rank_from_sharded_dir( + sharded_dir) model_state_dict_lst[pp_rank][tp_rank] = state_dict with ThreadPoolExecutor(max_workers=min(32, os.cpu_count())) as executor: - futures = [executor.submit(_process_one_megatron_shard, sharded_dir) for sharded_dir in sharded_dirs] - for future in tqdm(futures, desc=f"Loading {len(sharded_dirs)} Megatron shards", total=len(sharded_dirs)): + futures = [ + executor.submit(_process_one_megatron_shard, sharded_dir) + for sharded_dir in sharded_dirs + ] + for future in tqdm( + futures, + desc=f"Loading {len(sharded_dirs)} Megatron shards", + total=len(sharded_dirs), + ): future.result() return model_state_dict_lst @@ -568,8 +681,7 @@ def _check_megatron_state_key(self, key: str) -> bool: """ if key.startswith("model."): raise ValueError( - f"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder/embedding/output_layer' in TransformerLayer." - ) + f"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder/embedding/output_layer' in TransformerLayer.") skip_checking_keys = ["embedding.word_embeddings", "output_layer"] for skip_key in skip_checking_keys: @@ -580,8 +692,7 @@ def _check_megatron_state_key(self, key: str) -> bool: # Exclude extra state keys if not key.startswith("decoder"): raise ValueError( - f"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder' in TransformerLayer." - ) + f"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder' in TransformerLayer.") def _merge_state_dicts( self, model_state_dict_lst: list[list[dict]], tp_size: int, pp_size: int @@ -597,13 +708,18 @@ def _merge_state_dicts( for key in keys: if "extra_state" in key: continue - if self.config.tie_word_embedding and ("output_layer" in key): - print("skip lm_head and reward_head loading because of tie_word_embeddings") + if self.config.tie_word_embedding and ( + "output_layer" in key): + print( + "skip lm_head and reward_head loading because of tie_word_embeddings" + ) continue self._check_megatron_state_key(key) hf_name = self._replace_name(key, self.params_mapping) - assert hf_name is not None, f"Failed to convert layer name [{key}] from megatron to huggingface." + assert ( + hf_name is not None + ), f"Failed to convert layer name [{key}] from megatron to huggingface." if "model.layers." in hf_name: local_layer_no = int(hf_name.split(".")[2]) layers_handled = max(local_layer_no, layers_handled) @@ -612,10 +728,20 @@ def _merge_state_dicts( new_key_list[2] = str(global_layer_no) hf_name = ".".join(new_key_list) else: - warnings.warn(f"hf_name {hf_name} will not be fixed with layer number", stacklevel=2) + warnings.warn( + f"hf_name {hf_name} will not be fixed with layer number", stacklevel=2, ) - tp_data = [model_state_dict_lst[pp_rank][tp_rank][vpp_rank][key] for tp_rank in range(tp_size)] - merged = self._merge_across_tp(key, tp_data, self.model_config, tp_size, self.config.is_value_model) + tp_data = [ + model_state_dict_lst[pp_rank][tp_rank][vpp_rank][key] + for tp_rank in range(tp_size) + ] + merged = self._merge_across_tp( + key, + tp_data, + self.model_config, + tp_size, + self.config.is_value_model, + ) if not isinstance(merged, list): state_dict[hf_name] = merged @@ -625,11 +751,15 @@ def _merge_state_dicts( state_dict[hf_name.replace("qkv", n)] = d elif len(merged) == 2: # split gate up - state_dict[hf_name.replace("gate_up", "gate")] = merged[0] - state_dict[hf_name.replace("gate_up", "up")] = merged[1] + state_dict[hf_name.replace( + "gate_up", "gate")] = merged[0] + state_dict[hf_name.replace( + "gate_up", "up")] = merged[1] print( - f"converted {key} to {hf_name} with shape {merged.shape if isinstance(merged, torch.Tensor) else [t.shape for t in merged]}" - ) + f"converted {key} to {hf_name} with shape { + merged.shape if isinstance( + merged, torch.Tensor) else [ + t.shape for t in merged]}") layers_cum += layers_handled + 1 # zero based @@ -639,16 +769,25 @@ def merge_and_save(self): from verl.utils.megatron_utils import get_model_checkpoint_path model_ckpt_path = get_model_checkpoint_path(self.config.local_dir) - sharded_dirs, tp_size, pp_size = self._check_megatron_checkpoint_path(model_ckpt_path) - print(f"sharded_dirs: {sharded_dirs}, tp_size: {tp_size}, pp_size: {pp_size}, mp_size: {len(sharded_dirs)}") + sharded_dirs, tp_size, pp_size = self._check_megatron_checkpoint_path( + model_ckpt_path + ) + print( + f"sharded_dirs: {sharded_dirs}, tp_size: {tp_size}, pp_size: {pp_size}, mp_size: { + len(sharded_dirs)}") - model_state_dict_lst = self._load_state_dicts(model_ckpt_path, sharded_dirs, tp_size, pp_size) - merged_state_dict = self._merge_state_dicts(model_state_dict_lst, tp_size, pp_size) + model_state_dict_lst = self._load_state_dicts( + model_ckpt_path, sharded_dirs, tp_size, pp_size + ) + merged_state_dict = self._merge_state_dicts( + model_state_dict_lst, tp_size, pp_size + ) del model_state_dict_lst if self.config.operation == "test": if not self.config.test_hf_dir: - raise ValueError("test_hf_dir must be provided for test operation") + raise ValueError( + "test_hf_dir must be provided for test operation") self._test_state_dict(merged_state_dict) elif self.config.operation == "merge": self.save_hf_model_and_tokenizer(merged_state_dict) @@ -662,11 +801,15 @@ def _test_state_dict(self, state_dict: dict[str, torch.Tensor]): Compares the merged Megatron state_dict against a reference safetensors model. Applies necessary name mappings from Megatron to Hugging Face conventions using _replace_name. """ - ref_state_dict = load_file(Path(self.config.test_hf_dir) / "model.safetensors") + ref_state_dict = load_file( + Path( + self.config.test_hf_dir) / + "model.safetensors") for name, loaded_weight in state_dict.items(): # name = self._replace_name(original_name, self.params_mapping) - if not name or name.endswith(".bias") and name not in ref_state_dict: + if not name or name.endswith( + ".bias") and name not in ref_state_dict: continue if "rotary_emb.inv_freq" in name: continue @@ -676,9 +819,11 @@ def _test_state_dict(self, state_dict: dict[str, torch.Tensor]): raise RuntimeError(f"key: {name} not exist in state_dict") param = ref_state_dict[name] assert loaded_weight.dtype == param.dtype - torch.testing.assert_close(loaded_weight, param, atol=1e-2, rtol=5e-2) + torch.testing.assert_close( + loaded_weight, param, atol=1e-2, rtol=5e-2) - def _replace_name(self, megatron_name: str, name_mapping: dict[str, str]) -> str: + def _replace_name(self, megatron_name: str, + name_mapping: dict[str, str]) -> str: for m_name, v_name in name_mapping.items(): if m_name not in megatron_name: continue @@ -692,13 +837,25 @@ def _replace_name(self, megatron_name: str, name_mapping: dict[str, str]) -> str def main(): parser = argparse.ArgumentParser(description="verl model merger") - subparsers = parser.add_subparsers(dest="operation", required=True, help="Specify 'merge' or 'test' operation.") + subparsers = parser.add_subparsers( + dest="operation", + required=True, + help="Specify 'merge' or 'test' operation.") base_op_parser = argparse.ArgumentParser(add_help=False) base_op_parser.add_argument( - "--backend", type=str, required=True, choices=["fsdp", "megatron"], help="The backend of the model" + "--backend", + type=str, + required=True, + choices=["fsdp", "megatron"], + help="The backend of the model", + ) + base_op_parser.add_argument( + "--local_dir", + type=str, + required=True, + help="Path to the saved model checkpoints", ) - base_op_parser.add_argument("--local_dir", type=str, required=True, help="Path to the saved model checkpoints") base_op_parser.add_argument( "--hf_model_path", type=str, @@ -716,22 +873,38 @@ def main(): help="Whether the model is a value model (currently only Megatron supported)", ) - merge_parser = subparsers.add_parser("merge", parents=[base_op_parser], help="Merge model checkpoints and save.") + merge_parser = subparsers.add_parser( + "merge", + parents=[base_op_parser], + help="Merge model checkpoints and save.") merge_parser.add_argument( - "--target_dir", default="tmp", type=str, help="Directory to save the merged huggingface model" + "--target_dir", + default="tmp", + type=str, + help="Directory to save the merged huggingface model", ) merge_parser.add_argument( - "--hf_upload_path", default=None, type=str, help="Hugging Face repository ID to upload the model" + "--hf_upload_path", + default=None, + type=str, + help="Hugging Face repository ID to upload the model", ) merge_parser.add_argument( - "--private", action="store_true", help="Whether to upload the model to a private Hugging Face repository" + "--private", + action="store_true", + help="Whether to upload the model to a private Hugging Face repository", ) test_parser = subparsers.add_parser( - "test", parents=[base_op_parser], help="Test merged model against a reference Hugging Face model" + "test", + parents=[base_op_parser], + help="Test merged model against a reference Hugging Face model", ) test_parser.add_argument( - "--test_hf_dir", type=str, required=True, help="Path to the reference Hugging Face model directory for testing" + "--test_hf_dir", + type=str, + required=True, + help="Path to the reference Hugging Face model directory for testing", ) args = parser.parse_args() diff --git a/Agent0/executor_train/verl/setup.py b/Agent0/executor_train/verl/setup.py index a4caeba..625eb08 100644 --- a/Agent0/executor_train/verl/setup.py +++ b/Agent0/executor_train/verl/setup.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,7 +12,8 @@ # 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 +# setup.py is the fallback installation script when pyproject.toml does +# not work import os from pathlib import Path diff --git a/Agent0/executor_train/verl/tests/__init__.py b/Agent0/executor_train/verl/tests/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/tests/__init__.py +++ b/Agent0/executor_train/verl/tests/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/experimental/agent_loop/agent_utils.py b/Agent0/executor_train/verl/tests/experimental/agent_loop/agent_utils.py index 3c708c4..fa5f164 100644 --- a/Agent0/executor_train/verl/tests/experimental/agent_loop/agent_utils.py +++ b/Agent0/executor_train/verl/tests/experimental/agent_loop/agent_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -22,36 +22,49 @@ from verl.workers.fsdp_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker -def init_agent_loop_manager(config: DictConfig) -> AgentLoopManager | RayWorkerGroup: - # =========================== 1. Create hybrid ActorRollout workers =========================== +def init_agent_loop_manager( + config: DictConfig) -> AgentLoopManager | RayWorkerGroup: + # =========================== 1. Create hybrid ActorRollout workers ====== actor_rollout_cls = ( - AsyncActorRolloutRefWorker if config.actor_rollout_ref.rollout.mode == "async" else ActorRolloutRefWorker + AsyncActorRolloutRefWorker + if config.actor_rollout_ref.rollout.mode == "async" + else ActorRolloutRefWorker ) role_worker_mapping = { Role.ActorRollout: ray.remote(actor_rollout_cls), } global_pool_id = "global_pool" resource_pool_spec = { - global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + global_pool_id: [ + config.trainer.n_gpus_per_node] * + config.trainer.nnodes, } mapping = { Role.ActorRollout: global_pool_id, } - resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, mapping=mapping + ) resource_pool_manager.create_resource_pool() - resource_pool_to_cls = {pool: {} for pool in resource_pool_manager.resource_pool_dict.values()} + resource_pool_to_cls = { + pool: {} for pool in resource_pool_manager.resource_pool_dict.values() + } # create actor and rollout resource_pool = resource_pool_manager.get_resource_pool(Role.ActorRollout) actor_rollout_cls = RayClassWithInitArgs( - cls=role_worker_mapping[Role.ActorRollout], config=config.actor_rollout_ref, role="actor_rollout" + cls=role_worker_mapping[Role.ActorRollout], + config=config.actor_rollout_ref, + role="actor_rollout", ) resource_pool_to_cls[resource_pool]["actor_rollout"] = actor_rollout_cls all_wg = {} for resource_pool, class_dict in resource_pool_to_cls.items(): worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict) - wg_dict = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls) + wg_dict = RayWorkerGroup( + 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) actor_rollout_wg = all_wg["actor_rollout"] @@ -60,7 +73,7 @@ def init_agent_loop_manager(config: DictConfig) -> AgentLoopManager | RayWorkerG if config.actor_rollout_ref.rollout.mode == "sync": return actor_rollout_wg - # =========================== 2. Create AgentLoopManager =========================== + # =========================== 2. Create AgentLoopManager ================= agent_loop_manager = AgentLoopManager( config=config, worker_group=actor_rollout_wg, diff --git a/Agent0/executor_train/verl/tests/experimental/agent_loop/test_basic_agent_loop.py b/Agent0/executor_train/verl/tests/experimental/agent_loop/test_basic_agent_loop.py index 20936aa..6b7d72d 100644 --- a/Agent0/executor_train/verl/tests/experimental/agent_loop/test_basic_agent_loop.py +++ b/Agent0/executor_train/verl/tests/experimental/agent_loop/test_basic_agent_loop.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -71,7 +71,12 @@ def test_single_turn(init_config): "content": "Let's play a role playing game. Your name is Alice, your favorite color is blue.", } ], - [{"role": "user", "content": "Let's play a role playing game. Your name is Bob, your favorite color is red."}], + [ + { + "role": "user", + "content": "Let's play a role playing game. Your name is Bob, your favorite color is red.", + } + ], ] batch = DataProto( non_tensor_batch={ @@ -85,7 +90,8 @@ def test_single_turn(init_config): assert len(result) == len(raw_prompts) * n # check result - seq_len = result.batch["prompts"].size(1) + result.batch["responses"].size(1) + seq_len = result.batch["prompts"].size( + 1) + result.batch["responses"].size(1) assert result.batch["input_ids"].size(1) == seq_len assert result.batch["attention_mask"].size(1) == seq_len assert result.batch["position_ids"].size(1) == seq_len @@ -119,7 +125,9 @@ def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: schema = get_json_schema(self.get_current_temperature) return OpenAIFunctionToolSchema(**schema) - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: try: result = self.get_current_temperature(**parameters) return json.dumps(result), 0, {} @@ -132,7 +140,11 @@ def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: schema = get_json_schema(self.get_temperature_date) return OpenAIFunctionToolSchema(**schema) - def get_temperature_date(self, location: str, date: str, unit: str = "celsius"): + def get_temperature_date( + self, + location: str, + date: str, + unit: str = "celsius"): """Get temperature at a location and date. Args: @@ -150,7 +162,9 @@ def get_temperature_date(self, location: str, date: str, unit: str = "celsius"): "unit": unit, } - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: try: result = self.get_temperature_date(**parameters) return json.dumps(result), 0, {} @@ -170,7 +184,7 @@ def test_tool_agent(init_config): } ) - # =========================== 1. Init rollout manager =========================== + # =========================== 1. Init rollout manager ==================== tool_config = { "tools": [ { @@ -193,29 +207,30 @@ def test_tool_agent(init_config): init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 2 agent_loop_manager = init_agent_loop_manager(init_config) - # =========================== 2. Generate sequences =========================== - raw_prompts = [ - [ - {"role": "user", "content": "How are you?"}, - ], - [ - {"role": "user", "content": "What's the temperature in Los Angeles now?"}, - ], - [ - {"role": "user", "content": "What's the temperature in New York now?"}, - ], - [ - { - "role": "system", - "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.\n\n" - "Current Date: 2024-09-30", - }, - {"role": "user", "content": "What's the temperature in San Francisco now? How about tomorrow?"}, - ], - ] + # =========================== 2. Generate sequences ===================== + raw_prompts = [[{"role": "user", + "content": "How are you?"}, + ], + [{"role": "user", + "content": "What's the temperature in Los Angeles now?"}, + ], + [{"role": "user", + "content": "What's the temperature in New York now?"}, + ], + [{"role": "system", + "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.\n\n" + "Current Date: 2024-09-30", + }, + {"role": "user", + "content": "What's the temperature in San Francisco now? How about tomorrow?", + }, + ], + ] batch = DataProto( non_tensor_batch={ - "raw_prompt": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object), + "raw_prompt": np.array( + [np.array(prompt) for prompt in raw_prompts], dtype=object + ), "agent_name": np.array(["tool_agent"] * len(raw_prompts)), }, ) @@ -238,14 +253,20 @@ def test_tool_agent(init_config): tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path) responses = result.batch["responses"] response_mask = result.batch["response_mask"] - assert responses.size() == response_mask.size(), f"{responses.size()} != {response_mask.size()}" + assert ( + responses.size() == response_mask.size() + ), f"{responses.size()} != {response_mask.size()}" # Decode responses with response_mask for i in range(len(responses)): valid_tokens = responses[i][response_mask[i].bool()] response_str = tokenizer.decode(valid_tokens) - assert "" not in response_str, f"found in response: {response_str}" - assert "" not in response_str, f"found in response: {response_str}" + assert ( + "" not in response_str + ), f"found in response: {response_str}" + assert ( + "" not in response_str + ), f"found in response: {response_str}" print(f"response: {response_str}") print("Test passed!") diff --git a/Agent0/executor_train/verl/tests/interactions/__init__.py b/Agent0/executor_train/verl/tests/interactions/__init__.py index b6db0fc..084c798 100644 --- a/Agent0/executor_train/verl/tests/interactions/__init__.py +++ b/Agent0/executor_train/verl/tests/interactions/__init__.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/interactions/test_gsm8k_interaction.py b/Agent0/executor_train/verl/tests/interactions/test_gsm8k_interaction.py index bc16877..8235e4f 100644 --- a/Agent0/executor_train/verl/tests/interactions/test_gsm8k_interaction.py +++ b/Agent0/executor_train/verl/tests/interactions/test_gsm8k_interaction.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -41,12 +41,15 @@ async def test_start_interaction_with_instance_id(self): instance_id = "test_instance" ground_truth = "42" - result_id = await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + result_id = await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) assert result_id == instance_id assert instance_id in self.interaction._instance_dict assert self.interaction._instance_dict[instance_id]["response"] == "" - assert self.interaction._instance_dict[instance_id]["ground_truth"] == ground_truth + assert ( + self.interaction._instance_dict[instance_id]["ground_truth"] == ground_truth) assert self.interaction._instance_dict[instance_id]["reward"] == 0.0 @pytest.mark.asyncio @@ -59,7 +62,8 @@ async def test_start_interaction_without_instance_id(self): assert result_id is not None assert len(result_id) == 36 # UUID4 length assert result_id in self.interaction._instance_dict - assert self.interaction._instance_dict[result_id]["ground_truth"] == ground_truth + assert ( + self.interaction._instance_dict[result_id]["ground_truth"] == ground_truth) @pytest.mark.asyncio async def test_start_interaction_without_ground_truth(self): @@ -78,13 +82,15 @@ async def test_generate_response_correct_answer_with_prefix(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) messages = [{"role": "user", "content": "#### 42"}] with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=1.0): - should_terminate, response, reward, metadata = await self.interaction.generate_response( - instance_id, messages + should_terminate, response, reward, metadata = ( + await self.interaction.generate_response(instance_id, messages) ) assert should_terminate is True @@ -100,13 +106,15 @@ async def test_generate_response_correct_answer_without_prefix(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) messages = [{"role": "user", "content": "42"}] with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=1.0): - should_terminate, response, reward, metadata = await self.interaction.generate_response( - instance_id, messages + should_terminate, response, reward, metadata = ( + await self.interaction.generate_response(instance_id, messages) ) assert should_terminate is True @@ -121,17 +129,22 @@ async def test_generate_response_incorrect_answer(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) messages = [{"role": "user", "content": "24"}] with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=0.0): - should_terminate, response, reward, metadata = await self.interaction.generate_response( - instance_id, messages + should_terminate, response, reward, metadata = ( + await self.interaction.generate_response(instance_id, messages) ) assert should_terminate is False - assert response == "Your response is incorrect! You need to reflect on your answer and try again." + assert ( + response + == "Your response is incorrect! You need to reflect on your answer and try again." + ) assert reward == 0.0 assert self.interaction._instance_dict[instance_id]["response"] == "#### 24" @@ -142,7 +155,9 @@ async def test_generate_response_multiple_messages(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) messages = [ {"role": "user", "content": "What is 2+2?"}, @@ -151,8 +166,8 @@ async def test_generate_response_multiple_messages(self): ] with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=1.0): - should_terminate, response, reward, metadata = await self.interaction.generate_response( - instance_id, messages + should_terminate, response, reward, metadata = ( + await self.interaction.generate_response(instance_id, messages) ) assert should_terminate is True @@ -166,13 +181,15 @@ async def test_generate_response_no_user_message(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) messages = [{"role": "assistant", "content": "Hello!"}] with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=0.0): - should_terminate, response, reward, metadata = await self.interaction.generate_response( - instance_id, messages + should_terminate, response, reward, metadata = ( + await self.interaction.generate_response(instance_id, messages) ) assert should_terminate is False @@ -185,16 +202,22 @@ async def test_calculate_score_direct_call(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) # Set a response self.interaction._instance_dict[instance_id]["response"] = "#### 42" - with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=1.0) as mock_compute: + with patch( + "verl.utils.reward_score.gsm8k.compute_score", return_value=1.0 + ) as mock_compute: score = await self.interaction.calculate_score(instance_id) assert score == 1.0 - mock_compute.assert_called_once_with("#### 42", "42", method="flexible", format_score=0.0, score=1.0) + mock_compute.assert_called_once_with( + "#### 42", "42", method="flexible", format_score=0.0, score=1.0 + ) @pytest.mark.asyncio async def test_calculate_score_with_kwargs(self): @@ -203,16 +226,24 @@ async def test_calculate_score_with_kwargs(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) # Set a response self.interaction._instance_dict[instance_id]["response"] = "#### 24" - with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=0.0) as mock_compute: - score = await self.interaction.calculate_score(instance_id, extra_param="test") + with patch( + "verl.utils.reward_score.gsm8k.compute_score", return_value=0.0 + ) as mock_compute: + score = await self.interaction.calculate_score( + instance_id, extra_param="test" + ) assert score == 0.0 - mock_compute.assert_called_once_with("#### 24", "42", method="flexible", format_score=0.0, score=1.0) + mock_compute.assert_called_once_with( + "#### 24", "42", method="flexible", format_score=0.0, score=1.0 + ) @pytest.mark.asyncio async def test_finalize_interaction(self): @@ -221,7 +252,9 @@ async def test_finalize_interaction(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) assert instance_id in self.interaction._instance_dict @@ -236,7 +269,9 @@ async def test_finalize_interaction_with_kwargs(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) assert instance_id in self.interaction._instance_dict @@ -259,14 +294,16 @@ async def test_full_interaction_workflow_correct(self): ground_truth = "42" # Start interaction - instance_id = await self.interaction.start_interaction(ground_truth=ground_truth) + instance_id = await self.interaction.start_interaction( + ground_truth=ground_truth + ) # Generate response with correct answer messages = [{"role": "user", "content": "42"}] with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=1.0): - should_terminate, response, reward, metadata = await self.interaction.generate_response( - instance_id, messages + should_terminate, response, reward, metadata = ( + await self.interaction.generate_response(instance_id, messages) ) assert should_terminate is True @@ -282,14 +319,16 @@ async def test_full_interaction_workflow_incorrect(self): ground_truth = "42" # Start interaction - instance_id = await self.interaction.start_interaction(ground_truth=ground_truth) + instance_id = await self.interaction.start_interaction( + ground_truth=ground_truth + ) # Generate response with incorrect answer messages = [{"role": "user", "content": "24"}] with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=0.0): - should_terminate, response, reward, metadata = await self.interaction.generate_response( - instance_id, messages + should_terminate, response, reward, metadata = ( + await self.interaction.generate_response(instance_id, messages) ) assert should_terminate is False @@ -300,8 +339,8 @@ async def test_full_interaction_workflow_incorrect(self): messages.append({"role": "user", "content": "42"}) with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=1.0): - should_terminate, response, reward, metadata = await self.interaction.generate_response( - instance_id, messages + should_terminate, response, reward, metadata = ( + await self.interaction.generate_response(instance_id, messages) ) assert should_terminate is True @@ -318,8 +357,12 @@ async def test_multiple_concurrent_interactions(self): ground_truth_2 = "24" # Start multiple interactions - instance_id_1 = await self.interaction.start_interaction(ground_truth=ground_truth_1) - instance_id_2 = await self.interaction.start_interaction(ground_truth=ground_truth_2) + instance_id_1 = await self.interaction.start_interaction( + ground_truth=ground_truth_1 + ) + instance_id_2 = await self.interaction.start_interaction( + ground_truth=ground_truth_2 + ) assert len(self.interaction._instance_dict) == 2 assert instance_id_1 in self.interaction._instance_dict @@ -329,9 +372,15 @@ async def test_multiple_concurrent_interactions(self): messages_1 = [{"role": "user", "content": "42"}] messages_2 = [{"role": "user", "content": "24"}] - with patch("verl.utils.reward_score.gsm8k.compute_score", side_effect=[1.0, 1.0]): - should_terminate_1, _, reward_1, _ = await self.interaction.generate_response(instance_id_1, messages_1) - should_terminate_2, _, reward_2, _ = await self.interaction.generate_response(instance_id_2, messages_2) + with patch( + "verl.utils.reward_score.gsm8k.compute_score", side_effect=[1.0, 1.0] + ): + should_terminate_1, _, reward_1, _ = ( + await self.interaction.generate_response(instance_id_1, messages_1) + ) + should_terminate_2, _, reward_2, _ = ( + await self.interaction.generate_response(instance_id_2, messages_2) + ) assert should_terminate_1 is True assert should_terminate_2 is True @@ -351,13 +400,15 @@ async def test_edge_case_empty_messages(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) messages = [] with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=0.0): - should_terminate, response, reward, metadata = await self.interaction.generate_response( - instance_id, messages + should_terminate, response, reward, metadata = ( + await self.interaction.generate_response(instance_id, messages) ) assert should_terminate is False @@ -371,15 +422,15 @@ async def test_edge_case_message_without_content(self): ground_truth = "42" # Setup instance - await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth) + await self.interaction.start_interaction( + instance_id=instance_id, ground_truth=ground_truth + ) - messages = [ - {"role": "user"} # Missing content field - ] + messages = [{"role": "user"}] # Missing content field with patch("verl.utils.reward_score.gsm8k.compute_score", return_value=0.0): - should_terminate, response, reward, metadata = await self.interaction.generate_response( - instance_id, messages + should_terminate, response, reward, metadata = ( + await self.interaction.generate_response(instance_id, messages) ) assert should_terminate is False @@ -414,7 +465,9 @@ def test_name_attribute_initialization(self): # Test with default name when not provided in config config_without_name = {} interaction_without_name = Gsm8kInteraction(config_without_name) - assert interaction_without_name.name == "interaction_agent" # Default from BaseInteraction + assert ( + interaction_without_name.name == "interaction_agent" + ) # Default from BaseInteraction # Test that name is accessible as attribute assert hasattr(self.interaction, "name") diff --git a/Agent0/executor_train/verl/tests/interactions/test_interaction_registry.py b/Agent0/executor_train/verl/tests/interactions/test_interaction_registry.py index 7fe193b..289e29b 100644 --- a/Agent0/executor_train/verl/tests/interactions/test_interaction_registry.py +++ b/Agent0/executor_train/verl/tests/interactions/test_interaction_registry.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,11 +31,14 @@ class TestInteractionRegistry: def test_get_interaction_class(self): """Test getting interaction class by name.""" # Test getting base interaction class - base_cls = get_interaction_class("verl.interactions.base.BaseInteraction") + base_cls = get_interaction_class( + "verl.interactions.base.BaseInteraction") assert base_cls == BaseInteraction # Test getting gsm8k interaction class - gsm8k_cls = get_interaction_class("verl.interactions.gsm8k_interaction.Gsm8kInteraction") + gsm8k_cls = get_interaction_class( + "verl.interactions.gsm8k_interaction.Gsm8kInteraction" + ) assert gsm8k_cls == Gsm8kInteraction def test_initialize_single_interaction_from_config(self): @@ -47,16 +50,15 @@ def test_initialize_single_interaction_from_config(self): "name": "test_gsm8k", "class_name": "verl.interactions.gsm8k_interaction.Gsm8kInteraction", "config": {}, - } - ] - } + }]} with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: OmegaConf.save(config_content, f.name) temp_config_path = f.name try: - interaction_map = initialize_interactions_from_config(temp_config_path) + interaction_map = initialize_interactions_from_config( + temp_config_path) # Check that interaction was created assert len(interaction_map) == 1 @@ -78,17 +80,18 @@ def test_initialize_multiple_interactions_from_config(self): { "name": "base_agent", "class_name": "verl.interactions.base.BaseInteraction", - "config": {"custom_param": "test_value"}, + "config": { + "custom_param": "test_value"}, }, - ] - } + ]} with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: OmegaConf.save(config_content, f.name) temp_config_path = f.name try: - interaction_map = initialize_interactions_from_config(temp_config_path) + interaction_map = initialize_interactions_from_config( + temp_config_path) # Check that both interactions were created assert len(interaction_map) == 2 @@ -96,7 +99,9 @@ def test_initialize_multiple_interactions_from_config(self): assert "base_agent" in interaction_map # Check types - assert isinstance(interaction_map["gsm8k_solver"], Gsm8kInteraction) + assert isinstance( + interaction_map["gsm8k_solver"], + Gsm8kInteraction) assert isinstance(interaction_map["base_agent"], BaseInteraction) # Check names were injected @@ -104,14 +109,20 @@ def test_initialize_multiple_interactions_from_config(self): assert interaction_map["base_agent"].name == "base_agent" # Check custom config was passed - assert interaction_map["base_agent"].config.get("custom_param") == "test_value" + assert (interaction_map["base_agent"].config.get( + "custom_param") == "test_value") finally: os.unlink(temp_config_path) def test_initialize_interaction_without_explicit_name(self): """Test that interaction name is derived from class name when not specified.""" config_content = { - "interaction": [{"class_name": "verl.interactions.gsm8k_interaction.Gsm8kInteraction", "config": {}}] + "interaction": [ + { + "class_name": "verl.interactions.gsm8k_interaction.Gsm8kInteraction", + "config": {}, + } + ] } with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: @@ -119,11 +130,14 @@ def test_initialize_interaction_without_explicit_name(self): temp_config_path = f.name try: - interaction_map = initialize_interactions_from_config(temp_config_path) + interaction_map = initialize_interactions_from_config( + temp_config_path) # Check that interaction name was derived from class name assert len(interaction_map) == 1 - assert "gsm8k" in interaction_map # Should be "gsm8k" after removing "interaction" suffix + assert ( + "gsm8k" in interaction_map + ) # Should be "gsm8k" after removing "interaction" suffix assert isinstance(interaction_map["gsm8k"], Gsm8kInteraction) assert interaction_map["gsm8k"].name == "gsm8k" finally: @@ -138,7 +152,8 @@ def test_initialize_empty_config(self): temp_config_path = f.name try: - interaction_map = initialize_interactions_from_config(temp_config_path) + interaction_map = initialize_interactions_from_config( + temp_config_path) assert len(interaction_map) == 0 finally: os.unlink(temp_config_path) @@ -146,7 +161,13 @@ def test_initialize_empty_config(self): def test_invalid_class_name(self): """Test handling of invalid class name.""" config_content = { - "interaction": [{"name": "invalid", "class_name": "invalid.module.InvalidClass", "config": {}}] + "interaction": [ + { + "name": "invalid", + "class_name": "invalid.module.InvalidClass", + "config": {}, + } + ] } with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: @@ -163,21 +184,26 @@ def test_duplicate_interaction_names(self): """Test handling of duplicate interaction names.""" config_content = { "interaction": [ - {"name": "duplicate", "class_name": "verl.interactions.base.BaseInteraction", "config": {}}, + { + "name": "duplicate", + "class_name": "verl.interactions.base.BaseInteraction", + "config": {}, + }, { "name": "duplicate", "class_name": "verl.interactions.gsm8k_interaction.Gsm8kInteraction", "config": {}, }, - ] - } + ]} with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: OmegaConf.save(config_content, f.name) temp_config_path = f.name try: - with pytest.raises(ValueError, match="Duplicate interaction name 'duplicate' found"): + with pytest.raises( + ValueError, match="Duplicate interaction name 'duplicate' found" + ): initialize_interactions_from_config(temp_config_path) finally: os.unlink(temp_config_path) @@ -187,7 +213,10 @@ def test_auto_name_generation_edge_cases(self): config_content = { "interaction": [ {"class_name": "verl.interactions.base.BaseInteraction", "config": {}}, - {"class_name": "verl.interactions.gsm8k_interaction.Gsm8kInteraction", "config": {}}, + { + "class_name": "verl.interactions.gsm8k_interaction.Gsm8kInteraction", + "config": {}, + }, ] } @@ -196,7 +225,8 @@ def test_auto_name_generation_edge_cases(self): temp_config_path = f.name try: - interaction_map = initialize_interactions_from_config(temp_config_path) + interaction_map = initialize_interactions_from_config( + temp_config_path) # Check that names were generated correctly assert len(interaction_map) == 2 diff --git a/Agent0/executor_train/verl/tests/models/test_transformer.py b/Agent0/executor_train/verl/tests/models/test_transformer.py index 111230a..ddd5c39 100644 --- a/Agent0/executor_train/verl/tests/models/test_transformer.py +++ b/Agent0/executor_train/verl/tests/models/test_transformer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -45,10 +45,14 @@ def test_hf_casual_models(): # config = AutoConfig.from_pretrained(test_case) with torch.device("cuda"): model = AutoModelForCausalLM.from_config( - config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + config=config, + torch_dtype=torch.bfloat16, + attn_implementation="flash_attention_2", ) model = model.to(device="cuda") - input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device="cuda") + input_ids = torch.randint( + low=0, high=config.vocab_size, size=( + batch_size, seqlen), device="cuda") attention_mask = create_random_mask( input_ids=input_ids, max_ratio_of_left_padding=0.1, @@ -69,15 +73,21 @@ def test_hf_casual_models(): rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices ).transpose(0, 1) - # input with input_ids_rmpad and postition_ids to enable flash attention varlen + # input with input_ids_rmpad and postition_ids to enable flash + # attention varlen logits_rmpad = model( input_ids_rmpad, position_ids=position_ids_rmpad, use_cache=False ).logits # (1, total_nnz, vocab_size) origin_logits = model( - input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + use_cache=False, ).logits - origin_logits_rmpad, origin_logits_indices, *_ = unpad_input(origin_logits, attention_mask) + origin_logits_rmpad, origin_logits_indices, *_ = unpad_input( + origin_logits, attention_mask + ) logits_rmpad = logits_rmpad.squeeze(0) log_probs = log_probs_from_logits_all_rmpad( @@ -98,8 +108,8 @@ def test_hf_casual_models(): ) # (batch, seqlen) torch.testing.assert_close( - masked_mean(log_probs, attention_mask[:, -response_length - 1 : -1]), - masked_mean(origin_log_probs, attention_mask[:, -response_length - 1 : -1]), + masked_mean(log_probs, attention_mask[:, -response_length - 1: -1]), + masked_mean(origin_log_probs, attention_mask[:, -response_length - 1: -1]), atol=1e-2, rtol=1e-5, ) @@ -117,10 +127,14 @@ def test_hf_value_models(): config.hidden_dropout = 0 with torch.device("cuda"): model = AutoModelForTokenClassification.from_config( - config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + config=config, + torch_dtype=torch.bfloat16, + attn_implementation="flash_attention_2", ) model = model.to(device="cuda") - input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device="cuda") + input_ids = torch.randint( + low=0, high=config.vocab_size, size=( + batch_size, seqlen), device="cuda") attention_mask = create_random_mask( input_ids=input_ids, max_ratio_of_left_padding=0.1, @@ -142,15 +156,23 @@ def test_hf_value_models(): ).transpose(0, 1) origin_logits = model( - input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + use_cache=False, ).logits - # input with input_ids_rmpad and postition_ids to enable flash attention varlen + # input with input_ids_rmpad and postition_ids to enable flash + # attention varlen rmpad_logits = model( input_ids_rmpad, position_ids=position_ids_rmpad, use_cache=False ).logits # (1, total_nnz, 1) rmpad_logits = rmpad_logits.squeeze(0) - pad_logits = pad_input(rmpad_logits, indices, batch_size, seqlen=seqlen) + pad_logits = pad_input( + rmpad_logits, + indices, + batch_size, + seqlen=seqlen) torch.testing.assert_close( masked_mean(pad_logits, attention_mask[:, :, None]), diff --git a/Agent0/executor_train/verl/tests/models/test_transformers_ulysses.py b/Agent0/executor_train/verl/tests/models/test_transformers_ulysses.py index 233633f..00c40de 100644 --- a/Agent0/executor_train/verl/tests/models/test_transformers_ulysses.py +++ b/Agent0/executor_train/verl/tests/models/test_transformers_ulysses.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,7 +20,12 @@ import torch.distributed from flash_attn.bert_padding import index_first_axis, rearrange, unpad_input from torch.distributed import init_device_mesh -from transformers import AutoModelForCausalLM, LlamaConfig, PretrainedConfig, Qwen2Config +from transformers import ( + AutoModelForCausalLM, + LlamaConfig, + PretrainedConfig, + Qwen2Config, +) from verl.models.transformers.monkey_patch import apply_monkey_patch from verl.protocol import DataProto @@ -48,23 +53,48 @@ class SequenceParallelConfig: def test_configs(): return [ SequenceParallelConfig( - LlamaConfig(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=32), sp_size=8, is_valid=True + LlamaConfig( + num_hidden_layers=2, + num_attention_heads=32, + num_key_value_heads=32), + sp_size=8, + is_valid=True, ), SequenceParallelConfig( - Qwen2Config(num_hidden_layers=2, num_attention_heads=28, num_key_value_heads=4, hidden_size=3584), + Qwen2Config( + num_hidden_layers=2, + num_attention_heads=28, + num_key_value_heads=4, + hidden_size=3584, + ), sp_size=4, is_valid=True, ), SequenceParallelConfig( - Qwen2Config(num_hidden_layers=2, num_attention_heads=28, num_key_value_heads=4, hidden_size=3584), + Qwen2Config( + num_hidden_layers=2, + num_attention_heads=28, + num_key_value_heads=4, + hidden_size=3584, + ), sp_size=8, is_valid=False, ), SequenceParallelConfig( - Qwen2Config(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=4), sp_size=4, is_valid=True + Qwen2Config( + num_hidden_layers=2, + num_attention_heads=32, + num_key_value_heads=4), + sp_size=4, + is_valid=True, ), SequenceParallelConfig( - Qwen2Config(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=4), sp_size=8, is_valid=True + Qwen2Config( + num_hidden_layers=2, + num_attention_heads=32, + num_key_value_heads=4), + sp_size=8, + is_valid=True, ), ] @@ -80,10 +110,18 @@ def test_hf_casual_fwd_bwd(test_config): if not torch.distributed.is_initialized(): initialize_global_process_group() - context = contextlib.nullcontext() if test_config.is_valid else pytest.raises(AssertionError) + context = ( + contextlib.nullcontext() + if test_config.is_valid + else pytest.raises(AssertionError) + ) with context: world_size = torch.distributed.get_world_size() - _hf_casual_fwd_bwd(test_config.config, test_config.sp_size, world_size // test_config.sp_size) + _hf_casual_fwd_bwd( + test_config.config, + test_config.sp_size, + world_size // + test_config.sp_size) # TODO: seems not work, will cause `socketStartConnect: Connect to xxx failed : Software caused connection abort` # torch.distributed.destroy_process_group() @@ -93,8 +131,9 @@ def _hf_casual_fwd(config, sp_size, dp_size): assert torch.cuda.device_count() >= 2, "need at least 2 gpus for test" ulysses_device_mesh = init_device_mesh( - device_type="cuda", mesh_shape=(dp_size, sp_size), mesh_dim_names=("dp", "sp") - ) + device_type="cuda", mesh_shape=( + dp_size, sp_size), mesh_dim_names=( + "dp", "sp")) sharding_manager = FSDPUlyssesShardingManager(ulysses_device_mesh) batch_size = 1 @@ -104,16 +143,23 @@ def _hf_casual_fwd(config, sp_size, dp_size): # patch before load with torch.device("cuda"): model = AutoModelForCausalLM.from_config( - config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + config=config, + torch_dtype=torch.bfloat16, + attn_implementation="flash_attention_2", ) apply_monkey_patch(model, sp_size) model = model.to(device="cuda") sync_model_parameters_global(model) # different rank will generate different input_ids following fsdp - input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device="cuda") + input_ids = torch.randint( + low=0, high=config.vocab_size, size=(batch_size, seqlen), device="cuda" + ) attention_mask = create_random_mask( - input_ids=input_ids, max_ratio_of_left_padding=0, max_ratio_of_valid_token=0.9, min_ratio_of_valid_token=0.8 + input_ids=input_ids, + max_ratio_of_left_padding=0, + max_ratio_of_valid_token=0.9, + min_ratio_of_valid_token=0.8, ) position_ids = compute_position_id_with_mask( attention_mask @@ -145,17 +191,28 @@ def _hf_casual_fwd(config, sp_size, dp_size): # slice input tensor for ulysses # input_ids are padded and sliced # postition_ids are only padded but not sliced - input_ids_rmpad_sliced, position_ids_rmpad_padded, pad_size = ulysses_pad_and_slice_inputs( - input_ids_rmpad, position_ids_rmpad, sp_size=get_ulysses_sequence_parallel_world_size() + input_ids_rmpad_sliced, position_ids_rmpad_padded, pad_size = ( + ulysses_pad_and_slice_inputs( + input_ids_rmpad, + position_ids_rmpad, + sp_size=get_ulysses_sequence_parallel_world_size(), + ) ) - # input with input_ids_rmpad and postition_ids to enable flash attention varlen + # input with input_ids_rmpad and postition_ids to enable flash + # attention varlen logits_split_in_seq = model( - input_ids_rmpad_sliced, position_ids=position_ids_rmpad_padded, use_cache=False + input_ids_rmpad_sliced, + position_ids=position_ids_rmpad_padded, + use_cache=False, ).logits # (1, total_nnz/n, vocab_size) # all_gather output - logits_full = gather_outpus_and_unpad(logits_split_in_seq, gather_dim=1, unpad_dim=1, padding_size=pad_size) + logits_full = gather_outpus_and_unpad( + logits_split_in_seq, + gather_dim=1, + unpad_dim=1, + padding_size=pad_size) # 2. perform normal forward set_ulysses_sequence_parallel_group(None) @@ -172,8 +229,9 @@ def _hf_casual_fwd_bwd(config, sp_size, dp_size): assert torch.cuda.device_count() >= 2, "need at least 2 gpus for test" ulysses_device_mesh = init_device_mesh( - device_type="cuda", mesh_shape=(dp_size, sp_size), mesh_dim_names=("dp", "sp") - ) + device_type="cuda", mesh_shape=( + dp_size, sp_size), mesh_dim_names=( + "dp", "sp")) sharding_manager = FSDPUlyssesShardingManager(ulysses_device_mesh) batch_size = 1 @@ -183,16 +241,23 @@ def _hf_casual_fwd_bwd(config, sp_size, dp_size): # patch before load with torch.device("cuda"): model = AutoModelForCausalLM.from_config( - config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + config=config, + torch_dtype=torch.bfloat16, + attn_implementation="flash_attention_2", ) apply_monkey_patch(model, sp_size) model = model.to(device="cuda") sync_model_parameters_global(model) # different rank will generate different input_ids following fsdp - input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device="cuda") + input_ids = torch.randint( + low=0, high=config.vocab_size, size=(batch_size, seqlen), device="cuda" + ) attention_mask = create_random_mask( - input_ids=input_ids, max_ratio_of_left_padding=0, max_ratio_of_valid_token=0.9, min_ratio_of_valid_token=0.8 + input_ids=input_ids, + max_ratio_of_left_padding=0, + max_ratio_of_valid_token=0.9, + min_ratio_of_valid_token=0.8, ) position_ids = compute_position_id_with_mask( attention_mask @@ -224,17 +289,28 @@ def _hf_casual_fwd_bwd(config, sp_size, dp_size): # slice input tensor for ulysses # input_ids are padded and sliced # postition_ids are only padded but not sliced - input_ids_rmpad_sliced, position_ids_rmpad_padded, pad_size = ulysses_pad_and_slice_inputs( - input_ids_rmpad, position_ids_rmpad, sp_size=get_ulysses_sequence_parallel_world_size() + input_ids_rmpad_sliced, position_ids_rmpad_padded, pad_size = ( + ulysses_pad_and_slice_inputs( + input_ids_rmpad, + position_ids_rmpad, + sp_size=get_ulysses_sequence_parallel_world_size(), + ) ) - # input with input_ids_rmpad and postition_ids to enable flash attention varlen + # input with input_ids_rmpad and postition_ids to enable flash + # attention varlen logits_split_in_seq = model( - input_ids_rmpad_sliced, position_ids=position_ids_rmpad_padded, use_cache=False + input_ids_rmpad_sliced, + position_ids=position_ids_rmpad_padded, + use_cache=False, ).logits # (1, total_nnz/n, vocab_size) # all_gather output - logits_full = gather_outpus_and_unpad(logits_split_in_seq, gather_dim=1, unpad_dim=1, padding_size=pad_size) + logits_full = gather_outpus_and_unpad( + logits_split_in_seq, + gather_dim=1, + unpad_dim=1, + padding_size=pad_size) # 2. perform normal forward set_ulysses_sequence_parallel_group(None) diff --git a/Agent0/executor_train/verl/tests/single_controller/__init__.py b/Agent0/executor_train/verl/tests/single_controller/__init__.py index 1cd1e84..4597f19 100644 --- a/Agent0/executor_train/verl/tests/single_controller/__init__.py +++ b/Agent0/executor_train/verl/tests/single_controller/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/single_controller/base/test_decorator.py b/Agent0/executor_train/verl/tests/single_controller/base/test_decorator.py index 5447d65..ae77bf0 100644 --- a/Agent0/executor_train/verl/tests/single_controller/base/test_decorator.py +++ b/Agent0/executor_train/verl/tests/single_controller/base/test_decorator.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -72,5 +72,7 @@ def new_collect(worker_group, output): update_dispatch_mode(original_mode, new_dispatch, new_collect) # Verify update - assert get_predefined_dispatch_fn(original_mode)["dispatch_fn"] == new_dispatch - assert get_predefined_dispatch_fn(original_mode)["collect_fn"] == new_collect + assert get_predefined_dispatch_fn( + original_mode)["dispatch_fn"] == new_dispatch + assert get_predefined_dispatch_fn( + original_mode)["collect_fn"] == new_collect diff --git a/Agent0/executor_train/verl/tests/single_controller/check_worker_alive/main.py b/Agent0/executor_train/verl/tests/single_controller/check_worker_alive/main.py index cbdee9a..27fd125 100644 --- a/Agent0/executor_train/verl/tests/single_controller/check_worker_alive/main.py +++ b/Agent0/executor_train/verl/tests/single_controller/check_worker_alive/main.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,7 +20,11 @@ from verl.single_controller.base.decorator import Dispatch, register from verl.single_controller.base.worker import Worker -from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.single_controller.ray.base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) @ray.remote @@ -57,7 +61,9 @@ def foo(self, wait_time): print( time.time(), - f"wait 6x wait time {wait_time * 6} to let signal returned to process but still not exceed process wait time", + f"wait 6x wait time { + wait_time * + 6} to let signal returned to process but still not exceed process wait time", ) time.sleep(wait_time * 6) diff --git a/Agent0/executor_train/verl/tests/single_controller/detached_worker/client.py b/Agent0/executor_train/verl/tests/single_controller/detached_worker/client.py index 52f2c72..4a29339 100644 --- a/Agent0/executor_train/verl/tests/single_controller/detached_worker/client.py +++ b/Agent0/executor_train/verl/tests/single_controller/detached_worker/client.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -42,13 +42,23 @@ def compute_position_id_with_mask(mask): sequence_length = 1024 # give Trainer some data to train - input_ids = torch.randint(low=0, high=256, size=(batch_size, sequence_length), dtype=torch.int64, device="cuda") + input_ids = torch.randint( + low=0, + high=256, + size=(batch_size, sequence_length), + dtype=torch.int64, + device="cuda", + ) attention_mask = torch.ones_like(input_ids) position_ids = compute_position_id_with_mask(attention_mask) data = DataProto( batch=TensorDict( - {"input_ids": input_ids, "attention_mask": attention_mask, "position_ids": position_ids}, + { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + }, batch_size=batch_size, ), meta_info={}, diff --git a/Agent0/executor_train/verl/tests/single_controller/detached_worker/server.py b/Agent0/executor_train/verl/tests/single_controller/detached_worker/server.py index 57e555a..43b4da2 100644 --- a/Agent0/executor_train/verl/tests/single_controller/detached_worker/server.py +++ b/Agent0/executor_train/verl/tests/single_controller/detached_worker/server.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,31 +15,33 @@ Server starts a Trainer. Client sends data to the server to train. """ +from verl.utils.megatron_utils import ( + get_model, + init_megatron_optim_config, + mcore_model_parallel_config, +) +from verl.utils.megatron.optimizer import get_megatron_optimizer +from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup +from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool +from verl.single_controller.base.megatron.worker import MegatronWorker +from verl.single_controller.base.decorator import Dispatch, register +from verl.models.llama.megatron import ParallelLlamaForCausalLMRmPadPP +from verl import DataProto +from transformers import LlamaConfig +from torch import nn +from tensordict import TensorDict +from omegaconf import OmegaConf +from megatron.core.models.gpt.gpt_model import ModelType +from megatron.core import tensor_parallel +from megatron.core import parallel_state as mpu +import torch +import ray import os os.environ["MEGATRON_USE_CUDA_TIMER"] = "0" os.environ["MEGATRON_START_PROCESS_TIMER"] = "False" os.environ["NCCL_DEBUG"] = "WARN" -import ray -import torch -from megatron.core import parallel_state as mpu -from megatron.core import tensor_parallel -from megatron.core.models.gpt.gpt_model import ModelType -from omegaconf import OmegaConf -from tensordict import TensorDict -from torch import nn -from transformers import LlamaConfig - -from verl import DataProto -from verl.models.llama.megatron import ParallelLlamaForCausalLMRmPadPP -from verl.single_controller.base.decorator import Dispatch, register -from verl.single_controller.base.megatron.worker import MegatronWorker -from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool -from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup -from verl.utils.megatron.optimizer import get_megatron_optimizer -from verl.utils.megatron_utils import get_model, init_megatron_optim_config, mcore_model_parallel_config - @ray.remote class Trainer(MegatronWorker): @@ -75,7 +77,9 @@ def init_model(self): num_key_value_heads=16, ) - megatron_config = mcore_model_parallel_config(sequence_parallel=True, params_dtype=torch.bfloat16) + megatron_config = mcore_model_parallel_config( + sequence_parallel=True, params_dtype=torch.bfloat16 + ) self.megatron_config = megatron_config def megatron_actor_model_provider(pre_process, post_process): @@ -102,7 +106,9 @@ def megatron_actor_model_provider(pre_process, post_process): optim_config = init_megatron_optim_config(optim_config) self.optimizer_config = optim_config - actor_optimizer = get_megatron_optimizer(model=actor_module, config=optim_config) + actor_optimizer = get_megatron_optimizer( + model=actor_module, config=optim_config + ) self.model = actor_module[0] self.optimizer = actor_optimizer @@ -118,14 +124,19 @@ def train_model(self, data: DataProto) -> DataProto: zero_buffer=(not self.optimizer_config.use_distributed_optimizer) ) # use use_contiguous_buffers_in_local_ddp and no overlap_dp_param_comm # update for 1 iteration - output = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids).logits + output = self.model( + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + ).logits output.mean().backward() update_successful, grad_norm, num_zeros_in_grad = self.optimizer.step( self.megatron_config, self.megatron_config.timers ) - return DataProto(batch=TensorDict({"loss": output.detach()}, batch_size=output.shape[0])) + return DataProto(batch=TensorDict( + {"loss": output.detach()}, batch_size=output.shape[0])) if __name__ == "__main__": diff --git a/Agent0/executor_train/verl/tests/single_controller/test_auto_padding_on_cpu.py b/Agent0/executor_train/verl/tests/single_controller/test_auto_padding_on_cpu.py index f2c4412..98b2b0b 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_auto_padding_on_cpu.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_auto_padding_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,7 +20,11 @@ from verl.protocol import DataProtoConfig from verl.single_controller.base import Worker from verl.single_controller.base.decorator import Dispatch, register -from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.single_controller.ray.base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) # or set env var VERL_AUTO_PADDING = "1" / "true" DataProtoConfig.auto_padding = True @@ -42,41 +46,67 @@ def test_auto_padding(): chunk_size = 4 actor_cls = RayClassWithInitArgs(cls=Actor) - resource_pool = RayResourcePool(process_on_nodes=[chunk_size], use_gpu=False) - actor_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=actor_cls) + resource_pool = RayResourcePool( + process_on_nodes=[chunk_size], use_gpu=False) + actor_wg = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=actor_cls) # test locally first for test_size in range(4, 20): - local_data = DataProto.from_dict({"a": torch.zeros(test_size)}, {"na": np.zeros(test_size, dtype=object)}) + local_data = DataProto.from_dict({"a": torch.zeros(test_size)}, { + "na": np.zeros(test_size, dtype=object)}) # print(f"before padding, local_data = {local_data}") - padding_size = (chunk_size - (test_size % chunk_size)) if (test_size % chunk_size > 0) else 0 + padding_size = ( + (chunk_size - (test_size % chunk_size)) + if (test_size % chunk_size > 0) + else 0 + ) local_data.padding(padding_size) # print(f"after padding, local_data = {local_data}") - assert len(local_data) == len(local_data) + len(local_data) % chunk_size, ( - f"expecting padded length to be {len(local_data) + len(local_data) % chunk_size}, but got {len(local_data)}" - ) + assert ( + len(local_data) == len(local_data) + len(local_data) % + chunk_size), f"expecting padded length to be { + len(local_data) + len(local_data) % + chunk_size}, but got { + len(local_data)}" chunked = local_data.chunk(chunk_size) - assert len(chunked) == chunk_size, f"during test_size = {test_size}, expecting {chunk_size}, got {chunked}" + assert ( + len(chunked) == chunk_size + ), f"during test_size = {test_size}, expecting {chunk_size}, got {chunked}" for dp in chunked: - assert len(dp) == test_size // chunk_size + bool(test_size % chunk_size), ( - f"test size = {test_size}, expecting dp to be length of " - f"{test_size // chunk_size + bool(test_size % chunk_size)}, but got {len(dp)}: {dp} {chunked}" - ) - - # test with RayWorkerGroup method decorated as dispatch_mode=Dispatch.DP_COMPUTE_PROTO - data = DataProto.from_dict({"a": torch.zeros(10)}, {"na": np.array([str(i) for i in range(10)], dtype=object)}) + assert len(dp) == test_size // chunk_size + bool( + test_size % + chunk_size), (f"test size = {test_size}, expecting dp to be length of " f"{ + test_size // chunk_size + bool( + test_size % + chunk_size)}, but got { + len(dp)}: {dp} {chunked}") + + # test with RayWorkerGroup method decorated as + # dispatch_mode=Dispatch.DP_COMPUTE_PROTO + data = DataProto.from_dict( + {"a": torch.zeros(10)}, + {"na": np.array([str(i) for i in range(10)], dtype=object)}, + ) output = actor_wg.add(data) print(output.batch["a"]) assert len(output) == 10 - data = DataProto.from_dict({"a": torch.zeros(1)}, {"na": np.array([str(i) for i in range(1)], dtype=object)}) + data = DataProto.from_dict( + {"a": torch.zeros(1)}, + {"na": np.array([str(i) for i in range(1)], dtype=object)}, + ) output = actor_wg.add(data) print(output.batch["a"]) assert len(output) == 1 - data = DataProto.from_dict({"a": torch.zeros(8)}, {"na": np.array([str(i) for i in range(8)], dtype=object)}) + data = DataProto.from_dict( + {"a": torch.zeros(8)}, + {"na": np.array([str(i) for i in range(8)], dtype=object)}, + ) output = actor_wg.add(data) print(output.batch["a"]) @@ -86,21 +116,23 @@ def test_auto_padding(): DataProtoConfig.auto_padding = False data = DataProto.from_dict( - {"a": torch.zeros(10)}, {"na": np.array([str(i) for i in range(10)], dtype=object)}, auto_padding=True + {"a": torch.zeros(10)}, + {"na": np.array([str(i) for i in range(10)], dtype=object)}, + auto_padding=True, ) output = actor_wg.add(data) print(output.batch["a"]) assert len(output) == 10 - data = DataProto.from_single_dict( - {"a": torch.zeros(1), "na": np.array([str(i) for i in range(1)], dtype=object)}, auto_padding=True - ) + data = DataProto.from_single_dict({"a": torch.zeros(1), "na": np.array( + [str(i) for i in range(1)], dtype=object)}, auto_padding=True, ) output = actor_wg.add(data) print(output.batch["a"]) assert len(output) == 1 - data = DataProto.from_single_dict({"a": torch.zeros(8), "na": np.array([str(i) for i in range(8)], dtype=object)}) + data = DataProto.from_single_dict({"a": torch.zeros( + 8), "na": np.array([str(i) for i in range(8)], dtype=object)}) output = actor_wg.add(data) print(output.batch["a"]) diff --git a/Agent0/executor_train/verl/tests/single_controller/test_colocated_workers.py b/Agent0/executor_train/verl/tests/single_controller/test_colocated_workers.py index cdaa747..21f5517 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_colocated_workers.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_colocated_workers.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -59,8 +59,12 @@ def test_colocated_workers(): critic_cls = RayClassWithInitArgs(cls=Critic, config={"b": 10}) resource_pool = RayResourcePool(process_on_nodes=[2]) - actor_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=actor_cls) - critic_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=critic_cls) + actor_wg = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=actor_cls) + critic_wg = RayWorkerGroup( + resource_pool=resource_pool, ray_cls_with_init=critic_cls + ) expected_actor_output = actor_wg.add(data) expected_critic_output = critic_wg.sub(data) @@ -68,7 +72,9 @@ def test_colocated_workers(): # create colocated workers cls_dict = {"actor": actor_cls, "critic": critic_cls} ray_cls_with_init = create_colocated_worker_cls(cls_dict) - wg_dict = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init) + wg_dict = RayWorkerGroup( + resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init + ) spawn_wg = wg_dict.spawn(prefix_set=cls_dict.keys()) colocated_actor_wg = spawn_wg["actor"] @@ -77,7 +83,11 @@ def test_colocated_workers(): actor_output = colocated_actor_wg.add(data) critic_output = colocated_critic_wg.sub(data) - torch.testing.assert_close(expected_actor_output.batch, actor_output.batch, atol=0, rtol=0) - torch.testing.assert_close(expected_critic_output.batch, critic_output.batch, atol=0, rtol=0) + torch.testing.assert_close( + expected_actor_output.batch, actor_output.batch, atol=0, rtol=0 + ) + torch.testing.assert_close( + expected_critic_output.batch, critic_output.batch, atol=0, rtol=0 + ) ray.shutdown() diff --git a/Agent0/executor_train/verl/tests/single_controller/test_colocated_workers_fused.py b/Agent0/executor_train/verl/tests/single_controller/test_colocated_workers_fused.py index 93b1a72..a1e0943 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_colocated_workers_fused.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_colocated_workers_fused.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -59,8 +59,12 @@ def test_colocated_workers_fused(): critic_cls = RayClassWithInitArgs(cls=Critic, config={"b": 10}) resource_pool = RayResourcePool(process_on_nodes=[2]) - actor_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=actor_cls) - critic_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=critic_cls) + actor_wg = RayWorkerGroup( + resource_pool=resource_pool, + ray_cls_with_init=actor_cls) + critic_wg = RayWorkerGroup( + resource_pool=resource_pool, ray_cls_with_init=critic_cls + ) expected_actor_output = actor_wg.add(data) expected_critic_output = critic_wg.sub(data) @@ -68,7 +72,9 @@ def test_colocated_workers_fused(): # create colocated workers cls_dict = {"actor": actor_cls, "critic": critic_cls} ray_cls_with_init = create_colocated_worker_cls_fused(cls_dict) - wg_dict = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init) + wg_dict = RayWorkerGroup( + resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init + ) spawn_wg = wg_dict.spawn(prefix_set=cls_dict.keys()) colocated_actor_wg = spawn_wg["actor"] @@ -77,7 +83,11 @@ def test_colocated_workers_fused(): actor_output = colocated_actor_wg.add(data) critic_output = colocated_critic_wg.sub(data) - torch.testing.assert_close(expected_actor_output.batch, actor_output.batch, atol=0, rtol=0) - torch.testing.assert_close(expected_critic_output.batch, critic_output.batch, atol=0, rtol=0) + torch.testing.assert_close( + expected_actor_output.batch, actor_output.batch, atol=0, rtol=0 + ) + torch.testing.assert_close( + expected_critic_output.batch, critic_output.batch, atol=0, rtol=0 + ) ray.shutdown() diff --git a/Agent0/executor_train/verl/tests/single_controller/test_data_transfer.py b/Agent0/executor_train/verl/tests/single_controller/test_data_transfer.py index 13777b0..c5481f6 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_data_transfer.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_data_transfer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,7 +24,11 @@ from verl import DataProto from verl.single_controller.base import Worker from verl.single_controller.base.decorator import Dispatch, register -from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.single_controller.ray import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) from verl.utils.ray_utils import parallel_put @@ -98,7 +102,9 @@ def test_data_transfer(): for input_data, output_data in zip(data_list, output_lst, strict=True): for key in input_data.batch.keys(): - assert torch.all(torch.eq(input_data.batch[key] + 1, output_data.batch[key])), ( + assert torch.all( + torch.eq(input_data.batch[key] + 1, output_data.batch[key]) + ), ( input_data.batch[key], output_data.batch[key], key, diff --git a/Agent0/executor_train/verl/tests/single_controller/test_decorator_on_cpu.py b/Agent0/executor_train/verl/tests/single_controller/test_decorator_on_cpu.py index 4dfec63..fbb25e1 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_decorator_on_cpu.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_decorator_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,7 +23,11 @@ from verl.protocol import DataProto, DataProtoFuture from verl.single_controller.base.decorator import Dispatch, register from verl.single_controller.base.worker import Worker -from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.single_controller.ray import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) # Pytest fixture for Ray setup/teardown @@ -47,7 +51,11 @@ def __init__(self, initial_value=0): @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) def dp_compute(self, data: DataProto) -> DataProto: time.sleep(0.1) # Simulate work - rank_value = torch.tensor(self.rank, device=data.batch["input"].device, dtype=data.batch["input"].dtype) + rank_value = torch.tensor( + self.rank, + device=data.batch["input"].device, + dtype=data.batch["input"].dtype, + ) data.batch["output"] = data.batch["input"] + self.value + rank_value return data @@ -56,8 +64,13 @@ def dp_compute(self, data: DataProto) -> DataProto: async def async_dp_compute(self, data: DataProto) -> DataProto: # Simulate async work await asyncio.sleep(0.1) # Simulate async work - rank_value = torch.tensor(self.rank, device=data.batch["input"].device, dtype=data.batch["input"].dtype) - data.batch["output_async"] = data.batch["input"] * 2 + self.value + rank_value + rank_value = torch.tensor( + self.rank, + device=data.batch["input"].device, + dtype=data.batch["input"].dtype, + ) + data.batch["output_async"] = data.batch["input"] * \ + 2 + self.value + rank_value return data @@ -68,10 +81,15 @@ def test_decorator_dp_compute(ray_init_shutdown): Verifies the result correctness. """ num_workers = 2 - resource_pool = RayResourcePool([num_workers], use_gpu=False, max_colocate_count=1) # Use CPU for simplicity - cls_with_args = RayClassWithInitArgs(cls=DecoratorTestWorker, initial_value=10) + resource_pool = RayResourcePool( + [num_workers], use_gpu=False, max_colocate_count=1 + ) # Use CPU for simplicity + cls_with_args = RayClassWithInitArgs( + cls=DecoratorTestWorker, initial_value=10) worker_group = RayWorkerGroup( - resource_pool, cls_with_args, name_prefix=f"decorator_test_sync_dp_{int(time.time())}" + resource_pool, + cls_with_args, + name_prefix=f"decorator_test_sync_dp_{int(time.time())}", ) # Prepare input data (size 4, for 2 workers) @@ -94,7 +112,11 @@ def test_decorator_dp_compute(ray_init_shutdown): expected_output_part2 = torch.tensor([2, 3], dtype=torch.float32) + 10 + 1 expected_output = torch.cat([expected_output_part1, expected_output_part2]) - torch.testing.assert_close(output.batch["output"], expected_output, msg="Sync DP compute output data mismatch") + torch.testing.assert_close( + output.batch["output"], + expected_output, + msg="Sync DP compute output data mismatch", + ) # Test function for async def method with DP compute @@ -104,10 +126,16 @@ def test_decorator_async_function(ray_init_shutdown): Verifies that the call returns a future and the result is correct after .get(). """ num_workers = 2 - resource_pool = RayResourcePool([num_workers], use_gpu=False, max_colocate_count=1) - cls_with_args = RayClassWithInitArgs(cls=DecoratorTestWorker, initial_value=5) + resource_pool = RayResourcePool( + [num_workers], + use_gpu=False, + max_colocate_count=1) + cls_with_args = RayClassWithInitArgs( + cls=DecoratorTestWorker, initial_value=5) worker_group = RayWorkerGroup( - resource_pool, cls_with_args, name_prefix=f"decorator_test_async_dp_{int(time.time())}" + resource_pool, + cls_with_args, + name_prefix=f"decorator_test_async_dp_{int(time.time())}", ) # Prepare input data (size 4, for 2 workers) @@ -118,7 +146,9 @@ def test_decorator_async_function(ray_init_shutdown): future_output: DataProtoFuture = worker_group.async_dp_compute(data) # Assert that the call returned a future - assert isinstance(future_output, DataProtoFuture), "Expected DataProtoFuture for async def call" + assert isinstance( + future_output, DataProtoFuture + ), "Expected DataProtoFuture for async def call" # Get the result (this should block) result_data = future_output.get() @@ -126,16 +156,21 @@ def test_decorator_async_function(ray_init_shutdown): # Assert the result correctness assert isinstance(result_data, DataProto) assert "output_async" in result_data.batch.keys() - assert len(result_data) == len(data), "Output length should match input length" + assert len(result_data) == len( + data), "Output length should match input length" # Expected output calculation for DP_COMPUTE_PROTO with 2 workers # Worker 0 gets data[0:2], Worker 1 gets data[2:4] # Worker 0 calculates: input * 2 + initial_value(5) + rank(0) # Worker 1 calculates: input * 2 + initial_value(5) + rank(1) - expected_output_part1 = (torch.tensor([0, 1], dtype=torch.float32) * 2) + 5 + 0 - expected_output_part2 = (torch.tensor([2, 3], dtype=torch.float32) * 2) + 5 + 1 + expected_output_part1 = (torch.tensor( + [0, 1], dtype=torch.float32) * 2) + 5 + 0 + expected_output_part2 = (torch.tensor( + [2, 3], dtype=torch.float32) * 2) + 5 + 1 expected_output = torch.cat([expected_output_part1, expected_output_part2]) torch.testing.assert_close( - result_data.batch["output_async"], expected_output, msg="Async DP compute output data mismatch" + result_data.batch["output_async"], + expected_output, + msg="Async DP compute output data mismatch", ) diff --git a/Agent0/executor_train/verl/tests/single_controller/test_driverfunc_to_worker.py b/Agent0/executor_train/verl/tests/single_controller/test_driverfunc_to_worker.py index a38d790..cc6ec1d 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_driverfunc_to_worker.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_driverfunc_to_worker.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -45,7 +45,8 @@ def get_aux_metrics(self, test_proto): decode_count.append(len(sequence_ids[i].tolist())) ret_proto = DataProto( batch=TensorDict( - {"sequence_ids": sequence_ids, "decode_count": torch.tensor(decode_count)}, batch_size=sequence_ids.size(0) + {"sequence_ids": sequence_ids, "decode_count": torch.tensor(decode_count)}, + batch_size=sequence_ids.size(0), ) ) return ret_proto @@ -73,12 +74,15 @@ def test(): ) # Sharding among different ranks - ret_proto1 = shard_wg.execute_with_func_generator(get_aux_metrics, test_proto) + ret_proto1 = shard_wg.execute_with_func_generator( + get_aux_metrics, test_proto) # compare execute on driver hs = HackSelf() ret_proto2 = get_aux_metrics(hs, test_proto) - torch.testing.assert_close(ret_proto1.batch["decode_count"], ret_proto2.batch["decode_count"]) + torch.testing.assert_close( + ret_proto1.batch["decode_count"], ret_proto2.batch["decode_count"] + ) ray.shutdown() diff --git a/Agent0/executor_train/verl/tests/single_controller/test_fused_workers_on_cpu.py b/Agent0/executor_train/verl/tests/single_controller/test_fused_workers_on_cpu.py index 527ddc1..a4831fe 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_fused_workers_on_cpu.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_fused_workers_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -65,13 +65,16 @@ def test_fused_workers(): # create separate workers on the same resource pool process_on_nodes = [2] - resource_pool = RayResourcePool(process_on_nodes=process_on_nodes, use_gpu=False) + resource_pool = RayResourcePool( + process_on_nodes=process_on_nodes, use_gpu=False) # create colocated workers hybrid_cls_with_init = RayClassWithInitArgs(cls=HybridWorker) hybrid_cls_with_init.fused_worker_used = True - fused_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=hybrid_cls_with_init) + fused_wg = RayWorkerGroup( + resource_pool=resource_pool, ray_cls_with_init=hybrid_cls_with_init + ) fused_wg.fuse(cls_dict.keys()) x = fused_wg.actor.add(0.1) diff --git a/Agent0/executor_train/verl/tests/single_controller/test_high_level_scheduling_api.py b/Agent0/executor_train/verl/tests/single_controller/test_high_level_scheduling_api.py index 52cc7c7..5002094 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_high_level_scheduling_api.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_high_level_scheduling_api.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,7 +17,12 @@ import ray from verl.single_controller.base.worker import Worker -from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup, merge_resource_pool +from verl.single_controller.ray.base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, + merge_resource_pool, +) @ray.remote @@ -40,25 +45,42 @@ def test(): class_with_args = RayClassWithInitArgs(cls=TestActor) print("create actor worker group") - actor_wg = RayWorkerGroup(resource_pool, class_with_args, name_prefix="high_level_api_actor") + actor_wg = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="high_level_api_actor" + ) print("create critic worker group") - critic_wg = RayWorkerGroup(resource_pool, class_with_args, name_prefix="hight_level_api_critic") + critic_wg = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="hight_level_api_critic" + ) print("create rm worker group") - rm_wg = RayWorkerGroup(resource_pool, class_with_args, name_prefix="high_level_api_rm") + rm_wg = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="high_level_api_rm" + ) print("create ref worker group") - ref_wg = RayWorkerGroup(resource_pool, class_with_args, name_prefix="high_level_api_ref") - - assert actor_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(8)] - assert critic_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(8)] - assert rm_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(8)] - assert ref_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(8)] + ref_wg = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="high_level_api_ref" + ) + + assert actor_wg.execute_all_sync("get_cuda_visible_devices") == [ + str(i) for i in range(8) + ] + assert critic_wg.execute_all_sync("get_cuda_visible_devices") == [ + str(i) for i in range(8) + ] + assert rm_wg.execute_all_sync("get_cuda_visible_devices") == [ + str(i) for i in range(8) + ] + assert ref_wg.execute_all_sync("get_cuda_visible_devices") == [ + str(i) for i in range(8) + ] del actor_wg del critic_wg del rm_wg del ref_wg - [ray.util.remove_placement_group(pg) for pg in resource_pool.get_placement_groups()] + [ray.util.remove_placement_group(pg) + for pg in resource_pool.get_placement_groups()] print("wait 5s to remove placemeng_group") time.sleep(5) # test single-node-multi-partition @@ -66,20 +88,39 @@ def test(): print("test single-node-multi-partition") rm_resource_pool = RayResourcePool([4], use_gpu=True, name_prefix="rm") ref_resource_pool = RayResourcePool([4], use_gpu=True, name_prefix="ref") - total_resource_pool = merge_resource_pool(rm_resource_pool, ref_resource_pool) + total_resource_pool = merge_resource_pool( + rm_resource_pool, ref_resource_pool) assert rm_resource_pool.world_size == 4 assert ref_resource_pool.world_size == 4 assert total_resource_pool.world_size == 8 - actor_wg = RayWorkerGroup(total_resource_pool, class_with_args, name_prefix="high_level_api_actor") - critic_wg = RayWorkerGroup(total_resource_pool, class_with_args, name_prefix="high_level_api_critic") - rm_wg = RayWorkerGroup(rm_resource_pool, class_with_args, name_prefix="high_level_api_rm") - ref_wg = RayWorkerGroup(ref_resource_pool, class_with_args, name_prefix="high_level_api_ref") - - assert actor_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(8)] - assert critic_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(8)] - assert rm_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(4)] - assert ref_wg.execute_all_sync("get_cuda_visible_devices") == [str(i) for i in range(4, 8)] + actor_wg = RayWorkerGroup( + total_resource_pool, + class_with_args, + name_prefix="high_level_api_actor") + critic_wg = RayWorkerGroup( + total_resource_pool, + class_with_args, + name_prefix="high_level_api_critic") + rm_wg = RayWorkerGroup( + rm_resource_pool, class_with_args, name_prefix="high_level_api_rm" + ) + ref_wg = RayWorkerGroup( + ref_resource_pool, class_with_args, name_prefix="high_level_api_ref" + ) + + assert actor_wg.execute_all_sync("get_cuda_visible_devices") == [ + str(i) for i in range(8) + ] + assert critic_wg.execute_all_sync("get_cuda_visible_devices") == [ + str(i) for i in range(8) + ] + assert rm_wg.execute_all_sync("get_cuda_visible_devices") == [ + str(i) for i in range(4) + ] + assert ref_wg.execute_all_sync("get_cuda_visible_devices") == [ + str(i) for i in range(4, 8) + ] ray.shutdown() diff --git a/Agent0/executor_train/verl/tests/single_controller/test_ray_collectives.py b/Agent0/executor_train/verl/tests/single_controller/test_ray_collectives.py index 3722a8f..0804e67 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_ray_collectives.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_ray_collectives.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,7 +26,11 @@ from verl.single_controller.base import Worker from verl.single_controller.base.decorator import Dispatch, register -from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.single_controller.ray import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) @ray.remote @@ -35,11 +39,18 @@ class Actor(Worker): def init(self): remote_rank = self.rank // 2 self.group_name = f"A{self.rank}_R{remote_rank}" - collective.init_collective_group(world_size=2, rank=0, backend="nccl", group_name=self.group_name) + collective.init_collective_group( + world_size=2, rank=0, backend="nccl", group_name=self.group_name + ) @register(Dispatch.ONE_TO_ALL, blocking=False) def send_tensors(self): - tensor = torch.ones(size=(4,), dtype=torch.float32, device="cuda") * self.rank + tensor = torch.ones( + size=( + 4, + ), + dtype=torch.float32, + device="cuda") * self.rank collective.send(tensor=tensor, dst_rank=1, group_name=self.group_name) @@ -52,20 +63,39 @@ def init(self): self.first_group_name = f"A{self.remote_first_rank}_R{self.rank}" self.second_group_name = f"A{self.remote_second_rank}_R{self.rank}" - collective.init_collective_group(world_size=2, rank=1, backend="nccl", group_name=self.first_group_name) - collective.init_collective_group(world_size=2, rank=1, backend="nccl", group_name=self.second_group_name) + collective.init_collective_group( + world_size=2, + rank=1, + backend="nccl", + group_name=self.first_group_name) + collective.init_collective_group( + world_size=2, + rank=1, + backend="nccl", + group_name=self.second_group_name) @register(Dispatch.ONE_TO_ALL, blocking=False) def receive_tensors(self): - self.tensor1 = torch.randn(size=(4,), dtype=torch.float32, device="cuda") - self.tensor2 = torch.randn(size=(4,), dtype=torch.float32, device="cuda") - - collective.recv(self.tensor1, src_rank=0, group_name=self.first_group_name) - collective.recv(self.tensor2, src_rank=0, group_name=self.second_group_name) + self.tensor1 = torch.randn( + size=(4,), dtype=torch.float32, device="cuda") + self.tensor2 = torch.randn( + size=(4,), dtype=torch.float32, device="cuda") + + collective.recv( + self.tensor1, + src_rank=0, + group_name=self.first_group_name) + collective.recv( + self.tensor2, + src_rank=0, + group_name=self.second_group_name) @register(Dispatch.ONE_TO_ALL) def get_tensors(self): - return {f"src_{self.remote_first_rank}": self.tensor1, f"src_{self.remote_second_rank}": self.tensor2} + return { + f"src_{self.remote_first_rank}": self.tensor1, + f"src_{self.remote_second_rank}": self.tensor2, + } def test_ray_collective_group(): @@ -78,10 +108,14 @@ def test_ray_collective_group(): rollout_cls = RayClassWithInitArgs(cls=Rollout) actor_wg = RayWorkerGroup( - resource_pool=actor_resource_pool, ray_cls_with_init=actor_cls, name_prefix="collective_group_actor" + resource_pool=actor_resource_pool, + ray_cls_with_init=actor_cls, + name_prefix="collective_group_actor", ) rollout_wg = RayWorkerGroup( - resource_pool=rollout_resource_pool, ray_cls_with_init=rollout_cls, name_prefix="collective_group_rollout" + resource_pool=rollout_resource_pool, + ray_cls_with_init=rollout_cls, + name_prefix="collective_group_rollout", ) actor_wg.init() diff --git a/Agent0/executor_train/verl/tests/single_controller/test_ray_local_envs_on_cpu.py b/Agent0/executor_train/verl/tests/single_controller/test_ray_local_envs_on_cpu.py index ee6c0cb..2e2ecf3 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_ray_local_envs_on_cpu.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_ray_local_envs_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,7 +20,11 @@ import ray from verl.single_controller.base.worker import Worker -from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.single_controller.ray.base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) @ray.remote @@ -41,10 +45,13 @@ def test_basics(): class_with_args = RayClassWithInitArgs(cls=TestActor) worker_group = RayWorkerGroup( - resource_pool=resource_pool, ray_cls_with_init=class_with_args, name_prefix="worker_group_basic" + resource_pool=resource_pool, + ray_cls_with_init=class_with_args, + name_prefix="worker_group_basic", ) - output = worker_group.execute_all_sync("getenv", key="RAY_LOCAL_WORLD_SIZE") + output = worker_group.execute_all_sync( + "getenv", key="RAY_LOCAL_WORLD_SIZE") assert output == ["4", "4", "4", "4"] output = worker_group.execute_all_sync("getenv", key="RAY_LOCAL_RANK") diff --git a/Agent0/executor_train/verl/tests/single_controller/test_ray_utils_on_cpu.py b/Agent0/executor_train/verl/tests/single_controller/test_ray_utils_on_cpu.py index e36497d..8cc7e24 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_ray_utils_on_cpu.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_ray_utils_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/single_controller/test_rvdz.py b/Agent0/executor_train/verl/tests/single_controller/test_rvdz.py index 7dea12f..ba0591d 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_rvdz.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_rvdz.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,7 +26,9 @@ def __init__(self, rank, world_size, group_name): def init(self): from verl.utils.rendezvous.ray_backend import create_nccl_communicator_in_ray - self.communicator = create_nccl_communicator_in_ray(self.rank, self.world_size, self.group_name) + self.communicator = create_nccl_communicator_in_ray( + self.rank, self.world_size, self.group_name + ) def test(self): if self.communicator is None: @@ -40,7 +42,10 @@ def test_rvdz(): group_name = "test_group" world_size = 2 - workers = [TestWorker.options(num_gpus=1).remote(rank, world_size, group_name) for rank in range(world_size)] + workers = [ + TestWorker.options(num_gpus=1).remote(rank, world_size, group_name) + for rank in range(world_size) + ] ray.get([worker.init.remote() for worker in workers]) diff --git a/Agent0/executor_train/verl/tests/single_controller/test_worker_group_basics.py b/Agent0/executor_train/verl/tests/single_controller/test_worker_group_basics.py index 5c4823d..091efb1 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_worker_group_basics.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_worker_group_basics.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,9 +18,18 @@ import ray import torch -from verl.single_controller.base.decorator import Dispatch, Execute, collect_all_to_all, register +from verl.single_controller.base.decorator import ( + Dispatch, + Execute, + collect_all_to_all, + register, +) from verl.single_controller.base.worker import Worker -from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.single_controller.ray.base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) def two_to_all_dispatch_fn(worker_group, *args, **kwargs): @@ -48,7 +57,8 @@ def __init__(self, x) -> None: def foo(self, y): return self._x + y - @register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.RANK_ZERO) + @register(dispatch_mode=Dispatch.ALL_TO_ALL, + execute_mode=Execute.RANK_ZERO) def foo_rank_zero(self, x, y): return self._x + y + x @@ -60,7 +70,12 @@ def foo_one_to_all(self, x, y): def foo_all_to_all(self, x, y): return self._x + y + x - @register(dispatch_mode={"dispatch_fn": two_to_all_dispatch_fn, "collect_fn": collect_all_to_all}) + @register( + dispatch_mode={ + "dispatch_fn": two_to_all_dispatch_fn, + "collect_fn": collect_all_to_all, + } + ) def foo_custom(self, x, y): return self._x + y + x @@ -69,8 +84,9 @@ def foo_custom(self, x, y): def remote_call_wg(worker_names): class_with_args = RayClassWithInitArgs(cls=TestActor, x=2) worker_group = RayWorkerGroup.from_detached( - worker_names=worker_names, ray_cls_with_init=class_with_args, name_prefix=None - ) + worker_names=worker_names, + ray_cls_with_init=class_with_args, + name_prefix=None) print(worker_group.worker_names) output_ref = worker_group.foo_custom(x=[1, 2], y=[5, 6]) @@ -97,7 +113,9 @@ def test_basics(): class_with_args = RayClassWithInitArgs(cls=TestActor, x=2) worker_group = RayWorkerGroup( - resource_pool=resource_pool, ray_cls_with_init=class_with_args, name_prefix="worker_group_basic" + resource_pool=resource_pool, + ray_cls_with_init=class_with_args, + name_prefix="worker_group_basic", ) print(worker_group.worker_names) diff --git a/Agent0/executor_train/verl/tests/single_controller/test_worker_group_torch.py b/Agent0/executor_train/verl/tests/single_controller/test_worker_group_torch.py index a601c43..877a03f 100644 --- a/Agent0/executor_train/verl/tests/single_controller/test_worker_group_torch.py +++ b/Agent0/executor_train/verl/tests/single_controller/test_worker_group_torch.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,18 +12,20 @@ # See the License for the specific language governing permissions and # limitations under the License. +from verl.single_controller.ray.base import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) +from verl.single_controller.base.worker import Worker +import torch.distributed +import torch +import ray import os os.environ["RAY_DEDUP_LOGS"] = "0" os.environ["NCCL_DEBUG"] = "WARN" -import ray -import torch -import torch.distributed - -from verl.single_controller.base.worker import Worker -from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup - @ray.remote class TestAllGatherActor(Worker): @@ -33,15 +35,23 @@ def __init__(self, size) -> None: def init(self): torch.distributed.init_process_group() - self.tensor = torch.zeros(size=(self.size,), dtype=torch.int64, device="cuda") + self.tensor = torch.zeros( + size=( + self.size, + ), + dtype=torch.int64, + device="cuda") self.tensor += self.rank def all_gather(self): world_size = self._world_size output = torch.zeros( - size=(self.tensor.shape[0] * world_size,), dtype=self.tensor.dtype, device=self.tensor.device + size=(self.tensor.shape[0] * world_size,), + dtype=self.tensor.dtype, + device=self.tensor.device, ) - torch.distributed.all_gather_into_tensor(output, self.tensor, async_op=False) + torch.distributed.all_gather_into_tensor( + output, self.tensor, async_op=False) return output @@ -52,15 +62,23 @@ def __init__(self, size) -> None: self.size = size torch.distributed.init_process_group() - self.tensor = torch.zeros(size=(self.size,), dtype=torch.int64, device="cuda") + self.tensor = torch.zeros( + size=( + self.size, + ), + dtype=torch.int64, + device="cuda") self.tensor += self.rank def all_gather(self): world_size = self._world_size output = torch.zeros( - size=(self.tensor.shape[0] * world_size,), dtype=self.tensor.dtype, device=self.tensor.device + size=(self.tensor.shape[0] * world_size,), + dtype=self.tensor.dtype, + device=self.tensor.device, ) - torch.distributed.all_gather_into_tensor(output, self.tensor, async_op=False) + torch.distributed.all_gather_into_tensor( + output, self.tensor, async_op=False) return output @@ -74,7 +92,9 @@ def test_all_gather_torch(): resource_pool = RayResourcePool([4], use_gpu=True) class_with_args = RayClassWithInitArgs(cls=TestAllGatherActor, size=2) - worker_group = RayWorkerGroup(resource_pool, class_with_args, name_prefix="worker_group_torch") + worker_group = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="worker_group_torch" + ) worker_group.execute_all_sync("init") output = worker_group.execute_all_sync("all_gather") @@ -83,7 +103,9 @@ def test_all_gather_torch(): output = output[0].cpu() print(output) - assert torch.all(output == torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], dtype=torch.int64)) + assert torch.all( + output == torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], dtype=torch.int64) + ) ray.shutdown() @@ -98,7 +120,9 @@ def test_all_gather_torch_v2(): resource_pool = RayResourcePool([4], use_gpu=True) class_with_args = RayClassWithInitArgs(cls=TestAllGatherActorV2, size=2) - worker_group = RayWorkerGroup(resource_pool, class_with_args, name_prefix="worker_group_torch") + worker_group = RayWorkerGroup( + resource_pool, class_with_args, name_prefix="worker_group_torch" + ) output = worker_group.execute_all_sync("all_gather") for i in range(1, len(output)): @@ -106,6 +130,8 @@ def test_all_gather_torch_v2(): output = output[0].cpu() print(output) - assert torch.all(output == torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], dtype=torch.int64)) + assert torch.all( + output == torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], dtype=torch.int64) + ) ray.shutdown() diff --git a/Agent0/executor_train/verl/tests/special_distributed/test_fsdp_ckpt.py b/Agent0/executor_train/verl/tests/special_distributed/test_fsdp_ckpt.py index 49dceb7..1e1ce2c 100644 --- a/Agent0/executor_train/verl/tests/special_distributed/test_fsdp_ckpt.py +++ b/Agent0/executor_train/verl/tests/special_distributed/test_fsdp_ckpt.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -30,21 +30,27 @@ def test_fsdp_ckpt(strategy="fsdp"): assert torch.cuda.device_count() >= 2, "need at least 2 gpus for test" local_rank, rank, world_size = initialize_global_process_group() - device_mesh = init_device_mesh("cuda", mesh_shape=(world_size,), mesh_dim_names=("dp",)) + device_mesh = init_device_mesh( + "cuda", mesh_shape=(world_size,), mesh_dim_names=("dp",) + ) model_name = "Qwen/Qwen2.5-0.5B-Instruct" config = Qwen2Config(num_hidden_layers=1) with torch.device("cuda"): model = AutoModelForCausalLM.from_config( - config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + config=config, + torch_dtype=torch.bfloat16, + attn_implementation="flash_attention_2", ) model = model.to(device="cuda") # Wrap model with FSDP if strategy == "fsdp": mixed_precision = MixedPrecision( - param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32 + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, ) model = FSDP( @@ -57,7 +63,9 @@ def test_fsdp_ckpt(strategy="fsdp"): ) else: mp_policy = MixedPrecisionPolicy( - param_dtype=torch.bfloat16, reduce_dtype=torch.float32, cast_forward_inputs=True + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + cast_forward_inputs=True, ) fsdp_kwargs = { "mesh": device_mesh, @@ -66,24 +74,29 @@ def test_fsdp_ckpt(strategy="fsdp"): apply_fsdp2(model, fsdp_kwargs, {}) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) - lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9) + lr_scheduler = torch.optim.lr_scheduler.StepLR( + optimizer, step_size=1, gamma=0.9) # Create checkpoint manager tokenizer = AutoTokenizer.from_pretrained(model_name) checkpoint_manager = FSDPCheckpointManager( - model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, tokenizer=tokenizer - ) + model=model, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + tokenizer=tokenizer) # Generate sample input batch_size = 2 seq_len = 32 vocab_size = 32000 # First input for initial update - input_ids1 = torch.randint(0, vocab_size, (batch_size, seq_len), device="cuda") + input_ids1 = torch.randint( + 0, vocab_size, (batch_size, seq_len), device="cuda") attention_mask1 = torch.ones_like(input_ids1) # Second input for verification - input_ids2 = torch.randint(0, vocab_size, (batch_size, seq_len), device="cuda") + input_ids2 = torch.randint( + 0, vocab_size, (batch_size, seq_len), device="cuda") attention_mask2 = torch.ones_like(input_ids2) # Step 1: Initial update and save checkpoint @@ -97,7 +110,9 @@ def test_fsdp_ckpt(strategy="fsdp"): # Save checkpoint after first update temp_dir = tempfile.mkdtemp() checkpoint_path = os.path.join(temp_dir, "checkpoint") - checkpoint_manager.save_checkpoint(local_path=checkpoint_path, hdfs_path=None, global_step=0) + checkpoint_manager.save_checkpoint( + local_path=checkpoint_path, hdfs_path=None, global_step=0 + ) # Step 2: Second update and forward pass outputs2 = model(input_ids=input_ids2, attention_mask=attention_mask2) @@ -109,7 +124,9 @@ def test_fsdp_ckpt(strategy="fsdp"): # Record logits after second update with torch.no_grad(): - logits_before_load = model(input_ids=input_ids2, attention_mask=attention_mask2).logits + logits_before_load = model( + input_ids=input_ids2, attention_mask=attention_mask2 + ).logits # Step 3: Load checkpoint and repeat second update checkpoint_manager.load_checkpoint(checkpoint_path) @@ -124,10 +141,14 @@ def test_fsdp_ckpt(strategy="fsdp"): # Record logits after loaded checkpoint and update with torch.no_grad(): - logits_after_load = model(input_ids=input_ids2, attention_mask=attention_mask2).logits + logits_after_load = model( + input_ids=input_ids2, attention_mask=attention_mask2 + ).logits # Step 4: Verify outputs match - torch.testing.assert_close(logits_before_load, logits_after_load, atol=0.0, rtol=0.0) + torch.testing.assert_close( + logits_before_load, logits_after_load, atol=0.0, rtol=0.0 + ) print("Checkpoint save/load test passed!") # Cleanup diff --git a/Agent0/executor_train/verl/tests/special_distributed/test_tensor_dict.py b/Agent0/executor_train/verl/tests/special_distributed/test_tensor_dict.py index 0a7f803..6d31d9e 100644 --- a/Agent0/executor_train/verl/tests/special_distributed/test_tensor_dict.py +++ b/Agent0/executor_train/verl/tests/special_distributed/test_tensor_dict.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,42 +12,51 @@ # See the License for the specific language governing permissions and # limitations under the License. +from verl.utils.distributed import initialize_global_process_group +from verl.protocol import DataProto, all_gather_data_proto +import torch.distributed +import torch +import numpy as np import os os.environ["NCCL_DEBUG"] = "WARN" -import numpy as np -import torch -import torch.distributed - -from verl.protocol import DataProto, all_gather_data_proto -from verl.utils.distributed import initialize_global_process_group - def test_all_gather_data_proto(): - device_mesh = torch.distributed.device_mesh.init_device_mesh("cuda", mesh_shape=[2, 2], mesh_dim_names=["dp", "tp"]) + device_mesh = torch.distributed.device_mesh.init_device_mesh( + "cuda", mesh_shape=[2, 2], mesh_dim_names=["dp", "tp"] + ) global_rank = torch.distributed.get_rank() - obs = torch.tensor([[1 * global_rank, 2 * global_rank + 1], [3 * global_rank, 4 * global_rank + 1]]) + obs = torch.tensor([[1 * global_rank, 2 * global_rank + 1], + [3 * global_rank, 4 * global_rank + 1]]) labels = ["a", "b"] if global_rank % 2 == 0 else ["b", "a"] labels = np.array(labels, dtype=object) - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + data = DataProto.from_dict( + tensors={"obs": obs}, + non_tensors={"labels": labels}, + meta_info={"info": "test_info"}, + ) all_gather_data_proto(data=data, process_group=device_mesh.get_group("dp")) if global_rank == 0: - expected_obs = torch.tensor([[0, 1], [0, 1], [2, 5], [6, 9]], device="cuda") + expected_obs = torch.tensor( + [[0, 1], [0, 1], [2, 5], [6, 9]], device="cuda") expected_labels = ["a", "b", "a", "b"] elif global_rank == 1: - expected_obs = torch.tensor([[1, 3], [3, 5], [3, 7], [9, 13]], device="cuda") + expected_obs = torch.tensor( + [[1, 3], [3, 5], [3, 7], [9, 13]], device="cuda") expected_labels = ["b", "a", "b", "a"] elif global_rank == 2: - expected_obs = torch.tensor([[0, 1], [0, 1], [2, 5], [6, 9]], device="cuda") + expected_obs = torch.tensor( + [[0, 1], [0, 1], [2, 5], [6, 9]], device="cuda") expected_labels = ["a", "b", "a", "b"] elif global_rank == 3: - expected_obs = torch.tensor([[1, 3], [3, 5], [3, 7], [9, 13]], device="cuda") + expected_obs = torch.tensor( + [[1, 3], [3, 5], [3, 7], [9, 13]], device="cuda") expected_labels = ["b", "a", "b", "a"] torch.testing.assert_close(data.batch["obs"], expected_obs, atol=0, rtol=0) @@ -63,22 +72,36 @@ def test_vocab_parallel_entropy(): from verl.utils.torch_functional import entropy_from_logits mpu.initialize_model_parallel( - tensor_model_parallel_size=2, pipeline_model_parallel_size=1, virtual_pipeline_model_parallel_size=None + tensor_model_parallel_size=2, + pipeline_model_parallel_size=1, + virtual_pipeline_model_parallel_size=None, ) batch_size = 2 seqlen = 128 vocab_size = 155136 - logits = torch.randn(batch_size * seqlen, vocab_size, device="cuda", requires_grad=True) - target = torch.randint(low=0, high=vocab_size, size=(batch_size * seqlen,), device="cuda", dtype=torch.int64) + logits = torch.randn( + batch_size * seqlen, vocab_size, device="cuda", requires_grad=True + ) + target = torch.randint( + low=0, + high=vocab_size, + size=(batch_size * seqlen,), + device="cuda", + dtype=torch.int64, + ) # broadcast across tp torch.distributed.broadcast( - logits, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group() + logits, + mpu.get_tensor_model_parallel_src_rank(), + group=mpu.get_tensor_model_parallel_group(), ) torch.distributed.broadcast( - target, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group() + target, + mpu.get_tensor_model_parallel_src_rank(), + group=mpu.get_tensor_model_parallel_group(), ) tp_rank = mpu.get_tensor_model_parallel_rank() @@ -86,8 +109,13 @@ def test_vocab_parallel_entropy(): # get the local logits of each tp vocab_parallel_logits = ( - logits.clone().detach()[:, tp_rank * vocab_size_per_tp : (tp_rank + 1) * vocab_size_per_tp].requires_grad_() - ) + logits.clone() .detach()[ + :, + tp_rank * + vocab_size_per_tp: ( + tp_rank + + 1) * + vocab_size_per_tp] .requires_grad_()) logits.grad = None vocab_parallel_logits.grad = None @@ -102,11 +130,13 @@ def test_vocab_parallel_entropy(): torch.testing.assert_close(output_entropy, target_entropy) target_entropy.backward(grad_output) torch.testing.assert_close( - logits.grad[:, tp_rank * vocab_size_per_tp : (tp_rank + 1) * vocab_size_per_tp], vocab_parallel_logits.grad + logits.grad[:, tp_rank * vocab_size_per_tp: (tp_rank + 1) * vocab_size_per_tp], + vocab_parallel_logits.grad, ) # make sure logits is not altered torch.testing.assert_close( - logits[:, tp_rank * vocab_size_per_tp : (tp_rank + 1) * vocab_size_per_tp], vocab_parallel_logits + logits[:, tp_rank * vocab_size_per_tp: (tp_rank + 1) * vocab_size_per_tp], + vocab_parallel_logits, ) if mpu.get_tensor_model_parallel_rank() == 0: diff --git a/Agent0/executor_train/verl/tests/special_e2e/__init__.py b/Agent0/executor_train/verl/tests/special_e2e/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/tests/special_e2e/__init__.py +++ b/Agent0/executor_train/verl/tests/special_e2e/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/special_e2e/check_custom_rwd_fn.py b/Agent0/executor_train/verl/tests/special_e2e/check_custom_rwd_fn.py index 8d77a53..69e71a3 100644 --- a/Agent0/executor_train/verl/tests/special_e2e/check_custom_rwd_fn.py +++ b/Agent0/executor_train/verl/tests/special_e2e/check_custom_rwd_fn.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,8 +19,12 @@ def check_congratulations_in_file(output_file): with open(output_file) as f: output = f.read() - success_message = "Congratulations!!! You have called my_reward_function successfully!!!" - assert success_message in output, f"Success message of my_reward_function not found in {output_file}" + success_message = ( + "Congratulations!!! You have called my_reward_function successfully!!!" + ) + assert ( + success_message in output + ), f"Success message of my_reward_function not found in {output_file}" print("Check passes") diff --git a/Agent0/executor_train/verl/tests/special_e2e/check_results.py b/Agent0/executor_train/verl/tests/special_e2e/check_results.py index 9453282..77273ff 100644 --- a/Agent0/executor_train/verl/tests/special_e2e/check_results.py +++ b/Agent0/executor_train/verl/tests/special_e2e/check_results.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -34,7 +34,11 @@ def extract_reward_from_line(line): if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--output_file", required=True, type=str) - parser.add_argument("--target", type=float, default=0.2, help="target reward score") + parser.add_argument( + "--target", + type=float, + default=0.2, + help="target reward score") args = parser.parse_args() @@ -49,5 +53,7 @@ def extract_reward_from_line(line): best_reward = reward print(f"Best reward is {best_reward}") - assert best_reward > args.target, f"Best reward must be greater than {args.target}. best_reward: {best_reward}" + assert ( + best_reward > args.target), f"Best reward must be greater than { + args.target}. best_reward: {best_reward}" print("Check passes") diff --git a/Agent0/executor_train/verl/tests/special_e2e/envs/__init__.py b/Agent0/executor_train/verl/tests/special_e2e/envs/__init__.py index eb85e22..67a1448 100644 --- a/Agent0/executor_train/verl/tests/special_e2e/envs/__init__.py +++ b/Agent0/executor_train/verl/tests/special_e2e/envs/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/__init__.py b/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/__init__.py index 80893ae..71b5100 100644 --- a/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/__init__.py +++ b/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,4 +19,7 @@ AutoTokenizer.register(LlamaConfig, CharTokenizer, exist_ok=True) -__all__ = ["DigitCompletion", "generate_ground_truth_response", "CharTokenizer"] +__all__ = [ + "DigitCompletion", + "generate_ground_truth_response", + "CharTokenizer"] diff --git a/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/task.py b/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/task.py index c3643a8..433ea7b 100644 --- a/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/task.py +++ b/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/task.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,7 +32,9 @@ class DigitCompletion: Note that the tokenizer is char-level to increase the difficulty. """ - def __init__(self, max_number: int, max_diff: int, max_num_in_response: int, seed=0): + def __init__( + self, max_number: int, max_diff: int, max_num_in_response: int, seed=0 + ): """ Args: @@ -49,16 +51,20 @@ def __init__(self, max_number: int, max_diff: int, max_num_in_response: int, see assert self.max_diff > 0 self.max_number_length = len(str(max_number)) # {num1},{num2}:{max_num_in_response},{max_number} - self._prompt_length = self.max_number_length * 2 + 4 + self.max_number_length # no negative is allowed + self._prompt_length = ( + self.max_number_length * 2 + 4 + self.max_number_length + ) # no negative is allowed self.np_rng = np.random.default_rng(seed=seed) def __str__(self): return ( - f"Prompt length: {self.prompt_length}. Response length: {self.response_length}, " - f"Max number: {self.max_number}. Max diff: {self.max_diff}, " - f"Max number in response: {self.max_num_in_response}" - ) + f"Prompt length: { + self.prompt_length}. Response length: { + self.response_length}, " f"Max number: { + self.max_number}. Max diff: { + self.max_diff}, " f"Max number in response: { + self.max_num_in_response}") def get_state(self): return {"rng": self.np_rng} @@ -75,7 +81,11 @@ def prompt_length(self): def response_length(self): # number length + comma length + [EOS] # The actual number times 1.5 to allow 'U' - return (self.max_num_in_response * self.max_number_length + (self.max_num_in_response - 1) + 1) * 2 + return ( + self.max_num_in_response * self.max_number_length + + (self.max_num_in_response - 1) + + 1 + ) * 2 def add(self, a, b): return (a + b) % self.max_number @@ -86,7 +96,12 @@ def get_all_prompts(self): for diff in range(0, self.max_diff + 1): second_num = self.add(first_num, diff) for num_to_complete in range(self.max_num_in_response + 1): - prompt = str(first_num) + "," + str(second_num) + f":{self.max_number},{num_to_complete}" + prompt = ( + str(first_num) + + "," + + str(second_num) + + f":{self.max_number},{num_to_complete}" + ) all_prompts.append(prompt) return all_prompts @@ -96,7 +111,12 @@ def sample_str_prompts(self): diff = self.np_rng.integers(self.max_diff + 1) second_num = self.add(first_num, diff) num_to_complete = self.np_rng.integers(self.max_num_in_response + 1) - prompt = str(first_num) + "," + str(second_num) + f":{self.max_number},{num_to_complete}" + prompt = ( + str(first_num) + + "," + + str(second_num) + + f":{self.max_number},{num_to_complete}" + ) return prompt def sample_batch_str_prompts(self, batch_size): @@ -140,10 +160,14 @@ def compute_reward(prompt: str, response: str, sequence_reward=1.0): """We compute dense reward here so that we can directly train RL without SFT""" response_length = len(response) ground_truth_response = generate_ground_truth_response(prompt) - per_token_reward = sequence_reward / (len(ground_truth_response) + 1) # including [EOS] + per_token_reward = sequence_reward / ( + len(ground_truth_response) + 1 + ) # including [EOS] # pad - reward = np.zeros(response_length, dtype=np.float32) # this assumes that each char is a token + reward = np.zeros( + response_length, dtype=np.float32 + ) # this assumes that each char is a token # assign reward until mismatches ground_truth_idx = 0 for i in range(response_length): diff --git a/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/tokenizer.py b/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/tokenizer.py index 6ff4719..7fb42df 100644 --- a/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/tokenizer.py +++ b/Agent0/executor_train/verl/tests/special_e2e/envs/digit_completion/tokenizer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,7 +27,12 @@ class CharTokenizer(PreTrainedTokenizer): - def __init__(self, characters: Sequence[str], model_max_length: int, chat_template, **kwargs): + def __init__( + self, + characters: Sequence[str], + model_max_length: int, + chat_template, + **kwargs): """Character tokenizer for Hugging Face transformers. Args: @@ -65,7 +70,8 @@ def __init__(self, characters: Sequence[str], model_max_length: int, chat_templa unk_token_str: 3, **{ch: i + 4 for i, ch in enumerate(characters)}, } - self._vocab_int_to_str = {v: k for k, v in self._vocab_str_to_int.items()} + self._vocab_int_to_str = { + v: k for k, v in self._vocab_str_to_int.items()} super().__init__( eos_token=eos_token, diff --git a/Agent0/executor_train/verl/tests/special_e2e/sft/test_sp_loss_match.py b/Agent0/executor_train/verl/tests/special_e2e/sft/test_sp_loss_match.py index 4dc0cbd..0d49cd3 100644 --- a/Agent0/executor_train/verl/tests/special_e2e/sft/test_sp_loss_match.py +++ b/Agent0/executor_train/verl/tests/special_e2e/sft/test_sp_loss_match.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -21,7 +21,9 @@ from verl.utils.distributed import initialize_global_process_group -def test_trainer_forward_consistency(trainer: FSDPSFTTrainer, total_steps: int = 4): +def test_trainer_forward_consistency( + trainer: FSDPSFTTrainer, + total_steps: int = 4): """Test consistency between original forward pass and SP+rmpad forward passes. Args: @@ -29,8 +31,12 @@ def test_trainer_forward_consistency(trainer: FSDPSFTTrainer, total_steps: int = total_steps: Number of steps to test (default: 4) """ if trainer.device_mesh.get_rank() == 0: - print("\nStarting debug comparison between original and SP+rmpad forward passes...") - print(f"Sequence parallel size: {trainer.config.ulysses_sequence_parallel_size}") + print( + "\nStarting debug comparison between original and SP+rmpad forward passes..." + ) + print( + f"Sequence parallel size: { + trainer.config.ulysses_sequence_parallel_size}") print(f"Remove padding: {trainer.use_remove_padding}\n") steps_remaining = total_steps @@ -38,34 +44,48 @@ def test_trainer_forward_consistency(trainer: FSDPSFTTrainer, total_steps: int = for epoch in range(1): # Just one epoch for testing trainer.train_sampler.set_epoch(epoch=epoch) for data in trainer.train_dataloader: - data = TensorDict(data, batch_size=trainer.config.data.train_batch_size).cuda() + data = TensorDict( + data, batch_size=trainer.config.data.train_batch_size + ).cuda() trainer.fsdp_model.train() - micro_batches = data.split(trainer.config.data.micro_batch_size_per_gpu) + micro_batches = data.split( + trainer.config.data.micro_batch_size_per_gpu) for idx, micro_batch in enumerate(micro_batches): if trainer.device_mesh.get_rank() == 0: - print(f"\nProcessing micro batch {idx + 1}/{len(micro_batches)}") + print( + f"\nProcessing micro batch {idx + 1}/{len(micro_batches)}") # Compute losses using both methods # Disable SP and rmpad trainer.use_remove_padding = False old_sp = trainer.config.ulysses_sequence_parallel_size trainer.config.ulysses_sequence_parallel_size = 1 - loss_ref = trainer._compute_loss_and_backward(micro_batch.copy(), do_backward=False) + loss_ref = trainer._compute_loss_and_backward( + micro_batch.copy(), do_backward=False + ) # Do SP and rmpad trainer.config.ulysses_sequence_parallel_size = old_sp trainer.use_remove_padding = True - loss_sp = trainer._compute_loss_and_backward(micro_batch.copy(), do_backward=False) + loss_sp = trainer._compute_loss_and_backward( + micro_batch.copy(), do_backward=False + ) # Collect losses across all ranks loss_ref_all = loss_ref.clone() loss_sp_all = loss_sp.clone() - torch.distributed.all_reduce(loss_ref_all, op=torch.distributed.ReduceOp.AVG) - torch.distributed.all_reduce(loss_sp_all, op=torch.distributed.ReduceOp.AVG) + torch.distributed.all_reduce( + loss_ref_all, op=torch.distributed.ReduceOp.AVG + ) + torch.distributed.all_reduce( + loss_sp_all, op=torch.distributed.ReduceOp.AVG + ) # Calculate relative difference of averaged losses - rel_diff = torch.abs(loss_ref_all - loss_sp_all) / (torch.abs(loss_ref_all) + 1e-8) + rel_diff = torch.abs(loss_ref_all - loss_sp_all) / ( + torch.abs(loss_ref_all) + 1e-8 + ) if trainer.device_mesh.get_rank() == 0: print("\nComparison Results (Averaged across ranks):") @@ -73,7 +93,9 @@ def test_trainer_forward_consistency(trainer: FSDPSFTTrainer, total_steps: int = print(f"SP+rmpad Loss: {loss_sp_all.item():.6f}") print(f"Relative Difference: {rel_diff.item():.6f}") - assert rel_diff.item() < 1e-2, "Significant difference detected between averaged losses!" + assert ( + rel_diff.item() < 1e-2 + ), "Significant difference detected between averaged losses!" print("Loss difference is within the acceptable range.") steps_remaining -= 1 @@ -98,11 +120,15 @@ def create_trainer(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=("fsdp",)) + device_mesh = init_device_mesh( + device_type="cuda", mesh_shape=(world_size,), mesh_dim_names=("fsdp",) + ) dp_size = world_size // config.ulysses_sequence_parallel_size ulysses_device_mesh = init_device_mesh( - device_type="cuda", mesh_shape=(dp_size, config.ulysses_sequence_parallel_size), mesh_dim_names=("dp", "sp") + device_type="cuda", + mesh_shape=(dp_size, config.ulysses_sequence_parallel_size), + mesh_dim_names=("dp", "sp"), ) # build tokenizer and datasets first @@ -110,10 +136,15 @@ def create_trainer(config): from verl.utils import hf_tokenizer from verl.utils.fs import copy_to_local - local_model_path = copy_to_local(src=config.model.partial_pretrain, verbose=True) - tokenizer = hf_tokenizer(local_model_path, trust_remote_code=config.model.trust_remote_code) - train_dataset = create_sft_dataset(config.data.train_files, config.data, tokenizer) - val_dataset = create_sft_dataset(config.data.val_files, config.data, tokenizer) + local_model_path = copy_to_local( + src=config.model.partial_pretrain, verbose=True) + tokenizer = hf_tokenizer( + local_model_path, trust_remote_code=config.model.trust_remote_code + ) + train_dataset = create_sft_dataset( + config.data.train_files, config.data, tokenizer) + val_dataset = create_sft_dataset( + config.data.val_files, config.data, tokenizer) return FSDPSFTTrainer( config=config, @@ -139,7 +170,8 @@ def main(config): import hydra from omegaconf import DictConfig - @hydra.main(config_path="../../../verl/trainer/config", config_name="sft_trainer") + @hydra.main(config_path="../../../verl/trainer/config", + config_name="sft_trainer") def hydra_entry(cfg: DictConfig) -> None: main(cfg) diff --git a/Agent0/executor_train/verl/tests/special_sanity/check_api_docs.py b/Agent0/executor_train/verl/tests/special_sanity/check_api_docs.py index fa31ec8..ccc1e36 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/check_api_docs.py +++ b/Agent0/executor_train/verl/tests/special_sanity/check_api_docs.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -55,11 +55,16 @@ def iter_submodules(root: ModuleType) -> Iterable[ModuleType]: """Yield *root* and every sub-module inside it.""" yield root if getattr(root, "__path__", None): # only packages have __path__ - for mod_info in pkgutil.walk_packages(root.__path__, prefix=f"{root.__name__}."): + for mod_info in pkgutil.walk_packages( + root.__path__, prefix=f"{root.__name__}." + ): try: yield importlib.import_module(mod_info.name) except Exception as exc: # noqa: BLE001 - print(f"[warn] Skipping {mod_info.name!r}: {exc}", file=sys.stderr) + print( + f"[warn] Skipping { + mod_info.name!r}: {exc}", + file=sys.stderr) def names_missing_doc(mod: ModuleType) -> list[str]: @@ -116,7 +121,9 @@ def main() -> None: targets = args.modules or autodiscover_packages() if not targets: - raise ValueError("[error] No modules specified and none detected automatically.") + raise ValueError( + "[error] No modules specified and none detected automatically." + ) all_missing: list[str] = [] for modname in targets: @@ -126,7 +133,9 @@ def main() -> None: print("\nMissing docstrings:") for name in sorted(all_missing): print(f" - {name}") - raise ValueError("Missing docstrings detected. Please enhance them with docs accordingly.") + raise ValueError( + "Missing docstrings detected. Please enhance them with docs accordingly." + ) print("โœ… All exported functions/classes have docstrings.") diff --git a/Agent0/executor_train/verl/tests/special_sanity/check_device_api_usage.py b/Agent0/executor_train/verl/tests/special_sanity/check_device_api_usage.py index c8988db..94e5f67 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/check_device_api_usage.py +++ b/Agent0/executor_train/verl/tests/special_sanity/check_device_api_usage.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,8 +33,10 @@ "verl/utils/rendezvous/ray_backend.py", # appear in cupy importance "verl/single_controller/ray/base.py", # appear in default device_name "verl/trainer/ppo/ray_trainer.py", # appear in default device_name - "verl/utils/reward_score/sandbox_fusion/utils.py", # appear in sandbox language type - "verl/workers/reward_model/megatron/reward_model.py", # appear in default device_name + "verl/utils/reward_score/sandbox_fusion/utils.py", + # appear in sandbox language type + "verl/workers/reward_model/megatron/reward_model.py", + # appear in default device_name ] # directory or file path must contain keyword "nccl" @@ -65,7 +67,8 @@ # for easy debugging in non-linux system sw = sw.replace("/", os.sep) if sw in path_in_str: - print(f"[SKIP] File {path_in_str} is in device api usage check whitelist, checking is skipped.") + print( + f"[SKIP] File {path_in_str} is in device api usage check whitelist, checking is skipped.") path_in_whitelist = True break @@ -83,9 +86,8 @@ break print( - f"[CHECK] File {path_in_str} is detected for device api usage check, check result: " - f"{'success' if not find_invalid_device_management else f'failed, because detect {sk}'}." - ) + f"[CHECK] File {path_in_str} is detected for device api usage check, check result: " f"{ + 'success' if not find_invalid_device_management else f'failed, because detect {sk}'}.") assert not find_invalid_device_management, ( f'file {path_in_str} contains .cuda/"cuda"/"nccl" usage, please use api in ' diff --git a/Agent0/executor_train/verl/tests/special_sanity/check_docs_time_info.py b/Agent0/executor_train/verl/tests/special_sanity/check_docs_time_info.py index a54d1d5..435ce7c 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/check_docs_time_info.py +++ b/Agent0/executor_train/verl/tests/special_sanity/check_docs_time_info.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -51,7 +51,10 @@ def is_allowed(path: Path) -> bool: def main(): if not DOCS_DIR.exists(): - print(f"Error: Documentation directory '{DOCS_DIR}' does not exist.", file=sys.stderr) + print( + f"Error: Documentation directory '{DOCS_DIR}' does not exist.", + file=sys.stderr, + ) sys.exit(1) missing = [] @@ -74,8 +77,7 @@ def main(): print(f"\nTotal missing: {len(missing)}\n", file=sys.stderr) raise AssertionError( "Some documentation files lack a 'Last updated' line. Please include info such as " - "'Last updated: mm/dd/yyyy' to indicate the last update time of the document." - ) + "'Last updated: mm/dd/yyyy' to indicate the last update time of the document.") else: print("โœ… All checked files contain 'Last updated'.") diff --git a/Agent0/executor_train/verl/tests/special_sanity/check_docstrings.py b/Agent0/executor_train/verl/tests/special_sanity/check_docstrings.py index 7c5d8ed..4f45ebd 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/check_docstrings.py +++ b/Agent0/executor_train/verl/tests/special_sanity/check_docstrings.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -35,8 +35,13 @@ def visit_FunctionDef(self, node: ast.FunctionDef): """Visit function definitions and check for docstrings.""" if not node.name.startswith("_") and self.function_nesting_level == 0: if not self._has_docstring(node): - func_name = f"{self.current_class}.{node.name}" if self.current_class else node.name - self.missing_docstrings.append((func_name, self.filename, node.lineno)) + func_name = ( + f"{self.current_class}.{node.name}" + if self.current_class + else node.name + ) + self.missing_docstrings.append( + (func_name, self.filename, node.lineno)) self.function_nesting_level += 1 self.generic_visit(node) @@ -46,8 +51,13 @@ def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef): """Visit async function definitions and check for docstrings.""" if not node.name.startswith("_") and self.function_nesting_level == 0: if not self._has_docstring(node): - func_name = f"{self.current_class}.{node.name}" if self.current_class else node.name - self.missing_docstrings.append((func_name, self.filename, node.lineno)) + func_name = ( + f"{self.current_class}.{node.name}" + if self.current_class + else node.name + ) + self.missing_docstrings.append( + (func_name, self.filename, node.lineno)) self.function_nesting_level += 1 self.generic_visit(node) @@ -57,7 +67,8 @@ def visit_ClassDef(self, node: ast.ClassDef): """Visit class definitions and check for docstrings.""" if not node.name.startswith("_"): if not self._has_docstring(node): - self.missing_docstrings.append((node.name, self.filename, node.lineno)) + self.missing_docstrings.append( + (node.name, self.filename, node.lineno)) old_class = self.current_class self.current_class = node.name @@ -130,7 +141,9 @@ def main(): print("=" * 60) if all_missing_docstrings: - print(f"\nSUMMARY: Found {len(all_missing_docstrings)} functions/classes missing docstrings:") + print( + f"\nSUMMARY: Found { + len(all_missing_docstrings)} functions/classes missing docstrings:") print("-" * 60) by_file = {} @@ -146,7 +159,9 @@ def main(): print(f"\nTotal missing docstrings: {len(all_missing_docstrings)}") - raise Exception(f"Found {len(all_missing_docstrings)} functions/classes without proper docstrings!") + raise Exception( + f"Found { + len(all_missing_docstrings)} functions/classes without proper docstrings!") else: print("\nโœ… All functions and classes have proper docstrings!") diff --git a/Agent0/executor_train/verl/tests/special_sanity/check_license.py b/Agent0/executor_train/verl/tests/special_sanity/check_license.py index a02afeb..67bfd68 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/check_license.py +++ b/Agent0/executor_train/verl/tests/special_sanity/check_license.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,13 +14,13 @@ from argparse import ArgumentParser from pathlib import Path -license_head_bytedance = "Copyright 2024 Bytedance Ltd. and/or its affiliates" -license_head_bytedance_25 = "Copyright 2025 Bytedance Ltd. and/or its affiliates" +license_head_bytedance = "Copyright 2024-2026 Bytedance Ltd. and/or its affiliates" +license_head_bytedance_25 = "Copyright 2025-2026 Bytedance Ltd. and/or its affiliates" # Add custom license headers below license_head_prime = "Copyright 2024 PRIME team and/or its affiliates" -license_head_individual = "Copyright 2025 Individual Contributor:" -license_head_sglang = "Copyright 2023-2024 SGLang Team" -license_head_modelbest = "Copyright 2025 ModelBest Inc. and/or its affiliates" +license_head_individual = "Copyright 2025-2026 Individual Contributor:" +license_head_sglang = "Copyright 2023-2026 SGLang Team" +license_head_modelbest = "Copyright 2025-2026 ModelBest Inc. and/or its affiliates" license_head_amazon = "Copyright 2025 Amazon.com Inc and/or its affiliates" license_headers = [ license_head_bytedance, diff --git a/Agent0/executor_train/verl/tests/special_sanity/check_pr_description.py b/Agent0/executor_train/verl/tests/special_sanity/check_pr_description.py index 4ed4563..2c49b2f 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/check_pr_description.py +++ b/Agent0/executor_train/verl/tests/special_sanity/check_pr_description.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -34,7 +34,9 @@ class PRDescriptionError(Exception): # Path to the PR template file -template_file = os.path.join(os.getenv("GITHUB_WORKSPACE", "."), ".github", "PULL_REQUEST_TEMPLATE.md") +template_file = os.path.join( + os.getenv("GITHUB_WORKSPACE", "."), ".github", "PULL_REQUEST_TEMPLATE.md" +) def load_template(path): @@ -52,7 +54,9 @@ def load_template(path): lines.append(line.strip()) return lines except Exception as e: - raise TemplateFileError(f"Failed to read PR template (first {NUM_LINES} lines) at {path}: {e}") from e + raise TemplateFileError( + f"Failed to read PR template (first {NUM_LINES} lines) at {path}: {e}" + ) from e def load_pr_body(event_path): @@ -61,7 +65,8 @@ def load_pr_body(event_path): payload = json.load(f) return payload.get("pull_request", {}).get("body", "") or "" except Exception as e: - raise PRBodyLoadError(f"Failed to read PR body from {event_path}: {e}") from e + raise PRBodyLoadError( + f"Failed to read PR body from {event_path}: {e}") from e def check_pr_description(body, template_lines): @@ -74,8 +79,7 @@ def check_pr_description(body, template_lines): if pr_first == template_lines: raise PRDescriptionError( "It looks like you haven't updated the '### What does this PR do?' section. Please replace " - "the placeholder text with a concise description of what your PR does." - ) + "the placeholder text with a concise description of what your PR does.") else: print(pr_first) print(template_lines) diff --git a/Agent0/executor_train/verl/tests/special_sanity/check_pr_title.py b/Agent0/executor_train/verl/tests/special_sanity/check_pr_title.py index f4cbd52..819b212 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/check_pr_title.py +++ b/Agent0/executor_train/verl/tests/special_sanity/check_pr_title.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,9 +20,21 @@ # Define rules allowed_modules = ["fsdp", "megatron", "sglang", "vllm", "rollout", "trainer"] -allowed_modules += ["tests", "training_utils", "recipe", "hardware", "deployment"] -allowed_modules += ["ray", "worker", "single_controller", "misc", "docker", "ci"] -allowed_modules += ["perf", "model", "algo", "env", "tool", "ckpt", "doc", "data", "cfg"] +allowed_modules += ["tests", "training_utils", + "recipe", "hardware", "deployment"] +allowed_modules += ["ray", "worker", + "single_controller", "misc", "docker", "ci"] +allowed_modules += [ + "perf", + "model", + "algo", + "env", + "tool", + "ckpt", + "doc", + "data", + "cfg", +] allowed_types = ["feat", "fix", "refactor", "chore", "test"] # Check for [BREAKING] prefix and extract the rest of the title @@ -45,13 +57,17 @@ else: modules = re.findall(r"[a-z_]+", re_modules.group(1).lower()) if not all(module in allowed_modules for module in modules): - invalid_modules = [module for module in modules if module not in allowed_modules] + invalid_modules = [ + module for module in modules if module not in allowed_modules + ] print(f"โŒ Invalid modules: {', '.join(invalid_modules)}") print(f"Allowed modules: {', '.join(allowed_modules)}") raise Exception("Invalid PR title") types_pattern = "|".join(re.escape(t) for t in allowed_types) -re_types_pattern = re.compile(rf"^\[[a-z_,\s]+\]\s+({types_pattern}):\s+.+$", re.IGNORECASE) +re_types_pattern = re.compile( + rf"^\[[a-z_,\s]+\]\s+({types_pattern}):\s+.+$", re.IGNORECASE +) match = re_types_pattern.match(core_pr_title) if not match: @@ -64,4 +80,6 @@ # Build the success message breaking_info = " (BREAKING CHANGE)" if is_breaking else "" -print(f"โœ… PR title is valid: {pr_title}, modules: {modules}, type: {change_type}{breaking_info}") +print( + f"โœ… PR title is valid: {pr_title}, modules: {modules}, type: {change_type}{breaking_info}" +) diff --git a/Agent0/executor_train/verl/tests/special_sanity/test_config_docs.py b/Agent0/executor_train/verl/tests/special_sanity/test_config_docs.py index 2f260f1..36d3794 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/test_config_docs.py +++ b/Agent0/executor_train/verl/tests/special_sanity/test_config_docs.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -34,23 +34,35 @@ def validate_yaml_format(yaml_lines): if key_match: # Check if there's a comment above if i == 0 or not yaml_lines[i - 1].strip().startswith("#"): - errors.append(f"Missing comment above line {i + 1}: {line.strip()}") + errors.append( + f"Missing comment above line { + i + + 1}: { + line.strip()}") # Check for inline comment if "#" in line and not stripped.startswith("#"): comment_index = line.index("#") colon_index = line.index(":") if comment_index > colon_index: - errors.append(f"Inline comment found on line {i + 1}: {line.strip()}") + errors.append( + f"Inline comment found on line {i + 1}: {line.strip()}" + ) - # Check for blank line after this key line (unless next is a deeper indent) + # Check for blank line after this key line (unless next is a deeper + # indent) if i + 1 < len(yaml_lines): next_line = yaml_lines[i + 1] next_stripped = next_line.strip() - # If next is not empty and not a deeper nested line, enforce blank line + # If next is not empty and not a deeper nested line, enforce + # blank line if next_stripped != "": - errors.append(f"Missing blank line after line {i + 1}: {line.strip()}") + errors.append( + f"Missing blank line after line { + i + + 1}: { + line.strip()}") i += 1 @@ -76,7 +88,8 @@ def test_trainer_config_doc(): if validation_errors: success = False print("YAML documentation format check failed:") - print(f"Please read the top block of {yaml_to_inspect} to see format rules:\n") + print( + f"Please read the top block of {yaml_to_inspect} to see format rules:\n") for err in validation_errors: print(" -", err) diff --git a/Agent0/executor_train/verl/tests/special_sanity/test_import.py b/Agent0/executor_train/verl/tests/special_sanity/test_import.py index 4f8a918..848e656 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/test_import.py +++ b/Agent0/executor_train/verl/tests/special_sanity/test_import.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/special_sanity/type_coverage_check.py b/Agent0/executor_train/verl/tests/special_sanity/type_coverage_check.py index dc6dc7c..91a0959 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/type_coverage_check.py +++ b/Agent0/executor_train/verl/tests/special_sanity/type_coverage_check.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,7 +26,9 @@ def get_changed_files() -> list[Path]: result = subprocess.run( - ["git", "diff", "--name-only", "--diff-filter=AM", "origin/main...HEAD"], stdout=subprocess.PIPE, text=True + ["git", "diff", "--name-only", "--diff-filter=AM", "origin/main...HEAD"], + stdout=subprocess.PIPE, + text=True, ) return [Path(f) for f in result.stdout.splitlines() if f.endswith(".py")] @@ -70,14 +72,25 @@ def has_type_annotations(node: ast.AST, debug: bool = False) -> int: if isinstance(node, ast.FunctionDef): is_private = node.name.startswith("_") has_ann = ( - all(arg.annotation is not None for arg in node.args.args if should_check_type(arg.arg)) + all( + arg.annotation is not None + for arg in node.args.args + if should_check_type(arg.arg) + ) and node.returns is not None ) if has_ann or is_private: return CHECK_SUCCESS else: if debug: - print(node, [(arg.annotation, arg.arg) for arg in node.args.args if should_check_type(arg.arg)]) + print( + node, + [ + (arg.annotation, arg.arg) + for arg in node.args.args + if should_check_type(arg.arg) + ], + ) return CHECK_FAILURE return CHECK_SUCCESS @@ -102,10 +115,11 @@ def check_file( annotated += 1 if result == CHECK_WARNING: warning_lines.append( - (file_path, node.lineno, linecache.getline(str(file_path), node.lineno).strip()) - ) + (file_path, node.lineno, linecache.getline( + str(file_path), node.lineno).strip(), )) else: - source_line = linecache.getline(str(file_path), node.lineno).strip() + source_line = linecache.getline( + str(file_path), node.lineno).strip() failure_lines.append((file_path, node.lineno, source_line)) return annotated, total, warning_lines, failure_lines @@ -114,15 +128,26 @@ def check_file( def main() -> None: parser = argparse.ArgumentParser() parser.add_argument( - "--threshold", type=float, default=0.3, help="Minimum ratio of annotated lines required (0.0 - 1.0)" + "--threshold", + type=float, + default=0.3, + help="Minimum ratio of annotated lines required (0.0 - 1.0)", + ) + parser.add_argument( + "--target-file", + type=str, + default=None, + help="Path to the Python source file to analyse", ) - parser.add_argument("--target-file", type=str, default=None, help="Path to the Python source file to analyse") parser.add_argument( "--all-lines", action="store_true", help="Check all lines in the file instead of only changed lines based on git", ) - parser.add_argument("--debug", action="store_true", help="Add debugging logs") + parser.add_argument( + "--debug", + action="store_true", + help="Add debugging logs") args = parser.parse_args() total_changed = 0 @@ -130,15 +155,19 @@ def main() -> None: all_warnings: list[tuple[Path, int, str]] = [] all_failures: list[tuple[Path, int, str]] = [] - target_files = [args.target_file] if args.target_file is not None else get_changed_files() + target_files = ([args.target_file] + if args.target_file is not None else get_changed_files()) for fpath in target_files: if "tests/" in str(fpath): continue if args.all_lines: - changed_lines = [i + 1 for i in range(len(open(fpath).readlines()))] + changed_lines = [ + i + 1 for i in range(len(open(fpath).readlines()))] else: changed_lines = get_changed_lines(fpath) - annotated, total, warning_lines, failure_lines = check_file(fpath, changed_lines, args.debug) + annotated, total, warning_lines, failure_lines = check_file( + fpath, changed_lines, args.debug + ) total_annotated += annotated total_changed += total all_warnings.extend(warning_lines) @@ -152,7 +181,9 @@ def main() -> None: ) if all_warnings: - print("\nโš ๏ธ Suggest Improve: Lines missing type annotations for inputs and outputs:\n") + print( + "\nโš ๏ธ Suggest Improve: Lines missing type annotations for inputs and outputs:\n" + ) for fname, lineno, line in all_warnings: print(f"{fname}:{lineno}: {line}") @@ -163,13 +194,14 @@ def main() -> None: if ratio < args.threshold: print( - f"Please add type annotations for inputs and outputs to meet threshold {args.threshold}. " - f"Cases exempt from checking:" - ) + f"Please add type annotations for inputs and outputs to meet threshold { + args.threshold}. " f"Cases exempt from checking:") print("1. Private methods.") print("2. Args with name in ('self', 'cls'), or *args / **kwargs") print("3. Files under tests/") - raise Exception(f"\nโŒ Type coverage below threshold ({args.threshold:.0%}).") + raise Exception( + f"\nโŒ Type coverage below threshold ({ + args.threshold:.0%}).") else: if all_warnings or all_failures: print("") diff --git a/Agent0/executor_train/verl/tests/special_sanity/validate_imported_docs.py b/Agent0/executor_train/verl/tests/special_sanity/validate_imported_docs.py index b36a407..7f7b838 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/validate_imported_docs.py +++ b/Agent0/executor_train/verl/tests/special_sanity/validate_imported_docs.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -30,7 +30,9 @@ def _parse_args() -> argparse.Namespace: - p = argparse.ArgumentParser(description="Verify that imported functions/classes have docstrings.") + p = argparse.ArgumentParser( + description="Verify that imported functions/classes have docstrings." + ) p.add_argument( "--target-file", default="verl/trainer/ppo/ray_trainer.py", @@ -60,7 +62,9 @@ def _import_attr(module_name: str, attr_name: str): return getattr(module, attr_name) -def _check_file(py_file: pathlib.Path, project_root: pathlib.Path, allow_list: list[str]) -> list[str]: +def _check_file( + py_file: pathlib.Path, project_root: pathlib.Path, allow_list: list[str] +) -> list[str]: """Return a list of error strings (empty == success).""" # Ensure local packages resolve sys.path.insert(0, str(project_root.resolve())) @@ -77,8 +81,8 @@ def _check_file(py_file: pathlib.Path, project_root: pathlib.Path, allow_list: l for alias in node.names: if alias.name == "*": problems.append( - f"{py_file}:{node.lineno} - wildcard import `from {module_name} import *` cannot be verified." - ) + f"{py_file}:{ + node.lineno} - wildcard import `from {module_name} import *` cannot be verified.") continue imported_name = alias.name @@ -101,8 +105,8 @@ def _check_file(py_file: pathlib.Path, project_root: pathlib.Path, allow_list: l if not (doc and doc.strip()): kind = "class" if inspect.isclass(obj) else "function" problems.append( - f"{py_file}:{node.lineno} - {kind} `{module_name}.{imported_name}` is missing a docstring." - ) + f"{py_file}:{ + node.lineno} - {kind} `{module_name}.{imported_name}` is missing a docstring.") return problems @@ -123,7 +127,9 @@ def main() -> None: raise Exception("โŒ Docstring verification failed.") if not args.quiet: - print(f"โœ… All explicitly imported functions/classes in {target_path} have docstrings.") + print( + f"โœ… All explicitly imported functions/classes in {target_path} have docstrings." + ) if __name__ == "__main__": diff --git a/Agent0/executor_train/verl/tests/special_sanity/validate_structure.py b/Agent0/executor_train/verl/tests/special_sanity/validate_structure.py index a5390b1..2979aa5 100644 --- a/Agent0/executor_train/verl/tests/special_sanity/validate_structure.py +++ b/Agent0/executor_train/verl/tests/special_sanity/validate_structure.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -43,7 +43,9 @@ def discover_allowed_modules(impl_root: Path, extra: list[str]) -> set[str]: return allowed -def find_violations(tests_root: Path, allowed: set[str], allowed_files: list[str]) -> list[str]: +def find_violations( + tests_root: Path, allowed: set[str], allowed_files: list[str] +) -> list[str]: """Return a list of error strings for test files in the wrong place.""" errors: list[str] = [] for test_file in tests_root.rglob("test*.py"): @@ -51,7 +53,9 @@ def find_violations(tests_root: Path, allowed: set[str], allowed_files: list[str continue rel_parts = test_file.relative_to(tests_root).parts if len(rel_parts) < 2: - errors.append(f"{test_file}: must be inside one of {sorted(allowed)} (not at tests root)") + errors.append( + f"{test_file}: must be inside one of { + sorted(allowed)} (not at tests root)") continue first_folder = rel_parts[0] @@ -64,7 +68,9 @@ def find_violations(tests_root: Path, allowed: set[str], allowed_files: list[str def main() -> None: - parser = argparse.ArgumentParser(description="Check that test files follow tests//โ€ฆ layout.") + parser = argparse.ArgumentParser( + description="Check that test files follow tests//โ€ฆ layout." + ) parser.add_argument( "--impl-root", type=Path, @@ -80,19 +86,28 @@ def main() -> None: parser.add_argument( "--allow-dirs", nargs="*", - default=["special_e2e", "special_sanity", "special_standalone", "special_distributed"], + default=[ + "special_e2e", + "special_sanity", + "special_standalone", + "special_distributed", + ], help="Extra top-level test folders that are exempt from the rule", ) parser.add_argument( "--allow-files", nargs="*", - default=["tests/test_protocol_on_cpu.py", "tests/test_base_config_on_cpu.py"], + default=[ + "tests/test_protocol_on_cpu.py", + "tests/test_base_config_on_cpu.py"], help="Extra top-level test folders that are exempt from the rule", ) args = parser.parse_args() if not args.impl_root.is_dir(): - raise Exception(f"Implementation root '{args.impl_root}' does not exist.") + raise Exception( + f"Implementation root '{ + args.impl_root}' does not exist.") if not args.tests_root.is_dir(): raise Exception(f"Tests root '{args.tests_root}' does not exist.") diff --git a/Agent0/executor_train/verl/tests/special_standalone/test_memory_buffers.py b/Agent0/executor_train/verl/tests/special_standalone/test_memory_buffers.py index 7785153..6bce22d 100644 --- a/Agent0/executor_train/verl/tests/special_standalone/test_memory_buffers.py +++ b/Agent0/executor_train/verl/tests/special_standalone/test_memory_buffers.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -43,7 +43,9 @@ def test_memory_buffers(): r_before = torch.cuda.memory_reserved(0) / norm_factor a_before = torch.cuda.memory_allocated(0) / norm_factor - print(f"Before Total memory: {t_before} GB, reserved: {r_before} GB, allocated: {a_before} GB") + print( + f"Before Total memory: {t_before} GB, reserved: {r_before} GB, allocated: {a_before} GB" + ) t = torch.cuda.get_device_properties(0).total_memory / norm_factor r = torch.cuda.memory_reserved(0) / norm_factor @@ -55,11 +57,17 @@ def test_memory_buffers(): print(f"After Total memory: {t} GB, reserved: {r} GB, allocated: {a} GB") change_ratio = (a - a_before) / a_before - assert change_ratio < 0.01, f"make sure the allocated change is less than 1%, Got {change_ratio}" + assert ( + change_ratio < 0.01 + ), f"make sure the allocated change is less than 1%, Got {change_ratio}" - for (name1, param1), (name2, param2) in zip(model.named_parameters(), model_copy.named_parameters(), strict=True): + for (name1, param1), (name2, param2) in zip( + model.named_parameters(), model_copy.named_parameters(), strict=True + ): assert name1 == name2 - assert torch.eq(param1.data, param2.data).all(), f"{param1.data}, {param2.data}, {name1}" + assert torch.eq( + param1.data, param2.data + ).all(), f"{param1.data}, {param2.data}, {name1}" if __name__ == "__main__": diff --git a/Agent0/executor_train/verl/tests/test_base_config_on_cpu.py b/Agent0/executor_train/verl/tests/test_base_config_on_cpu.py index 9a50235..df465c9 100644 --- a/Agent0/executor_train/verl/tests/test_base_config_on_cpu.py +++ b/Agent0/executor_train/verl/tests/test_base_config_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/test_protocol_on_cpu.py b/Agent0/executor_train/verl/tests/test_protocol_on_cpu.py index 2052635..066b70a 100644 --- a/Agent0/executor_train/verl/tests/test_protocol_on_cpu.py +++ b/Agent0/executor_train/verl/tests/test_protocol_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,12 +26,13 @@ def test_union_tensor_dict(): obs = torch.randn(100, 10) - data1 = TensorDict({"obs": obs, "act": torch.randn(100, 3)}, batch_size=[100]) - data2 = TensorDict({"obs": obs, "next_obs": torch.randn(100, 10), "rew": torch.randn(100)}, batch_size=[100]) + data1 = TensorDict( + {"obs": obs, "act": torch.randn(100, 3)}, batch_size=[100]) + data2 = TensorDict({"obs": obs, "next_obs": torch.randn( + 100, 10), "rew": torch.randn(100)}, batch_size=[100], ) - data_with_copied_obs = TensorDict( - {"obs": obs.clone(), "next_obs": torch.randn(100, 10), "rew": torch.randn(100)}, batch_size=[100] - ) + data_with_copied_obs = TensorDict({"obs": obs.clone(), "next_obs": torch.randn( + 100, 10), "rew": torch.randn(100)}, batch_size=[100], ) data = union_tensor_dict(data1, data2) with pytest.raises(AssertionError): @@ -58,16 +59,27 @@ def test_tensor_dict_constructor(): assert data.batch.batch_size == torch.Size([100]) with pytest.raises(AssertionError): - data = DataProto.from_dict(tensors={"obs": obs, "act": act}, num_batch_dims=2) + data = DataProto.from_dict( + tensors={ + "obs": obs, + "act": act}, + num_batch_dims=2) with pytest.raises(AssertionError): - data = DataProto.from_dict(tensors={"obs": obs, "act": act}, num_batch_dims=3) + data = DataProto.from_dict( + tensors={ + "obs": obs, + "act": act}, + num_batch_dims=3) def test_tensor_dict_make_iterator(): obs = torch.randn(100, 10) labels = [random.choice(["abc", "cde"]) for _ in range(100)] - dataset = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}) + dataset = DataProto.from_dict( + tensors={ + "obs": obs}, non_tensors={ + "labels": labels}) data_iter_1 = dataset.make_iterator(mini_batch_size=10, epochs=2, seed=1) data_list_1 = [] @@ -87,7 +99,10 @@ def test_tensor_dict_make_iterator(): print(data1.batch["obs"]) print(data2.batch["obs"]) raise AssertionError() - non_tensor_result = np.all(np.equal(data1.non_tensor_batch["labels"], data2.non_tensor_batch["labels"])) + non_tensor_result = np.all( + np.equal( + data1.non_tensor_batch["labels"], + data2.non_tensor_batch["labels"])) if not non_tensor_result.item(): print(data1.non_tensor_batch["labels"]) print(data2.non_tensor_batch["labels"]) @@ -96,42 +111,58 @@ def test_tensor_dict_make_iterator(): def test_reorder(): obs = torch.tensor([1, 2, 3, 4, 5, 6]) labels = ["a", "b", "c", "d", "e", "f"] - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"name": "abdce"}) + data = DataProto.from_dict( + tensors={"obs": obs}, + non_tensors={"labels": labels}, + meta_info={"name": "abdce"}, + ) data.reorder(torch.tensor([3, 4, 2, 0, 1, 5])) - assert torch.all(torch.eq(data.batch["obs"], torch.tensor([4, 5, 3, 1, 2, 6]))) - assert np.all(data.non_tensor_batch["labels"] == np.array(["d", "e", "c", "a", "b", "f"])) + assert torch.all( + torch.eq(data.batch["obs"], torch.tensor([4, 5, 3, 1, 2, 6]))) + assert np.all(data.non_tensor_batch["labels"] == np.array( + ["d", "e", "c", "a", "b", "f"])) assert data.meta_info == {"name": "abdce"} def test_chunk_concat(): obs = torch.tensor([1, 2, 3, 4, 5, 6]) labels = ["a", "b", "c", "d", "e", "f"] - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"name": "abdce"}) + data = DataProto.from_dict( + tensors={"obs": obs}, + non_tensors={"labels": labels}, + meta_info={"name": "abdce"}, + ) with pytest.raises(AssertionError): data.chunk(5) data_split = data.chunk(2) assert len(data_split) == 2 - assert torch.all(torch.eq(data_split[0].batch["obs"], torch.tensor([1, 2, 3]))) - assert np.all(data_split[0].non_tensor_batch["labels"] == np.array(["a", "b", "c"])) + assert torch.all( + torch.eq(data_split[0].batch["obs"], torch.tensor([1, 2, 3]))) + assert np.all(data_split[0].non_tensor_batch["labels"] + == np.array(["a", "b", "c"])) assert data_split[0].meta_info == {"name": "abdce"} - assert torch.all(torch.eq(data_split[1].batch["obs"], torch.tensor([4, 5, 6]))) - assert np.all(data_split[1].non_tensor_batch["labels"] == np.array(["d", "e", "f"])) + assert torch.all( + torch.eq(data_split[1].batch["obs"], torch.tensor([4, 5, 6]))) + assert np.all(data_split[1].non_tensor_batch["labels"] + == np.array(["d", "e", "f"])) assert data_split[1].meta_info == {"name": "abdce"} concat_data = DataProto.concat(data_split) assert torch.all(torch.eq(concat_data.batch["obs"], data.batch["obs"])) - assert np.all(concat_data.non_tensor_batch["labels"] == data.non_tensor_batch["labels"]) + assert np.all( + concat_data.non_tensor_batch["labels"] == data.non_tensor_batch["labels"]) assert concat_data.meta_info == data.meta_info def test_pop(): obs = torch.randn(100, 10) act = torch.randn(100, 3) - dataset = DataProto.from_dict({"obs": obs, "act": act}, meta_info={"2": 2, "1": 1}) + dataset = DataProto.from_dict( + {"obs": obs, "act": act}, meta_info={"2": 2, "1": 1}) poped_dataset = dataset.pop(batch_keys=["obs"], meta_info_keys=["2"]) assert poped_dataset.batch.keys() == {"obs"} @@ -145,31 +176,55 @@ def test_repeat(): # Create a DataProto object with some batch and non-tensor data obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) labels = ["a", "b", "c"] - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + data = DataProto.from_dict( + tensors={"obs": obs}, + non_tensors={"labels": labels}, + meta_info={"info": "test_info"}, + ) # Test interleave=True repeated_data_interleave = data.repeat(repeat_times=2, interleave=True) - expected_obs_interleave = torch.tensor([[1, 2], [1, 2], [3, 4], [3, 4], [5, 6], [5, 6]]) + expected_obs_interleave = torch.tensor( + [[1, 2], [1, 2], [3, 4], [3, 4], [5, 6], [5, 6]] + ) expected_labels_interleave = ["a", "a", "b", "b", "c", "c"] - assert torch.all(torch.eq(repeated_data_interleave.batch["obs"], expected_obs_interleave)) - assert (repeated_data_interleave.non_tensor_batch["labels"] == expected_labels_interleave).all() + assert torch.all( + torch.eq( + repeated_data_interleave.batch["obs"], + expected_obs_interleave)) + assert ( + repeated_data_interleave.non_tensor_batch["labels"] + == expected_labels_interleave + ).all() assert repeated_data_interleave.meta_info == {"info": "test_info"} # Test interleave=False repeated_data_no_interleave = data.repeat(repeat_times=2, interleave=False) - expected_obs_no_interleave = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6]]) + expected_obs_no_interleave = torch.tensor( + [[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6]] + ) expected_labels_no_interleave = ["a", "b", "c", "a", "b", "c"] - assert torch.all(torch.eq(repeated_data_no_interleave.batch["obs"], expected_obs_no_interleave)) - assert (repeated_data_no_interleave.non_tensor_batch["labels"] == expected_labels_no_interleave).all() + assert torch.all( + torch.eq( + repeated_data_no_interleave.batch["obs"], + expected_obs_no_interleave)) + assert ( + repeated_data_no_interleave.non_tensor_batch["labels"] + == expected_labels_no_interleave + ).all() assert repeated_data_no_interleave.meta_info == {"info": "test_info"} def test_dataproto_pad_unpad(): obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) labels = ["a", "b", "c"] - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + data = DataProto.from_dict( + tensors={"obs": obs}, + non_tensors={"labels": labels}, + meta_info={"info": "test_info"}, + ) from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto @@ -206,7 +261,9 @@ def test_dataproto_pad_unpad(): padded_data, pad_size = pad_dataproto_to_divisor(data, size_divisor=7) assert pad_size == 4 - expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6], [1, 2]]) + expected_obs = torch.tensor( + [[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6], [1, 2]] + ) expected_labels = ["a", "b", "c", "a", "b", "c", "a"] assert torch.all(torch.eq(padded_data.batch["obs"], expected_obs)) assert (padded_data.non_tensor_batch["labels"] == expected_labels).all() @@ -223,33 +280,50 @@ def test_dataproto_fold_unfold(): obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) labels = ["a", "b", "c"] - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + data = DataProto.from_dict( + tensors={"obs": obs}, + non_tensors={"labels": labels}, + meta_info={"info": "test_info"}, + ) data1 = data.repeat(repeat_times=2, interleave=True) data2 = fold_batch_dim(data1, new_batch_size=3) - torch.testing.assert_close(data2.batch["obs"], torch.tensor([[[1, 2], [1, 2]], [[3, 4], [3, 4]], [[5, 6], [5, 6]]])) - assert (data2.non_tensor_batch["labels"] == [["a", "a"], ["b", "b"], ["c", "c"]]).all() + torch.testing.assert_close( + data2.batch["obs"], + torch.tensor([[[1, 2], [1, 2]], [[3, 4], [3, 4]], [[5, 6], [5, 6]]]), + ) + assert (data2.non_tensor_batch["labels"] == [ + ["a", "a"], ["b", "b"], ["c", "c"]]).all() data2.reorder(indices=torch.tensor([1, 2, 0])) data3 = unfold_batch_dim(data2, batch_dims=2) - torch.testing.assert_close(data3.batch["obs"], torch.tensor([[3, 4], [3, 4], [5, 6], [5, 6], [1, 2], [1, 2]])) - assert (data3.non_tensor_batch["labels"] == ["b", "b", "c", "c", "a", "a"]).all() + torch.testing.assert_close( + data3.batch["obs"], + torch.tensor([[3, 4], [3, 4], [5, 6], [5, 6], [1, 2], [1, 2]]), + ) + assert (data3.non_tensor_batch["labels"] == [ + "b", "b", "c", "c", "a", "a"]).all() assert data3.meta_info == {"info": "test_info"} def test_torch_save_data_proto(): obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) labels = ["a", "b", "c"] - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + data = DataProto.from_dict( + tensors={"obs": obs}, + non_tensors={"labels": labels}, + meta_info={"info": "test_info"}, + ) data.save_to_disk("test_data.pt") loaded_data = DataProto.load_from_disk("test_data.pt") assert torch.all(torch.eq(loaded_data.batch["obs"], data.batch["obs"])) - assert (loaded_data.non_tensor_batch["labels"] == data.non_tensor_batch["labels"]).all() + assert (loaded_data.non_tensor_batch["labels"] + == data.non_tensor_batch["labels"]).all() assert loaded_data.meta_info == data.meta_info import os @@ -260,19 +334,34 @@ def test_torch_save_data_proto(): def test_len(): obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) labels = np.array(["a", "b", "c"], dtype=object) - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + data = DataProto.from_dict( + tensors={"obs": obs}, + non_tensors={"labels": labels}, + meta_info={"info": "test_info"}, + ) assert len(data) == 3 - data = DataProto(batch=None, non_tensor_batch={"labels": labels}, meta_info={"info": "test_info"}) + data = DataProto( + batch=None, non_tensor_batch={ + "labels": labels}, meta_info={ + "info": "test_info"}) assert len(data) == 3 - data = DataProto(batch=None, non_tensor_batch={}, meta_info={"info": "test_info"}) + data = DataProto( + batch=None, + non_tensor_batch={}, + meta_info={ + "info": "test_info"}) assert len(data) == 0 - data = DataProto(batch=None, non_tensor_batch=None, meta_info={"info": "test_info"}) + data = DataProto( + batch=None, + non_tensor_batch=None, + meta_info={ + "info": "test_info"}) assert len(data) == 0 @@ -283,7 +372,10 @@ def test_dataproto_index(): obs = torch.randn(data_len, 10) labels = [random.choice(["abc", "cde"]) for _ in range(data_len)] - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}) + data = DataProto.from_dict( + tensors={ + "obs": obs}, non_tensors={ + "labels": labels}) labels_np = np.array(labels) idx_np_int = np.random.randint(0, data_len, size=(idx_num,)) @@ -292,8 +384,12 @@ def test_dataproto_index(): assert result_np_int.non_tensor_batch.keys() == data.non_tensor_batch.keys() assert result_np_int.batch["obs"].shape[0] == idx_num assert result_np_int.non_tensor_batch["labels"].shape[0] == idx_num - assert np.array_equal(result_np_int.batch["obs"].cpu().numpy(), obs[idx_np_int].numpy()) - assert np.array_equal(result_np_int.non_tensor_batch["labels"], labels_np[idx_np_int]) + assert np.array_equal( + result_np_int.batch["obs"].cpu().numpy(), obs[idx_np_int].numpy() + ) + assert np.array_equal( + result_np_int.non_tensor_batch["labels"], labels_np[idx_np_int] + ) idx_torch_int = torch.randint(0, data_len, size=(idx_num,)) result_torch_int = data[idx_torch_int] @@ -301,8 +397,13 @@ def test_dataproto_index(): assert result_torch_int.non_tensor_batch.keys() == data.non_tensor_batch.keys() assert result_torch_int.batch["obs"].shape[0] == idx_num assert result_torch_int.non_tensor_batch["labels"].shape[0] == idx_num - assert np.array_equal(result_torch_int.batch["obs"].cpu().numpy(), obs[idx_torch_int].cpu().numpy()) - assert np.array_equal(result_torch_int.non_tensor_batch["labels"], labels_np[idx_torch_int.cpu().numpy()]) + assert np.array_equal( + result_torch_int.batch["obs"].cpu().numpy(), + obs[idx_torch_int].cpu().numpy()) + assert np.array_equal( + result_torch_int.non_tensor_batch["labels"], + labels_np[idx_torch_int.cpu().numpy()], + ) idx_list_int = [np.random.randint(0, data_len) for _ in range(idx_num)] result_list_int = data[idx_list_int] @@ -310,35 +411,58 @@ def test_dataproto_index(): assert result_list_int.non_tensor_batch.keys() == data.non_tensor_batch.keys() assert result_list_int.batch["obs"].shape[0] == idx_num assert result_list_int.non_tensor_batch["labels"].shape[0] == idx_num - assert np.array_equal(result_list_int.batch["obs"].cpu().numpy(), obs[idx_list_int].cpu().numpy()) - assert np.array_equal(result_list_int.non_tensor_batch["labels"], labels_np[idx_list_int]) + assert np.array_equal( + result_list_int.batch["obs"].cpu().numpy(), + obs[idx_list_int].cpu().numpy()) + assert np.array_equal( + result_list_int.non_tensor_batch["labels"], labels_np[idx_list_int] + ) idx_np_bool = np.random.randint(0, 2, size=(data_len,), dtype=bool) result_np_bool = data[idx_np_bool] assert result_np_bool.batch.keys() == data.batch.keys() assert result_np_bool.non_tensor_batch.keys() == data.non_tensor_batch.keys() assert result_np_bool.batch["obs"].shape[0] == idx_np_bool.sum() - assert result_np_bool.non_tensor_batch["labels"].shape[0] == idx_np_bool.sum() - assert np.array_equal(result_np_bool.batch["obs"].cpu().numpy(), obs[idx_np_bool].cpu().numpy()) - assert np.array_equal(result_np_bool.non_tensor_batch["labels"], labels_np[idx_np_bool]) + assert result_np_bool.non_tensor_batch["labels"].shape[0] == idx_np_bool.sum( + ) + assert np.array_equal( + result_np_bool.batch["obs"].cpu().numpy(), + obs[idx_np_bool].cpu().numpy()) + assert np.array_equal( + result_np_bool.non_tensor_batch["labels"], labels_np[idx_np_bool] + ) idx_torch_bool = torch.randint(0, 2, size=(data_len,), dtype=torch.bool) result_torch_bool = data[idx_torch_bool] assert result_torch_bool.batch.keys() == data.batch.keys() assert result_torch_bool.non_tensor_batch.keys() == data.non_tensor_batch.keys() - assert result_torch_bool.batch["obs"].shape[0] == idx_torch_bool.sum().item() - assert result_torch_bool.non_tensor_batch["labels"].shape[0] == idx_torch_bool.sum().item() - assert np.array_equal(result_torch_bool.batch["obs"].cpu().numpy(), obs[idx_torch_bool].cpu().numpy()) - assert np.array_equal(result_torch_bool.non_tensor_batch["labels"], labels_np[idx_torch_bool]) + assert result_torch_bool.batch["obs"].shape[0] == idx_torch_bool.sum( + ).item() + assert ( + result_torch_bool.non_tensor_batch["labels"].shape[0] + == idx_torch_bool.sum().item() + ) + assert np.array_equal( + result_torch_bool.batch["obs"].cpu().numpy(), + obs[idx_torch_bool].cpu().numpy()) + assert np.array_equal( + result_torch_bool.non_tensor_batch["labels"], labels_np[idx_torch_bool] + ) - idx_list_bool = [np.random.randint(0, 2, dtype=bool) for _ in range(data_len)] + idx_list_bool = [np.random.randint(0, 2, dtype=bool) + for _ in range(data_len)] result_list_bool = data[idx_list_bool] assert result_list_bool.batch.keys() == data.batch.keys() assert result_list_bool.non_tensor_batch.keys() == data.non_tensor_batch.keys() assert result_list_bool.batch["obs"].shape[0] == sum(idx_list_bool) - assert result_list_bool.non_tensor_batch["labels"].shape[0] == sum(idx_list_bool) - assert np.array_equal(result_list_bool.batch["obs"].cpu().numpy(), obs[idx_list_bool].cpu().numpy()) - assert np.array_equal(result_list_bool.non_tensor_batch["labels"], labels_np[idx_list_bool]) + assert result_list_bool.non_tensor_batch["labels"].shape[0] == sum( + idx_list_bool) + assert np.array_equal( + result_list_bool.batch["obs"].cpu().numpy(), + obs[idx_list_bool].cpu().numpy()) + assert np.array_equal( + result_list_bool.non_tensor_batch["labels"], labels_np[idx_list_bool] + ) def test_old_vs_new_from_single_dict(): @@ -380,7 +504,9 @@ def from_single_dict(cls, data, meta_info=None, auto_padding=False): def test_dataproto_no_batch(): labels = ["a", "b", "c"] - data = DataProto.from_dict(non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + data = DataProto.from_dict( + non_tensors={"labels": labels}, meta_info={"info": "test_info"} + ) selected = data.select(non_tensor_batch_keys=["labels"]) assert (selected.non_tensor_batch["labels"] == labels).all() pop_data = data.pop(non_tensor_batch_keys=["labels"]) @@ -392,24 +518,46 @@ def test_sample_level_repeat(): # Create a DataProto object with some batch and non-tensor data obs = torch.tensor([[1, 2], [3, 4], [5, 6]]) labels = ["a", "b", "c"] - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"info": "test_info"}) + data = DataProto.from_dict( + tensors={"obs": obs}, + non_tensors={"labels": labels}, + meta_info={"info": "test_info"}, + ) # list repeated_data_interleave = data.sample_level_repeat(repeat_times=[3, 1, 2]) - expected_obs_interleave = torch.tensor([[1, 2], [1, 2], [1, 2], [3, 4], [5, 6], [5, 6]]) + expected_obs_interleave = torch.tensor( + [[1, 2], [1, 2], [1, 2], [3, 4], [5, 6], [5, 6]] + ) expected_labels_interleave = ["a", "a", "a", "b", "c", "c"] - assert torch.all(torch.eq(repeated_data_interleave.batch["obs"], expected_obs_interleave)) - assert (repeated_data_interleave.non_tensor_batch["labels"] == expected_labels_interleave).all() + assert torch.all( + torch.eq( + repeated_data_interleave.batch["obs"], + expected_obs_interleave)) + assert ( + repeated_data_interleave.non_tensor_batch["labels"] + == expected_labels_interleave + ).all() assert repeated_data_interleave.meta_info == {"info": "test_info"} # torch.tensor - repeated_data_no_interleave = data.sample_level_repeat(repeat_times=torch.tensor([1, 2, 3])) - expected_obs_no_interleave = torch.tensor([[1, 2], [3, 4], [3, 4], [5, 6], [5, 6], [5, 6]]) + repeated_data_no_interleave = data.sample_level_repeat( + repeat_times=torch.tensor([1, 2, 3]) + ) + expected_obs_no_interleave = torch.tensor( + [[1, 2], [3, 4], [3, 4], [5, 6], [5, 6], [5, 6]] + ) expected_labels_no_interleave = ["a", "b", "b", "c", "c", "c"] - assert torch.all(torch.eq(repeated_data_no_interleave.batch["obs"], expected_obs_no_interleave)) - assert (repeated_data_no_interleave.non_tensor_batch["labels"] == expected_labels_no_interleave).all() + assert torch.all( + torch.eq( + repeated_data_no_interleave.batch["obs"], + expected_obs_no_interleave)) + assert ( + repeated_data_no_interleave.non_tensor_batch["labels"] + == expected_labels_no_interleave + ).all() assert repeated_data_no_interleave.meta_info == {"info": "test_info"} @@ -419,12 +567,16 @@ def test_dataproto_unfold_column_chunks(): labels = ["a", "b", "c"] data = DataProto.from_dict( - tensors={"obs1": obs1, "obs2": obs2}, non_tensors={"labels": labels}, meta_info={"name": "abc"} + tensors={"obs1": obs1, "obs2": obs2}, + non_tensors={"labels": labels}, + meta_info={"name": "abc"}, ) ret = data.unfold_column_chunks(2, split_keys=["obs1"]) - expect_obs1 = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]]) - expect_obs2 = torch.tensor([[1, 2], [1, 2], [5, 6], [5, 6], [9, 10], [9, 10]]) + expect_obs1 = torch.tensor( + [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]]) + expect_obs2 = torch.tensor( + [[1, 2], [1, 2], [5, 6], [5, 6], [9, 10], [9, 10]]) expect_labels = ["a", "a", "b", "b", "c", "c"] assert torch.all(torch.eq(ret.batch["obs1"], expect_obs1)) assert torch.all(torch.eq(ret.batch["obs2"], expect_obs2)) @@ -436,12 +588,16 @@ def test_dataproto_unfold_column_chunks(): labels = [["a1", "a2"], ["b1", "b2"], ["c1", "c2"]] data = DataProto.from_dict( - tensors={"obs1": obs1, "obs2": obs2}, non_tensors={"labels": labels}, meta_info={"name": "abc"} + tensors={"obs1": obs1, "obs2": obs2}, + non_tensors={"labels": labels}, + meta_info={"name": "abc"}, ) ret = data.unfold_column_chunks(2, split_keys=["obs1", "labels"]) - expect_obs1 = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]]) - expect_obs2 = torch.tensor([[1, 2], [1, 2], [5, 6], [5, 6], [9, 10], [9, 10]]) + expect_obs1 = torch.tensor( + [[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]]) + expect_obs2 = torch.tensor( + [[1, 2], [1, 2], [5, 6], [5, 6], [9, 10], [9, 10]]) expect_labels = [["a1"], ["a2"], ["b1"], ["b2"], ["c1"], ["c2"]] assert torch.all(torch.eq(ret.batch["obs1"], expect_obs1)) assert torch.all(torch.eq(ret.batch["obs2"], expect_obs2)) @@ -449,13 +605,20 @@ def test_dataproto_unfold_column_chunks(): assert ret.meta_info == {"name": "abc"} obs1 = torch.tensor( - [[[1, 1], [2, 2], [3, 3], [4, 4]], [[5, 5], [6, 6], [7, 7], [8, 8]], [[9, 9], [10, 10], [11, 11], [12, 12]]] + [ + [[1, 1], [2, 2], [3, 3], [4, 4]], + [[5, 5], [6, 6], [7, 7], [8, 8]], + [[9, 9], [10, 10], [11, 11], [12, 12]], + ] ) - obs2 = torch.tensor([[[1, 1], [2, 2]], [[5, 5], [6, 6]], [[9, 9], [10, 10]]]) + obs2 = torch.tensor( + [[[1, 1], [2, 2]], [[5, 5], [6, 6]], [[9, 9], [10, 10]]]) labels = ["a", "b", "c"] data = DataProto.from_dict( - tensors={"obs1": obs1, "obs2": obs2}, non_tensors={"labels": labels}, meta_info={"name": "abc"} + tensors={"obs1": obs1, "obs2": obs2}, + non_tensors={"labels": labels}, + meta_info={"name": "abc"}, ) ret = data.unfold_column_chunks(2, split_keys=["obs1"]) @@ -470,7 +633,14 @@ def test_dataproto_unfold_column_chunks(): ] ) expect_obs2 = torch.tensor( - [[[1, 1], [2, 2]], [[1, 1], [2, 2]], [[5, 5], [6, 6]], [[5, 5], [6, 6]], [[9, 9], [10, 10]], [[9, 9], [10, 10]]] + [ + [[1, 1], [2, 2]], + [[1, 1], [2, 2]], + [[5, 5], [6, 6]], + [[5, 5], [6, 6]], + [[9, 9], [10, 10]], + [[9, 9], [10, 10]], + ] ) expect_labels = ["a", "a", "b", "b", "c", "c"] assert torch.all(torch.eq(ret.batch["obs1"], expect_obs1)) @@ -483,13 +653,19 @@ def test_dataproto_chunk_after_index(): data_len = 4 obs = torch.randn(data_len, 4) labels = [f"label_{i}" for i in range(data_len)] - data = DataProto.from_dict(tensors={"obs": obs}, non_tensors={"labels": labels}, meta_info={"name": "abc"}) + data = DataProto.from_dict( + tensors={ + "obs": obs}, non_tensors={ + "labels": labels}, meta_info={ + "name": "abc"}) # Test with boolean numpy array bool_mask = np.array([True, False, True, False]) selected = data[bool_mask] assert isinstance(selected.batch.batch_size, torch.Size) - assert all(isinstance(d, int) for d in selected.batch.batch_size) # int or List[int] + assert all( + isinstance(d, int) for d in selected.batch.batch_size + ) # int or List[int] # Test with integer numpy array int_mask = np.array([0, 2]) diff --git a/Agent0/executor_train/verl/tests/tools/test_base_tool_on_cpu.py b/Agent0/executor_train/verl/tests/tools/test_base_tool_on_cpu.py index 63a2bbb..b647c2a 100644 --- a/Agent0/executor_train/verl/tests/tools/test_base_tool_on_cpu.py +++ b/Agent0/executor_train/verl/tests/tools/test_base_tool_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -44,7 +44,9 @@ def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: schema = get_json_schema(self.get_current_temperature) return OpenAIFunctionToolSchema(**schema) - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: try: result = self.get_current_temperature(**parameters) return json.dumps(result), 0, {} @@ -57,7 +59,11 @@ def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: schema = get_json_schema(self.get_temperature_date) return OpenAIFunctionToolSchema(**schema) - def get_temperature_date(self, location: str, date: str, unit: str = "celsius"): + def get_temperature_date( + self, + location: str, + date: str, + unit: str = "celsius"): """Get temperature at a location and date. Args: @@ -75,7 +81,9 @@ def get_temperature_date(self, location: str, date: str, unit: str = "celsius"): "unit": unit, } - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: try: result = self.get_temperature_date(**parameters) return json.dumps(result), 0, {} @@ -131,7 +139,8 @@ def test_initialize_tools_from_fake_config(create_fake_tool_config): tool_config_path = create_fake_tool_config # Use pytest.raises to check if an exception is raised when calling initialize_tools_from_config. - # Since the tool configuration uses fake paths, an exception is expected during the tool initialization process. + # Since the tool configuration uses fake paths, an exception is expected + # during the tool initialization process. with pytest.raises(ModuleNotFoundError): _ = initialize_tools_from_config(tool_config_path) @@ -146,13 +155,17 @@ def test_initialize_tools_from_local_config(create_local_tool_config): and returns its path. After the test is completed, the fixture will clean up the configuration file. """ - # Retrieve the path of the local tool configuration file generated by the fixture + # Retrieve the path of the local tool configuration file generated by the + # fixture tool_config_path = create_local_tool_config tools = initialize_tools_from_config(tool_config_path) assert len(tools) == 2 - from tests.tools.test_base_tool_on_cpu import WeatherToolForTest, WeatherToolWithDataForTest + from tests.tools.test_base_tool_on_cpu import ( + WeatherToolForTest, + WeatherToolWithDataForTest, + ) assert isinstance(tools[0], WeatherToolForTest) assert isinstance(tools[1], WeatherToolWithDataForTest) diff --git a/Agent0/executor_train/verl/tests/trainer/__init__.py b/Agent0/executor_train/verl/tests/trainer/__init__.py index 6f79d47..c4f217c 100644 --- a/Agent0/executor_train/verl/tests/trainer/__init__.py +++ b/Agent0/executor_train/verl/tests/trainer/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/trainer/config/__init__.py b/Agent0/executor_train/verl/tests/trainer/config/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/tests/trainer/config/__init__.py +++ b/Agent0/executor_train/verl/tests/trainer/config/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/trainer/config/test_algo_config_on_cpu.py b/Agent0/executor_train/verl/tests/trainer/config/test_algo_config_on_cpu.py index 848a3ff..adbf66d 100644 --- a/Agent0/executor_train/verl/tests/trainer/config/test_algo_config_on_cpu.py +++ b/Agent0/executor_train/verl/tests/trainer/config/test_algo_config_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,7 +32,8 @@ class TestAlgoConfig(unittest.TestCase): def setUp(self): """Set up test fixtures.""" - # Create a sample algorithm config as DictConfig (similar to what comes from YAML) + # Create a sample algorithm config as DictConfig (similar to what comes + # from YAML) self.config_dict = { "_target_": "verl.trainer.config.AlgoConfig", "gamma": 0.99, @@ -49,7 +50,11 @@ def setUp(self): "target_kl": 0.05, }, "use_pf_ppo": True, - "pf_ppo": {"_target_": "verl.trainer.config.PFPPOConfig", "reweight_method": "max_min", "weight_pow": 3.0}, + "pf_ppo": { + "_target_": "verl.trainer.config.PFPPOConfig", + "reweight_method": "max_min", + "weight_pow": 3.0, + }, } self.omega_config = OmegaConf.create(self.config_dict) @@ -151,7 +156,9 @@ def setUp(self): norm_adv_by_std_in_grpo=True, use_kl_in_reward=True, kl_penalty="kl", - kl_ctrl=KLControlConfig(type="adaptive", kl_coef=0.002, horizon=5000, target_kl=0.05), + kl_ctrl=KLControlConfig( + type="adaptive", kl_coef=0.002, horizon=5000, target_kl=0.05 + ), use_pf_ppo=True, pf_ppo=PFPPOConfig(reweight_method="max_min", weight_pow=3.0), ) @@ -183,11 +190,15 @@ def test_advantage_estimator_with_cfg(self): def test_grpo_advantage_estimator_with_cfg(self): """Test integration with GRPO advantage estimator.""" - grpo_config = AlgoConfig(adv_estimator="grpo", norm_adv_by_std_in_grpo=True) + grpo_config = AlgoConfig( + adv_estimator="grpo", + norm_adv_by_std_in_grpo=True) # Test GRPO advantage computation batch_size, seq_len = 4, 3 - token_level_rewards = torch.tensor([[1.0, 0.5, 0.0], [2.0, 1.0, 0.0], [0.5, 0.2, 0.0], [1.5, 0.8, 0.0]]) + token_level_rewards = torch.tensor( + [[1.0, 0.5, 0.0], [2.0, 1.0, 0.0], [0.5, 0.2, 0.0], [1.5, 0.8, 0.0]] + ) response_mask = torch.ones(batch_size, seq_len) index = np.array([0, 0, 1, 1]) # Two groups diff --git a/Agent0/executor_train/verl/tests/trainer/config/test_legacy_config_on_cpu.py b/Agent0/executor_train/verl/tests/trainer/config/test_legacy_config_on_cpu.py index 39862aa..a7a8ea1 100644 --- a/Agent0/executor_train/verl/tests/trainer/config/test_legacy_config_on_cpu.py +++ b/Agent0/executor_train/verl/tests/trainer/config/test_legacy_config_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,14 +23,20 @@ class TestConfigComparison(unittest.TestCase): """Test that current configs match their legacy counterparts exactly.""" - def _compare_configs_recursively(self, current_config, legacy_config, path="", legacy_allow_missing=True): + def _compare_configs_recursively( + self, current_config, legacy_config, path="", legacy_allow_missing=True + ): """Recursively compare two OmegaConf configs and assert they are identical. Args: legacy_allow_missing (bool): sometimes the legacy megatron config contains fewer keys and we allow that to happen """ - if isinstance(current_config, dict) and isinstance(legacy_config, dict): + if isinstance( + current_config, + dict) and isinstance( + legacy_config, + dict): current_keys = set(current_config.keys()) legacy_keys = set(legacy_config.keys()) @@ -38,7 +44,8 @@ def _compare_configs_recursively(self, current_config, legacy_config, path="", l missing_in_legacy = current_keys - legacy_keys if missing_in_current: - self.fail(f"Keys missing in current config at {path}: {missing_in_current}") + self.fail( + f"Keys missing in current config at {path}: {missing_in_current}") if missing_in_legacy: # if the legacy msg = f"Keys missing in legacy config at {path}: {missing_in_legacy}" @@ -50,15 +57,23 @@ def _compare_configs_recursively(self, current_config, legacy_config, path="", l for key in current_keys: current_path = f"{path}.{key}" if path else key if key in legacy_config: - self._compare_configs_recursively(current_config[key], legacy_config[key], current_path) + self._compare_configs_recursively( + current_config[key], legacy_config[key], current_path + ) elif isinstance(current_config, list) and isinstance(legacy_config, list): self.assertEqual( len(current_config), len(legacy_config), - f"List lengths differ at {path}: current={len(current_config)}, legacy={len(legacy_config)}", + f"List lengths differ at {path}: current={ + len(current_config)}, legacy={ + len(legacy_config)}", ) - for i, (current_item, legacy_item) in enumerate(zip(current_config, legacy_config, strict=True)): - self._compare_configs_recursively(current_item, legacy_item, f"{path}[{i}]") + for i, (current_item, legacy_item) in enumerate( + zip(current_config, legacy_config, strict=True) + ): + self._compare_configs_recursively( + current_item, legacy_item, f"{path}[{i}]" + ) else: self.assertEqual( current_config, @@ -76,10 +91,14 @@ def test_ppo_trainer_config_matches_legacy(self): GlobalHydra.instance().clear() try: - with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + with initialize_config_dir( + config_dir=os.path.abspath("verl/trainer/config") + ): current_config = compose(config_name="ppo_trainer") - legacy_config = OmegaConf.load("tests/trainer/config/legacy_ppo_trainer.yaml") + legacy_config = OmegaConf.load( + "tests/trainer/config/legacy_ppo_trainer.yaml" + ) current_dict = OmegaConf.to_container(current_config, resolve=True) legacy_dict = OmegaConf.to_container(legacy_config, resolve=True) @@ -96,17 +115,23 @@ def test_ppo_megatron_trainer_config_matches_legacy(self): GlobalHydra.instance().clear() try: - with initialize_config_dir(config_dir=os.path.abspath("verl/trainer/config")): + with initialize_config_dir( + config_dir=os.path.abspath("verl/trainer/config") + ): current_config = compose(config_name="ppo_megatron_trainer") - legacy_config = OmegaConf.load("tests/trainer/config/legacy_ppo_megatron_trainer.yaml") + legacy_config = OmegaConf.load( + "tests/trainer/config/legacy_ppo_megatron_trainer.yaml" + ) current_dict = OmegaConf.to_container(current_config, resolve=True) legacy_dict = OmegaConf.to_container(legacy_config, resolve=True) if "defaults" in current_dict: del current_dict["defaults"] - self._compare_configs_recursively(current_dict, legacy_dict, legacy_allow_missing=True) + self._compare_configs_recursively( + current_dict, legacy_dict, legacy_allow_missing=True + ) finally: GlobalHydra.instance().clear() diff --git a/Agent0/executor_train/verl/tests/trainer/ppo/__init__.py b/Agent0/executor_train/verl/tests/trainer/ppo/__init__.py index 26d7c04..9c5db9e 100644 --- a/Agent0/executor_train/verl/tests/trainer/ppo/__init__.py +++ b/Agent0/executor_train/verl/tests/trainer/ppo/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/trainer/ppo/test_core_algos_on_cpu.py b/Agent0/executor_train/verl/tests/trainer/ppo/test_core_algos_on_cpu.py index 087a0d2..cced9d7 100644 --- a/Agent0/executor_train/verl/tests/trainer/ppo/test_core_algos_on_cpu.py +++ b/Agent0/executor_train/verl/tests/trainer/ppo/test_core_algos_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,7 +19,11 @@ import torch import verl.trainer.ppo.core_algos -from verl.trainer.ppo.core_algos import compute_gae_advantage_return, get_adv_estimator_fn, register_adv_est +from verl.trainer.ppo.core_algos import ( + compute_gae_advantage_return, + get_adv_estimator_fn, + register_adv_est, +) def mock_test_fn(): @@ -48,7 +52,9 @@ def test_fn(): pass self.assertIn("test_estimator", self.ADV_ESTIMATOR_REGISTRY) - self.assertEqual(self.ADV_ESTIMATOR_REGISTRY["test_estimator"], test_fn) + self.assertEqual( + self.ADV_ESTIMATOR_REGISTRY["test_estimator"], + test_fn) def test_register_with_enum(self): """Test registering with an enum value (assuming AdvantageEstimator exists)""" @@ -62,14 +68,18 @@ def test_fn(): pass self.assertIn("test_enum_estimator", self.ADV_ESTIMATOR_REGISTRY) - self.assertEqual(self.ADV_ESTIMATOR_REGISTRY["test_enum_estimator"], test_fn) + self.assertEqual( + self.ADV_ESTIMATOR_REGISTRY["test_enum_estimator"], + test_fn) def test_duplicate_registration_same_function(self): """Test that registering the same function twice doesn't raise an error""" register_adv_est("duplicate_test")(mock_test_fn) register_adv_est("duplicate_test")(mock_test_fn) - self.assertEqual(self.ADV_ESTIMATOR_REGISTRY["duplicate_test"], mock_test_fn) + self.assertEqual( + self.ADV_ESTIMATOR_REGISTRY["duplicate_test"], + mock_test_fn) def test_duplicate_registration_different_function(self): """Test that registering different functions with same name raises ValueError""" @@ -123,7 +133,8 @@ def test_get_adv_estimator_fn_invalid_name(self): """Test that invalid names raise ValueError.""" with pytest.raises(ValueError) as excinfo: get_adv_estimator_fn("invalid_name") - assert "Unknown advantage estimator simply: invalid_name" in str(excinfo.value) + assert "Unknown advantage estimator simply: invalid_name" in str( + excinfo.value) def test_get_adv_estimator_fn_case_sensitive(self): """Test that name lookup is case-sensitive.""" @@ -136,7 +147,9 @@ def test_multi_turn_compute_gae_advantage_return(): gamma = random.uniform(0.0, 1.0) lam = random.uniform(0.0, 1.0) - rewards = torch.tensor([[0.0, 0.0, 0.1, 0.1, 0.1, 0.0, 0.0, 0.1, 1.0, 0.0, 0.0]], dtype=torch.float) + rewards = torch.tensor( + [[0.0, 0.0, 0.1, 0.1, 0.1, 0.0, 0.0, 0.1, 1.0, 0.0, 0.0]], dtype=torch.float + ) values1 = torch.tensor( [ @@ -176,10 +189,15 @@ def test_multi_turn_compute_gae_advantage_return(): dtype=torch.float, ) - response_mask = torch.tensor([[0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0]], dtype=torch.float) + response_mask = torch.tensor( + [[0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0]], dtype=torch.float) - adv1, ret1 = compute_gae_advantage_return(rewards, values1, response_mask, gamma, lam) - adv2, ret2 = compute_gae_advantage_return(rewards, values2, response_mask, gamma, lam) + adv1, ret1 = compute_gae_advantage_return( + rewards, values1, response_mask, gamma, lam + ) + adv2, ret2 = compute_gae_advantage_return( + rewards, values2, response_mask, gamma, lam + ) ret1 *= response_mask ret2 *= response_mask diff --git a/Agent0/executor_train/verl/tests/trainer/ppo/test_metric_utils_on_cpu.py b/Agent0/executor_train/verl/tests/trainer/ppo/test_metric_utils_on_cpu.py index 50fe952..db3cb66 100644 --- a/Agent0/executor_train/verl/tests/trainer/ppo/test_metric_utils_on_cpu.py +++ b/Agent0/executor_train/verl/tests/trainer/ppo/test_metric_utils_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -110,8 +110,12 @@ def test_compute_data_metrics_with_critic(self): self.assertIn("prompt_length/mean", metrics) # Check some specific values - self.assertAlmostEqual(metrics["critic/score/mean"], 5.0) # Sum of token_level_scores - self.assertAlmostEqual(metrics["critic/rewards/mean"], 2.5) # Sum of token_level_rewards + self.assertAlmostEqual( + metrics["critic/score/mean"], 5.0 + ) # Sum of token_level_scores + self.assertAlmostEqual( + metrics["critic/rewards/mean"], 2.5 + ) # Sum of token_level_rewards def test_compute_data_metrics_without_critic(self): """Test compute_data_metrics with critic disabled.""" @@ -135,7 +139,8 @@ def setUp(self): # Create a mock DataProto object self.batch = MagicMock() self.batch.batch = { - "responses": torch.zeros((2, 3)), # 2 samples, 3 response tokens each + # 2 samples, 3 response tokens each + "responses": torch.zeros((2, 3)), "attention_mask": torch.tensor( [ [1, 1, 1, 1, 1, 1], # 3 prompt tokens, 3 response tokens @@ -171,11 +176,17 @@ def test_compute_timing_metrics(self, mock_compute_response_info): # Check per-token timing metrics # gen uses only response tokens (6 tokens) - self.assertAlmostEqual(metrics["timing_per_token_ms/gen"], 0.5 * 1000 / 6, places=5) + self.assertAlmostEqual( + metrics["timing_per_token_ms/gen"], 0.5 * 1000 / 6, places=5 + ) # ref and values use all tokens (12 tokens) - self.assertAlmostEqual(metrics["timing_per_token_ms/ref"], 0.3 * 1000 / 12, places=5) - self.assertAlmostEqual(metrics["timing_per_token_ms/values"], 0.2 * 1000 / 12, places=5) + self.assertAlmostEqual( + metrics["timing_per_token_ms/ref"], 0.3 * 1000 / 12, places=5 + ) + self.assertAlmostEqual( + metrics["timing_per_token_ms/values"], 0.2 * 1000 / 12, places=5 + ) class TestComputeThroughputMetrics(unittest.TestCase): @@ -200,14 +211,18 @@ def test_compute_throughout_metrics(self): self.assertEqual(metrics["perf/total_num_tokens"], 600) self.assertEqual(metrics["perf/time_per_step"], 2.0) - self.assertEqual(metrics["perf/throughput"], 600 / 2.0) # 300 tokens/sec + self.assertEqual( + metrics["perf/throughput"], + 600 / 2.0) # 300 tokens/sec # Test with 2 GPUs metrics = compute_throughout_metrics(self.batch, timing_raw, n_gpus=2) self.assertEqual(metrics["perf/total_num_tokens"], 600) self.assertEqual(metrics["perf/time_per_step"], 2.0) - self.assertEqual(metrics["perf/throughput"], 600 / (2.0 * 2)) # 150 tokens/sec/GPU + self.assertEqual( + metrics["perf/throughput"], 600 / (2.0 * 2) + ) # 150 tokens/sec/GPU class TestBootstrapMetric(unittest.TestCase): @@ -219,7 +234,12 @@ def test_bootstrap_metric_basic(self): reduce_fns = [np.mean, np.max] # Use a fixed seed for reproducibility - result = bootstrap_metric(data, subset_size=3, reduce_fns=reduce_fns, n_bootstrap=100, seed=42) + result = bootstrap_metric( + data, + subset_size=3, + reduce_fns=reduce_fns, + n_bootstrap=100, + seed=42) # Check that we get two results (one for each reduce_fn) self.assertEqual(len(result), 2) @@ -233,7 +253,8 @@ def test_bootstrap_metric_basic(self): self.assertAlmostEqual(mean_result[0], 3.0, delta=0.3) # The mean of maxes should be close to the expected value for samples of size 3 - # For samples of size 3 from [1,2,3,4,5], the expected max is around 4.0-4.5 + # For samples of size 3 from [1,2,3,4,5], the expected max is around + # 4.0-4.5 self.assertGreater(max_result[0], 3.5) self.assertLess(max_result[0], 5.0) @@ -287,7 +308,9 @@ def test_process_validation_metrics_basic(self): "score": [0.8, 0.9, 0.7], } - result = process_validation_metrics(data_sources, sample_inputs, infos_dict, seed=42) + result = process_validation_metrics( + data_sources, sample_inputs, infos_dict, seed=42 + ) # Check the structure of the result self.assertIn("source1", result) @@ -311,7 +334,9 @@ def test_process_validation_metrics_with_pred(self): "pred": ["A", "B", "A"], } - result = process_validation_metrics(data_sources, sample_inputs, infos_dict, seed=42) + result = process_validation_metrics( + data_sources, sample_inputs, infos_dict, seed=42 + ) # Check that majority voting metrics are present self.assertIn("maj@2/mean", result["source1"]["score"]) diff --git a/Agent0/executor_train/verl/tests/utils/_test_module.py b/Agent0/executor_train/verl/tests/utils/_test_module.py index ec3d5fb..ac00f27 100644 --- a/Agent0/executor_train/verl/tests/utils/_test_module.py +++ b/Agent0/executor_train/verl/tests/utils/_test_module.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/utils/ckpt/test_esi_save_ckpt_on_cpu.py b/Agent0/executor_train/verl/tests/utils/ckpt/test_esi_save_ckpt_on_cpu.py index 203494b..692665b 100644 --- a/Agent0/executor_train/verl/tests/utils/ckpt/test_esi_save_ckpt_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/ckpt/test_esi_save_ckpt_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,19 +23,25 @@ class TestShouldSaveCkptEsi(TestCase): def test_no_expiration_timestamp(self): """Test case when no expiration timestamp is set""" os.environ.pop("MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP", None) - os.environ.pop("SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP", None) + os.environ.pop( + "SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP", None) self.assertFalse(should_save_ckpt_esi(100)) def test_mlp_expiration_valid(self): """Test valid MLP expiration timestamp requiring save""" current_time = time.time() - os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(current_time + 90) - self.assertTrue(should_save_ckpt_esi(30)) # max_steps_duration=30 seconds + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str( + current_time + 90 + ) + # max_steps_duration=30 seconds + self.assertTrue(should_save_ckpt_esi(30)) def test_mlp_expiration_passed(self): """Test expired MLP timestamp""" current_time = time.time() - os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(current_time - 10) + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str( + current_time - 10 + ) self.assertFalse(should_save_ckpt_esi(30)) def test_mlp_invalid_timestamp(self): @@ -46,25 +52,33 @@ def test_mlp_invalid_timestamp(self): def test_mlp_expiration_not_reached(self): """Test MLP expiration timestamp with insufficient remaining time""" current_time = time.time() - os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(current_time + 200) + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str( + current_time + 200 + ) self.assertFalse(should_save_ckpt_esi(30)) # max_steps_duration=30 def test_aws_expiration_not_reached(self): """Test AWS expiration timestamp with sufficient remaining time""" now = datetime.now() - expiration = now + timedelta(minutes=100) # Exceeds 90-minute threshold - os.environ["SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(int(expiration.timestamp())) + # Exceeds 90-minute threshold + expiration = now + timedelta(minutes=100) + os.environ["SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str( + int(expiration.timestamp())) self.assertFalse(should_save_ckpt_esi(30 * 60)) def test_redundant_time(self): """Test redundant_time parameter effect""" current_time = time.time() # Total required: 60+30+30=120 seconds - os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(current_time + 120) + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str( + current_time + 120 + ) self.assertTrue(should_save_ckpt_esi(30, redundant_time=30)) def test_zero_max_steps_duration(self): """Test zero max_steps_duration""" current_time = time.time() - os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str(current_time + 60) + os.environ["MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP"] = str( + current_time + 60 + ) self.assertFalse(should_save_ckpt_esi(0)) diff --git a/Agent0/executor_train/verl/tests/utils/dataset/test_create_rl_sampler_on_cpu.py b/Agent0/executor_train/verl/tests/utils/dataset/test_create_rl_sampler_on_cpu.py index 35bf5a3..8f1aa3b 100644 --- a/Agent0/executor_train/verl/tests/utils/dataset/test_create_rl_sampler_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/dataset/test_create_rl_sampler_on_cpu.py @@ -82,8 +82,7 @@ def test_create_custom_curriculum_samper(): "class_path": "pkg://tests.utils.dataset.test_create_rl_sampler_on_cpu", "class_name": "RandomCurriculumSampler", }, - } - ) + }) dataset = MockChatDataset() @@ -97,12 +96,11 @@ def test_create_custom_curriculum_samper_wrong_class(): "sampler": { "class_path": "pkg://tests.utils.dataset.test_create_rl_sampler_on_cpu", "class_name": "MockIncorrectSampler", - } - } - ) + }}) dataset = MockChatDataset() - # MockIncorrectSampler is not an instance of AbstractCurriculumSampler, so raises + # MockIncorrectSampler is not an instance of AbstractCurriculumSampler, so + # raises with pytest.raises(AssertionError): create_rl_sampler(data_config, dataset) diff --git a/Agent0/executor_train/verl/tests/utils/dataset/test_multiturn_sft_dataset_on_cpu.py b/Agent0/executor_train/verl/tests/utils/dataset/test_multiturn_sft_dataset_on_cpu.py index 8028d44..6a5201a 100644 --- a/Agent0/executor_train/verl/tests/utils/dataset/test_multiturn_sft_dataset_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/dataset/test_multiturn_sft_dataset_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -56,8 +56,14 @@ def test_multiturn_sft_dataset(): # Initialize tokenizer and dataset tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-Coder-7B-Instruct") - config = {"max_length": 512, "truncation": "error", "multiturn": {"messages_key": "messages"}} - dataset = MultiTurnSFTDataset(parquet_files=test_file, tokenizer=tokenizer, config=config) + config = { + "max_length": 512, + "truncation": "error", + "multiturn": {"messages_key": "messages"}, + } + dataset = MultiTurnSFTDataset( + parquet_files=test_file, tokenizer=tokenizer, config=config + ) # Test 1: Dataset Length assert len(dataset) == 2, f"Expected dataset length 2, got {len(dataset)}" @@ -67,18 +73,29 @@ def test_multiturn_sft_dataset(): item1 = dataset[1] # Joke conversation # Test 2: Required Keys and Types - required_keys = ["input_ids", "attention_mask", "position_ids", "loss_mask"] + required_keys = [ + "input_ids", + "attention_mask", + "position_ids", + "loss_mask"] for key in required_keys: assert key in item0, f"Missing key {key} in dataset item" - assert isinstance(item0[key], torch.Tensor), f"Expected torch.Tensor for {key}" - assert item0[key].dtype == torch.long, f"Expected torch.long for {key}, got {item0[key].dtype}" + assert isinstance( + item0[key], torch.Tensor), f"Expected torch.Tensor for {key}" + assert ( + item0[key].dtype == torch.long + ), f"Expected torch.long for {key}, got {item0[key].dtype}" # Test 3: Shape Consistency - assert item0["loss_mask"].shape == item0["input_ids"].shape, "Loss mask shape doesn't match input_ids shape" - assert item0["attention_mask"].shape == item0["input_ids"].shape, ( - "Attention mask shape doesn't match input_ids shape" - ) - assert item0["position_ids"].shape == item0["input_ids"].shape, "Position IDs shape doesn't match input_ids shape" + assert ( + item0["loss_mask"].shape == item0["input_ids"].shape + ), "Loss mask shape doesn't match input_ids shape" + assert ( + item0["attention_mask"].shape == item0["input_ids"].shape + ), "Attention mask shape doesn't match input_ids shape" + assert ( + item0["position_ids"].shape == item0["input_ids"].shape + ), "Position IDs shape doesn't match input_ids shape" # Test 4: Loss Mask Pattern - Math Conversation loss_mask0 = item0["loss_mask"] @@ -86,7 +103,8 @@ def test_multiturn_sft_dataset(): # Find assistant response positions assistant_positions0 = torch.where(loss_mask0 == 1)[0] - assert len(assistant_positions0) > 0, "No assistant positions found in loss mask" + assert len( + assistant_positions0) > 0, "No assistant positions found in loss mask" # Decode and verify assistant responses assistant_text0 = tokenizer.decode(input_ids0[loss_mask0 == 1]) @@ -100,29 +118,38 @@ def test_multiturn_sft_dataset(): # Find assistant response positions assistant_positions1 = torch.where(loss_mask1 == 1)[0] - assert len(assistant_positions1) > 0, "No assistant positions found in loss mask" + assert len( + assistant_positions1) > 0, "No assistant positions found in loss mask" # Decode and verify assistant responses assistant_text1 = tokenizer.decode(input_ids1[loss_mask1 == 1]) print(f"Joke conversation assistant text: {assistant_text1}") - assert "chicken cross the road" in assistant_text1, "First assistant response not found" + assert ( + "chicken cross the road" in assistant_text1 + ), "First assistant response not found" assert "other side" in assistant_text1, "Second assistant response not found" # Test 6: Attention Mask Pattern attention_mask0 = item0["attention_mask"] sequence_length = torch.sum(attention_mask0) assert sequence_length > 0, "No tokens marked as attended in attention mask" - assert torch.all(attention_mask0[:sequence_length] == 1), "Incorrect attention mask pattern" + assert torch.all( + attention_mask0[:sequence_length] == 1 + ), "Incorrect attention mask pattern" if sequence_length < len(attention_mask0): - assert torch.all(attention_mask0[sequence_length:] == 0), "Padding not properly masked" + assert torch.all( + attention_mask0[sequence_length:] == 0 + ), "Padding not properly masked" # Test 7: Position IDs Pattern position_ids0 = item0["position_ids"] - assert torch.equal(position_ids0[:sequence_length], torch.arange(sequence_length)), ( - "Position IDs not sequential for non-padded tokens" - ) + assert torch.equal( + position_ids0[:sequence_length], torch.arange(sequence_length) + ), "Position IDs not sequential for non-padded tokens" if sequence_length < len(position_ids0): - assert torch.all(position_ids0[sequence_length:] == 0), "Padding position IDs not zero" + assert torch.all( + position_ids0[sequence_length:] == 0 + ), "Padding position IDs not zero" # Test 8: Verify loss mask for assistant responses # Get the full conversation text @@ -137,13 +164,15 @@ def test_multiturn_sft_dataset(): for msg in test_data["messages"][0]: # First conversation if msg["role"] == "assistant": # The content should appear in the masked text - assert msg["content"] in assistant_text, f"Assistant message '{msg['content']}' not found in masked text" + assert ( + msg["content"] in assistant_text + ), f"Assistant message '{msg['content']}' not found in masked text" # The content should NOT appear in the non-masked text non_assistant_text = tokenizer.decode(input_ids0[loss_mask0 == 0]) - assert msg["content"] not in non_assistant_text, ( - f"Assistant message '{msg['content']}' found in non-assistant text" - ) + assert ( + msg["content"] not in non_assistant_text), f"Assistant message '{ + msg['content']}' found in non-assistant text" # Test 9: Verify non-assistant parts have loss_mask=0 # Get non-assistant text @@ -153,29 +182,41 @@ def test_multiturn_sft_dataset(): # Verify that system and user messages are in the non-assistant text for msg in test_data["messages"][0]: # First conversation if msg["role"] in ["system", "user"]: - assert msg["content"] in non_assistant_text, ( - f"{msg['role'].title()} message '{msg['content']}' not found in non-assistant text" - ) + assert ( + msg["content"] in non_assistant_text), f"{ + msg['role'].title()} message '{ + msg['content']}' not found in non-assistant text" # And verify they're NOT in the assistant text - assert msg["content"] not in assistant_text, ( - f"{msg['role'].title()} message '{msg['content']}' found in assistant text" - ) + assert ( + msg["content"] not in assistant_text), f"{ + msg['role'].title()} message '{ + msg['content']}' found in assistant text" # Test 10: Verify padding behavior - padding_config = {"max_length": 1024, "truncation": "error", "multiturn": {"messages_key": "messages"}} - small_dataset = MultiTurnSFTDataset(parquet_files=test_file, tokenizer=tokenizer, config=padding_config) + padding_config = { + "max_length": 1024, + "truncation": "error", + "multiturn": {"messages_key": "messages"}, + } + small_dataset = MultiTurnSFTDataset( + parquet_files=test_file, tokenizer=tokenizer, config=padding_config + ) padded_item = small_dataset[0] # Get actual sequence length (before padding) actual_length = torch.sum(padded_item["attention_mask"]) # Verify padding tokens - assert torch.all(padded_item["input_ids"][actual_length:] == tokenizer.pad_token_id), ( - "Padding tokens not set correctly" - ) - assert torch.all(padded_item["attention_mask"][actual_length:] == 0), "Attention mask not set correctly for padding" - assert torch.all(padded_item["loss_mask"][actual_length:] == 0), "Loss mask not set correctly for padding" + assert torch.all( + padded_item["input_ids"][actual_length:] == tokenizer.pad_token_id + ), "Padding tokens not set correctly" + assert torch.all( + padded_item["attention_mask"][actual_length:] == 0 + ), "Attention mask not set correctly for padding" + assert torch.all( + padded_item["loss_mask"][actual_length:] == 0 + ), "Loss mask not set correctly for padding" print("All tests passed!") print("Starting test...") diff --git a/Agent0/executor_train/verl/tests/utils/dataset/test_rl_dataset_on_cpu.py b/Agent0/executor_train/verl/tests/utils/dataset/test_rl_dataset_on_cpu.py index 2afc3ef..970aa14 100644 --- a/Agent0/executor_train/verl/tests/utils/dataset/test_rl_dataset_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/dataset/test_rl_dataset_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -40,9 +40,18 @@ def test_rl_dataset(): "filter_overlong_prompts_workers": 2, } ) - dataset = RLHFDataset(data_files=local_path, tokenizer=tokenizer, config=config) - - dataloader = DataLoader(dataset=dataset, batch_size=16, shuffle=True, drop_last=True, collate_fn=collate_fn) + dataset = RLHFDataset( + data_files=local_path, + tokenizer=tokenizer, + config=config) + + dataloader = DataLoader( + dataset=dataset, + batch_size=16, + shuffle=True, + drop_last=True, + collate_fn=collate_fn, + ) a = next(iter(dataloader)) @@ -87,7 +96,13 @@ def test_image_rl_data(): processor=processor, ) - dataloader = DataLoader(dataset=dataset, batch_size=16, shuffle=True, drop_last=True, collate_fn=collate_fn) + dataloader = DataLoader( + dataset=dataset, + batch_size=16, + shuffle=True, + drop_last=True, + collate_fn=collate_fn, + ) a = next(iter(dataloader)) diff --git a/Agent0/executor_train/verl/tests/utils/dataset/test_sft_dataset_on_cpu.py b/Agent0/executor_train/verl/tests/utils/dataset/test_sft_dataset_on_cpu.py index 680fce4..3dd8021 100644 --- a/Agent0/executor_train/verl/tests/utils/dataset/test_sft_dataset_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/dataset/test_sft_dataset_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/utils/megatron/test_pipeline_parallel.py b/Agent0/executor_train/verl/tests/utils/megatron/test_pipeline_parallel.py index cf442a0..595e6b7 100644 --- a/Agent0/executor_train/verl/tests/utils/megatron/test_pipeline_parallel.py +++ b/Agent0/executor_train/verl/tests/utils/megatron/test_pipeline_parallel.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/tests/utils/reward_score/reward_score/test_sandbox_fusion_on_cpu.py b/Agent0/executor_train/verl/tests/utils/reward_score/reward_score/test_sandbox_fusion_on_cpu.py index 997cb8a..07a779d 100644 --- a/Agent0/executor_train/verl/tests/utils/reward_score/reward_score/test_sandbox_fusion_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/reward_score/reward_score/test_sandbox_fusion_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -63,11 +63,18 @@ """ # --- Test input/output data --- -INPUT_OUTPUT_VALID = {"inputs": ["input1", "input2"], "outputs": ["output1\n", "output2\n"]} +INPUT_OUTPUT_VALID = { + "inputs": ["input1", "input2"], + "outputs": ["output1\n", "output2\n"], +} INPUT_OUTPUT_SINGLE = {"inputs": ["input1"], "outputs": ["output1\n"]} -INPUT_OUTPUT_MISMATCH = {"inputs": ["input1"], "outputs": ["output1\n", "output2\n"]} +INPUT_OUTPUT_MISMATCH = { + "inputs": ["input1"], + "outputs": [ + "output1\n", + "output2\n"]} INPUT_OUTPUT_INVALID_MISSING_KEY = {"inputs": ["input1"]} @@ -77,7 +84,9 @@ @pytest.mark.skipif(skip_condition, reason=skip_reason) def test_integration_success_correct(): """Integration test: Code is correct, output is correct""" - results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_SUCCESS) + results, metadata_list = check_correctness( + SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_SUCCESS + ) assert results == [True, True] assert metadata_list[0]["status"] == "success" assert metadata_list[0]["stdout"] == "output1\n" @@ -88,7 +97,9 @@ def test_integration_success_correct(): @pytest.mark.skipif(skip_condition, reason=skip_reason) def test_integration_success_wrong_output(): """Integration test: Code runs successfully, but output is wrong""" - results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_WRONG_OUTPUT) + results, metadata_list = check_correctness( + SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_WRONG_OUTPUT + ) assert results == [False, False] assert metadata_list[0]["status"] == "wrong_answer" assert metadata_list[0]["stdout"] == "wrong_output\n" @@ -98,7 +109,9 @@ def test_integration_success_wrong_output(): @pytest.mark.skipif(skip_condition, reason=skip_reason) def test_integration_compile_error(): """Integration test: Code causes compile error""" - results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_COMPILE_ERROR, language="cpp") + results, metadata_list = check_correctness( + SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_COMPILE_ERROR, language="cpp" + ) assert results == [-4, -4] assert metadata_list[0]["status"] == "compile_error" assert metadata_list[1]["status"] == "compile_error" @@ -107,20 +120,26 @@ def test_integration_compile_error(): @pytest.mark.skipif(skip_condition, reason=skip_reason) def test_integration_runtime_error(): """Integration test: Code causes runtime error""" - results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_SINGLE, CODE_RUNTIME_ERROR) + results, metadata_list = check_correctness( + SANDBOX_URL, INPUT_OUTPUT_SINGLE, CODE_RUNTIME_ERROR + ) assert results == [-2] assert metadata_list[0]["status"] == "runtime_error" - # More assertions can be added based on the actual API response, e.g., exit_code, stderr + # More assertions can be added based on the actual API response, e.g., + # exit_code, stderr @pytest.mark.skipif(skip_condition, reason=skip_reason) def test_integration_runtime_timeout(): """Integration test: Code causes runtime timeout""" test_timeout = 5 # Set a timeout shorter than the sleep time in CODE_TIMEOUT - results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_SINGLE, CODE_TIMEOUT, timeout=test_timeout) + results, metadata_list = check_correctness( + SANDBOX_URL, INPUT_OUTPUT_SINGLE, CODE_TIMEOUT, timeout=test_timeout + ) assert results == [-3] assert metadata_list[0]["status"] == "timeout" - # More assertions can be added based on the actual API response, e.g., run_status + # More assertions can be added based on the actual API response, e.g., + # run_status @pytest.mark.skipif(skip_condition, reason=skip_reason) @@ -155,7 +174,9 @@ def test_integration_concurrency_high_load(): high_load_outputs.append(f"output_{i}\n") expected_results_map[i] = True # Expect success - high_load_in_outs = {"inputs": high_load_inputs, "outputs": high_load_outputs} + high_load_in_outs = { + "inputs": high_load_inputs, + "outputs": high_load_outputs} # Code that handles normal inputs, and sleeps on specific "timeout" inputs code_mixed_concurrent = """ @@ -170,8 +191,11 @@ def test_integration_concurrency_high_load(): else: print("unknown_input\\n", end='') """ - # Set a reasonable timeout per case (must be less than the sleep time in the code) - test_timeout = 15 # Allow slightly more time due to potential API load, but less than 20s sleep + # Set a reasonable timeout per case (must be less than the sleep time in + # the code) + # Allow slightly more time due to potential API load, but less than 20s + # sleep + test_timeout = 15 start_time = time.time() results, metadata_list = check_correctness( @@ -183,12 +207,15 @@ def test_integration_concurrency_high_load(): end_time = time.time() duration = end_time - start_time print( - f"\nHigh concurrency test ({concurrency_level} cases with {len(wrong_answer_indices)} wrong answers, " - f"{len(timeout_indices)} timeouts) duration: {duration:.2f} seconds" - ) + f"\nHigh concurrency test ({concurrency_level} cases with { + len(wrong_answer_indices)} wrong answers, " f"{ + len(timeout_indices)} timeouts) duration: { + duration:.2f} seconds") # Verify results against the expected map - assert len(results) == concurrency_level, f"Expected {concurrency_level} results, got {len(results)}" + assert ( + len(results) == concurrency_level + ), f"Expected {concurrency_level} results, got {len(results)}" correct_count = 0 wrong_count = 0 @@ -210,35 +237,57 @@ def test_integration_concurrency_high_load(): f"Correct results (True): {correct_count}/" f"{concurrency_level - len(wrong_answer_indices) - len(timeout_indices)}" ) - print(f"Expected wrong answers (False, correctly identified): {wrong_count}/{len(wrong_answer_indices)}") - print(f"Expected timeouts (-3, correctly identified): {timeout_count}/{len(timeout_indices)}") + print( + f"Expected wrong answers (False, correctly identified): {wrong_count}/{ + len(wrong_answer_indices)}") + print( + f"Expected timeouts (-3, correctly identified): {timeout_count}/{len(timeout_indices)}" + ) if unexpected_results: print("Unexpected results found:") - for idx, res, expected_str in unexpected_results[:10]: # Print first 10 unexpected - print(f" Index {idx}: Got {res}, {expected_str}. Metadata: {metadata_list[idx]}") - raise AssertionError(f"Found {len(unexpected_results)} unexpected results.") - - assert correct_count == concurrency_level - len(wrong_answer_indices) - len(timeout_indices), ( - "Incorrect number of successful results" - ) - assert wrong_count == len(wrong_answer_indices), "Incorrect number of identified wrong answers" - assert timeout_count == len(timeout_indices), "Incorrect number of identified timeouts" + for idx, res, expected_str in unexpected_results[ + :10 + ]: # Print first 10 unexpected + print( + f" Index {idx}: Got {res}, {expected_str}. Metadata: { + metadata_list[idx]}") + raise AssertionError( + f"Found { + len(unexpected_results)} unexpected results.") + + assert correct_count == concurrency_level - len(wrong_answer_indices) - len( + timeout_indices + ), "Incorrect number of successful results" + assert wrong_count == len( + wrong_answer_indices + ), "Incorrect number of identified wrong answers" + assert timeout_count == len( + timeout_indices + ), "Incorrect number of identified timeouts" # Verify metadata count and basic status of one of each type assert len(metadata_list) == concurrency_level # Find the first correct index first_correct_index = next( - i for i in range(concurrency_level) if i not in wrong_answer_indices and i not in timeout_indices + i + for i in range(concurrency_level) + if i not in wrong_answer_indices and i not in timeout_indices ) assert metadata_list[first_correct_index]["status"] == "success" - assert metadata_list[first_correct_index]["stdout"] == f"output_{first_correct_index}\n" + assert ( + metadata_list[first_correct_index]["stdout"] + == f"output_{first_correct_index}\n" + ) # Check the status of the first intentionally wrong case first_wrong_index = min(wrong_answer_indices) assert metadata_list[first_wrong_index]["status"] == "wrong_answer" assert metadata_list[first_wrong_index]["stdout"] == f"output_{first_wrong_index}\n" - assert metadata_list[first_wrong_index]["expected_output"] == f"wrong_output_{first_wrong_index}\n" + assert ( + metadata_list[first_wrong_index]["expected_output"] + == f"wrong_output_{first_wrong_index}\n" + ) # Check the status of the first intentionally timeout case first_timeout_index = min(timeout_indices) @@ -256,24 +305,48 @@ def test_unit_concurrency_order(mock_call_sandbox_api): generation = "print(input())" language = "python" timeout = 5 - in_outs = {"inputs": ["input1", "input2", "input3"], "outputs": ["output1", "output2", "output3"]} + in_outs = { + "inputs": ["input1", "input2", "input3"], + "outputs": ["output1", "output2", "output3"], + } def side_effect(*args, **kwargs): stdin = kwargs.get("stdin") if stdin == "input1": return ( - {"status": "Success", "run_result": {"status": "Finished", "stdout": "output1", "return_code": 0}}, + { + "status": "Success", + "run_result": { + "status": "Finished", + "stdout": "output1", + "return_code": 0, + }, + }, None, ) elif stdin == "input2": time.sleep(0.1) return ( - {"status": "Success", "run_result": {"status": "Finished", "stdout": "output2", "return_code": 0}}, + { + "status": "Success", + "run_result": { + "status": "Finished", + "stdout": "output2", + "return_code": 0, + }, + }, None, ) elif stdin == "input3": return ( - {"status": "Success", "run_result": {"status": "Finished", "stdout": "output3", "return_code": 0}}, + { + "status": "Success", + "run_result": { + "status": "Finished", + "stdout": "output3", + "return_code": 0, + }, + }, None, ) else: @@ -281,7 +354,9 @@ def side_effect(*args, **kwargs): mock_call_sandbox_api.side_effect = side_effect - results, metadata_list = check_correctness(sandbox_url, in_outs, generation, timeout, language) + results, metadata_list = check_correctness( + sandbox_url, in_outs, generation, timeout, language + ) assert results == [True, True, True] assert len(metadata_list) == 3 @@ -300,7 +375,10 @@ def test_unit_api_timeout_error_concurrent(mock_call_sandbox_api): generation = "print(input())" language = "python" timeout = 5 - in_outs = {"inputs": ["input1", "input2_timeout", "input3"], "outputs": ["output1", "output2", "output3"]} + in_outs = { + "inputs": ["input1", "input2_timeout", "input3"], + "outputs": ["output1", "output2", "output3"], + } api_error_message = "API Call Failed: Gateway Timeout (504) on attempt 3/3" @@ -308,14 +386,28 @@ def side_effect(*args, **kwargs): stdin = kwargs.get("stdin") if stdin == "input1": return ( - {"status": "Success", "run_result": {"status": "Finished", "stdout": "output1", "return_code": 0}}, + { + "status": "Success", + "run_result": { + "status": "Finished", + "stdout": "output1", + "return_code": 0, + }, + }, None, ) elif stdin == "input2_timeout": return (None, api_error_message) elif stdin == "input3": return ( - {"status": "Success", "run_result": {"status": "Finished", "stdout": "output3", "return_code": 0}}, + { + "status": "Success", + "run_result": { + "status": "Finished", + "stdout": "output3", + "return_code": 0, + }, + }, None, ) else: @@ -323,7 +415,9 @@ def side_effect(*args, **kwargs): mock_call_sandbox_api.side_effect = side_effect - results, metadata_list = check_correctness(sandbox_url, in_outs, generation, timeout, language) + results, metadata_list = check_correctness( + sandbox_url, in_outs, generation, timeout, language + ) assert results == [True, -1, True] assert len(metadata_list) == 3 @@ -347,7 +441,8 @@ def side_effect(*args, **kwargs): # --- Mock API call function for concurrency tracking --- -# This function will replace the real call_sandbox_api and use shared variables to track concurrency +# This function will replace the real call_sandbox_api and use shared +# variables to track concurrency def _mock_api_call_for_concurrency_tracking( active_calls_counter, # multiprocessing.Value max_calls_tracker, # multiprocessing.Value @@ -368,21 +463,28 @@ def _mock_api_call_for_concurrency_tracking( max_calls_tracker.value = active_calls_counter.value # Optional debug log: # print(f"[PID:{os.getpid()}-TID:{threading.get_ident()}] API Call Start. Active: " - # f"{active_calls_counter.value}, Max Observed: {max_calls_tracker.value}, Input: {stdin}") + # f"{active_calls_counter.value}, Max Observed: + # {max_calls_tracker.value}, Input: {stdin}") - time.sleep(SIMULATED_API_CALL_DURATION_TEST) # Simulate actual work duration + # Simulate actual work duration + time.sleep(SIMULATED_API_CALL_DURATION_TEST) # exit_time = time.time() # For detailed logging with call_lock: active_calls_counter.value -= 1 # Optional debug log: # print(f"[PID:{os.getpid()}-TID:{threading.get_ident()}] API Call End. Active: " - # f"{active_calls_counter.value}, Input: {stdin}, Duration: {exit_time - entry_time:.2f}s") + # f"{active_calls_counter.value}, Input: {stdin}, Duration: {exit_time + # - entry_time:.2f}s") # Return a simulated successful API response return { "status": "Success", - "run_result": {"status": "Finished", "stdout": f"mock_output_for_{stdin}", "return_code": 0}, + "run_result": { + "status": "Finished", + "stdout": f"mock_output_for_{stdin}", + "return_code": 0, + }, }, None @@ -400,23 +502,28 @@ def _process_pool_worker_for_concurrency_test( max_calls_tracker, call_lock, ): - # Corrected lambda to accept keyword arguments matching call_sandbox_api's usage - curried_mock_api_call = ( - lambda sandbox_fusion_url, code, stdin, compile_timeout, run_timeout, memory_limit_mb, language: ( - _mock_api_call_for_concurrency_tracking( - active_calls_counter, - max_calls_tracker, - call_lock, - sandbox_fusion_url, - code, - stdin, - compile_timeout, - run_timeout, - memory_limit_mb, - language, - ) - ) - ) + # Corrected lambda to accept keyword arguments matching call_sandbox_api's + # usage + def curried_mock_api_call( + sandbox_fusion_url, + code, + stdin, + compile_timeout, + run_timeout, + memory_limit_mb, + language): return ( + _mock_api_call_for_concurrency_tracking( + active_calls_counter, + max_calls_tracker, + call_lock, + sandbox_fusion_url, + code, + stdin, + compile_timeout, + run_timeout, + memory_limit_mb, + language, + )) # ---- START DEBUG PRINTS ---- import os @@ -431,7 +538,8 @@ def _process_pool_worker_for_concurrency_test( # ---- END DEBUG PRINTS ---- with patch( - "verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api", side_effect=curried_mock_api_call + "verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api", + side_effect=curried_mock_api_call, ) as mock_obj: # ---- START DEBUG PRINTS ---- print( @@ -439,7 +547,10 @@ def _process_pool_worker_for_concurrency_test( f"{verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api}", flush=True, ) - print(f"[Worker PID:{os.getpid()}] Mock object: {mock_obj}", flush=True) + print( + f"[Worker PID:{ + os.getpid()}] Mock object: {mock_obj}", + flush=True) # ---- END DEBUG PRINTS ---- results, metadata_list = check_correctness( sandbox_fusion_url=sandbox_url, @@ -448,10 +559,12 @@ def _process_pool_worker_for_concurrency_test( timeout=timeout, memory_limit_mb=memory_limit_mb, language=language, - concurrent_semaphore=mp_semaphore_for_check_correctness, # Pass multiprocessing.Semaphore + # Pass multiprocessing.Semaphore + concurrent_semaphore=mp_semaphore_for_check_correctness, ) # print(f"Process {os.getpid()} finished check_correctness. Processed {len(results)} tasks.") - return len(results) # Return the number of processed tasks for basic validation + # Return the number of processed tasks for basic validation + return len(results) # --- The actual test case for multiprocess concurrency control --- @@ -463,26 +576,33 @@ def test_multiprocess_global_concurrency_limit_with_semaphore(): via check_correctness's internal ThreadPoolExecutor. """ manager = multiprocessing.Manager() - active_calls_counter = manager.Value("i", 0) # Current active mock API calls - max_calls_tracker = manager.Value("i", 0) # Observed maximum concurrent mock API calls + active_calls_counter = manager.Value( + "i", 0) # Current active mock API calls + max_calls_tracker = manager.Value( + "i", 0 + ) # Observed maximum concurrent mock API calls call_lock = manager.Lock() # Lock to protect counters # Create a multiprocessing.Semaphore instance, this is the global semaphore we are testing. - # It will be passed to check_correctness and used by _process_single_case to limit calls to call_sandbox_api. + # It will be passed to check_correctness and used by _process_single_case + # to limit calls to call_sandbox_api. global_mp_semaphore = manager.Semaphore(MAX_GLOBAL_CONCURRENCY_LIMIT_TEST) mock_sandbox_url = "mock_url_for_concurrency_test" - mock_generation = "pass" # Specific code content is not important as API call is mocked + mock_generation = ( + "pass" # Specific code content is not important as API call is mocked + ) mock_memory_limit_mb = 1024 mock_language = "python" mock_timeout = 5 # Timeout setting, not critical for mock calls # Input/output data for each process - # NUM_TASKS_PER_PROCESS_TEST tasks will be handled by check_correctness's internal ThreadPoolExecutor + # NUM_TASKS_PER_PROCESS_TEST tasks will be handled by check_correctness's + # internal ThreadPoolExecutor process_in_outs = { - "inputs": [f"task_input_{i}" for i in range(NUM_TASKS_PER_PROCESS_TEST)], - "outputs": [f"task_output_{i}" for i in range(NUM_TASKS_PER_PROCESS_TEST)], - } + "inputs": [ + f"task_input_{i}" for i in range(NUM_TASKS_PER_PROCESS_TEST)], "outputs": [ + f"task_output_{i}" for i in range(NUM_TASKS_PER_PROCESS_TEST)], } futures = [] total_tasks_expected_to_run = NUM_PROCESSES_TEST * NUM_TASKS_PER_PROCESS_TEST @@ -513,9 +633,13 @@ def test_multiprocess_global_concurrency_limit_with_semaphore(): # Print some test statistics for debugging and validation print("\n--- Global Concurrency Test Stats ---") - print(f"Semaphore Limit (MAX_GLOBAL_CONCURRENCY_LIMIT_TEST): {MAX_GLOBAL_CONCURRENCY_LIMIT_TEST}") + print( + f"Semaphore Limit (MAX_GLOBAL_CONCURRENCY_LIMIT_TEST): {MAX_GLOBAL_CONCURRENCY_LIMIT_TEST}" + ) print(f"Number of Processes (NUM_PROCESSES_TEST): {NUM_PROCESSES_TEST}") - print(f"Tasks per Process (NUM_TASKS_PER_PROCESS_TEST): {NUM_TASKS_PER_PROCESS_TEST}") + print( + f"Tasks per Process (NUM_TASKS_PER_PROCESS_TEST): {NUM_TASKS_PER_PROCESS_TEST}" + ) print(f"Total Tasks Submitted: {total_tasks_expected_to_run}") print(f"Simulated API Call Duration: {SIMULATED_API_CALL_DURATION_TEST}s") print(f"Total Test Execution Time: {total_execution_time:.2f}s") @@ -523,14 +647,17 @@ def test_multiprocess_global_concurrency_limit_with_semaphore(): # print(f"Tasks processed per worker: {num_tasks_processed_per_worker}") # Verify that all submitted tasks have been processed - assert sum(num_tasks_processed_per_worker) == total_tasks_expected_to_run, ( - "Mismatch in the number of tasks processed." - ) + assert ( + sum(num_tasks_processed_per_worker) == total_tasks_expected_to_run + ), "Mismatch in the number of tasks processed." # Verify that the mock API was called at least once - assert max_calls_tracker.value > 0, "The mocked API call_sandbox_api was not called." + assert ( + max_calls_tracker.value > 0 + ), "The mocked API call_sandbox_api was not called." - # Core assertion: Observed maximum concurrent calls should not exceed the semaphore's limit + # Core assertion: Observed maximum concurrent calls should not exceed the + # semaphore's limit assert max_calls_tracker.value <= MAX_GLOBAL_CONCURRENCY_LIMIT_TEST, ( f"Observed concurrency ({max_calls_tracker.value}) exceeded semaphore limit " f"({MAX_GLOBAL_CONCURRENCY_LIMIT_TEST})." @@ -563,7 +690,9 @@ def test_unit_invalid_input_format(): assert results == [-1] assert metadata_list[0]["error"] == "Invalid input/output data" - results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_INVALID_MISSING_KEY, CODE_SUCCESS) + results, metadata_list = check_correctness( + SANDBOX_URL, INPUT_OUTPUT_INVALID_MISSING_KEY, CODE_SUCCESS + ) assert results == [-1] assert metadata_list[0]["error"] == "Invalid input/output data" @@ -571,7 +700,9 @@ def test_unit_invalid_input_format(): @pytest.mark.skipif(skip_condition, reason=skip_reason) def test_unit_input_output_mismatch(): """Unit test: Mismatch between the number of inputs and outputs""" - results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_MISMATCH, CODE_SUCCESS) + results, metadata_list = check_correctness( + SANDBOX_URL, INPUT_OUTPUT_MISMATCH, CODE_SUCCESS + ) assert results == [-1] assert len(metadata_list) == 1 assert metadata_list[0]["error"] == "Input/output count mismatch" @@ -608,13 +739,19 @@ def solve(): test_timeout = 10 # Set a timeout value start_time = time.time() - results, metadata_list = check_correctness(SANDBOX_URL, timeout_in_outs, code_infinite_loop, timeout=test_timeout) + results, metadata_list = check_correctness( + SANDBOX_URL, timeout_in_outs, code_infinite_loop, timeout=test_timeout + ) end_time = time.time() duration = end_time - start_time - print(f"\nHigh concurrency all timeout test ({concurrency_level} cases) duration: {duration:.2f} seconds") + print( + f"\nHigh concurrency all timeout test ({concurrency_level} cases) duration: { + duration:.2f} seconds") # Verify all results are -3 (timeout) - assert len(results) == concurrency_level, f"Expected {concurrency_level} results, got {len(results)}" + assert ( + len(results) == concurrency_level + ), f"Expected {concurrency_level} results, got {len(results)}" all_timed_out = all(r == -3 for r in results) if not all_timed_out: non_timeout_indices = [i for i, r in enumerate(results) if r != -3] @@ -622,7 +759,9 @@ def solve(): # Print metadata for the first few non-timeout cases for debugging for i in non_timeout_indices[:5]: print(f"Metadata for non-timeout case {i}: {metadata_list[i]}") - assert all_timed_out, f"Not all {concurrency_level} concurrent tests resulted in timeout (-3). Results: {results}" + assert ( + all_timed_out + ), f"Not all {concurrency_level} concurrent tests resulted in timeout (-3). Results: {results}" # Verify metadata count and status of the first case assert len(metadata_list) == concurrency_level @@ -657,10 +796,12 @@ def occurrencesOfElement(self, nums: List[int], queries: List[int], x: int) -> L } # Use a short timeout for fast tests - results, metadata_list = check_correctness(SANDBOX_URL, in_outs, generation_code, timeout=5) + results, metadata_list = check_correctness( + SANDBOX_URL, in_outs, generation_code, timeout=5 + ) # from verl.utils.reward_score.prime_code import apps_check_correctness # results, metadata_list = apps_check_correctness(in_outs=in_outs, generation=generation_code, - # timeout=50000, debug=True) + # timeout=50000, debug=True) assert results == [True, True] assert "error" not in metadata_list[0] diff --git a/Agent0/executor_train/verl/tests/utils/reward_score/test_sandbox_on_cpu.py b/Agent0/executor_train/verl/tests/utils/reward_score/test_sandbox_on_cpu.py index ff40732..125de2b 100644 --- a/Agent0/executor_train/verl/tests/utils/reward_score/test_sandbox_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/reward_score/test_sandbox_on_cpu.py @@ -33,17 +33,18 @@ """(x + 2)^2 + (y - 3)^2 """, # symbolic test ] -prime_code_answers = [ - """import sys +prime_code_answers = ( + [ + """import sys from collections import deque def main(): data = sys.stdin.read().split() it = iter(data) - + # Read start and target positions x0, y0, x1, y1 = int(next(it)), int(next(it)), int(next(it)), int(next(it)) - + n = int(next(it)) allowed = set() # The total number of allowed cells is at most 10^5. @@ -53,21 +54,21 @@ def main(): b = int(next(it)) for c in range(a, b + 1): allowed.add((r, c)) - + # Directions for the king (8 neighboring cells) directions = [(-1, -1), (-1, 0), (-1, 1), (0, -1), (0, 1), (1, -1), (1, 0), (1, 1)] - + start = (x0, y0) target = (x1, y1) - + # BFS initialization queue = deque() queue.append((x0, y0, 0)) # Mark the starting cell as visited by removing it from allowed set. allowed.discard(start) - + while queue: x, y, moves = queue.popleft() if (x, y) == target: @@ -78,13 +79,15 @@ def main(): if (nx, ny) in allowed: allowed.remove((nx, ny)) queue.append((nx, ny, moves + 1)) - + print(-1) if __name__ == '__main__': main() """ -] * 2 + ] + * 2 +) prime_code_gts = [ """{\n \"inputs\": [\n \"5 7 6 11\\n3\\n5 3 8\\n6 7 11\\n5 2 5\\n\",\n \"3 4 3 10\\n3\\n3 1 4\\n4 5 9\\n3 10 10\\n\",\n \"1 1 2 10\\n2\\n1 1 3\\n2 6 10\\n\",\n \"9 8 7 8\\n9\\n10 6 6\\n10 6 6\\n7 7 8\\n9 5 6\\n8 9 9\\n9 5 5\\n9 8 8\\n8 5 6\\n9 10 10\\n\",\n \"6 15 7 15\\n9\\n6 15 15\\n7 14 14\\n6 15 15\\n9 14 14\\n7 14 16\\n6 15 15\\n6 15 15\\n7 14 14\\n8 15 15\\n\",\n \"13 16 20 10\\n18\\n13 16 16\\n20 10 10\\n19 10 10\\n12 15 15\\n20 10 10\\n18 11 11\\n19 10 10\\n19 10 10\\n20 10 10\\n19 10 10\\n20 10 10\\n20 10 10\\n19 10 10\\n18 11 11\\n13 16 16\\n12 15 15\\n19 10 10\\n19 10 10\\n\",\n \"89 29 88 30\\n16\\n87 31 31\\n14 95 95\\n98 88 89\\n96 88 88\\n14 97 97\\n13 97 98\\n100 88 88\\n88 32 32\\n99 88 89\\n90 29 29\\n87 31 31\\n15 94 96\\n89 29 29\\n88 32 32\\n97 89 89\\n88 29 30\\n\",\n \"30 14 39 19\\n31\\n35 7 11\\n37 11 12\\n32 13 13\\n37 5 6\\n46 13 13\\n37 14 14\\n31 13 13\\n43 13 19\\n45 15 19\\n46 13 13\\n32 17 17\\n41 14 19\\n30 14 14\\n43 13 17\\n34 16 18\\n44 11 19\\n38 13 13\\n40 12 20\\n37 16 18\\n46 16 18\\n34 10 14\\n36 9 10\\n36 15 19\\n38 15 19\\n42 13 19\\n33 14 15\\n35 15 19\\n33 17 18\\n39 12 20\\n36 5 7\\n45 12 12\\n\",\n \"2 1 1 1\\n2\\n1 1 2\\n2 1 2\\n\",\n \"1 1 1 2\\n5\\n1000000000 1 10000\\n19920401 1188 5566\\n1000000000 1 10000\\n1 1 10000\\n5 100 200\\n\",\n \"1 1 1000000000 2\\n5\\n1000000000 1 10000\\n19920401 1188 5566\\n1000000000 1 10000\\n1 1 10000\\n5 100 200\\n\"\n ],\n \"outputs\": [\n \"4\\n\",\n \"6\\n\",\n \"-1\\n\",\n \"2\\n\",\n \"1\\n\",\n \"-1\\n\",\n \"1\\n\",\n \"9\\n\",\n \"1\\n\",\n \"1\\n\",\n \"-1\\n\"\n ]\n}""", # A correct sample # noqa: E501 """{\n \"inputs\": [\n \"5 7 6 11\\n3\\n5 3 8\\n6 7 11\\n5 2 5\\n\",\n \"3 4 3 10\\n3\\n3 1 4\\n4 5 9\\n3 10 10\\n\",\n \"1 1 2 10\\n2\\n1 1 3\\n2 6 10\\n\",\n \"9 8 7 8\\n9\\n10 6 6\\n10 6 6\\n7 7 8\\n9 5 6\\n8 9 9\\n9 5 5\\n9 8 8\\n8 5 6\\n9 10 10\\n\",\n \"6 15 7 15\\n9\\n6 15 15\\n7 14 14\\n6 15 15\\n9 14 14\\n7 14 16\\n6 15 15\\n6 15 15\\n7 14 14\\n8 15 15\\n\",\n \"13 16 20 10\\n18\\n13 16 16\\n20 10 10\\n19 10 10\\n12 15 15\\n20 10 10\\n18 11 11\\n19 10 10\\n19 10 10\\n20 10 10\\n19 10 10\\n20 10 10\\n20 10 10\\n19 10 10\\n18 11 11\\n13 16 16\\n12 15 15\\n19 10 10\\n19 10 10\\n\",\n \"89 29 88 30\\n16\\n87 31 31\\n14 95 95\\n98 88 89\\n96 88 88\\n14 97 97\\n13 97 98\\n100 88 88\\n88 32 32\\n99 88 89\\n90 29 29\\n87 31 31\\n15 94 96\\n89 29 29\\n88 32 32\\n97 89 89\\n88 29 30\\n\",\n \"30 14 39 19\\n31\\n35 7 11\\n37 11 12\\n32 13 13\\n37 5 6\\n46 13 13\\n37 14 14\\n31 13 13\\n43 13 19\\n45 15 19\\n46 13 13\\n32 17 17\\n41 14 19\\n30 14 14\\n43 13 17\\n34 16 18\\n44 11 19\\n38 13 13\\n40 12 20\\n37 16 18\\n46 16 18\\n34 10 14\\n36 9 10\\n36 15 19\\n38 15 19\\n42 13 19\\n33 14 15\\n35 15 19\\n33 17 18\\n39 12 20\\n36 5 7\\n45 12 12\\n\",\n \"2 1 1 1\\n2\\n1 1 2\\n2 1 2\\n\",\n \"1 1 1 2\\n5\\n1000000000 1 10000\\n19920401 1188 5566\\n1000000000 1 10000\\n1 1 10000\\n5 100 200\\n\",\n \"1 1 1000000000 2\\n5\\n1000000000 1 10000\\n19920401 1188 5566\\n1000000000 1 10000\\n1 1 10000\\n5 100 200\\n\"\n ],\n \"outputs\": [\n \"4\\n\",\n \"6\\n\",\n \"-1\\n\",\n \"-1\\n\",\n \"1\\n\",\n \"-1\\n\",\n \"1\\n\",\n \"9\\n\",\n \"1\\n\",\n \"1\\n\",\n \"-1\\n\"\n ]\n}""", # noqa: E501 @@ -110,7 +113,13 @@ def test_parallelism(): data_sources.extend(["numina_aops_forum"] * len(prime_math_answers)) scores = asyncio.run( - parallel_compute_score_async(default_compute_score, sequences_str, ground_truth, data_sources, num_processes=16) + parallel_compute_score_async( + default_compute_score, + sequences_str, + ground_truth, + data_sources, + num_processes=16, + ) ) print(scores) @@ -120,47 +129,69 @@ def test_prime_code(): Test PRIME code sandbox. """ data_source = "codecontests" - for completion, ground_truth, score_ in zip(prime_code_answers, prime_code_gts, prime_code_scores, strict=True): + for completion, ground_truth, score_ in zip( + prime_code_answers, prime_code_gts, prime_code_scores, strict=True + ): score = default_compute_score(data_source, completion, ground_truth) assert float(score) == score_ # Use the pytest.mark.skipif decorator to skip the test -@pytest.mark.skipif(not os.environ.get("SANDBOX_FUSION_URL"), reason="SANDBOX_FUSION_URL environment variable not set") +@pytest.mark.skipif( + not os.environ.get("SANDBOX_FUSION_URL"), + reason="SANDBOX_FUSION_URL environment variable not set", +) def test_prime_code_sandbox_fusion(): """ Test PRIME code on sandbox fusion. Skips if SANDBOX_FUSION_URL is not set. """ data_source = "codecontests" - # Get the URL from the environment variable, as skipif ensures it is set at this point + # Get the URL from the environment variable, as skipif ensures it is set + # at this point sandbox_fusion_url = os.environ.get("SANDBOX_FUSION_URL") # Removed the previous 'if not sandbox_url' check block - for completion, ground_truth, score_ in zip(prime_code_answers, prime_code_gts, prime_code_scores, strict=True): + for completion, ground_truth, score_ in zip( + prime_code_answers, prime_code_gts, prime_code_scores, strict=True + ): score = default_compute_score( - data_source, completion, ground_truth, extra_info={"sandbox_fusion_url": sandbox_fusion_url} + data_source, + completion, + ground_truth, + extra_info={"sandbox_fusion_url": sandbox_fusion_url}, ) # <-- Use the URL obtained from the environment variable assert float(score) == score_ -@pytest.mark.skipif(not os.environ.get("SANDBOX_FUSION_URL"), reason="SANDBOX_FUSION_URL environment variable not set") +@pytest.mark.skipif( + not os.environ.get("SANDBOX_FUSION_URL"), + reason="SANDBOX_FUSION_URL environment variable not set", +) def test_continuous_score_consistency(): """ Verify that continuous score calculation is consistent between prime_code and sandbox_fusion. Uses a test case where the first 9 out of 11 sub-cases pass (expected score 0.9). """ completion = prime_code_answers[1] # Use the second sample - ground_truth = prime_code_gts[1] # Use the second sample (9/11 pass, first 9 pass) + # Use the second sample (9/11 pass, first 9 pass) + ground_truth = prime_code_gts[1] expected_continuous_score = 0.9 # 1. Calculate score using prime_code (default) with continuous=True prime_score, _ = sandbox_fusion.compute_score( - os.environ.get("SANDBOX_FUSION_URL"), None, completion, ground_truth, continuous=True + os.environ.get("SANDBOX_FUSION_URL"), + None, + completion, + ground_truth, + continuous=True, ) # 2. Calculate score using sandbox_fusion with continuous=True - # Ensure the extra_info key triggers the sandbox_fusion path in default_compute_score - fusion_score, _ = prime_code.compute_score(completion, ground_truth, continuous=True) + # Ensure the extra_info key triggers the sandbox_fusion path in + # default_compute_score + fusion_score, _ = prime_code.compute_score( + completion, ground_truth, continuous=True + ) # 3. Assert scores are equal (using pytest.approx for float comparison) assert float(prime_score) == pytest.approx(expected_continuous_score) @@ -173,13 +204,19 @@ def test_continuous_score_consistency(): def test_check_correctness(): completion = prime_code_answers[0] ground_truth = json.loads(prime_code_gts[0]) - ground_truth_single = {"inputs": ground_truth["inputs"][:1], "outputs": ground_truth["outputs"][:1]} - res, meta = apps_check_correctness(in_outs=ground_truth_single, generation=completion, timeout=5, debug=False) + ground_truth_single = { + "inputs": ground_truth["inputs"][:1], + "outputs": ground_truth["outputs"][:1], + } + res, meta = apps_check_correctness( + in_outs=ground_truth_single, generation=completion, timeout=5, debug=False) print(res, meta) def test_prime_math(): data_source = "numina_aops_forum" - for completion, ground_truth in zip(prime_math_answers, prime_math_gts, strict=True): + for completion, ground_truth in zip( + prime_math_answers, prime_math_gts, strict=True + ): score = default_compute_score(data_source, completion, ground_truth) assert float(score) == 1.0 diff --git a/Agent0/executor_train/verl/tests/utils/test_activation_offload.py b/Agent0/executor_train/verl/tests/utils/test_activation_offload.py index 2393d79..0e22d60 100644 --- a/Agent0/executor_train/verl/tests/utils/test_activation_offload.py +++ b/Agent0/executor_train/verl/tests/utils/test_activation_offload.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,10 +26,16 @@ from verl.utils.activation_offload import enable_activation_offloading from verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager -from verl.utils.fsdp_utils import MixedPrecisionPolicy, apply_fsdp2, get_fsdp_wrap_policy +from verl.utils.fsdp_utils import ( + MixedPrecisionPolicy, + apply_fsdp2, + get_fsdp_wrap_policy, +) -def _fsdp_activation_offloading_test(rank, world_size, rendezvous_file, strategy="fsdp"): +def _fsdp_activation_offloading_test( + rank, world_size, rendezvous_file, strategy="fsdp" +): torch.cuda.set_device(rank) torch.distributed.init_process_group( backend="nccl", @@ -37,19 +43,27 @@ def _fsdp_activation_offloading_test(rank, world_size, rendezvous_file, strategy rank=rank, world_size=world_size, ) - device_mesh = init_device_mesh("cuda", mesh_shape=(world_size,), mesh_dim_names=("dp",)) + device_mesh = init_device_mesh( + "cuda", mesh_shape=(world_size,), mesh_dim_names=("dp",) + ) model_name = "Qwen/Qwen2.5-0.5B-Instruct" config = Qwen2Config(num_hidden_layers=4) with torch.device("cuda"): model = AutoModelForCausalLM.from_config( - config=config, torch_dtype=torch.bfloat16, attn_implementation="flash_attention_2" + config=config, + torch_dtype=torch.bfloat16, + attn_implementation="flash_attention_2", ) model = model.to(device="cuda") # Wrap model with FSDP - mixed_precision = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32) + mixed_precision = MixedPrecision( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + ) if strategy == "fsdp": model = FSDP( @@ -63,7 +77,9 @@ def _fsdp_activation_offloading_test(rank, world_size, rendezvous_file, strategy ) else: mp_policy = MixedPrecisionPolicy( - param_dtype=torch.bfloat16, reduce_dtype=torch.float32, cast_forward_inputs=True + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + cast_forward_inputs=True, ) fsdp_kwargs = { "mesh": device_mesh, @@ -72,24 +88,29 @@ def _fsdp_activation_offloading_test(rank, world_size, rendezvous_file, strategy apply_fsdp2(model, fsdp_kwargs, {}) optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4) - lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9) + lr_scheduler = torch.optim.lr_scheduler.StepLR( + optimizer, step_size=1, gamma=0.9) # Create checkpoint manager tokenizer = AutoTokenizer.from_pretrained(model_name) checkpoint_manager = FSDPCheckpointManager( - model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, tokenizer=tokenizer - ) + model=model, + optimizer=optimizer, + lr_scheduler=lr_scheduler, + tokenizer=tokenizer) # Generate sample input batch_size = 2 seq_len = 32 vocab_size = 32000 # First input for initial update - input_ids1 = torch.randint(0, vocab_size, (batch_size, seq_len), device="cuda") + input_ids1 = torch.randint( + 0, vocab_size, (batch_size, seq_len), device="cuda") attention_mask1 = torch.ones_like(input_ids1) # Second input for verification - input_ids2 = torch.randint(0, vocab_size, (batch_size, seq_len), device="cuda") + input_ids2 = torch.randint( + 0, vocab_size, (batch_size, seq_len), device="cuda") attention_mask2 = torch.ones_like(input_ids2) # Step 1: Initial update and save checkpoint @@ -103,7 +124,9 @@ def _fsdp_activation_offloading_test(rank, world_size, rendezvous_file, strategy # Save checkpoint after first update temp_dir = tempfile.mkdtemp() checkpoint_path = os.path.join(temp_dir, "checkpoint") - checkpoint_manager.save_checkpoint(local_path=checkpoint_path, hdfs_path=None, global_step=0) + checkpoint_manager.save_checkpoint( + local_path=checkpoint_path, hdfs_path=None, global_step=0 + ) # Step 2: Second update and forward pass outputs2 = model(input_ids=input_ids2, attention_mask=attention_mask2) @@ -115,7 +138,9 @@ def _fsdp_activation_offloading_test(rank, world_size, rendezvous_file, strategy # Record logits after second update with torch.no_grad(): - logits_without_offloading = model(input_ids=input_ids2, attention_mask=attention_mask2).logits + logits_without_offloading = model( + input_ids=input_ids2, attention_mask=attention_mask2 + ).logits # Step 3: wrap module with activation offloading and load checkpoint enable_activation_offloading(model, "fsdp") @@ -131,11 +156,16 @@ def _fsdp_activation_offloading_test(rank, world_size, rendezvous_file, strategy # Record logits after loaded checkpoint and update with torch.no_grad(): - logits_with_offloading = model(input_ids=input_ids2, attention_mask=attention_mask2).logits + logits_with_offloading = model( + input_ids=input_ids2, attention_mask=attention_mask2 + ).logits # Step 4: Verify outputs match - torch.testing.assert_close(logits_without_offloading, logits_with_offloading, atol=0.0, rtol=0.0) - print(f"Activaiton offloading for {strategy} test passed on {world_size} GPUs!") + torch.testing.assert_close( + logits_without_offloading, logits_with_offloading, atol=0.0, rtol=0.0 + ) + print( + f"Activaiton offloading for {strategy} test passed on {world_size} GPUs!") # Cleanup shutil.rmtree(temp_dir) diff --git a/Agent0/executor_train/verl/tests/utils/test_config_on_cpu.py b/Agent0/executor_train/verl/tests/utils/test_config_on_cpu.py index 03d952c..0523648 100644 --- a/Agent0/executor_train/verl/tests/utils/test_config_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/test_config_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -59,7 +59,8 @@ def test_omega_conf_to_dataclass(self): assert isinstance(cfg, TestDataclass) def test_nested_omega_conf_to_dataclass(self): - cfg = omega_conf_to_dataclass(self.config.train_config, TestTrainConfig) + cfg = omega_conf_to_dataclass( + self.config.train_config, TestTrainConfig) self.assertEqual(cfg.batch_size, 32) self.assertEqual(cfg.model.hidden_size, 768) self.assertEqual(cfg.model.activation, "relu") diff --git a/Agent0/executor_train/verl/tests/utils/test_flops_counter.py b/Agent0/executor_train/verl/tests/utils/test_flops_counter.py index 0b8889b..fa14bc2 100644 --- a/Agent0/executor_train/verl/tests/utils/test_flops_counter.py +++ b/Agent0/executor_train/verl/tests/utils/test_flops_counter.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -147,11 +147,15 @@ def test_flops_counter(config_type: str): config = Config(test_config["config"]) flops_counter = FlopsCounter(config) for batch_seqlens, expected_flops in zip( - test_config["batch_seqlens_tuple"], test_config["expected_flops_tuple"], strict=True + test_config["batch_seqlens_tuple"], + test_config["expected_flops_tuple"], + strict=True, ): # set delta time to 1 to get the flops counted_flops, _ = flops_counter.estimate_flops(batch_seqlens, 1) - print(f"Expect flops for {test_config['config']} is {expected_flops}, but get {counted_flops}") - assert math.isclose(counted_flops, expected_flops), ( - f"Expect flops for {test_config['config']} is {expected_flops}, but get {counted_flops}" - ) + print( + f"Expect flops for { + test_config['config']} is {expected_flops}, but get {counted_flops}") + assert math.isclose( + counted_flops, expected_flops), f"Expect flops for { + test_config['config']} is {expected_flops}, but get {counted_flops}" diff --git a/Agent0/executor_train/verl/tests/utils/test_fs_on_cpu.py b/Agent0/executor_train/verl/tests/utils/test_fs_on_cpu.py index 7ae85e0..d286804 100644 --- a/Agent0/executor_train/verl/tests/utils/test_fs_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/test_fs_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -58,7 +58,9 @@ def fake_copy(src: str, dst: str, *args, **kwargs): # Test initial copy local_path = fs.copy_to_local(hdfs_path, cache_dir=test_cache) - expected_path = os.path.join(test_cache, fs.md5_encode(hdfs_path), os.path.basename(hdfs_path)) + expected_path = os.path.join( + test_cache, fs.md5_encode(hdfs_path), os.path.basename(hdfs_path) + ) assert local_path == expected_path assert os.path.exists(local_path) diff --git a/Agent0/executor_train/verl/tests/utils/test_import_utils_on_cpu.py b/Agent0/executor_train/verl/tests/utils/test_import_utils_on_cpu.py index 59709b8..afaaf5a 100644 --- a/Agent0/executor_train/verl/tests/utils/test_import_utils_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/test_import_utils_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -84,7 +84,9 @@ def test_load_extern_type_invalid_module(): # Create a temporary file with syntax errors import tempfile - with tempfile.NamedTemporaryFile(suffix=".py", mode="w+", delete=False) as temp_file: + with tempfile.NamedTemporaryFile( + suffix=".py", mode="w+", delete=False + ) as temp_file: temp_file.write("This is not valid Python syntax :") temp_path = temp_file.name diff --git a/Agent0/executor_train/verl/tests/utils/test_linear_cross_entropy.py b/Agent0/executor_train/verl/tests/utils/test_linear_cross_entropy.py index 0512d13..d68d848 100644 --- a/Agent0/executor_train/verl/tests/utils/test_linear_cross_entropy.py +++ b/Agent0/executor_train/verl/tests/utils/test_linear_cross_entropy.py @@ -15,7 +15,7 @@ # limitations under the License. # -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -38,7 +38,8 @@ from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy from verl.utils.torch_functional import logprobs_from_logits -compute_entropy_from_logits = torch.compile(verl_F.entropy_from_logits, dynamic=True) +compute_entropy_from_logits = torch.compile( + verl_F.entropy_from_logits, dynamic=True) fused_linear_for_ppo = FusedLinearForPPO() fused_linear_for_ppo.compile(dynamic=True) @@ -46,17 +47,24 @@ def run_torch_entropy( - hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, temperature: float, reduction="none" + hidden: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + temperature: float, + reduction="none", ) -> list[torch.Tensor]: hidden = hidden.squeeze(0).to(torch.float32) weight = weight.transpose(0, 1).to(torch.float32) logits = torch.matmul(hidden, weight) # [num_tokens, vocab_size] logits /= temperature - pd = torch.nn.functional.softmax(logits, dim=-1) # [num_tokens, vocab_size] + pd = torch.nn.functional.softmax( + logits, dim=-1) # [num_tokens, vocab_size] entropy_a = torch.logsumexp(logits, dim=-1) # [num_tokens] entropy_b = torch.sum(pd * logits, dim=-1) # [num_tokens] entropy = entropy_a - entropy_b - logprobs = torch.nn.functional.cross_entropy(logits, labels.squeeze(0), reduction=reduction) # [num_tokens] + logprobs = torch.nn.functional.cross_entropy( + logits, labels.squeeze(0), reduction=reduction + ) # [num_tokens] logprobs = torch.neg(logprobs) return logprobs, entropy @@ -74,7 +82,9 @@ def run_verl_original_entropy( # compute entropy entropy = compute_entropy_from_logits(logits) # ((total_nnz / sp) + pad) # if use_sp: ((total_nnz / sp) + pad) ; if not use_sp: (batch, seqlen) - logprobs = logprobs_from_logits(logits=logits, labels=labels, inplace_backward=False) + logprobs = logprobs_from_logits( + logits=logits, labels=labels, inplace_backward=False + ) return logprobs, entropy @@ -144,21 +154,34 @@ def generate_hyper(self): def generate_forward_inputs(self): hidden = ( - torch.empty((self.batch_size, self.num_tokens, self.hidden_size), dtype=self.dtype, device="cuda") - .uniform_(-0.5, 0.5) - .requires_grad_() - ) - weight = ( - torch.empty((self.vocab_size, self.hidden_size), dtype=self.dtype, device="cuda") + torch.empty( + (self.batch_size, self.num_tokens, self.hidden_size), + dtype=self.dtype, + device="cuda", + ) .uniform_(-0.5, 0.5) .requires_grad_() ) - labels = torch.randint(0, self.vocab_size, (self.batch_size, self.num_tokens), device="cuda") + weight = (torch.empty((self.vocab_size, + self.hidden_size), + dtype=self.dtype, + device="cuda") .uniform_(-0.5, + 0.5) .requires_grad_()) + labels = torch.randint( + 0, + self.vocab_size, + (self.batch_size, + self.num_tokens), + device="cuda") return hidden, weight, labels def generate_backward_inputs(self): - g_entropy = torch.empty((self.num_tokens,), dtype=self.dtype, device="cuda").uniform_(-0.5, 0.5) - g_logprobs = torch.empty((self.num_tokens,), dtype=self.dtype, device="cuda").uniform_(-1, 1) + g_entropy = torch.empty( + (self.num_tokens,), dtype=self.dtype, device="cuda" + ).uniform_(-0.5, 0.5) + g_logprobs = torch.empty( + (self.num_tokens,), dtype=self.dtype, device="cuda" + ).uniform_(-1, 1) return g_entropy, g_logprobs def verify_correctness(self, iterations=5): @@ -182,52 +205,85 @@ def verify_correctness(self, iterations=5): hidden, weight, labels = self.generate_forward_inputs() start_event.record() - (torch_logprobs, torch_entropy) = run_torch_entropy(hidden, weight, labels, self.temperature) + (torch_logprobs, torch_entropy) = run_torch_entropy( + hidden, weight, labels, self.temperature + ) end_event.record() torch.cuda.synchronize() torch_forward_latency.append(start_event.elapsed_time(end_event)) start_event.record() - (verl_logprobs, verl_entropy) = run_verl_original_entropy(hidden, weight, labels, self.temperature) + (verl_logprobs, verl_entropy) = run_verl_original_entropy( + hidden, weight, labels, self.temperature + ) end_event.record() torch.cuda.synchronize() verl_forward_latency.append(start_event.elapsed_time(end_event)) start_event.record() (verl_fused_logprobs, verl_fused_entropy) = run_verl_torch_fused_entropy( - hidden, weight, labels, self.temperature - ) + hidden, weight, labels, self.temperature) end_event.record() torch.cuda.synchronize() - verl_fused_forward_latency.append(start_event.elapsed_time(end_event)) + verl_fused_forward_latency.append( + start_event.elapsed_time(end_event)) start_event.record() - (kernel_logprobs, kernel_entropy) = linear_cross_entropy(hidden, weight, labels, self.temperature) + (kernel_logprobs, kernel_entropy) = linear_cross_entropy( + hidden, weight, labels, self.temperature + ) end_event.record() torch.cuda.synchronize() kernel_forward_latency.append(start_event.elapsed_time(end_event)) - torch.testing.assert_close(torch_logprobs, verl_logprobs, atol=1e-4, rtol=1e-4) - torch.testing.assert_close(torch_entropy, verl_entropy, atol=1e-4, rtol=1e-4) + torch.testing.assert_close( + torch_logprobs, verl_logprobs, atol=1e-4, rtol=1e-4 + ) + torch.testing.assert_close( + torch_entropy, verl_entropy, atol=1e-4, rtol=1e-4 + ) - torch.testing.assert_close(torch_logprobs, verl_fused_logprobs, atol=1e-4, rtol=1e-4) - torch.testing.assert_close(torch_entropy, verl_fused_entropy, atol=1e-4, rtol=1e-4) - torch.testing.assert_close(verl_logprobs, verl_fused_logprobs, atol=1e-4, rtol=1e-4) - torch.testing.assert_close(verl_entropy, verl_fused_entropy, atol=1e-4, rtol=1e-4) + torch.testing.assert_close( + torch_logprobs, verl_fused_logprobs, atol=1e-4, rtol=1e-4 + ) + torch.testing.assert_close( + torch_entropy, verl_fused_entropy, atol=1e-4, rtol=1e-4 + ) + torch.testing.assert_close( + verl_logprobs, verl_fused_logprobs, atol=1e-4, rtol=1e-4 + ) + torch.testing.assert_close( + verl_entropy, verl_fused_entropy, atol=1e-4, rtol=1e-4 + ) - torch.testing.assert_close(torch_logprobs, kernel_logprobs, atol=1e-3, rtol=2e-4) - torch.testing.assert_close(torch_entropy, kernel_entropy, atol=5e-3, rtol=5e-4) - torch.testing.assert_close(verl_logprobs, kernel_logprobs, atol=1e-3, rtol=2e-4) - torch.testing.assert_close(verl_entropy, kernel_entropy, atol=5e-3, rtol=5e-4) - torch.testing.assert_close(verl_fused_logprobs, kernel_logprobs, atol=1e-3, rtol=2e-4) - torch.testing.assert_close(verl_fused_entropy, kernel_entropy, atol=5e-3, rtol=5e-4) + torch.testing.assert_close( + torch_logprobs, kernel_logprobs, atol=1e-3, rtol=2e-4 + ) + torch.testing.assert_close( + torch_entropy, kernel_entropy, atol=5e-3, rtol=5e-4 + ) + torch.testing.assert_close( + verl_logprobs, kernel_logprobs, atol=1e-3, rtol=2e-4 + ) + torch.testing.assert_close( + verl_entropy, kernel_entropy, atol=5e-3, rtol=5e-4 + ) + torch.testing.assert_close( + verl_fused_logprobs, kernel_logprobs, atol=1e-3, rtol=2e-4 + ) + torch.testing.assert_close( + verl_fused_entropy, kernel_entropy, atol=5e-3, rtol=5e-4 + ) # backward g_entropy, g_logprobs = self.generate_backward_inputs() start_event.record() (d_torch_hidden, d_torch_weight) = torch.autograd.grad( - (torch_entropy, torch_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False + (torch_entropy, torch_logprobs), + (hidden, weight), + (g_entropy, g_logprobs), + retain_graph=False, ) end_event.record() torch.cuda.synchronize() @@ -235,7 +291,10 @@ def verify_correctness(self, iterations=5): start_event.record() (d_verl_hidden, d_verl_weight) = torch.autograd.grad( - (verl_entropy, verl_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False + (verl_entropy, verl_logprobs), + (hidden, weight), + (g_entropy, g_logprobs), + retain_graph=False, ) end_event.record() torch.cuda.synchronize() @@ -243,36 +302,71 @@ def verify_correctness(self, iterations=5): start_event.record() (d_verl_fused_hidden, d_verl_fused_weight) = torch.autograd.grad( - (verl_fused_entropy, verl_fused_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False + (verl_fused_entropy, verl_fused_logprobs), + (hidden, weight), + (g_entropy, g_logprobs), + retain_graph=False, ) end_event.record() torch.cuda.synchronize() - verl_fused_backward_latency.append(start_event.elapsed_time(end_event)) + verl_fused_backward_latency.append( + start_event.elapsed_time(end_event)) start_event.record() (d_kernel_hidden, d_kernel_weight) = torch.autograd.grad( - (kernel_entropy, kernel_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False + (kernel_entropy, kernel_logprobs), + (hidden, weight), + (g_entropy, g_logprobs), + retain_graph=False, ) end_event.record() torch.cuda.synchronize() kernel_backward_latency.append(start_event.elapsed_time(end_event)) - torch.testing.assert_close(d_torch_hidden, d_verl_hidden, atol=1e-2, rtol=1e-4) - torch.testing.assert_close(d_torch_weight, d_verl_weight, atol=1e-2, rtol=1e-4) + torch.testing.assert_close( + d_torch_hidden, d_verl_hidden, atol=1e-2, rtol=1e-4 + ) + torch.testing.assert_close( + d_torch_weight, d_verl_weight, atol=1e-2, rtol=1e-4 + ) - torch.testing.assert_close(d_torch_hidden, d_verl_fused_hidden, atol=1e-2, rtol=1e-4) - torch.testing.assert_close(d_torch_weight, d_verl_fused_weight, atol=1e-2, rtol=1e-4) - torch.testing.assert_close(d_verl_hidden, d_verl_fused_hidden, atol=1e-2, rtol=1e-4) - torch.testing.assert_close(d_verl_weight, d_verl_fused_weight, atol=1e-2, rtol=1e-4) - torch.testing.assert_close(d_torch_hidden, d_verl_hidden, atol=1e-2, rtol=1e-4) - torch.testing.assert_close(d_torch_weight, d_verl_weight, atol=1e-2, rtol=1e-4) + torch.testing.assert_close( + d_torch_hidden, d_verl_fused_hidden, atol=1e-2, rtol=1e-4 + ) + torch.testing.assert_close( + d_torch_weight, d_verl_fused_weight, atol=1e-2, rtol=1e-4 + ) + torch.testing.assert_close( + d_verl_hidden, d_verl_fused_hidden, atol=1e-2, rtol=1e-4 + ) + torch.testing.assert_close( + d_verl_weight, d_verl_fused_weight, atol=1e-2, rtol=1e-4 + ) + torch.testing.assert_close( + d_torch_hidden, d_verl_hidden, atol=1e-2, rtol=1e-4 + ) + torch.testing.assert_close( + d_torch_weight, d_verl_weight, atol=1e-2, rtol=1e-4 + ) - torch.testing.assert_close(d_torch_hidden, d_kernel_hidden, atol=2e-2, rtol=4e-2) - torch.testing.assert_close(d_torch_weight, d_kernel_weight, atol=2e-2, rtol=4e-2) - torch.testing.assert_close(d_verl_hidden, d_kernel_hidden, atol=2e-2, rtol=4e-2) - torch.testing.assert_close(d_verl_weight, d_kernel_weight, atol=2e-2, rtol=4e-2) - torch.testing.assert_close(d_verl_fused_hidden, d_kernel_hidden, atol=2e-2, rtol=4e-2) - torch.testing.assert_close(d_verl_fused_weight, d_kernel_weight, atol=2e-2, rtol=4e-2) + torch.testing.assert_close( + d_torch_hidden, d_kernel_hidden, atol=2e-2, rtol=4e-2 + ) + torch.testing.assert_close( + d_torch_weight, d_kernel_weight, atol=2e-2, rtol=4e-2 + ) + torch.testing.assert_close( + d_verl_hidden, d_kernel_hidden, atol=2e-2, rtol=4e-2 + ) + torch.testing.assert_close( + d_verl_weight, d_kernel_weight, atol=2e-2, rtol=4e-2 + ) + torch.testing.assert_close( + d_verl_fused_hidden, d_kernel_hidden, atol=2e-2, rtol=4e-2 + ) + torch.testing.assert_close( + d_verl_fused_weight, d_kernel_weight, atol=2e-2, rtol=4e-2 + ) # remove first latency torch_forward_latency = torch_forward_latency[1:] @@ -291,9 +385,9 @@ def verify_correctness(self, iterations=5): f"{sum(torch_forward_latency) / len(torch_forward_latency):.2f} ms" ) print( - f"[INFO]: Backward pass: torch implementation average time: " - f"{sum(torch_backward_latency) / len(torch_backward_latency):.2f} ms" - ) + f"[INFO]: Backward pass: torch implementation average time: " f"{ + sum(torch_backward_latency) / + len(torch_backward_latency):.2f} ms") print( f"[INFO]: Forward pass: VeRL implementation average time: " f"{sum(verl_forward_latency) / len(verl_forward_latency):.2f} ms" @@ -303,21 +397,21 @@ def verify_correctness(self, iterations=5): f"{sum(verl_backward_latency) / len(verl_backward_latency):.2f} ms" ) print( - f"[INFO]: Forward pass: VeRL Fused Entropy implementation average time: " - f"{sum(verl_fused_forward_latency) / len(verl_fused_forward_latency):.2f} ms" - ) + f"[INFO]: Forward pass: VeRL Fused Entropy implementation average time: " f"{ + sum(verl_fused_forward_latency) / + len(verl_fused_forward_latency):.2f} ms") print( - f"[INFO]: Backward pass: VeRL Fused Entropy implementation average time: " - f"{sum(verl_fused_backward_latency) / len(verl_fused_backward_latency):.2f} ms" - ) + f"[INFO]: Backward pass: VeRL Fused Entropy implementation average time: " f"{ + sum(verl_fused_backward_latency) / + len(verl_fused_backward_latency):.2f} ms") print( - f"[INFO]: Forward pass: Kernel implementation average time: " - f"{sum(kernel_forward_latency) / len(kernel_forward_latency):.2f} ms" - ) + f"[INFO]: Forward pass: Kernel implementation average time: " f"{ + sum(kernel_forward_latency) / + len(kernel_forward_latency):.2f} ms") print( - f"[INFO]: Backward pass: kernel implementation average time: " - f"{sum(kernel_backward_latency) / len(kernel_backward_latency):.2f} ms" - ) + f"[INFO]: Backward pass: kernel implementation average time: " f"{ + sum(kernel_backward_latency) / + len(kernel_backward_latency):.2f} ms") def check_storage(self, method_name, run_forward): self.cleanup() @@ -326,20 +420,28 @@ def check_storage(self, method_name, run_forward): hidden, weight, labels = self.generate_forward_inputs() torch.cuda.reset_peak_memory_stats() - (logprobs, entropy) = run_forward(hidden, weight, labels, self.temperature) + (logprobs, entropy) = run_forward( + hidden, weight, labels, self.temperature) torch.cuda.synchronize() torch_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024 - print(f"[INFO]: {method_name} Forward pass peak memory: {torch_max_memory:.2f} MB") + print( + f"[INFO]: {method_name} Forward pass peak memory: { + torch_max_memory:.2f} MB") g_entropy, g_logprobs = self.generate_backward_inputs() torch.cuda.reset_peak_memory_stats() (d_torch_hidden, d_torch_weight) = torch.autograd.grad( - (entropy, logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False + (entropy, logprobs), + (hidden, weight), + (g_entropy, g_logprobs), + retain_graph=False, ) torch.cuda.synchronize() torch_backward_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024 - print(f"[INFO]: {method_name} Backward pass peak memory: {torch_backward_max_memory:.2f} MB") + print( + f"[INFO]: {method_name} Backward pass peak memory: { + torch_backward_max_memory:.2f} MB") def check_storage_all(self): self.check_storage("Torch", run_torch_entropy) diff --git a/Agent0/executor_train/verl/tests/utils/test_linear_cross_entropy_tp.py b/Agent0/executor_train/verl/tests/utils/test_linear_cross_entropy_tp.py index 9c1f868..2f88993 100644 --- a/Agent0/executor_train/verl/tests/utils/test_linear_cross_entropy_tp.py +++ b/Agent0/executor_train/verl/tests/utils/test_linear_cross_entropy_tp.py @@ -15,7 +15,7 @@ # limitations under the License. # -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -40,13 +40,18 @@ # FIXME: remove these manually included paths import sys - sys.path.append(os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../"))) + sys.path.append( + os.path.abspath( + os.path.join(os.path.dirname(os.path.abspath(__file__)), "../../") + ) + ) finally: from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy import verl.utils.torch_functional as verl_F -compute_entropy_from_logits = torch.compile(verl_F.entropy_from_logits, dynamic=True) +compute_entropy_from_logits = torch.compile( + verl_F.entropy_from_logits, dynamic=True) MAX_TEST_CASES = os.environ.get("MAX_TEST_CASES", 5) VERIFY_TORCH_SELF = os.environ.get("VERIFY_TORCH_SELF", False) @@ -55,7 +60,11 @@ def run_torch_entropy( - hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, temperature: float, reduction="none" + hidden: torch.Tensor, + weight: torch.Tensor, + labels: torch.Tensor, + temperature: float, + reduction="none", ) -> list[torch.Tensor]: # [num_tokens, vocab_size] if len(hidden.shape) > 2: @@ -64,14 +73,21 @@ def run_torch_entropy( labels = labels.view(-1) logits = torch.matmul( hidden.to(torch.float32), - weight.to(torch.float32) if weight.size(0) == hidden.size(1) else weight.T.to(torch.float32), + ( + weight.to(torch.float32) + if weight.size(0) == hidden.size(1) + else weight.T.to(torch.float32) + ), ) logits /= temperature - pd = torch.nn.functional.softmax(logits, dim=-1) # [num_tokens, vocab_size] + pd = torch.nn.functional.softmax( + logits, dim=-1) # [num_tokens, vocab_size] entropy_a = torch.logsumexp(logits, dim=-1) # [num_tokens] entropy_b = torch.sum(pd * logits, dim=-1) # [num_tokens] entropy = entropy_a - entropy_b - logprobs = torch.nn.functional.cross_entropy(logits, labels, reduction=reduction) # [num_tokens] + logprobs = torch.nn.functional.cross_entropy( + logits, labels, reduction=reduction + ) # [num_tokens] logprobs = torch.neg(logprobs) return logprobs, entropy @@ -91,22 +107,29 @@ def forward( temperature: float, dist_process_group: torch.distributed.ProcessGroup, ): - # weight has shape [vocab_size, hidden_size], hidden has shape [num_tokens, hidden_size] + # weight has shape [vocab_size, hidden_size], hidden has shape + # [num_tokens, hidden_size] ctx.original_hidden_shape = hidden.shape if len(hidden.shape) > 2: - hidden = hidden.view(-1, hidden.shape[-1]) # [num_tokens, hidden_size] + # [num_tokens, hidden_size] + hidden = hidden.view(-1, hidden.shape[-1]) if len(labels.shape) > 1: labels = labels.view(-1) - logits = torch.matmul(hidden.to(torch.float32), weight.to(torch.float32).T) # [num_tokens, vocab_size] + logits = torch.matmul( + hidden.to(torch.float32), weight.to(torch.float32).T + ) # [num_tokens, vocab_size] logits /= temperature whole_logits = torch.empty( - (logits.shape[0], logits.shape[1] * dist.get_world_size(dist_process_group)), + ( + logits.shape[0], + logits.shape[1] * dist.get_world_size(dist_process_group), + ), dtype=logits.dtype, device=logits.device, ) whole_logits_ref = [ - whole_logits[:, i * logits.shape[1] : (i + 1) * logits.shape[1]] + whole_logits[:, i * logits.shape[1]: (i + 1) * logits.shape[1]] for i in range(dist.get_world_size(dist_process_group)) ] dist.all_gather(whole_logits_ref, logits, group=dist_process_group) @@ -116,7 +139,9 @@ def forward( entropy_b = torch.sum(pd * whole_logits, dim=-1) # [num_tokens] entropy = entropy_a - entropy_b - logprobs = torch.nn.functional.cross_entropy(whole_logits, labels, reduction="none") + logprobs = torch.nn.functional.cross_entropy( + whole_logits, labels, reduction="none" + ) logprobs = torch.neg(logprobs) ctx.save_for_backward(hidden, weight, labels, whole_logits, entropy_b) @@ -148,7 +173,9 @@ def backward(ctx, g_logprobs: torch.Tensor, g_entropy: torch.Tensor): # d_entropy/d_logits = d_entropy_a - d_entropy_b # d_entropy/d_logits = pd - pd * (logits - b.unsqueeze(1) + 1) # d_entropy/d_logits = -pd * (logits - b.unsqueeze(1)) - d_logits_entropy = g_entropy.unsqueeze(1) * (-pd * (whole_logits - entropy_b.unsqueeze(1))) + d_logits_entropy = g_entropy.unsqueeze(1) * ( + -pd * (whole_logits - entropy_b.unsqueeze(1)) + ) # Gradient for logprobs # logprobs = -cross_entropy = -log(pd[labels]) @@ -165,7 +192,8 @@ def backward(ctx, g_logprobs: torch.Tensor, g_entropy: torch.Tensor): d_logits /= temperature # Get local slice of gradients - local_d_logits = d_logits[:, rank * vocab_size : (rank + 1) * vocab_size] + local_d_logits = d_logits[:, rank * + vocab_size: (rank + 1) * vocab_size] # Compute gradients for hidden and weight d_hidden = torch.matmul(local_d_logits, weight.to(torch.float32)) @@ -187,7 +215,10 @@ def __init__(self): self.world_size = dist.get_world_size(self.group) device = torch.device(f"cuda:{self.local_rank}") torch.cuda.set_device(device) - print(f"[INFO]: Local rank: {self.local_rank}, World size: {self.world_size}") + print( + f"[INFO]: Local rank: { + self.local_rank}, World size: { + self.world_size}") def initialize(self, test_case_idx: int, temperature: float = 1.5): self.test_case_idx = test_case_idx @@ -241,21 +272,34 @@ def generate_hyper(self): def generate_forward_inputs(self): hidden = ( - torch.empty((self.batch_size, self.num_tokens, self.hidden_size), dtype=self.dtype, device="cuda") - .uniform_(-0.5, 0.5) - .requires_grad_() - ) - weight = ( - torch.empty((self.vocab_size, self.hidden_size), dtype=self.dtype, device="cuda") + torch.empty( + (self.batch_size, self.num_tokens, self.hidden_size), + dtype=self.dtype, + device="cuda", + ) .uniform_(-0.5, 0.5) .requires_grad_() ) - labels = torch.randint(0, self.vocab_size, (self.batch_size, self.num_tokens), device="cuda") + weight = (torch.empty((self.vocab_size, + self.hidden_size), + dtype=self.dtype, + device="cuda") .uniform_(-0.5, + 0.5) .requires_grad_()) + labels = torch.randint( + 0, + self.vocab_size, + (self.batch_size, + self.num_tokens), + device="cuda") return hidden, weight, labels def generate_backward_inputs(self): - g_entropy = torch.empty((self.num_tokens,), dtype=self.dtype, device="cuda").uniform_(-0.5, 0.5) - g_logprobs = torch.empty((self.num_tokens,), dtype=self.dtype, device="cuda").uniform_(-1, 1) + g_entropy = torch.empty( + (self.num_tokens,), dtype=self.dtype, device="cuda" + ).uniform_(-0.5, 0.5) + g_logprobs = torch.empty( + (self.num_tokens,), dtype=self.dtype, device="cuda" + ).uniform_(-1, 1) return g_entropy, g_logprobs def verify_torch_itself(self, iterations: int = 5): @@ -265,23 +309,28 @@ def verify_torch_itself(self, iterations: int = 5): for i in range(iterations): hidden, weight, labels = self.generate_forward_inputs() - # NOTE: we need to manually synchronize hidden and labels among Process Group + # NOTE: we need to manually synchronize hidden and labels among + # Process Group dist.broadcast(hidden, src=0, group=self.group) dist.broadcast(labels, src=0, group=self.group) # forward pass # Create a tensor to hold the gathered weights from all ranks # weight has shape [vocab_size, hidden_size] - # We want to gather along the first dimension to get [vocab_size * world_size, hidden_size] + # We want to gather along the first dimension to get [vocab_size * + # world_size, hidden_size] # Create a single contiguous tensor to hold all gathered weights whole_weight = torch.empty( - (self.vocab_size * self.world_size, self.hidden_size), dtype=weight.dtype, device=weight.device + (self.vocab_size * self.world_size, self.hidden_size), + dtype=weight.dtype, + device=weight.device, ) # Create views into the tensor for each rank's portion whole_weight_views = [ - whole_weight[i * self.vocab_size : (i + 1) * self.vocab_size] for i in range(self.world_size) + whole_weight[i * self.vocab_size: (i + 1) * self.vocab_size] + for i in range(self.world_size) ] # Perform all_gather operation using the views @@ -290,36 +339,59 @@ def verify_torch_itself(self, iterations: int = 5): # Set requires_grad for autograd whole_weight.requires_grad_() - (single_logprobs, single_entropy) = run_torch_entropy(hidden, whole_weight, labels, self.temperature) + (single_logprobs, single_entropy) = run_torch_entropy( + hidden, whole_weight, labels, self.temperature + ) - (tp_logprobs, tp_entropy) = run_torch_entropy_tp(hidden, weight, labels, self.temperature, self.group) + (tp_logprobs, tp_entropy) = run_torch_entropy_tp( + hidden, weight, labels, self.temperature, self.group + ) - torch.testing.assert_close(single_logprobs, tp_logprobs, atol=1e-4, rtol=1e-4) - torch.testing.assert_close(single_entropy, tp_entropy, atol=1e-4, rtol=1e-4) + torch.testing.assert_close( + single_logprobs, tp_logprobs, atol=1e-4, rtol=1e-4 + ) + torch.testing.assert_close( + single_entropy, tp_entropy, atol=1e-4, rtol=1e-4) # backward pass g_entropy, g_logprobs = self.generate_backward_inputs() - # NOTE: we need to manually synchronize g_entropy and g_logprobs among Process Group + # NOTE: we need to manually synchronize g_entropy and g_logprobs + # among Process Group dist.broadcast(g_entropy, src=0, group=self.group) dist.broadcast(g_logprobs, src=0, group=self.group) (single_d_hidden, single_d_weight) = torch.autograd.grad( - (single_entropy, single_logprobs), (hidden, whole_weight), (g_entropy, g_logprobs), retain_graph=False + (single_entropy, single_logprobs), + (hidden, whole_weight), + (g_entropy, g_logprobs), + retain_graph=False, ) (tp_d_hidden, tp_d_weight) = torch.autograd.grad( - (tp_entropy, tp_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False + (tp_entropy, tp_logprobs), + (hidden, weight), + (g_entropy, g_logprobs), + retain_graph=False, ) # NOTE: all-reduce on hidden is conducted outside the kernel - dist.all_reduce(tp_d_hidden, op=dist.ReduceOp.SUM, group=self.group) + dist.all_reduce( + tp_d_hidden, + op=dist.ReduceOp.SUM, + group=self.group) - torch.testing.assert_close(tp_d_hidden, single_d_hidden, atol=1e-2, rtol=1e-4) + torch.testing.assert_close( + tp_d_hidden, single_d_hidden, atol=1e-2, rtol=1e-4 + ) # Extract the corresponding slice from single_d_weight for comparison # tp_d_weight has shape [vocab_size, hidden_size] # single_d_weight has shape [vocab_size * world_size, hidden_size] torch.testing.assert_close( tp_d_weight, - single_d_weight[self.local_rank * self.vocab_size : (self.local_rank + 1) * self.vocab_size], + single_d_weight[ + self.local_rank + * self.vocab_size: (self.local_rank + 1) + * self.vocab_size + ], atol=1e-2, rtol=1e-4, ) @@ -334,23 +406,30 @@ def check_torch_storage(self): hidden, weight, labels = self.generate_forward_inputs() - # NOTE: we need to manually synchronize hidden and labels among Process Group + # NOTE: we need to manually synchronize hidden and labels among Process + # Group dist.broadcast(hidden, src=0, group=self.group) dist.broadcast(labels, src=0, group=self.group) torch.cuda.reset_peak_memory_stats() - (tp_logprobs, tp_entropy) = run_torch_entropy_tp(hidden, weight, labels, self.temperature, self.group) + (tp_logprobs, tp_entropy) = run_torch_entropy_tp( + hidden, weight, labels, self.temperature, self.group + ) torch.cuda.synchronize() forward_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024 g_entropy, g_logprobs = self.generate_backward_inputs() - # NOTE: we need to manually synchronize g_entropy and g_logprobs among Process Group + # NOTE: we need to manually synchronize g_entropy and g_logprobs among + # Process Group dist.broadcast(g_entropy, src=0, group=self.group) dist.broadcast(g_logprobs, src=0, group=self.group) torch.cuda.reset_peak_memory_stats() (d_tp_hidden, d_tp_weight) = torch.autograd.grad( - (tp_entropy, tp_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False + (tp_entropy, tp_logprobs), + (hidden, weight), + (g_entropy, g_logprobs), + retain_graph=False, ) torch.cuda.synchronize() backward_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024 @@ -358,8 +437,12 @@ def check_torch_storage(self): dist.all_reduce(d_tp_hidden, op=dist.ReduceOp.SUM, group=self.group) if self.local_rank == 0: - print(f"[INFO]: Torch Forward pass peak memory: {forward_max_memory:.2f} MB") - print(f"[INFO]: Torch Backward pass peak memory: {backward_max_memory:.2f} MB") + print( + f"[INFO]: Torch Forward pass peak memory: { + forward_max_memory:.2f} MB") + print( + f"[INFO]: Torch Backward pass peak memory: { + backward_max_memory:.2f} MB") def verify_kernel_correctness(self, iterations: int = 5): self.cleanup() @@ -376,12 +459,15 @@ def verify_kernel_correctness(self, iterations: int = 5): for i in range(iterations): hidden, weight, labels = self.generate_forward_inputs() - # NOTE: we need to manually synchronize hidden and labels among Process Group + # NOTE: we need to manually synchronize hidden and labels among + # Process Group dist.broadcast(hidden, src=0, group=self.group) dist.broadcast(labels, src=0, group=self.group) start_event.record() - (torch_logprobs, torch_entropy) = run_torch_entropy_tp(hidden, weight, labels, self.temperature, self.group) + (torch_logprobs, torch_entropy) = run_torch_entropy_tp( + hidden, weight, labels, self.temperature, self.group + ) end_event.record() torch.cuda.synchronize() torch_forward_latency.append(start_event.elapsed_time(end_event)) @@ -394,37 +480,58 @@ def verify_kernel_correctness(self, iterations: int = 5): torch.cuda.synchronize() kernel_forward_latency.append(start_event.elapsed_time(end_event)) - torch.testing.assert_close(torch_logprobs, kernel_logprobs, atol=1e-1, rtol=1e-2) - torch.testing.assert_close(torch_entropy, kernel_entropy, atol=1e-1, rtol=1e-2) + torch.testing.assert_close( + torch_logprobs, kernel_logprobs, atol=1e-1, rtol=1e-2 + ) + torch.testing.assert_close( + torch_entropy, kernel_entropy, atol=1e-1, rtol=1e-2 + ) # backward pass g_entropy, g_logprobs = self.generate_backward_inputs() - # NOTE: we need to manually synchronize g_entropy and g_logprobs among Process Group + # NOTE: we need to manually synchronize g_entropy and g_logprobs + # among Process Group dist.broadcast(g_entropy, src=0, group=self.group) dist.broadcast(g_logprobs, src=0, group=self.group) start_event.record() (torch_d_hidden, torch_d_weight) = torch.autograd.grad( - (torch_entropy, torch_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False + (torch_entropy, torch_logprobs), + (hidden, weight), + (g_entropy, g_logprobs), + retain_graph=False, ) end_event.record() torch.cuda.synchronize() torch_backward_latency.append(start_event.elapsed_time(end_event)) # NOTE: all-reduce on hidden is conducted outside the kernel - dist.all_reduce(torch_d_hidden, op=dist.ReduceOp.SUM, group=self.group) + dist.all_reduce( + torch_d_hidden, + op=dist.ReduceOp.SUM, + group=self.group) start_event.record() (kernel_d_hidden, kernel_d_weight) = torch.autograd.grad( - (kernel_entropy, kernel_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False + (kernel_entropy, kernel_logprobs), + (hidden, weight), + (g_entropy, g_logprobs), + retain_graph=False, ) end_event.record() torch.cuda.synchronize() kernel_backward_latency.append(start_event.elapsed_time(end_event)) # NOTE: all-reduce on hidden is conducted outside the kernel - dist.all_reduce(kernel_d_hidden, op=dist.ReduceOp.SUM, group=self.group) + dist.all_reduce( + kernel_d_hidden, + op=dist.ReduceOp.SUM, + group=self.group) - torch.testing.assert_close(torch_d_hidden, kernel_d_hidden, atol=2e-2, rtol=4e-2) - torch.testing.assert_close(torch_d_weight, kernel_d_weight, atol=2e-2, rtol=4e-2) + torch.testing.assert_close( + torch_d_hidden, kernel_d_hidden, atol=2e-2, rtol=4e-2 + ) + torch.testing.assert_close( + torch_d_weight, kernel_d_weight, atol=2e-2, rtol=4e-2 + ) # remove first latency torch_forward_latency = torch_forward_latency[1:] @@ -436,21 +543,21 @@ def verify_kernel_correctness(self, iterations: int = 5): print("\n[PASS]: Verified kernel forward & backward correctness.") print( - f"[INFO]: Forward pass: Torch implementation average time: " - f"{sum(torch_forward_latency) / len(torch_forward_latency):.2f} ms" - ) + f"[INFO]: Forward pass: Torch implementation average time: " f"{ + sum(torch_forward_latency) / + len(torch_forward_latency):.2f} ms") print( - f"[INFO]: Backward pass: torch implementation average time: " - f"{sum(torch_backward_latency) / len(torch_backward_latency):.2f} ms" - ) + f"[INFO]: Backward pass: torch implementation average time: " f"{ + sum(torch_backward_latency) / + len(torch_backward_latency):.2f} ms") print( - f"[INFO]: Forward pass: Kernel implementation average time: " - f"{sum(kernel_forward_latency) / len(kernel_forward_latency):.2f} ms" - ) + f"[INFO]: Forward pass: Kernel implementation average time: " f"{ + sum(kernel_forward_latency) / + len(kernel_forward_latency):.2f} ms") print( - f"[INFO]: Backward pass: kernel implementation average time: " - f"{sum(kernel_backward_latency) / len(kernel_backward_latency):.2f} ms" - ) + f"[INFO]: Backward pass: kernel implementation average time: " f"{ + sum(kernel_backward_latency) / + len(kernel_backward_latency):.2f} ms") def check_kernel_storage(self): self.cleanup() @@ -458,7 +565,8 @@ def check_kernel_storage(self): hidden, weight, labels = self.generate_forward_inputs() - # NOTE: we need to manually synchronize hidden and labels among Process Group + # NOTE: we need to manually synchronize hidden and labels among Process + # Group dist.broadcast(hidden, src=0, group=self.group) dist.broadcast(labels, src=0, group=self.group) @@ -470,26 +578,38 @@ def check_kernel_storage(self): kernel_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024 g_entropy, g_logprobs = self.generate_backward_inputs() - # NOTE: we need to manually synchronize g_entropy and g_logprobs among Process Group + # NOTE: we need to manually synchronize g_entropy and g_logprobs among + # Process Group dist.broadcast(g_entropy, src=0, group=self.group) dist.broadcast(g_logprobs, src=0, group=self.group) torch.cuda.reset_peak_memory_stats() (d_kernel_hidden, d_kernel_weight) = torch.autograd.grad( - (kernel_entropy, kernel_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False + (kernel_entropy, kernel_logprobs), + (hidden, weight), + (g_entropy, g_logprobs), + retain_graph=False, ) torch.cuda.synchronize() kernel_backward_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024 # NOTE: all-reduce on hidden is conducted outside the kernel - dist.all_reduce(d_kernel_hidden, op=dist.ReduceOp.SUM, group=self.group) + dist.all_reduce( + d_kernel_hidden, + op=dist.ReduceOp.SUM, + group=self.group) if self.local_rank == 0: - print(f"[INFO]: Kernel Forward pass peak memory: {kernel_max_memory:.2f} MB") - print(f"[INFO]: Kernel Backward pass peak memory: {kernel_backward_max_memory:.2f} MB") + print( + f"[INFO]: Kernel Forward pass peak memory: { + kernel_max_memory:.2f} MB") + print( + f"[INFO]: Kernel Backward pass peak memory: { + kernel_backward_max_memory:.2f} MB") if __name__ == "__main__": - # TP command: torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/kernels/test_linear_cross_entropy_tp.py + # TP command: torchrun --standalone --nnodes=1 --nproc-per-node=2 + # tests/kernels/test_linear_cross_entropy_tp.py # Check if running with torchrun (distributed mode) assert int(os.environ["WORLD_SIZE"]) > 1, ( diff --git a/Agent0/executor_train/verl/tests/utils/test_model_on_cpu.py b/Agent0/executor_train/verl/tests/utils/test_model_on_cpu.py index 8b1416c..bee88df 100644 --- a/Agent0/executor_train/verl/tests/utils/test_model_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/test_model_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,7 +24,10 @@ "override_kwargs", [ {"param_a": 5, "new_param": "plain_added"}, - {"param_a": 2, "nested_params": {"sub_param_x": "updated_x", "sub_param_z": True}}, + { + "param_a": 2, + "nested_params": {"sub_param_x": "updated_x", "sub_param_z": True}, + }, ], ) def test_update_model_config(override_kwargs): @@ -33,20 +36,31 @@ def test_update_model_config(override_kwargs): handling both plain and nested overrides via parametrization. """ # Create a fresh mock config object for each test case - mock_config = SimpleNamespace( - param_a=1, nested_params=SimpleNamespace(sub_param_x="original_x", sub_param_y=100), other_param="keep_me" - ) + mock_config = SimpleNamespace(param_a=1, nested_params=SimpleNamespace( + sub_param_x="original_x", sub_param_y=100), other_param="keep_me", ) # Apply the updates using the parametrized override_kwargs update_model_config(mock_config, override_kwargs) # Assertions to check if the config was updated correctly if "nested_params" in override_kwargs: # Case 2: Nested override override_nested = override_kwargs["nested_params"] - assert mock_config.nested_params.sub_param_x == override_nested["sub_param_x"], "Nested sub_param_x mismatch" - assert mock_config.nested_params.sub_param_y == 100, "Nested sub_param_y should be unchanged" - assert hasattr(mock_config.nested_params, "sub_param_z"), "Expected nested sub_param_z to be added" - assert mock_config.nested_params.sub_param_z == override_nested["sub_param_z"], "Value of sub_param_z mismatch" + assert (mock_config.nested_params.sub_param_x == + override_nested["sub_param_x"]), "Nested sub_param_x mismatch" + assert ( + mock_config.nested_params.sub_param_y == 100 + ), "Nested sub_param_y should be unchanged" + assert hasattr( + mock_config.nested_params, "sub_param_z" + ), "Expected nested sub_param_z to be added" + assert (mock_config.nested_params.sub_param_z == + override_nested["sub_param_z"]), "Value of sub_param_z mismatch" else: # Case 1: Plain override (nested params untouched) - assert mock_config.nested_params.sub_param_x == "original_x", "Nested sub_param_x should be unchanged" - assert mock_config.nested_params.sub_param_y == 100, "Nested sub_param_y should be unchanged" - assert not hasattr(mock_config.nested_params, "sub_param_z"), "Nested sub_param_z should not exist" + assert ( + mock_config.nested_params.sub_param_x == "original_x" + ), "Nested sub_param_x should be unchanged" + assert ( + mock_config.nested_params.sub_param_y == 100 + ), "Nested sub_param_y should be unchanged" + assert not hasattr( + mock_config.nested_params, "sub_param_z" + ), "Nested sub_param_z should not exist" diff --git a/Agent0/executor_train/verl/tests/utils/test_nvtx_profile.py b/Agent0/executor_train/verl/tests/utils/test_nvtx_profile.py index 3450260..914d728 100644 --- a/Agent0/executor_train/verl/tests/utils/test_nvtx_profile.py +++ b/Agent0/executor_train/verl/tests/utils/test_nvtx_profile.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -42,8 +42,12 @@ def test_config_init(self): assert isinstance(profiler_config, ProfilerConfig) with self.assertRaises(AttributeError): _ = profiler_config.non_existing_key - assert config.get("non_existing_key") == profiler_config.get("non_existing_key") - assert config.get("non_existing_key", 1) == profiler_config.get("non_existing_key", 1) + assert config.get("non_existing_key") == profiler_config.get( + "non_existing_key" + ) + assert config.get("non_existing_key", 1) == profiler_config.get( + "non_existing_key", 1 + ) assert config["discrete"] == profiler_config["discrete"] from dataclasses import FrozenInstanceError @@ -73,7 +77,10 @@ def test_initialization(self): self.assertEqual(self.profiler.discrete, False) def test_start_stop_profiling(self): - with patch("torch.cuda.profiler.start") as mock_start, patch("torch.cuda.profiler.stop") as mock_stop: + with ( + patch("torch.cuda.profiler.start") as mock_start, + patch("torch.cuda.profiler.stop") as mock_stop, + ): # Test start self.profiler.start() self.assertTrue(self.profiler.this_step) @@ -88,7 +95,10 @@ def test_discrete_profiling(self): discrete_config = ProfilerConfig(discrete=True, all_ranks=True) profiler = NsightSystemsProfiler(self.rank, discrete_config) - with patch("torch.cuda.profiler.start") as mock_start, patch("torch.cuda.profiler.stop") as mock_stop: + with ( + patch("torch.cuda.profiler.start") as mock_start, + patch("torch.cuda.profiler.stop") as mock_stop, + ): profiler.start() self.assertTrue(profiler.this_step) mock_start.assert_not_called() # Shouldn't start immediately in discrete mode @@ -109,7 +119,9 @@ def test_func(self, *args, **kwargs): with ( patch("torch.cuda.profiler.start") as mock_start, patch("torch.cuda.profiler.stop") as mock_stop, - patch("verl.utils.profiler.nvtx_profile.mark_start_range") as mock_start_range, + patch( + "verl.utils.profiler.nvtx_profile.mark_start_range" + ) as mock_start_range, patch("verl.utils.profiler.nvtx_profile.mark_end_range") as mock_end_range, ): result = test_func(mock_self) @@ -133,7 +145,9 @@ def test_func(self, *args, **kwargs): with ( patch("torch.cuda.profiler.start") as mock_start, patch("torch.cuda.profiler.stop") as mock_stop, - patch("verl.utils.profiler.nvtx_profile.mark_start_range") as mock_start_range, + patch( + "verl.utils.profiler.nvtx_profile.mark_start_range" + ) as mock_start_range, patch("verl.utils.profiler.nvtx_profile.mark_end_range") as mock_end_range, ): result = test_func(mock_self) diff --git a/Agent0/executor_train/verl/tests/utils/test_rollout_trace_on_cpu.py b/Agent0/executor_train/verl/tests/utils/test_rollout_trace_on_cpu.py index e9358c1..9a5a521 100644 --- a/Agent0/executor_train/verl/tests/utils/test_rollout_trace_on_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/test_rollout_trace_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,7 +18,11 @@ import pytest -from verl.utils.rollout_trace import RolloutTraceConfig, rollout_trace_attr, rollout_trace_op +from verl.utils.rollout_trace import ( + RolloutTraceConfig, + rollout_trace_attr, + rollout_trace_op, +) @pytest.fixture(autouse=True) @@ -39,7 +43,10 @@ def mock_weave_client(): # Also mock the call_context if it's used internally by the decorator mock_weave.trace.context.call_context.return_value = MagicMock() - with patch.dict(sys.modules, {"weave": mock_weave, "weave.trace.context": mock_weave.trace.context}): + with patch.dict( + sys.modules, + {"weave": mock_weave, "weave.trace.context": mock_weave.trace.context}, + ): yield mock_client @@ -78,7 +85,10 @@ async def test_rollout_trace_on_untraced_class(): async def test_rollout_trace_with_tracer(mock_weave_client): """Tests that the decorator calls the tracer's methods correctly.""" - RolloutTraceConfig.init(project_name="my-project", experiment_name="my-experiment", backend="weave") + RolloutTraceConfig.init( + project_name="my-project", + experiment_name="my-experiment", + backend="weave") instance = TracedClass() assert RolloutTraceConfig.get_client() is mock_weave_client @@ -92,12 +102,16 @@ async def test_rollout_trace_with_tracer(mock_weave_client): assert call_kwargs["inputs"] == expected_inputs mock_call = mock_weave_client.create_call.return_value - mock_weave_client.finish_call.assert_called_once_with(mock_call, output=result) + mock_weave_client.finish_call.assert_called_once_with( + mock_call, output=result) async def test_rollout_trace_with_exception(mock_weave_client): """Tests that `finish` is called with the exception when one is raised.""" - RolloutTraceConfig.init(project_name="my-project", experiment_name="my-experiment", backend="weave") + RolloutTraceConfig.init( + project_name="my-project", + experiment_name="my-experiment", + backend="weave") instance = TracedClass() with pytest.raises(ValueError, match="Test Exception"): @@ -116,7 +130,10 @@ async def test_rollout_trace_with_exception(mock_weave_client): async def test_rollout_trace_with_dummy_backend(mock_weave_client): """Tests that the tracer is not called when the backend is 'dummy'.""" - RolloutTraceConfig.init(project_name="my-project", experiment_name="my-experiment", backend="dummy") + RolloutTraceConfig.init( + project_name="my-project", + experiment_name="my-experiment", + backend="dummy") instance = TracedClass() await instance.my_method("test_a") @@ -132,7 +149,10 @@ async def test_rollout_trace_with_real_weave_backend(): """Integration test with a real weave backend.""" # This assumes that the weave environment (e.g., project) is configured - RolloutTraceConfig.init(project_name="my-project", experiment_name="my-experiment", backend="weave") + RolloutTraceConfig.init( + project_name="my-project", + experiment_name="my-experiment", + backend="weave") instance = TracedClass() @@ -142,4 +162,6 @@ async def test_rollout_trace_with_real_weave_backend(): with pytest.raises(ValueError, match="Test Exception"): await instance.my_method_with_exception() - print("\nWeave integration test ran successfully. Check your weave project for the trace.") + print( + "\nWeave integration test ran successfully. Check your weave project for the trace." + ) diff --git a/Agent0/executor_train/verl/tests/utils/test_seqlen_balancing.py b/Agent0/executor_train/verl/tests/utils/test_seqlen_balancing.py index df7760b..452343f 100644 --- a/Agent0/executor_train/verl/tests/utils/test_seqlen_balancing.py +++ b/Agent0/executor_train/verl/tests/utils/test_seqlen_balancing.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,18 +18,27 @@ from verl import DataProto from verl.utils.model import create_random_mask -from verl.utils.seqlen_balancing import ceildiv, get_reverse_idx, rearrange_micro_batches +from verl.utils.seqlen_balancing import ( + ceildiv, + get_reverse_idx, + rearrange_micro_batches, +) def test_seqlen_balancing(): input_ids = torch.randint(low=0, high=10, size=(20, 100)) attention_mask = create_random_mask( - input_ids=input_ids, max_ratio_of_left_padding=0.1, max_ratio_of_valid_token=0.9, min_ratio_of_valid_token=0.5 + input_ids=input_ids, + max_ratio_of_left_padding=0.1, + max_ratio_of_valid_token=0.9, + min_ratio_of_valid_token=0.5, ) data = {"input_ids": input_ids, "attention_mask": attention_mask} dataproto = DataProto.from_single_dict(data) - micro_batches, micro_bsz_idx_lst = rearrange_micro_batches(dataproto.batch, max_token_len=300) + micro_batches, micro_bsz_idx_lst = rearrange_micro_batches( + dataproto.batch, max_token_len=300 + ) batch = torch.cat(micro_batches) micro_bsz_idx = [] for idx in micro_bsz_idx_lst: @@ -50,9 +59,11 @@ def _worker(rank, world_size, init_method, max_token_len, use_same_dp, min_mb): rank=rank, ) - # 2) build a small random batch (each rank different length to force mismatch) + # 2) build a small random batch (each rank different length to force + # mismatch) torch.manual_seed(42 + rank) - input_ids = torch.randint(0, 10, (20 + rank * 5, 100), device=f"cuda:{rank}") + input_ids = torch.randint( + 0, 10, (20 + rank * 5, 100), device=f"cuda:{rank}") attention_mask = create_random_mask( input_ids=input_ids, max_ratio_of_left_padding=0.1, @@ -82,7 +93,8 @@ def _worker(rank, world_size, init_method, max_token_len, use_same_dp, min_mb): assert len(micros) == expected if use_same_dp: # gather all local_counts - counts = [torch.zeros(1, device=f"cuda:{rank}") for _ in range(world_size)] + counts = [torch.zeros(1, device=f"cuda:{rank}") + for _ in range(world_size)] counts[rank].fill_(local) dist.all_gather(counts, counts[rank]) expected = max(int(c.item()) for c in counts) diff --git a/Agent0/executor_train/verl/tests/utils/test_timeout_decorator_cpu.py b/Agent0/executor_train/verl/tests/utils/test_timeout_decorator_cpu.py index 3417469..f5b22c5 100644 --- a/Agent0/executor_train/verl/tests/utils/test_timeout_decorator_cpu.py +++ b/Agent0/executor_train/verl/tests/utils/test_timeout_decorator_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,7 +23,8 @@ # --- Test Task Functions --- TEST_TIMEOUT_SECONDS = 1.5 # Timeout duration for tests -LONG_TASK_DURATION = TEST_TIMEOUT_SECONDS + 0.5 # Duration slightly longer than timeout +LONG_TASK_DURATION = TEST_TIMEOUT_SECONDS + \ + 0.5 # Duration slightly longer than timeout @timeout(seconds=TEST_TIMEOUT_SECONDS) # Keep global decorator for mp tests @@ -85,12 +86,14 @@ def set_macos_start_method(): # Force fork method on macOS to avoid pickling issues with globally decorated functions # when running tests via pytest discovery. current_method = multiprocessing.get_start_method(allow_none=True) - # Only set if not already set or if set to something else (less likely in test run) + # Only set if not already set or if set to something else (less likely + # in test run) if current_method is None or current_method != "fork": try: multiprocessing.set_start_method("fork", force=True) except RuntimeError: - # Might fail if context is already started, ignore in that case. + # Might fail if context is already started, ignore in that + # case. pass @@ -107,16 +110,22 @@ def test_slow_task_timeout(): # Renamed from test_multiprocessing_slow_task_tim with pytest.raises(TimeoutError) as excinfo: # Use pytest.raises slow_task(1) # Check the error message from the multiprocessing implementation - assert f"timed out after {TEST_TIMEOUT_SECONDS} seconds" in str(excinfo.value) # Use pytest assert + assert f"timed out after {TEST_TIMEOUT_SECONDS} seconds" in str( + excinfo.value + ) # Use pytest assert def test_internal_exception(): # Renamed from test_multiprocessing_internal_exception """Tests timeout correctly propagates internal exceptions.""" - # Apply the default timeout decorator dynamically to the undecorated function - decorated_task = timeout(seconds=TEST_TIMEOUT_SECONDS)(task_raises_value_error) # Apply decorator dynamically + # Apply the default timeout decorator dynamically to the undecorated + # function + decorated_task = timeout(seconds=TEST_TIMEOUT_SECONDS)( + task_raises_value_error + ) # Apply decorator dynamically with pytest.raises(ValueError) as excinfo: # Use pytest.raises decorated_task() # Call the dynamically decorated function - assert str(excinfo.value) == "Specific value error from task" # Use pytest assert + # Use pytest assert + assert str(excinfo.value) == "Specific value error from task" # --- Test the signal implementation (use_signals=True) --- @@ -132,7 +141,9 @@ def plain_quick_task_logic(): time.sleep(0.1) return "quick_ok_signal" - decorated_task = timeout(seconds=TEST_TIMEOUT_SECONDS, use_signals=True)(plain_quick_task_logic) + decorated_task = timeout(seconds=TEST_TIMEOUT_SECONDS, use_signals=True)( + plain_quick_task_logic + ) assert decorated_task() == "quick_ok_signal" # Use pytest assert @@ -144,14 +155,20 @@ def plain_slow_task_logic(): time.sleep(LONG_TASK_DURATION) return "slow_finished_signal" - decorated_task = timeout(seconds=TEST_TIMEOUT_SECONDS, use_signals=True)(plain_slow_task_logic) + decorated_task = timeout(seconds=TEST_TIMEOUT_SECONDS, use_signals=True)( + plain_slow_task_logic + ) with pytest.raises(TimeoutError) as excinfo: # Use pytest.raises decorated_task() # Check the error message (falls back to multiprocessing message on POSIX) - assert f"timed out after {TEST_TIMEOUT_SECONDS} seconds" in str(excinfo.value) # Use pytest assert + assert f"timed out after {TEST_TIMEOUT_SECONDS} seconds" in str( + excinfo.value + ) # Use pytest assert -@pytest.mark.skip(reason="this test won't pass. Just to show why use_signals should not be used") +@pytest.mark.skip( + reason="this test won't pass. Just to show why use_signals should not be used" +) def test_signal_in_thread_does_not_timeout(): """ Tests that signal-based timeout does NOT work reliably in a child thread. diff --git a/Agent0/executor_train/verl/tests/utils/test_torch_functional.py b/Agent0/executor_train/verl/tests/utils/test_torch_functional.py index 900cb5d..35c6e19 100644 --- a/Agent0/executor_train/verl/tests/utils/test_torch_functional.py +++ b/Agent0/executor_train/verl/tests/utils/test_torch_functional.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,7 +19,11 @@ import torch.distributed as dist import torch.multiprocessing as mp -from verl.utils.torch_functional import distributed_masked_mean, distributed_mean_max_min_std, masked_mean +from verl.utils.torch_functional import ( + distributed_masked_mean, + distributed_mean_max_min_std, + masked_mean, +) def _worker_mean(rank: int, world_size: int, rendezvous_file: str): @@ -34,7 +38,8 @@ def _worker_mean(rank: int, world_size: int, rendezvous_file: str): # each rank holds tensor [rank+1] local = torch.tensor([float(rank + 1)], device=f"cuda:{rank}") - mean, gmax, gmin, gstd = distributed_mean_max_min_std(local, True, True, True) + mean, gmax, gmin, gstd = distributed_mean_max_min_std( + local, True, True, True) values = [float(i + 1) for i in range(world_size)] exp_mean = sum(values) / len(values) @@ -89,7 +94,8 @@ def _worker_mask(rank: int, world_size: int, rendezvous_file: str): ) # build perโ€rank tensor and mask - local_tensor = torch.tensor([rank * 2 + 1.0, rank * 2 + 2.0], device=f"cuda:{rank}") + local_tensor = torch.tensor( + [rank * 2 + 1.0, rank * 2 + 2.0], device=f"cuda:{rank}") if rank == 0: mask = torch.tensor([1, 0], device=f"cuda:{rank}", dtype=torch.float32) else: @@ -99,7 +105,9 @@ def _worker_mask(rank: int, world_size: int, rendezvous_file: str): valid_values = [1.0] + [2 * i + 2.0 for i in range(1, world_size)] expected_mean = sum(valid_values) / len(valid_values) - assert torch.allclose(gmean.cpu(), torch.tensor(expected_mean)), f"masked_mean@{rank}" + assert torch.allclose( + gmean.cpu(), torch.tensor(expected_mean) + ), f"masked_mean@{rank}" dist.destroy_process_group() diff --git a/Agent0/executor_train/verl/tests/workers/reward_manager/test_registry_on_cpu.py b/Agent0/executor_train/verl/tests/workers/reward_manager/test_registry_on_cpu.py index 9932ae8..7542a3f 100644 --- a/Agent0/executor_train/verl/tests/workers/reward_manager/test_registry_on_cpu.py +++ b/Agent0/executor_train/verl/tests/workers/reward_manager/test_registry_on_cpu.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,14 +15,20 @@ import pytest # Assuming REWARD_MANAGER_REGISTRY is defined somewhere in the module -from verl.workers.reward_manager.registry import REWARD_MANAGER_REGISTRY, get_reward_manager_cls, register +from verl.workers.reward_manager.registry import ( + REWARD_MANAGER_REGISTRY, + get_reward_manager_cls, + register, +) @pytest.fixture def setup(): """Setup test cases with a mock registry.""" REWARD_MANAGER_REGISTRY.clear() - REWARD_MANAGER_REGISTRY.update({"manager1": "Manager1Class", "manager2": "Manager2Class"}) + REWARD_MANAGER_REGISTRY.update( + {"manager1": "Manager1Class", "manager2": "Manager2Class"} + ) return REWARD_MANAGER_REGISTRY diff --git a/Agent0/executor_train/verl/tests/workers/rollout/async_rollout_utils.py b/Agent0/executor_train/verl/tests/workers/rollout/async_rollout_utils.py index 22f2029..aab8688 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/async_rollout_utils.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/async_rollout_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -23,38 +23,48 @@ def init_async_rollout_manager(config: DictConfig) -> AsyncLLMServerManager: - # =========================== 1. Create hybrid ActorRollout workers =========================== + # =========================== 1. Create hybrid ActorRollout workers ====== role_worker_mapping = { Role.ActorRollout: ray.remote(AsyncActorRolloutRefWorker), } global_pool_id = "global_pool" resource_pool_spec = { - global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + global_pool_id: [ + config.trainer.n_gpus_per_node] * + config.trainer.nnodes, } mapping = { Role.ActorRollout: global_pool_id, } - resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, mapping=mapping + ) resource_pool_manager.create_resource_pool() - resource_pool_to_cls = {pool: {} for pool in resource_pool_manager.resource_pool_dict.values()} + resource_pool_to_cls = { + pool: {} for pool in resource_pool_manager.resource_pool_dict.values() + } # create actor and rollout resource_pool = resource_pool_manager.get_resource_pool(Role.ActorRollout) actor_rollout_cls = RayClassWithInitArgs( - cls=role_worker_mapping[Role.ActorRollout], config=config.actor_rollout_ref, role="actor_rollout" + cls=role_worker_mapping[Role.ActorRollout], + config=config.actor_rollout_ref, + role="actor_rollout", ) resource_pool_to_cls[resource_pool]["actor_rollout"] = actor_rollout_cls all_wg = {} for resource_pool, class_dict in resource_pool_to_cls.items(): worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict) - wg_dict = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls) + wg_dict = RayWorkerGroup( + 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) actor_rollout_wg = all_wg["actor_rollout"] actor_rollout_wg.init_model() - # =========================== 2. Create AsyncLLMServerManager =========================== + # =========================== 2. Create AsyncLLMServerManager =========== async_rollout_manager = AsyncLLMServerManager( config=config, worker_group=actor_rollout_wg, diff --git a/Agent0/executor_train/verl/tests/workers/rollout/perf/vllm_async_rollout.py b/Agent0/executor_train/verl/tests/workers/rollout/perf/vllm_async_rollout.py index dbcd255..316d3c2 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/perf/vllm_async_rollout.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/perf/vllm_async_rollout.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -38,7 +38,11 @@ from torch.utils.data import SequentialSampler from torchdata.stateful_dataloader import StatefulDataLoader -from tests.experimental.agent_loop.agent_utils import AgentLoopManager, RayWorkerGroup, init_agent_loop_manager +from tests.experimental.agent_loop.agent_utils import ( + AgentLoopManager, + RayWorkerGroup, + init_agent_loop_manager, +) from verl.protocol import DataProto from verl.utils import hf_tokenizer from verl.utils.dataset import RLHFDataset @@ -71,7 +75,9 @@ def init_config(n_gpus_per_node) -> DictConfig: return config -def initialize(config, backend) -> tuple[AgentLoopManager | RayWorkerGroup, StatefulDataLoader]: +def initialize( + config, backend +) -> tuple[AgentLoopManager | RayWorkerGroup, StatefulDataLoader]: env_vars = { "NCCL_DEBUG": "WARN", "VLLM_USE_V1": "1", @@ -91,8 +97,12 @@ def initialize(config, backend) -> tuple[AgentLoopManager | RayWorkerGroup, Stat ) dataloader = StatefulDataLoader( dataset=dataset, - batch_size=config.data.get("gen_batch_size", config.data.train_batch_size), - num_workers=config.data.get("dataloader_num_workers", 8), + batch_size=config.data.get( + "gen_batch_size", + config.data.train_batch_size), + num_workers=config.data.get( + "dataloader_num_workers", + 8), drop_last=True, collate_fn=default_collate_fn, sampler=SequentialSampler(dataset), @@ -116,9 +126,9 @@ def perf_rollout(mode, backend, n_gpus_per_node, num_steps): gen_batch = agent_loop_manager.generate_sequences(batch) t_end = time.time() print( - f"[DEBUG] backend: {backend}, n_gpus_per_node: {n_gpus_per_node}, batch_size: {len(gen_batch)}, " - f"step: {step}, step_time: {t_end - t_start:.2f} secs" - ) + f"[DEBUG] backend: {backend}, n_gpus_per_node: {n_gpus_per_node}, batch_size: { + len(gen_batch)}, " f"step: {step}, step_time: { + t_end - t_start:.2f} secs") if step + 1 >= num_steps: break @@ -132,4 +142,9 @@ def perf_rollout(mode, backend, n_gpus_per_node, num_steps): # test_cases = [("sync", "sync"), ("async", "zeromq"), ("async", "ray")] test_cases = [("async", "zeromq"), ("async", "ray")] for mode, backend in test_cases: - perf_rollout(mode=mode, backend=backend, n_gpus_per_node=n_gpus_per_node, num_steps=num_steps) + perf_rollout( + mode=mode, + backend=backend, + n_gpus_per_node=n_gpus_per_node, + num_steps=num_steps, + ) diff --git a/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/run_fsdp_vllm.py b/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/run_fsdp_vllm.py index 6922389..784c76b 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/run_fsdp_vllm.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/run_fsdp_vllm.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,7 +19,11 @@ import torch.distributed as dist from torch.distributed.fsdp import CPUOffload, MixedPrecision from torch.distributed.fsdp import FullyShardedDataParallel as FSDP -from torch.distributed.fsdp.api import ShardedStateDictConfig, ShardingStrategy, StateDictType +from torch.distributed.fsdp.api import ( + ShardedStateDictConfig, + ShardingStrategy, + StateDictType, +) from transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer from vllm import SamplingParams @@ -38,10 +42,15 @@ def main(): from verl.utils.fs import copy_to_local local_model_path = copy_to_local(src=hdfs_path, cache_dir=local_cache_path) - tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True) - actor_model_config = AutoConfig.from_pretrained(local_model_path, trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained( + local_model_path, trust_remote_code=True) + actor_model_config = AutoConfig.from_pretrained( + local_model_path, trust_remote_code=True + ) with torch.device("cuda"): - actor_model = AutoModelForCausalLM.from_pretrained(local_model_path, trust_remote_code=True) + actor_model = AutoModelForCausalLM.from_pretrained( + local_model_path, trust_remote_code=True + ) actor_model.to(torch.bfloat16) max_prompt_length = 16 @@ -57,8 +66,12 @@ def main(): attention_mask = prompts["attention_mask"] from verl.utils.torch_functional import pad_sequence_to_length - input_ids = pad_sequence_to_length(input_ids, max_prompt_length, tokenizer.pad_token_id, left_pad=True).cuda() - attention_mask = pad_sequence_to_length(attention_mask, max_prompt_length, 0, left_pad=True).cuda() + input_ids = pad_sequence_to_length( + input_ids, max_prompt_length, tokenizer.pad_token_id, left_pad=True + ).cuda() + attention_mask = pad_sequence_to_length( + attention_mask, max_prompt_length, 0, left_pad=True + ).cuda() from transformers import GenerationConfig @@ -85,9 +98,15 @@ def main(): tensor_model_parallel_size = 4 from torch.distributed.device_mesh import init_device_mesh - device_mesh = init_device_mesh("cuda", mesh_shape=(world_size,), mesh_dim_names=["fsdp"]) + device_mesh = init_device_mesh( + "cuda", mesh_shape=(world_size,), mesh_dim_names=["fsdp"] + ) - mixed_precision = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32) + mixed_precision = MixedPrecision( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + ) fsdp_model = FSDP( actor_model, use_orig_params=True, @@ -101,13 +120,21 @@ def main(): ) FSDP.set_state_dict_type( - fsdp_model, state_dict_type=StateDictType.SHARDED_STATE_DICT, state_dict_config=ShardedStateDictConfig() + fsdp_model, + state_dict_type=StateDictType.SHARDED_STATE_DICT, + state_dict_config=ShardedStateDictConfig(), ) state_dict = fsdp_model.state_dict() sampling_params = SamplingParams( - temperature=0, top_p=1, n=1, max_tokens=response_length, logprobs=1, ignore_eos=True, detokenize=False + temperature=0, + top_p=1, + n=1, + max_tokens=response_length, + logprobs=1, + ignore_eos=True, + detokenize=False, ) print(actor_model_config) @@ -145,13 +172,20 @@ def main(): idx_list = [] batch_size = input_ids.shape[0] - pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + pad_token_id = ( + tokenizer.pad_token_id + if tokenizer.pad_token_id is not None + else tokenizer.eos_token_id + ) from verl.workers.rollout.vllm_rollout.vllm_rollout_spmd import _pre_process_inputs for i in range(batch_size): idx_list.append(_pre_process_inputs(pad_token_id, input_ids[i])) print("start generation") - outputs = llm.generate(prompt_token_ids=idx_list, sampling_params=sampling_params, use_tqdm=False) + outputs = llm.generate( + prompt_token_ids=idx_list, + sampling_params=sampling_params, + use_tqdm=False) vllm_output = outputs[0].cuda() if torch.distributed.get_rank() == 0: print(f"hf response: {tokenizer.batch_decode(response)}") diff --git a/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_chat_scheduler.py b/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_chat_scheduler.py index 93aca6a..3dee73e 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_chat_scheduler.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_chat_scheduler.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -60,14 +60,14 @@ def test_vllm_async_rollout_without_tool_calls(init_config): } ) - # =========================== 1. Init rollout manager =========================== + # =========================== 1. Init rollout manager ==================== async_rollout_manager = init_async_rollout_manager(init_config) # test sleep and wake_up async_rollout_manager.sleep() async_rollout_manager.wake_up() - # =========================== 2. Generate sequences =========================== + # =========================== 2. Generate sequences ===================== raw_prompts = [ [ { @@ -75,7 +75,12 @@ def test_vllm_async_rollout_without_tool_calls(init_config): "content": "Let's play a role playing game. Your name is Alice, your favorite color is blue.", } ], - [{"role": "user", "content": "Let's play a role playing game. Your name is Bob, your favorite color is red."}], + [ + { + "role": "user", + "content": "Let's play a role playing game. Your name is Bob, your favorite color is red.", + } + ], ] batch = DataProto( non_tensor_batch={ @@ -85,7 +90,8 @@ def test_vllm_async_rollout_without_tool_calls(init_config): result = async_rollout_manager.generate_sequences(prompts=batch) # check result - seq_len = result.batch["prompts"].size(1) + result.batch["responses"].size(1) + seq_len = result.batch["prompts"].size( + 1) + result.batch["responses"].size(1) assert len(result) == 2 assert result.batch["input_ids"].size(1) == seq_len assert result.batch["attention_mask"].size(1) == seq_len @@ -120,7 +126,9 @@ def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: schema = get_json_schema(self.get_current_temperature) return OpenAIFunctionToolSchema(**schema) - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: try: result = self.get_current_temperature(**parameters) return json.dumps(result), 0, {} @@ -133,7 +141,11 @@ def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: schema = get_json_schema(self.get_temperature_date) return OpenAIFunctionToolSchema(**schema) - def get_temperature_date(self, location: str, date: str, unit: str = "celsius"): + def get_temperature_date( + self, + location: str, + date: str, + unit: str = "celsius"): """Get temperature at a location and date. Args: @@ -151,7 +163,9 @@ def get_temperature_date(self, location: str, date: str, unit: str = "celsius"): "unit": unit, } - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: try: result = self.get_temperature_date(**parameters) return json.dumps(result), 0, {} @@ -171,7 +185,7 @@ def test_vllm_async_rollout_with_tool_calls(init_config): } ) - # =========================== 1. Init rollout manager =========================== + # =========================== 1. Init rollout manager ==================== tool_config = { "tools": [ { @@ -191,26 +205,27 @@ def test_vllm_async_rollout_with_tool_calls(init_config): init_config.actor_rollout_ref.rollout.multi_turn.tool_config_path = tool_config_path async_rollout_manager = init_async_rollout_manager(init_config) - # =========================== 2. Generate sequences =========================== - raw_prompts = [ - [ - {"role": "user", "content": "How are you?"}, - ], - [ - {"role": "user", "content": "What's the temperature in Los Angeles now?"}, - ], - [ - { - "role": "system", - "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.\n\n" - "Current Date: 2024-09-30", - }, - {"role": "user", "content": "What's the temperature in San Francisco now? How about tomorrow?"}, - ], - ] + # =========================== 2. Generate sequences ===================== + raw_prompts = [[{"role": "user", + "content": "How are you?"}, + ], + [{"role": "user", + "content": "What's the temperature in Los Angeles now?"}, + ], + [{"role": "system", + "content": "You are Qwen, created by Alibaba Cloud. You are a helpful assistant.\n\n" + "Current Date: 2024-09-30", + }, + {"role": "user", + "content": "What's the temperature in San Francisco now? How about tomorrow?", + }, + ], + ] batch = DataProto( non_tensor_batch={ - "raw_prompt": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object), + "raw_prompt": np.array( + [np.array(prompt) for prompt in raw_prompts], dtype=object + ), }, ) result = async_rollout_manager.generate_sequences(prompts=batch) @@ -228,14 +243,20 @@ def test_vllm_async_rollout_with_tool_calls(init_config): tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path) responses = result.batch["responses"] response_mask = result.batch["response_mask"] - assert responses.size() == response_mask.size(), f"{responses.size()} != {response_mask.size()}" + assert ( + responses.size() == response_mask.size() + ), f"{responses.size()} != {response_mask.size()}" # Decode responses with response_mask for i in range(len(responses)): valid_tokens = responses[i][response_mask[i].bool()] response_str = tokenizer.decode(valid_tokens) - assert "" not in response_str, f"found in response: {response_str}" - assert "" not in response_str, f"found in response: {response_str}" + assert ( + "" not in response_str + ), f"found in response: {response_str}" + assert ( + "" not in response_str + ), f"found in response: {response_str}" print(f"response: {response_str}") print("Test passed!") diff --git a/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_model_rope_scaling.py b/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_model_rope_scaling.py index 30c9ae2..3054885 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_model_rope_scaling.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_model_rope_scaling.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -61,12 +61,16 @@ def test_vllm_rollout_with_yarn_position_embeddings(): } ) - tokenizer = AutoTokenizer.from_pretrained(config.model_path, trust_remote_code=True, padding_side="left") + tokenizer = AutoTokenizer.from_pretrained( + config.model_path, trust_remote_code=True, padding_side="left" + ) tokenizer.pad_token = tokenizer.eos_token model_hf_config = AutoConfig.from_pretrained(config.model_path) # do_sample=False for temperate=0 deterministic - input_dataproto = prepare_input_dataproto(tokenizer, config, validate=True, do_sample=False) + input_dataproto = prepare_input_dataproto( + tokenizer, config, validate=True, do_sample=False + ) vllm_rollout = vLLMRollout( model_path=config.model_path, @@ -80,11 +84,14 @@ def test_vllm_rollout_with_yarn_position_embeddings(): ) if rank == 0: print("VLLM Rollout Outputs:") - print(tokenizer.batch_decode(rollout_response.batch["responses"][:], skip_special_tokens=False)) + print( + tokenizer.batch_decode( + rollout_response.batch["responses"][:], + skip_special_tokens=False)) for response in rollout_response.batch["responses"]: - assert "<|im_end|>" in tokenizer.decode(response, skip_special_tokens=False), ( - "Response should contain <|im_end|> token" - ) + assert "<|im_end|>" in tokenizer.decode( + response, skip_special_tokens=False + ), "Response should contain <|im_end|> token" print("Checks passed.") del vllm_rollout @@ -99,20 +106,33 @@ def prepare_input_dataproto(tokenizer, config, validate, do_sample=False): base_phrase = "Roses are red, sky is blue. " * 4096 preencode_prompts = [ # 32810 tokens > 32768 tokens - [{"role": "user", "content": base_phrase + "Who won the Champions League in 2019?"}], + [ + { + "role": "user", + "content": base_phrase + "Who won the Champions League in 2019?", + } + ], [{"role": "user", "content": base_phrase + "The founder of Apple is"}], [{"role": "user", "content": base_phrase + "What's your name"}], ] formatted_prompts = [ - tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True) + tokenizer.apply_chat_template( + conversation, tokenize=False, add_generation_prompt=True + ) for conversation in preencode_prompts ] - prompts = tokenizer(formatted_prompts, return_tensors="pt", padding="max_length", max_length=config.prompt_length) + prompts = tokenizer( + formatted_prompts, + return_tensors="pt", + padding="max_length", + max_length=config.prompt_length, + ) input_dataproto = DataProto.from_dict( { "input_ids": prompts["input_ids"], "attention_mask": prompts["attention_mask"], - "position_ids": compute_position_id_with_mask(prompts["attention_mask"]), + "position_ids": compute_position_id_with_mask( + prompts["attention_mask"]), }, meta_info={ "bos_token_id": tokenizer.bos_token_id, diff --git a/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_spmd.py b/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_spmd.py index c2b8f51..e497a1c 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_spmd.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/rollout_vllm/test_vllm_spmd.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,7 +18,11 @@ import torch from torch.distributed.fsdp import CPUOffload, MixedPrecision from torch.distributed.fsdp import FullyShardedDataParallel as FSDP -from torch.distributed.fsdp.api import ShardedStateDictConfig, ShardingStrategy, StateDictType +from torch.distributed.fsdp.api import ( + ShardedStateDictConfig, + ShardingStrategy, + StateDictType, +) from transformers import AutoModelForCausalLM, AutoTokenizer from vllm import LLM, SamplingParams @@ -38,7 +42,8 @@ def levenshtein(s1, s2): # Compute the Levenshtein distance matrix for i in range(1, m + 1): for j in range(1, n + 1): - cost = 0 if s1[i - 1] == s2[j - 1] else 1 # No cost if characters match + # No cost if characters match + cost = 0 if s1[i - 1] == s2[j - 1] else 1 dp[i][j] = min( dp[i - 1][j] + 1, # Deletion dp[i][j - 1] + 1, # Insertion @@ -60,7 +65,8 @@ def are_lists_similar(a, b): total_length += max_len diff = levenshtein(s1, s2) total_diff += diff - print(f"Comparing strings:\n{s1}\n{s2}\nDifference: {diff} characters\n") + print( + f"Comparing strings:\n{s1}\n{s2}\nDifference: {diff} characters\n") percentage_difference = (total_diff / total_length) * 100 print(f"Total difference: {percentage_difference:.2f}%") @@ -70,7 +76,9 @@ def are_lists_similar(a, b): @pytest.mark.skip("https://github.com/vllm-project/vllm/issues/16993") def test_vllm_spmd(): - assert torch.cuda.device_count() >= 2, "At least 2 GPUs is required to run tp+dp tests." + assert ( + torch.cuda.device_count() >= 2 + ), "At least 2 GPUs is required to run tp+dp tests." local_rank, rank, world_size = initialize_global_process_group() # Initialize model and token @@ -80,9 +88,13 @@ def test_vllm_spmd(): from verl.utils.fs import copy_to_local local_model_path = copy_to_local(src=hdfs_path, cache_dir=local_cache_path) - tokenizer = AutoTokenizer.from_pretrained(local_model_path, padding_side="left", trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained( + local_model_path, padding_side="left", trust_remote_code=True + ) - actor_model = AutoModelForCausalLM.from_pretrained(local_model_path, trust_remote_code=True) + actor_model = AutoModelForCausalLM.from_pretrained( + local_model_path, trust_remote_code=True + ) actor_model.to(torch.bfloat16) # fill rollout config @@ -98,8 +110,12 @@ def test_vllm_spmd(): input_ids = prompts["input_ids"] attention_mask = prompts["attention_mask"] - input_ids = pad_sequence_to_length(input_ids, max_prompt_length, tokenizer.pad_token_id, left_pad=True) - attention_mask = pad_sequence_to_length(attention_mask, max_prompt_length, 0, left_pad=True) + input_ids = pad_sequence_to_length( + input_ids, max_prompt_length, tokenizer.pad_token_id, left_pad=True + ) + attention_mask = pad_sequence_to_length( + attention_mask, max_prompt_length, 0, left_pad=True + ) print("start generation") input_ids = input_ids.cuda() @@ -108,16 +124,27 @@ def test_vllm_spmd(): temperature = 0 top_p = 1 kwargs = dict( - n=1, temperature=temperature, top_p=top_p, max_tokens=max_response_length, logprobs=1, ignore_eos=True + n=1, + temperature=temperature, + top_p=top_p, + max_tokens=max_response_length, + logprobs=1, + ignore_eos=True, ) tensor_parallel_size = 4 from torch.distributed.device_mesh import init_device_mesh - device_mesh = init_device_mesh("cuda", mesh_shape=(world_size,), mesh_dim_names=["fsdp"]) + device_mesh = init_device_mesh( + "cuda", mesh_shape=(world_size,), mesh_dim_names=["fsdp"] + ) - mixed_precision = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32) + mixed_precision = MixedPrecision( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + ) fsdp_model = FSDP( actor_model, @@ -132,7 +159,9 @@ def test_vllm_spmd(): ) FSDP.set_state_dict_type( - fsdp_model, state_dict_type=StateDictType.SHARDED_STATE_DICT, state_dict_config=ShardedStateDictConfig() + fsdp_model, + state_dict_type=StateDictType.SHARDED_STATE_DICT, + state_dict_config=ShardedStateDictConfig(), ) state_dict = fsdp_model.state_dict() @@ -153,7 +182,9 @@ def test_vllm_spmd(): seed=1, ) - outputs = llm.generate(preencode_prompts, sampling_params=sampling_params, use_tqdm=False) + outputs = llm.generate( + preencode_prompts, sampling_params=sampling_params, use_tqdm=False + ) vllm_response_tokens = [] for output in outputs: generated_text = output.outputs[0].text @@ -162,10 +193,15 @@ def test_vllm_spmd(): world_size = torch.distributed.get_world_size() model = llm.llm_engine.model_executor.driver_worker.worker.model_runner.model model.load_weights( - ((name, param.full_tensor() if world_size != 1 else param) for name, param in state_dict.items()) + ( + (name, param.full_tensor() if world_size != 1 else param) + for name, param in state_dict.items() + ) ) - outputs = llm.generate(preencode_prompts, sampling_params=sampling_params, use_tqdm=False) + outputs = llm.generate( + preencode_prompts, sampling_params=sampling_params, use_tqdm=False + ) verl_vllm_response_tokens = [] for output in outputs: generated_text = output.outputs[0].text @@ -174,7 +210,9 @@ def test_vllm_spmd(): if torch.distributed.get_rank() == 0: print(f"vllm response: {vllm_response_tokens}") print(f"verl-vllm response: {verl_vllm_response_tokens}") - assert are_lists_similar(vllm_response_tokens, verl_vllm_response_tokens), "Strings differ more than 10%:\n" + assert are_lists_similar( + vllm_response_tokens, verl_vllm_response_tokens + ), "Strings differ more than 10%:\n" print("Check Pass") torch.distributed.destroy_process_group() diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_async_sglang_server.py b/Agent0/executor_train/verl/tests/workers/rollout/test_async_sglang_server.py index 0b4e914..8ee5fd1 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_async_sglang_server.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_async_sglang_server.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -21,7 +21,9 @@ @patch.dict( "sys.modules", { - "verl.workers.rollout.sglang_rollout.sglang_rollout": MagicMock(SGLangRollout=MagicMock()), + "verl.workers.rollout.sglang_rollout.sglang_rollout": MagicMock( + SGLangRollout=MagicMock() + ), }, ) class TestAsyncSglangServer: @@ -30,10 +32,19 @@ def server_config(self): return DictConfig({"rollout": {"tensor_model_parallel_size": 2}}) @pytest.mark.asyncio - @patch("verl.workers.rollout.sglang_rollout.async_sglang_server.ray.util.list_named_actors") - @patch("verl.workers.rollout.async_server.AsyncServerBase._start_fastapi_server", new_callable=AsyncMock) - @pytest.mark.filterwarnings("ignore:Ray state API is no longer experimental:DeprecationWarning") - async def test_init_engine(self, mock_start_fastapi_server, mock_list_actors, server_config): + @patch( + "verl.workers.rollout.sglang_rollout.async_sglang_server.ray.util.list_named_actors" + ) + @patch( + "verl.workers.rollout.async_server.AsyncServerBase._start_fastapi_server", + new_callable=AsyncMock, + ) + @pytest.mark.filterwarnings( + "ignore:Ray state API is no longer experimental:DeprecationWarning" + ) + async def test_init_engine( + self, mock_start_fastapi_server, mock_list_actors, server_config + ): mock_list_actors.return_value = [ {"name": "test_prefixWorkerDict_1:0", "namespace": "test"}, {"name": "test_prefixWorkerDict_1:1", "namespace": "test"}, @@ -44,7 +55,9 @@ async def test_init_engine(self, mock_start_fastapi_server, mock_list_actors, se {"name": "test_prefixWorkerDict_0:2", "namespace": "test"}, {"name": "test_prefixWorkerDict_0:3", "namespace": "test"}, ] - from verl.workers.rollout.sglang_rollout.async_sglang_server import AsyncSglangServer + from verl.workers.rollout.sglang_rollout.async_sglang_server import ( + AsyncSglangServer, + ) ActualClassToInstantiate = AsyncSglangServer if hasattr(AsyncSglangServer, "__ray_metadata__") and hasattr( @@ -64,7 +77,8 @@ def getitem_mock(key): if key == "name": return name # Use 'name' here # For other keys, return a new MagicMock to mimic default behavior or raise KeyError - # Returning a MagicMock is consistent with the original error's cause for unmocked keys + # Returning a MagicMock is consistent with the original error's + # cause for unmocked keys return MagicMock(name=f"mock.__getitem__('{key}')") actor_mock.__getitem__.side_effect = getitem_mock @@ -77,7 +91,8 @@ def getitem_mock(key): side_effect=mock_get_actor_side_effect, ): # Instance 1 - instance = ActualClassToInstantiate(server_config, 4, 0, "test_prefix") + instance = ActualClassToInstantiate( + server_config, 4, 0, "test_prefix") await instance.init_engine() assert len(instance.workers) == 2 @@ -86,7 +101,8 @@ def getitem_mock(key): assert instance.workers[1].name == "test_prefixWorkerDict_0:1" # Instance 2 - instance = ActualClassToInstantiate(server_config, 4, 1, "test_prefix") + instance = ActualClassToInstantiate( + server_config, 4, 1, "test_prefix") await instance.init_engine() assert len(instance.workers) == 2 @@ -95,7 +111,8 @@ def getitem_mock(key): assert instance.workers[1].name == "test_prefixWorkerDict_0:3" # Instance 3 - instance = ActualClassToInstantiate(server_config, 4, 3, "test_prefix") + instance = ActualClassToInstantiate( + server_config, 4, 3, "test_prefix") await instance.init_engine() assert len(instance.workers) == 2 diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_custom_completion_callback.py b/Agent0/executor_train/verl/tests/workers/rollout/test_custom_completion_callback.py index 495bce9..afe9717 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_custom_completion_callback.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_custom_completion_callback.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -35,7 +35,10 @@ from verl.protocol import DataProto from verl.utils import hf_tokenizer from verl.utils.reward_score.sandbox_fusion.utils import _process_single_case -from verl.workers.rollout.chat_scheduler import ChatCompletionScheduler, ToolCompletionCallback +from verl.workers.rollout.chat_scheduler import ( + ChatCompletionScheduler, + ToolCompletionCallback, +) def _get_free_port(): @@ -63,13 +66,18 @@ async def code_execution(self, request: Request): code = request_json["code"] print(f"execute code:\n{code}") - _, temp_file = tempfile.mkstemp(suffix=".py", prefix="temp_code", dir=None, text=True) + _, temp_file = tempfile.mkstemp( + suffix=".py", prefix="temp_code", dir=None, text=True + ) with open(temp_file, "w") as f: f.write(code) try: process = await asyncio.create_subprocess_exec( - sys.executable, temp_file, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE + sys.executable, + temp_file, + stdout=asyncio.subprocess.PIPE, + stderr=asyncio.subprocess.PIPE, ) stdout, stderr = await process.communicate() @@ -97,14 +105,21 @@ async def lifespan(app: fastapi.FastAPI): self.server_ready.set() yield - print("FastAPI shutdown, maybe address already in use, exit process immediately.") + print( + "FastAPI shutdown, maybe address already in use, exit process immediately." + ) os._exit(-1) app = fastapi.FastAPI(lifespan=lifespan) - app.router.add_api_route("/run_code", self.code_execution, methods=["POST"]) + app.router.add_api_route( + "/run_code", + self.code_execution, + methods=["POST"]) self.port = _get_free_port() - config = uvicorn.Config(app, host=["::", "0.0.0.0"], port=self.port, log_level="warning") + config = uvicorn.Config( + app, host=["::", "0.0.0.0"], port=self.port, log_level="warning" + ) server = uvicorn.Server(config) await server.serve() @@ -120,13 +135,17 @@ def __init__(self, config: DictConfig, scheduler: ChatCompletionScheduler): self.max_assistant_turns = 16 self.answer_pattern = re.compile(r"(.*?)", re.DOTALL) - self.code_pattern = re.compile(r"\s*```python(.*?)```\s*", re.DOTALL) + self.code_pattern = re.compile( + r"\s*```python(.*?)```\s*", re.DOTALL + ) self.sandbox_fusion_url = config.reward_model.sandbox_fusion.url self.default_timeout = 10 self.memory_limit_mb = config.reward_model.sandbox_fusion.memory_limit_mb # TODO: support asyncio executor - self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=max(32, os.cpu_count() * 5)) + self.executor = concurrent.futures.ThreadPoolExecutor( + max_workers=max(32, os.cpu_count() * 5) + ) async def sandbox_code_execution(self, code: str) -> dict[str, Any]: loop = asyncio.get_running_loop() @@ -153,7 +172,12 @@ def extra_body(self): } return extra - async def __call__(self, messages: list[dict[str, str]], completions: ChatCompletion, info: dict[str, Any]): + async def __call__( + self, + messages: list[dict[str, str]], + completions: ChatCompletion, + info: dict[str, Any], + ): role, content, finish_reason = ( completions.choices[0].message.role, completions.choices[0].message.content, @@ -164,24 +188,33 @@ async def __call__(self, messages: list[dict[str, str]], completions: ChatComple # STEP 0: check if we reach max turns if len(messages) >= self.max_assistant_turns: - print(f"[id={completions.id},turn={turn},finish_reason={finish_reason}] Reach max turns, done!") + print( + f"[id={ + completions.id},turn={turn},finish_reason={finish_reason}] Reach max turns, done!") return # STEP 1: check if we reach max tokens if finish_reason == "length": - print(f"[id={completions.id},turn={turn},finish_reason={finish_reason}] Reach max tokens, done!") + print( + f"[id={ + completions.id},turn={turn},finish_reason={finish_reason}] Reach max tokens, done!") return # STEP 2: check if we got answer matches = self.answer_pattern.findall(content) if matches: - print(f"[id={completions.id},turn={turn},finish_reason={finish_reason}] Got answer: {matches[0]}, done!") + print( + f"[id={ + completions.id},turn={turn},finish_reason={finish_reason}] Got answer: { + matches[0]}, done!") return # STEP 3: check if we got code block matches = self.code_pattern.findall(content) if not matches: - print(f"[id={completions.id},turn={turn},finish_reason={finish_reason}] No code block found, done!") + print( + f"[id={ + completions.id},turn={turn},finish_reason={finish_reason}] No code block found, done!") return # STEP 4: execute code block in sandbox @@ -195,8 +228,12 @@ async def __call__(self, messages: list[dict[str, str]], completions: ChatComple return stdout, stderr = metadata["stdout"], metadata["stderr"] - messages.append({"role": "tool", "content": f"{stdout}{stderr}"}) - print(f"[id={completions.id},turn={turn},finish_reason={finish_reason}] Code block executed, continue...") + messages.append( + {"role": "tool", "content": f"{stdout}{stderr}"} + ) + print( + f"[id={ + completions.id},turn={turn},finish_reason={finish_reason}] Code block executed, continue...") # STEP 5: resubmit chat completions with code block output self.scheduler.submit_chat_completions( @@ -218,10 +255,10 @@ async def __call__(self, messages: list[dict[str, str]], completions: ChatComple ``` -The code must explictly print necessary output to stdout. Remember stop generation at immediately and +The code must explictly print necessary output to stdout. Remember stop generation at immediately and return the code. 2. User will send the python code to a external sandbox to execute and get output from stdout. -3. User will send the output in format output to you, and you should use the +3. User will send the output in format output to you, and you should use the output to answer the question. The answer format must be: \\boxed{'The final answer goes here.'} @@ -273,7 +310,14 @@ async def __call__(self, messages: list[dict[str, str]], completions: ChatComple non_tensor_batch={ "raw_prompt": np.array( [ - [{"role": "user", "content": user_prompt_template.replace("{question}", problem)}] + [ + { + "role": "user", + "content": user_prompt_template.replace( + "{question}", problem + ), + } + ] for problem in dataset["Problem"] ] ), @@ -292,14 +336,20 @@ async def __call__(self, messages: list[dict[str, str]], completions: ChatComple tokenizer = hf_tokenizer(config.actor_rollout_ref.model.path) responses = result.batch["responses"] response_mask = result.batch["response_mask"] - assert responses.size() == response_mask.size(), f"{responses.size()} != {response_mask.size()}" + assert ( + responses.size() == response_mask.size() + ), f"{responses.size()} != {response_mask.size()}" # Decode responses with response_mask for i in range(len(responses)): valid_tokens = responses[i][response_mask[i].bool()] response_str = tokenizer.decode(valid_tokens) - assert "" not in response_str, f"found in response: {response_str}" - assert "" not in response_str, f"found in response: {response_str}" + assert ( + "" not in response_str + ), f"found in response: {response_str}" + assert ( + "" not in response_str + ), f"found in response: {response_str}" print(f"response: {response_str}") print("Test passed!") diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_hf_rollout.py b/Agent0/executor_train/verl/tests/workers/rollout/test_hf_rollout.py index 3eb6f4b..51b240a 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_hf_rollout.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_hf_rollout.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,7 +18,11 @@ from omegaconf import OmegaConf from torch.distributed.fsdp import CPUOffload, MixedPrecision from torch.distributed.fsdp import FullyShardedDataParallel as FSDP -from torch.distributed.fsdp.api import ShardedStateDictConfig, ShardingStrategy, StateDictType +from torch.distributed.fsdp.api import ( + ShardedStateDictConfig, + ShardingStrategy, + StateDictType, +) from transformers import AutoModelForCausalLM, AutoTokenizer from verl import DataProto @@ -52,15 +56,23 @@ def prepare_input_dataproto(tokenizer, config, validate): [{"role": "user", "content": "What's your name"}], ] formatted_prompts = [ - tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True) + tokenizer.apply_chat_template( + conversation, tokenize=False, add_generation_prompt=True + ) for conversation in preencode_prompts ] - prompts = tokenizer(formatted_prompts, return_tensors="pt", padding="max_length", max_length=config.prompt_length) + prompts = tokenizer( + formatted_prompts, + return_tensors="pt", + padding="max_length", + max_length=config.prompt_length, + ) input_dataproto = DataProto.from_dict( { "input_ids": prompts["input_ids"], "attention_mask": prompts["attention_mask"], - "position_ids": compute_position_id_with_mask(prompts["attention_mask"]), + "position_ids": compute_position_id_with_mask( + prompts["attention_mask"]), }, meta_info={ "bos_token_id": tokenizer.bos_token_id, @@ -75,9 +87,15 @@ def prepare_input_dataproto(tokenizer, config, validate): def prepare_fsdp_model(model, world_size): from torch.distributed.device_mesh import init_device_mesh - device_mesh = init_device_mesh("cuda", mesh_shape=(world_size,), mesh_dim_names=["fsdp"]) + device_mesh = init_device_mesh( + "cuda", mesh_shape=(world_size,), mesh_dim_names=["fsdp"] + ) - mixed_precision = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32) + mixed_precision = MixedPrecision( + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, + ) fsdp_model = FSDP( model, @@ -92,16 +110,23 @@ def prepare_fsdp_model(model, world_size): ) FSDP.set_state_dict_type( - fsdp_model, state_dict_type=StateDictType.SHARDED_STATE_DICT, state_dict_config=ShardedStateDictConfig() + fsdp_model, + state_dict_type=StateDictType.SHARDED_STATE_DICT, + state_dict_config=ShardedStateDictConfig(), ) return fsdp_model -def test_hf_rollout(n: int = 1, do_sample: bool = True, validate: bool = False): +def test_hf_rollout( + n: int = 1, + do_sample: bool = True, + validate: bool = False): config = OmegaConf.create(BASE_HF_ROLLOUT_CONFIG) config.update({"n": n, "do_sample": do_sample}) - assert torch.cuda.device_count() >= 2, "At least 2 GPUs is required to run tp+dp tests." + assert ( + torch.cuda.device_count() >= 2 + ), "At least 2 GPUs is required to run tp+dp tests." local_rank, rank, world_size = initialize_global_process_group() # Initialize model and tokenizer @@ -109,17 +134,23 @@ def test_hf_rollout(n: int = 1, do_sample: bool = True, validate: bool = False): local_cache_path = os.path.expanduser(local_cache_path) hdfs_path = "Qwen/Qwen2-7B-Instruct" local_model_path = copy_to_local(src=hdfs_path, cache_dir=local_cache_path) - tokenizer = AutoTokenizer.from_pretrained(local_model_path, padding_side="left", trust_remote_code=True) + tokenizer = AutoTokenizer.from_pretrained( + local_model_path, padding_side="left", trust_remote_code=True + ) tokenizer.pad_token = tokenizer.eos_token # Initialize FSDP model - actor_model = AutoModelForCausalLM.from_pretrained(local_model_path, trust_remote_code=True) + actor_model = AutoModelForCausalLM.from_pretrained( + local_model_path, trust_remote_code=True + ) actor_model.to(torch.bfloat16) fsdp_model = prepare_fsdp_model(actor_model, world_size) # Initialize HFRollout and start generate hf_rollout = HFRollout(fsdp_model, OmegaConf.create(config)) - input = prepare_input_dataproto(tokenizer, config, validate).to(torch.cuda.current_device()) + input = prepare_input_dataproto(tokenizer, config, validate).to( + torch.cuda.current_device() + ) outputs = hf_rollout.generate_sequences(input) # check generated batch size is expected @@ -130,12 +161,14 @@ def test_hf_rollout(n: int = 1, do_sample: bool = True, validate: bool = False): prompt_tokens = outputs.batch["prompts"][i] prompt_mask = prompt_tokens != tokenizer.pad_token_id prompt_tokens = prompt_tokens[prompt_mask] - decoded_prompt = tokenizer.decode(prompt_tokens, skip_special_tokens=False) + decoded_prompt = tokenizer.decode( + prompt_tokens, skip_special_tokens=False) response_tokens = outputs.batch["responses"][i] response_mask = response_tokens != tokenizer.pad_token_id response_tokens = response_tokens[response_mask] - decoded_response = tokenizer.decode(response_tokens, skip_special_tokens=False) + decoded_response = tokenizer.decode( + response_tokens, skip_special_tokens=False) attention_mask = outputs.batch["attention_mask"][i] position_ids = outputs.batch["position_ids"][i] @@ -147,16 +180,22 @@ def test_hf_rollout(n: int = 1, do_sample: bool = True, validate: bool = False): # check response attention mask is expected response_attention = attention_mask[prompt_length:] - eos_positions = (outputs.batch["responses"][i] == tokenizer.pad_token_id).nonzero(as_tuple=True)[0] + eos_positions = ( + outputs.batch["responses"][i] == tokenizer.pad_token_id + ).nonzero(as_tuple=True)[0] if len(eos_positions) > 0: first_eos_pos = eos_positions[0].item() - assert response_attention[: first_eos_pos + 1].all(), "Response attention mask should be 1 until EOS" + assert response_attention[ + : first_eos_pos + 1 + ].all(), "Response attention mask should be 1 until EOS" if first_eos_pos + 1 < response_length: - assert not response_attention[first_eos_pos + 1 :].any(), ( - "Response attention mask should be 0 after EOS" - ) + assert not response_attention[ + first_eos_pos + 1: + ].any(), "Response attention mask should be 0 after EOS" else: - assert response_attention.all(), "Response attention mask should be all 1 if no EOS token" + assert ( + response_attention.all() + ), "Response attention mask should be all 1 if no EOS token" # check response position ids is expected prompt_positions = position_ids[:prompt_length] diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_mcp_tools.py b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_mcp_tools.py index 387de16..176a170 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_mcp_tools.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_mcp_tools.py @@ -1,6 +1,6 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -29,7 +29,11 @@ from verl.protocol import DataProto from verl.tools.mcp_search_tool import MCPSearchTool from verl.tools.utils.mcp_clients.McpClientManager import MCPClientManager -from verl.workers.rollout.schemas import AsyncRolloutRequest, AsyncRolloutRequestStateEnum, Message +from verl.workers.rollout.schemas import ( + AsyncRolloutRequest, + AsyncRolloutRequestStateEnum, + Message, +) from verl.workers.rollout.sglang_rollout.sglang_rollout import SGLangRollout DEFAULT_USER_CONTENT_PREFIX = ( @@ -40,9 +44,9 @@ ". You can search as many times as your want. If you find no " "further external knowledge needed, you can directly provide the answer inside " " and , without detailed illustrations. For example, " - " Beijing . Question: " -) -user_content = DEFAULT_USER_CONTENT_PREFIX.rstrip("\n") + "How's the weather lately?" + " Beijing . Question: ") +user_content = DEFAULT_USER_CONTENT_PREFIX.rstrip( + "\n") + "How's the weather lately?" def get_search_messages(): @@ -100,14 +104,20 @@ def get_search_messages(): } # Mock search tool responses - tool_return_0_msg = {"role": "tool", "content": [{"type": "text", "text": "Today's weather in Beijing is sunny."}]} + tool_return_0_msg = {"role": "tool", "content": [ + {"type": "text", "text": "Today's weather in Beijing is sunny."}], } tool_return_1_msg = { "role": "tool", - "content": [{"type": "text", "text": "Tomorrow's weather in Beijing is cloudy."}], + "content": [ + {"type": "text", "text": "Tomorrow's weather in Beijing is cloudy."} + ], } user_prompts = [user_prompt] - expect_turn_array = [expect_turn_0_msg, expect_turn_1_msg, expect_turn_2_msg] + expect_turn_array = [ + expect_turn_0_msg, + expect_turn_1_msg, + expect_turn_2_msg] tool_return_array = [tool_return_0_msg, tool_return_1_msg] return user_prompts, expect_turn_array, tool_return_array @@ -117,7 +127,8 @@ class TestRolloutWithMCPSearchTools: @pytest.fixture def qwen_tokenizer(self): local_model_path = "Qwen/Qwen2.5-0.5B" - tokenizer = AutoTokenizer.from_pretrained(local_model_path, padding_side="left") + tokenizer = AutoTokenizer.from_pretrained( + local_model_path, padding_side="left") tokenizer.pad_token = tokenizer.eos_token return tokenizer @@ -133,11 +144,15 @@ def search_data(self, qwen_tokenizer): user_prompt, expect_turn_array, tool_return_array = get_search_messages() prompts = [[message] for message in user_prompt] preencode_turn_array = [ - qwen_tokenizer.apply_chat_template([turn], tokenize=False, add_generation_prompt=False) + qwen_tokenizer.apply_chat_template( + [turn], tokenize=False, add_generation_prompt=False + ) for turn in expect_turn_array ] preencode_tool_return_array = [ - qwen_tokenizer.apply_chat_template([turn], tokenize=False, add_generation_prompt=True) + qwen_tokenizer.apply_chat_template( + [turn], tokenize=False, add_generation_prompt=True + ) for turn in tool_return_array ] return prompts, preencode_turn_array, preencode_tool_return_array @@ -150,7 +165,11 @@ def search_rollout_config(self): tensor_parallel_size = 1 tool_path = "./resource/tool_configs/mcp_tool_config" rollout_config = get_rollout_config( - max_response_length, max_prompt_length, dtype, tensor_parallel_size, tool_path + max_response_length, + max_prompt_length, + dtype, + tensor_parallel_size, + tool_path, ) return rollout_config @@ -158,10 +177,14 @@ def search_rollout_config(self): def search_data_proto(self, search_data, qwen_tokenizer): preencode_prompts, _, _ = search_data prompts = [ - qwen_tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True) + qwen_tokenizer.apply_chat_template( + message, tokenize=False, add_generation_prompt=True + ) for message in preencode_prompts ] - input_ids, attention_mask, position_ids = prepare_inputs(qwen_tokenizer, prompts, 1000) + input_ids, attention_mask, position_ids = prepare_inputs( + qwen_tokenizer, prompts, 1000 + ) prompt_dict = TensorDict( { "input_ids": input_ids, @@ -176,7 +199,9 @@ def search_data_proto(self, search_data, qwen_tokenizer): [ { "tavily_search_tool": { - "create_kwargs": {"ground_truth": "Today is sunny and tomorrow will be cloudy in Beijing."}, + "create_kwargs": { + "ground_truth": "Today is sunny and tomorrow will be cloudy in Beijing." + }, }, } ], @@ -184,12 +209,21 @@ def search_data_proto(self, search_data, qwen_tokenizer): ) index = np.array([0], dtype=object) prompts = DataProto( - batch=prompt_dict, non_tensor_batch={"raw_prompt": messages, "tools_kwargs": tools_kwargs, "index": index} + batch=prompt_dict, + non_tensor_batch={ + "raw_prompt": messages, + "tools_kwargs": tools_kwargs, + "index": index, + }, ) return prompts @pytest.fixture - def mock_rollout(self, search_rollout_config, qwen_tokenizer, qwen_model_config): + def mock_rollout( + self, + search_rollout_config, + qwen_tokenizer, + qwen_model_config): """Mock the rollout instance with sampling_params initialized.""" tool_schema = [ { @@ -263,7 +297,9 @@ def mock_rollout(self, search_rollout_config, qwen_tokenizer, qwen_model_config) } ] with ( - patch.object(MCPClientManager, "fetch_tool_schemas", return_value=tool_schema), + patch.object( + MCPClientManager, "fetch_tool_schemas", return_value=tool_schema + ), patch.object(SGLangRollout, "_init_distributed_env", return_value=None), patch.object(SGLangRollout, "_init_inference_engine", return_value=None), patch.object(SGLangRollout, "_init_sampling_params", return_value=None), @@ -288,25 +324,36 @@ def test_tools_registration(self, mock_rollout): assert "tavily_search_tool" in mock_rollout._tool_map.keys() from verl.tools.mcp_search_tool import MCPSearchTool - assert isinstance(mock_rollout._tool_map["tavily_search_tool"], MCPSearchTool) + assert isinstance( + mock_rollout._tool_map["tavily_search_tool"], + MCPSearchTool) # depend on the tokenizer assert mock_rollout._tool_call_parser_type == "qwen25" def test_rollout_req_creation(self, mock_rollout, search_data_proto): - req_list = mock_rollout._preprocess_prompt_to_async_rollout_requests(search_data_proto, n=1) + req_list = mock_rollout._preprocess_prompt_to_async_rollout_requests( + search_data_proto, n=1 + ) assert len(req_list) == 1 assert req_list[0].state == AsyncRolloutRequestStateEnum.PENDING assert len(req_list[0].tool_schemas) == 1 - def test_over_size_case(self, mock_rollout, search_data_proto, search_data): + def test_over_size_case( + self, + mock_rollout, + search_data_proto, + search_data): mock_rollout.config.multi_turn.max_assistant_turns = 1 - req = mock_rollout._preprocess_prompt_to_async_rollout_requests(search_data_proto, n=1)[0] + req = mock_rollout._preprocess_prompt_to_async_rollout_requests( + search_data_proto, n=1 + )[0] req = MagicMock(wraps=req, spec=AsyncRolloutRequest) req.finalize = MagicMock() req_list = [req] _, expect_turn_array, _ = search_data - # here we mock a meta info with 'length'. indicate the response is truncate + # here we mock a meta info with 'length'. indicate the response is + # truncate mock_rollout._handle_engine_call = MagicMock() future = asyncio.Future() future.set_result( @@ -327,7 +374,10 @@ def test_over_size_case(self, mock_rollout, search_data_proto, search_data): loop = asyncio.get_event_loop() output_req_list = loop.run_until_complete( asyncio.gather( - *[mock_rollout._async_rollout_a_request(req, True, False) for req in req_list], + *[ + mock_rollout._async_rollout_a_request(req, True, False) + for req in req_list + ], ) ) assert len(output_req_list) == 1 @@ -343,26 +393,39 @@ def test_over_size_case(self, mock_rollout, search_data_proto, search_data): ) @patch.object(MCPSearchTool, "execute", new_callable=AsyncMock) - def test_tool_call_basic_case(self, mock_execute, mock_rollout, search_data_proto, search_data): + def test_tool_call_basic_case( + self, mock_execute, mock_rollout, search_data_proto, search_data + ): _, expect_turn_array, tool_return_array = search_data # Mock search tool execution to return predefined responses - mock_execute.side_effect = [(msg, 0.0, {"status": "success"}) for msg in tool_return_array] + mock_execute.side_effect = [ + (msg, 0.0, {"status": "success"}) for msg in tool_return_array + ] mock_rollout.config.multi_turn.max_assistant_turns = 10 - req = mock_rollout._preprocess_prompt_to_async_rollout_requests(search_data_proto, n=1)[0] + req = mock_rollout._preprocess_prompt_to_async_rollout_requests( + search_data_proto, n=1 + )[0] req = MagicMock(wraps=req, spec=AsyncRolloutRequest) req.finalize = MagicMock() req_list = [req] mock_rollout._handle_engine_call = MagicMock() futures = [asyncio.Future() for i in expect_turn_array] - for idx, (i, turn) in enumerate(zip(futures, expect_turn_array, strict=True)): + for idx, (i, turn) in enumerate( + zip(futures, expect_turn_array, strict=True)): i.set_result( { "text": turn, "meta_info": { "id": "d1188d81cba840359df5b352b344bc8e", - "finish_reason": {"type": "tool_calls" if idx < len(expect_turn_array) - 1 else "stop"}, + "finish_reason": { + "type": ( + "tool_calls" + if idx < len(expect_turn_array) - 1 + else "stop" + ) + }, "prompt_tokens": len(turn), "completion_tokens": 100, "cached_tokens": 0, @@ -372,14 +435,20 @@ def test_tool_call_basic_case(self, mock_execute, mock_rollout, search_data_prot ) if idx < len(expect_turn_array) - 1: assert mock_rollout._function_call_parser.has_tool_call(turn) - assert mock_rollout._function_call_parser.parse_non_stream(turn) + assert mock_rollout._function_call_parser.parse_non_stream( + turn) mock_rollout._handle_engine_call.side_effect = futures mock_rollout._tp_rank = 0 loop = asyncio.get_event_loop() output_req_list = loop.run_until_complete( - asyncio.gather(*[mock_rollout._async_rollout_a_request(req, True, False) for req in req_list]) + asyncio.gather( + *[ + mock_rollout._async_rollout_a_request(req, True, False) + for req in req_list + ] + ) ) # Verify conversation completed successfully with proper tool usage @@ -398,7 +467,9 @@ def test_tool_call_basic_case(self, mock_execute, mock_rollout, search_data_prot assert search_counter == 2 @patch.object(MCPSearchTool, "execute", new_callable=AsyncMock) - def test_tool_call_batch_case(self, mock_execute, mock_rollout, search_data_proto, search_data): + def test_tool_call_batch_case( + self, mock_execute, mock_rollout, search_data_proto, search_data + ): _, expect_turn_array, tool_return_array = search_data # Mock tool execution for large batch (100 requests * 2 calls each) mock_execute.side_effect = [ @@ -407,7 +478,9 @@ def test_tool_call_batch_case(self, mock_execute, mock_rollout, search_data_prot ] * 100 mock_rollout.config.multi_turn.max_assistant_turns = 10 - base_req = mock_rollout._preprocess_prompt_to_async_rollout_requests(search_data_proto, n=1)[0] + base_req = mock_rollout._preprocess_prompt_to_async_rollout_requests( + search_data_proto, n=1 + )[0] req_nums = 100 req_list = [] @@ -421,13 +494,21 @@ def test_tool_call_batch_case(self, mock_execute, mock_rollout, search_data_prot req_list.append(MagicMock(wraps=tmp_req, spec=AsyncRolloutRequest)) futures = [asyncio.Future() for _ in expect_turn_array] - for idx, (fut, turn) in enumerate(zip(futures, expect_turn_array, strict=True)): + for idx, (fut, turn) in enumerate( + zip(futures, expect_turn_array, strict=True) + ): fut.set_result( { "text": turn, "meta_info": { "id": "dummy", - "finish_reason": {"type": "tool_calls" if idx < len(expect_turn_array) - 1 else "stop"}, + "finish_reason": { + "type": ( + "tool_calls" + if idx < len(expect_turn_array) - 1 + else "stop" + ) + }, "prompt_tokens": len(turn), "completion_tokens": 100, }, @@ -436,16 +517,27 @@ def test_tool_call_batch_case(self, mock_execute, mock_rollout, search_data_prot req_turns_map[i] = futures req_turns_counter[i] = 0 - async def hacked_handle_engine_call(self, _req: AsyncRolloutRequest, *_args, **_kwargs): - fut = req_turns_map[_req.batch_data_id][req_turns_counter[_req.batch_data_id]] + async def hacked_handle_engine_call( + self, _req: AsyncRolloutRequest, *_args, **_kwargs + ): + fut = req_turns_map[_req.batch_data_id][ + req_turns_counter[_req.batch_data_id] + ] req_turns_counter[_req.batch_data_id] += 1 return await fut - with patch.object(SGLangRollout, "_handle_engine_call", new=hacked_handle_engine_call): + with patch.object( + SGLangRollout, "_handle_engine_call", new=hacked_handle_engine_call + ): mock_rollout._tp_rank = 0 loop = asyncio.get_event_loop() output_req_list = loop.run_until_complete( - asyncio.gather(*[mock_rollout._async_rollout_a_request(r, True, False) for r in req_list]) + asyncio.gather( + *[ + mock_rollout._async_rollout_a_request(r, True, False) + for r in req_list + ] + ) ) # Verify all requests completed successfully diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_multimodal_delta.py b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_multimodal_delta.py index 47fefca..e0e0c6c 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_multimodal_delta.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_multimodal_delta.py @@ -1,5 +1,5 @@ # Copyright 2025 Amazon.com, Inc. or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -25,7 +25,9 @@ ) -def _test_add_tool_response_messages_image_delta(processor, image_list, description_list, resize_image=False): +def _test_add_tool_response_messages_image_delta( + processor, image_list, description_list, resize_image=False +): assert len(image_list) == len(description_list) # Get the smallest dimensions across all images processed_images = [] @@ -45,9 +47,7 @@ def _test_add_tool_response_messages_image_delta(processor, image_list, descript processed_images = processed_images_resized # Initial message history - system_prompt = ( - "You will be provided with an image. Describe this image and then generate a new image for the next round" - ) + system_prompt = "You will be provided with an image. Describe this image and then generate a new image for the next round" messages = [ { "role": "system", @@ -109,7 +109,10 @@ def _test_add_tool_response_messages_image_delta(processor, image_list, descript _ = req.get_generation_prompt_ids(processor) req.add_assistant_message(processor, content=description_list[idx - 1]) before_tool_call_len = req.input_ids.shape[-1] - req.add_tool_response_messages(processor, [{"image": [img], "text": "Here is the new image you requested: "}]) + req.add_tool_response_messages( + processor, + [{"image": [img], "text": "Here is the new image you requested: "}], + ) after_tool_call_len = req.input_ids.shape[-1] if prev_generated_len == 0: prev_generated_len = after_tool_call_len - before_tool_call_len @@ -122,7 +125,8 @@ def _test_add_tool_response_messages_image_delta(processor, image_list, descript req.add_assistant_message(processor, content=description_list[-1]) messages = [msg.model_dump() for msg in req.messages] - tools = [tool.model_dump() for tool in req.tool_schemas] if req.tool_schemas else None + tools = ([tool.model_dump() + for tool in req.tool_schemas] if req.tool_schemas else None) full_prompt_info = req._handle_apply_chat_template( processor, messages, @@ -142,46 +146,67 @@ def _test_add_tool_response_messages_image_delta(processor, image_list, descript full_prompt_multi_modal_inputs.pop("attention_mask", None) for key in full_prompt_multi_modal_inputs: - assert full_prompt_multi_modal_inputs[key].eq(req.multi_modal_inputs[key]).all() + assert full_prompt_multi_modal_inputs[key].eq( + req.multi_modal_inputs[key]).all() @pytest.mark.skipif( - hf_processor("Qwen/Qwen2.5-VL-3B-Instruct") is None, reason="Processor not available for Qwen/Qwen2.5-VL-B-Instruct" + hf_processor("Qwen/Qwen2.5-VL-3B-Instruct") is None, + reason="Processor not available for Qwen/Qwen2.5-VL-B-Instruct", ) def test_add_tool_response_messages_image_delta(): processor = hf_processor("Qwen/Qwen2.5-VL-3B-Instruct") # From Qwen2.5-VL-3B-Instruct HF example - img_1_url = {"image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"} + img_1_url = { + "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" + } img_1_description = "A woman sits on the beach at sunset, smiling as she shares a high five with her large dog." # GitHub Logo - img_2_url = {"image": "https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png"} + img_2_url = { + "image": "https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png"} img_2_description = "A GitHub Logo image" # Octocat - img_3_url = {"image": "https://octodex.github.com/images/orderedlistocat.png"} + img_3_url = { + "image": "https://octodex.github.com/images/orderedlistocat.png"} img_3_description = "An Octocat image" image_list = [img_1_url, img_2_url, img_3_url] - description_list = [img_1_description, img_2_description, img_3_description] - _test_add_tool_response_messages_image_delta(processor, image_list, description_list, resize_image=False) + description_list = [ + img_1_description, + img_2_description, + img_3_description] + _test_add_tool_response_messages_image_delta( + processor, image_list, description_list, resize_image=False + ) @pytest.mark.skipif( - hf_processor("Qwen/Qwen2.5-VL-3B-Instruct") is None, reason="Processor not available for Qwen/Qwen2.5-VL-B-Instruct" + hf_processor("Qwen/Qwen2.5-VL-3B-Instruct") is None, + reason="Processor not available for Qwen/Qwen2.5-VL-B-Instruct", ) def test_add_tool_response_messages_image_delta_resize_image(): processor = hf_processor("Qwen/Qwen2.5-VL-3B-Instruct") # From Qwen2.5-VL-3B-Instruct HF example - img_1_url = {"image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg"} + img_1_url = { + "image": "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg" + } img_1_description = "A woman sits on the beach at sunset, smiling as she shares a high five with her large dog." # GitHub Logo - img_2_url = {"image": "https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png"} + img_2_url = { + "image": "https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png"} img_2_description = "A GitHub Logo image" # Octocat - img_3_url = {"image": "https://octodex.github.com/images/orderedlistocat.png"} + img_3_url = { + "image": "https://octodex.github.com/images/orderedlistocat.png"} img_3_description = "An Octocat image" image_list = [img_1_url, img_2_url, img_3_url] - description_list = [img_1_description, img_2_description, img_3_description] - _test_add_tool_response_messages_image_delta(processor, image_list, description_list, resize_image=True) + description_list = [ + img_1_description, + img_2_description, + img_3_description] + _test_add_tool_response_messages_image_delta( + processor, image_list, description_list, resize_image=True + ) diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_search_tools.py b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_search_tools.py index 2400d5c..680b8b9 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_search_tools.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_search_tools.py @@ -1,5 +1,5 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,7 +33,11 @@ OpenAIFunctionToolSchema, ) from verl.tools.search_tool import SearchTool -from verl.workers.rollout.schemas import AsyncRolloutRequest, AsyncRolloutRequestStateEnum, Message +from verl.workers.rollout.schemas import ( + AsyncRolloutRequest, + AsyncRolloutRequestStateEnum, + Message, +) from verl.workers.rollout.sglang_rollout.sglang_rollout import SGLangRollout DEFAULT_USER_CONTENT_PREFIX = ( @@ -44,9 +48,9 @@ ". You can search as many times as your want. If you find no " "further external knowledge needed, you can directly provide the answer inside " " and , without detailed illustrations. For example, " - " Beijing . Question: " -) -user_content = DEFAULT_USER_CONTENT_PREFIX.rstrip("\n") + "How's the weather lately?" + " Beijing . Question: ") +user_content = DEFAULT_USER_CONTENT_PREFIX.rstrip( + "\n") + "How's the weather lately?" def get_search_messages(): @@ -58,14 +62,28 @@ def get_search_messages(): expect_turn_0_msg = { "role": "assistant", "content": "Let me search the web.", - "tool_calls": [{"type": "function", "function": {"name": "search", "arguments": {"query": "today's weather"}}}], + "tool_calls": [ + { + "type": "function", + "function": { + "name": "search", + "arguments": {"query": "today's weather"}, + }, + } + ], } expect_turn_1_msg = { "role": "assistant", "content": "Let me search again.", "tool_calls": [ - {"type": "function", "function": {"name": "search", "arguments": {"query": "tomorrow's weather"}}} + { + "type": "function", + "function": { + "name": "search", + "arguments": {"query": "tomorrow's weather"}, + }, + } ], } @@ -75,11 +93,20 @@ def get_search_messages(): } # Mock search tool responses - tool_return_0_msg = {"role": "tool", "content": "Today's weather in Beijing is sunny."} - tool_return_1_msg = {"role": "tool", "content": "Tomorrow's weather in Beijing is cloudy."} + tool_return_0_msg = { + "role": "tool", + "content": "Today's weather in Beijing is sunny.", + } + tool_return_1_msg = { + "role": "tool", + "content": "Tomorrow's weather in Beijing is cloudy.", + } user_prompts = [user_prompt] - expect_turn_array = [expect_turn_0_msg, expect_turn_1_msg, expect_turn_2_msg] + expect_turn_array = [ + expect_turn_0_msg, + expect_turn_1_msg, + expect_turn_2_msg] tool_return_array = [tool_return_0_msg, tool_return_1_msg] return user_prompts, expect_turn_array, tool_return_array @@ -89,7 +116,8 @@ class TestRolloutWithSearchTools: @pytest.fixture def qwen_tokenizer(self): local_model_path = "Qwen/Qwen2.5-0.5B" - tokenizer = AutoTokenizer.from_pretrained(local_model_path, padding_side="left") + tokenizer = AutoTokenizer.from_pretrained( + local_model_path, padding_side="left") tokenizer.pad_token = tokenizer.eos_token return tokenizer @@ -105,11 +133,15 @@ def search_data(self, qwen_tokenizer): user_prompt, expect_turn_array, tool_return_array = get_search_messages() prompts = [[message] for message in user_prompt] preencode_turn_array = [ - qwen_tokenizer.apply_chat_template([turn], tokenize=False, add_generation_prompt=False) + qwen_tokenizer.apply_chat_template( + [turn], tokenize=False, add_generation_prompt=False + ) for turn in expect_turn_array ] preencode_tool_return_array = [ - qwen_tokenizer.apply_chat_template([turn], tokenize=False, add_generation_prompt=True) + qwen_tokenizer.apply_chat_template( + [turn], tokenize=False, add_generation_prompt=True + ) for turn in tool_return_array ] return prompts, preencode_turn_array, preencode_tool_return_array @@ -122,7 +154,11 @@ def search_rollout_config(self): tensor_parallel_size = 1 tool_path = "./resource/tool_configs/search_tool_config" rollout_config = get_rollout_config( - max_response_length, max_prompt_length, dtype, tensor_parallel_size, tool_path + max_response_length, + max_prompt_length, + dtype, + tensor_parallel_size, + tool_path, ) return rollout_config @@ -130,10 +166,14 @@ def search_rollout_config(self): def search_data_proto(self, search_data, qwen_tokenizer): preencode_prompts, _, _ = search_data prompts = [ - qwen_tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True) + qwen_tokenizer.apply_chat_template( + message, tokenize=False, add_generation_prompt=True + ) for message in preencode_prompts ] - input_ids, attention_mask, position_ids = prepare_inputs(qwen_tokenizer, prompts, 1000) + input_ids, attention_mask, position_ids = prepare_inputs( + qwen_tokenizer, prompts, 1000 + ) prompt_dict = TensorDict( { "input_ids": input_ids, @@ -153,18 +193,26 @@ def search_data_proto(self, search_data, qwen_tokenizer): "data_source": "searchR1_nq", }, }, - } - ], + }], dtype=object, ) index = np.array([0], dtype=object) prompts = DataProto( - batch=prompt_dict, non_tensor_batch={"raw_prompt": messages, "tools_kwargs": tools_kwargs, "index": index} + batch=prompt_dict, + non_tensor_batch={ + "raw_prompt": messages, + "tools_kwargs": tools_kwargs, + "index": index, + }, ) return prompts @pytest.fixture - def mock_rollout(self, search_rollout_config, qwen_tokenizer, qwen_model_config): + def mock_rollout( + self, + search_rollout_config, + qwen_tokenizer, + qwen_model_config): """Mock the rollout instance with sampling_params initialized.""" with ( patch.object(SGLangRollout, "_init_distributed_env", return_value=None), @@ -190,7 +238,13 @@ def mock_rollout(self, search_rollout_config, qwen_tokenizer, qwen_model_config) @patch.object(SGLangRollout, "_init_inference_engine", return_value=None) @patch.object(SGLangRollout, "_init_sampling_params", return_value=None) def test_tools_registration( - self, mock_env, mock_engine, mock_sampling, search_rollout_config, qwen_tokenizer, qwen_model_config + self, + mock_env, + mock_engine, + mock_sampling, + search_rollout_config, + qwen_tokenizer, + qwen_model_config, ): rollout = SGLangRollout( actor_module="", @@ -225,7 +279,9 @@ def test_rollout_req_creation( processing_class=qwen_tokenizer, model_hf_config=qwen_model_config, ) - req_list = rollout._preprocess_prompt_to_async_rollout_requests(search_data_proto, n=1) + req_list = rollout._preprocess_prompt_to_async_rollout_requests( + search_data_proto, n=1 + ) assert len(req_list) == 1 assert req_list[0].state == AsyncRolloutRequestStateEnum.PENDING assert len(req_list[0].tool_schemas) == 1 @@ -242,18 +298,24 @@ def test_rollout_req_creation( type="array", description="A list of fully-formed semantic queries. The tool will return search " "results for each query.", - items={"type": "string"}, - ) - }, + items={ + "type": "string"}, + )}, required=["query_list"], ), strict=False, ), ) - def test_over_size_case(self, mock_rollout, search_data_proto, search_data): + def test_over_size_case( + self, + mock_rollout, + search_data_proto, + search_data): mock_rollout.config.multi_turn.max_assistant_turns = 1 - req = mock_rollout._preprocess_prompt_to_async_rollout_requests(search_data_proto, n=1)[0] + req = mock_rollout._preprocess_prompt_to_async_rollout_requests( + search_data_proto, n=1 + )[0] req = MagicMock(wraps=req, spec=AsyncRolloutRequest) req.finalize = MagicMock() req_list = [req] @@ -279,7 +341,10 @@ def test_over_size_case(self, mock_rollout, search_data_proto, search_data): loop = asyncio.get_event_loop() output_req_list = loop.run_until_complete( asyncio.gather( - *[mock_rollout._async_rollout_a_request(req, True, False) for req in req_list], + *[ + mock_rollout._async_rollout_a_request(req, True, False) + for req in req_list + ], ) ) assert len(output_req_list) == 1 @@ -294,29 +359,42 @@ def test_over_size_case(self, mock_rollout, search_data_proto, search_data): ) @patch.object(SearchTool, "execute", new_callable=AsyncMock) - def test_tool_call_basic_case(self, mock_execute, mock_rollout, search_data_proto, search_data): + def test_tool_call_basic_case( + self, mock_execute, mock_rollout, search_data_proto, search_data + ): _, expect_turn_array, tool_return_array = search_data # Mock search tool execution to return predefined responses - mock_execute.side_effect = [(msg, 0.0, {"status": "success"}) for msg in tool_return_array] + mock_execute.side_effect = [ + (msg, 0.0, {"status": "success"}) for msg in tool_return_array + ] mock_rollout.config.multi_turn.max_assistant_turns = 10 mock_rollout._tool_map["search"].retrieval_service_url = "mock://dummy" - req = mock_rollout._preprocess_prompt_to_async_rollout_requests(search_data_proto, n=1)[0] + req = mock_rollout._preprocess_prompt_to_async_rollout_requests( + search_data_proto, n=1 + )[0] req = MagicMock(wraps=req, spec=AsyncRolloutRequest) req.finalize = MagicMock() req_list = [req] mock_rollout._handle_engine_call = MagicMock() futures = [asyncio.Future() for i in expect_turn_array] - for idx, (i, turn) in enumerate(zip(futures, expect_turn_array, strict=True)): + for idx, (i, turn) in enumerate( + zip(futures, expect_turn_array, strict=True)): i.set_result( { "text": turn, "meta_info": { "id": "d1188d81cba840359df5b352b344bc8e", - "finish_reason": {"type": "tool_calls" if idx < len(expect_turn_array) - 1 else "stop"}, + "finish_reason": { + "type": ( + "tool_calls" + if idx < len(expect_turn_array) - 1 + else "stop" + ) + }, "prompt_tokens": len(turn), "completion_tokens": 100, "cached_tokens": 0, @@ -326,14 +404,20 @@ def test_tool_call_basic_case(self, mock_execute, mock_rollout, search_data_prot ) if idx < len(expect_turn_array) - 1: assert mock_rollout._function_call_parser.has_tool_call(turn) - assert mock_rollout._function_call_parser.parse_non_stream(turn) + assert mock_rollout._function_call_parser.parse_non_stream( + turn) mock_rollout._handle_engine_call.side_effect = futures mock_rollout._tp_rank = 0 loop = asyncio.get_event_loop() output_req_list = loop.run_until_complete( - asyncio.gather(*[mock_rollout._async_rollout_a_request(req, True, False) for req in req_list]) + asyncio.gather( + *[ + mock_rollout._async_rollout_a_request(req, True, False) + for req in req_list + ] + ) ) # Verify conversation completed successfully with proper tool usage @@ -342,7 +426,8 @@ def test_tool_call_basic_case(self, mock_execute, mock_rollout, search_data_prot assert "search" in output_req.metrics assert output_req.metrics["search"][0]["status"] == "success" assert mock_execute.await_count == 2 - assert len(output_req.messages) == 6 # user + 3*assistant + 2*tool_call + # user + 3*assistant + 2*tool_call + assert len(output_req.messages) == 6 # Verify tool response messages contain expected content search_counter = 0 for msg in output_req.messages: @@ -352,7 +437,9 @@ def test_tool_call_basic_case(self, mock_execute, mock_rollout, search_data_prot assert search_counter == 2 @patch.object(SearchTool, "execute", new_callable=AsyncMock) - def test_tool_call_batch_case(self, mock_execute, mock_rollout, search_data_proto, search_data): + def test_tool_call_batch_case( + self, mock_execute, mock_rollout, search_data_proto, search_data + ): _, expect_turn_array, tool_return_array = search_data # Mock tool execution for large batch (100 requests * 2 calls each) @@ -364,7 +451,9 @@ def test_tool_call_batch_case(self, mock_execute, mock_rollout, search_data_prot mock_rollout.config.multi_turn.max_assistant_turns = 10 mock_rollout._tool_map["search"].retrieval_service_url = "mock://dummy" - base_req = mock_rollout._preprocess_prompt_to_async_rollout_requests(search_data_proto, n=1)[0] + base_req = mock_rollout._preprocess_prompt_to_async_rollout_requests( + search_data_proto, n=1 + )[0] req_nums = 100 req_list = [] @@ -378,13 +467,21 @@ def test_tool_call_batch_case(self, mock_execute, mock_rollout, search_data_prot req_list.append(MagicMock(wraps=tmp_req, spec=AsyncRolloutRequest)) futures = [asyncio.Future() for _ in expect_turn_array] - for idx, (fut, turn) in enumerate(zip(futures, expect_turn_array, strict=True)): + for idx, (fut, turn) in enumerate( + zip(futures, expect_turn_array, strict=True) + ): fut.set_result( { "text": turn, "meta_info": { "id": "dummy", - "finish_reason": {"type": "tool_calls" if idx < len(expect_turn_array) - 1 else "stop"}, + "finish_reason": { + "type": ( + "tool_calls" + if idx < len(expect_turn_array) - 1 + else "stop" + ) + }, "prompt_tokens": len(turn), "completion_tokens": 100, }, @@ -393,16 +490,27 @@ def test_tool_call_batch_case(self, mock_execute, mock_rollout, search_data_prot req_turns_map[i] = futures req_turns_counter[i] = 0 - async def hacked_handle_engine_call(self, _req: AsyncRolloutRequest, *_args, **_kwargs): - fut = req_turns_map[_req.batch_data_id][req_turns_counter[_req.batch_data_id]] + async def hacked_handle_engine_call( + self, _req: AsyncRolloutRequest, *_args, **_kwargs + ): + fut = req_turns_map[_req.batch_data_id][ + req_turns_counter[_req.batch_data_id] + ] req_turns_counter[_req.batch_data_id] += 1 return await fut - with patch.object(SGLangRollout, "_handle_engine_call", new=hacked_handle_engine_call): + with patch.object( + SGLangRollout, "_handle_engine_call", new=hacked_handle_engine_call + ): mock_rollout._tp_rank = 0 loop = asyncio.get_event_loop() output_req_list = loop.run_until_complete( - asyncio.gather(*[mock_rollout._async_rollout_a_request(r, True, False) for r in req_list]) + asyncio.gather( + *[ + mock_rollout._async_rollout_a_request(r, True, False) + for r in req_list + ] + ) ) # Verify all requests completed successfully diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_sf_tools.py b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_sf_tools.py index 3f30929..9848e4d 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_sf_tools.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_sf_tools.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -38,7 +38,11 @@ OpenAIFunctionSchema, OpenAIFunctionToolSchema, ) -from verl.workers.rollout.schemas import AsyncRolloutRequest, AsyncRolloutRequestStateEnum, Message +from verl.workers.rollout.schemas import ( + AsyncRolloutRequest, + AsyncRolloutRequestStateEnum, + Message, +) from verl.workers.rollout.sglang_rollout.sglang_rollout import SGLangRollout sandbox_url = "" @@ -48,29 +52,29 @@ def get_sandbox_fusion_messages(): user_prompt = { "role": "user", "content": """ - Solve the following problem step by step. You now have the ability to selectively - write executable Python code to enhance your reasoning process. \n\n**user question:**\nThere - are 152 students at Dala High School. Assume the following: \n- 100 students take a Math class \n- 94 - students take a Science class \n- 57 students take an English class \n- 73 students take a Math class - and a Science class \n- 24 students take a Math class and an English class \n- 27 students take a Science - class and an English class \n- 22 students take a Math class and a Science class and an English class\n \nHow - many students take neither a Math class nor a Science class nor an Eglish class?\n\nRemember to place the final + Solve the following problem step by step. You now have the ability to selectively + write executable Python code to enhance your reasoning process. \n\n**user question:**\nThere + are 152 students at Dala High School. Assume the following: \n- 100 students take a Math class \n- 94 + students take a Science class \n- 57 students take an English class \n- 73 students take a Math class + and a Science class \n- 24 students take a Math class and an English class \n- 27 students take a Science + class and an English class \n- 22 students take a Math class and a Science class and an English class\n \nHow + many students take neither a Math class nor a Science class nor an Eglish class?\n\nRemember to place the final answer in the last part using the format: \n\n\boxed{'The final answer goes here.'}\n """, } expect_turn_0_msg = { "role": "assistant", "content": """ - Okay, so I need to find out how many students at Dala High School are not taking any of the three classes: Math, - Science, or English. The total number of students is 152. Let me see... I remember this is a problem about sets - and maybe using the principle of inclusion-exclusion. Let me recall how that works.\n\nFirst, the inclusion-exclusion - principle for three sets says that the total number of students taking at least one of the classes is equal to the - sum of the numbers in each individual class, minus the sum of the numbers in each pair of classes, plus the number in - all three classes. Then, subtract that total from the overall number of students to get those not taking any of the - three. \n\nLet me write that down step by step. Let M be the set of students taking Math, S for Science, and E for English. - Then:\n\nTotal in at least one class = |M โˆช S โˆช E| = |M| + |S| + |E| - |M โˆฉ S| - |M โˆฉ E| - |S โˆฉ E| + |M โˆฉ S โˆฉ E|\n\nGiven the - numbers:\n\n|M| = 100\n\n|S| = 94\n\n|E| = 57\n\n|M โˆฉ S| = 73\n\n|M โˆฉ E| = 24\n\n|S โˆฉ E| = 27\n\n|M โˆฉ S โˆฉ E| = 22\n\nSo plugging - these into the formula:\n\nTotal = 100 + 94 + 57 - 73 - 24 - 27 + 22\n\nLet me compute that step by step using code to ensure + Okay, so I need to find out how many students at Dala High School are not taking any of the three classes: Math, + Science, or English. The total number of students is 152. Let me see... I remember this is a problem about sets + and maybe using the principle of inclusion-exclusion. Let me recall how that works.\n\nFirst, the inclusion-exclusion + principle for three sets says that the total number of students taking at least one of the classes is equal to the + sum of the numbers in each individual class, minus the sum of the numbers in each pair of classes, plus the number in + all three classes. Then, subtract that total from the overall number of students to get those not taking any of the + three. \n\nLet me write that down step by step. Let M be the set of students taking Math, S for Science, and E for English. + Then:\n\nTotal in at least one class = |M โˆช S โˆช E| = |M| + |S| + |E| - |M โˆฉ S| - |M โˆฉ E| - |S โˆฉ E| + |M โˆฉ S โˆฉ E|\n\nGiven the + numbers:\n\n|M| = 100\n\n|S| = 94\n\n|E| = 57\n\n|M โˆฉ S| = 73\n\n|M โˆฉ E| = 24\n\n|S โˆฉ E| = 27\n\n|M โˆฉ S โˆฉ E| = 22\n\nSo plugging + these into the formula:\n\nTotal = 100 + 94 + 57 - 73 - 24 - 27 + 22\n\nLet me compute that step by step using code to ensure accuracy.\n """, "tool_calls": [ @@ -79,12 +83,11 @@ class and an English class \n- 22 students take a Math class and a Science class "function": { "name": "code_interpreter", "arguments": { - "code": """M = 100\\nS = 94\\nE = 57\\nM_S = 73\\nM_E = 24\\nS_E = 27\\nM_S_E = + "code": """M = 100\\nS = 94\\nE = 57\\nM_S = 73\\nM_E = 24\\nS_E = 27\\nM_S_E = 22\\n\\ntotal_in_any = M + S + E - M_S - M_E - S_E + M_S_E\\nstudents_neither = 152 - total_in_any\\nprint(students_neither)""", }, }, - } - ], + }], } tool_return_0_msg = { @@ -93,27 +96,18 @@ class and an English class \n- 22 students take a Math class and a Science class } expect_turn_1_msg = { - "role": "assistant", - "content": """ + "role": "assistant", "content": """ The code returns 3, indicating that 3 students take none of the classes. To verify, I'll break down each category to ensure the total matches:\n\n- Only Math and Science: 73 - 22 = 51\n- Only Math and English: 24 - 22 = 2\n- Only Science and English: 27 - 22 = 5\n- Only Math: 100 - 51 - 2 - 22 = 25\n- Only Science: 94 - 51 - 5 - 22 = 16\n- Only English: 57 - 2 - 5 - 22 = 28\n\nSumming all categories:\n - """, - "tool_calls": [ + """, "tool_calls": [ { - "type": "function", - "function": { - "name": "code_interpreter", - "arguments": { + "type": "function", "function": { + "name": "code_interpreter", "arguments": { "code": """only_M_S = 73 - 22\\nonly_M_E = 24 - 22\\nonly_S_E = 27 - 22\\n\\nonly_M = 100 - only_M_S - only_M_E - 22\\nonly_S = 94 - only_M_S - only_S_E - 22\\nonly_E = 57 - only_M_E - only_S_E - 22\\n\\ntotal_verify - = only_M + only_S + only_E + only_M_S + only_M_E + only_S_E + 22\\nprint(total_verify)""", - }, - }, - } - ], - } + = only_M + only_S + only_E + only_M_S + only_M_E + only_S_E + 22\\nprint(total_verify)""", }, }, }], } tool_return_1_msg = { "role": "tool", @@ -128,7 +122,10 @@ class and an English class \n- 22 students take a Math class and a Science class } user_prompts = [user_prompt] - expect_turn_array = [expect_turn_0_msg, expect_turn_1_msg, expect_turn_2_msg] + expect_turn_array = [ + expect_turn_0_msg, + expect_turn_1_msg, + expect_turn_2_msg] tool_return_array = [tool_return_0_msg, tool_return_1_msg] return user_prompts, expect_turn_array, tool_return_array @@ -150,7 +147,8 @@ class TestRolloutWithTools: @pytest.fixture def qwen_tokenizer(self): local_model_path = "Qwen/Qwen2.5-0.5B" - tokenizer = AutoTokenizer.from_pretrained(local_model_path, padding_side="left") + tokenizer = AutoTokenizer.from_pretrained( + local_model_path, padding_side="left") tokenizer.pad_token = tokenizer.eos_token return tokenizer @@ -163,14 +161,20 @@ def qwen_model_config(self): @pytest.fixture def sandbox_fusion_data(self, qwen_tokenizer): - user_prompt, expect_turn_array, tool_return_array = get_sandbox_fusion_messages() + user_prompt, expect_turn_array, tool_return_array = ( + get_sandbox_fusion_messages() + ) prompts = [[message] for message in user_prompt] preencode_turn_array = [ - qwen_tokenizer.apply_chat_template([turn], tokenize=False, add_generation_prompt=False) + qwen_tokenizer.apply_chat_template( + [turn], tokenize=False, add_generation_prompt=False + ) for turn in expect_turn_array ] preencode_tool_return_array = [ - qwen_tokenizer.apply_chat_template([turn], tokenize=False, add_generation_prompt=True) + qwen_tokenizer.apply_chat_template( + [turn], tokenize=False, add_generation_prompt=True + ) for turn in tool_return_array ] return prompts, preencode_turn_array, preencode_tool_return_array @@ -183,7 +187,11 @@ def sandbox_fusion_rollout_config(self): tensor_parallel_size = 1 tool_path = "./resource/tool_configs/sandbox_fusion_tool_config" rollout_config = get_rollout_config( - max_response_length, max_prompt_length, dtype, tensor_parallel_size, tool_path + max_response_length, + max_prompt_length, + dtype, + tensor_parallel_size, + tool_path, ) return rollout_config @@ -191,10 +199,14 @@ def sandbox_fusion_rollout_config(self): def sandbox_data_proto(self, sandbox_fusion_data, qwen_tokenizer): preencode_prompts, _, _ = sandbox_fusion_data prompts = [ - qwen_tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True) + qwen_tokenizer.apply_chat_template( + message, tokenize=False, add_generation_prompt=True + ) for message in preencode_prompts ] - input_ids, attention_mask, position_ids = prepare_inputs(qwen_tokenizer, prompts, 1000) + input_ids, attention_mask, position_ids = prepare_inputs( + qwen_tokenizer, prompts, 1000 + ) prompt_dict = TensorDict( { "input_ids": input_ids, @@ -216,16 +228,27 @@ def sandbox_data_proto(self, sandbox_fusion_data, qwen_tokenizer): ) index = np.array([0], dtype=object) prompts = DataProto( - batch=prompt_dict, non_tensor_batch={"raw_prompt": messages, "tools_kwargs": tools_kwargs, "index": index} + batch=prompt_dict, + non_tensor_batch={ + "raw_prompt": messages, + "tools_kwargs": tools_kwargs, + "index": index, + }, ) return prompts @pytest.fixture - def mock_rollout(self, sandbox_fusion_rollout_config, qwen_tokenizer, qwen_model_config): + def mock_rollout( + self, sandbox_fusion_rollout_config, qwen_tokenizer, qwen_model_config + ): """Mock the rollout instance""" - with patch.object(SGLangRollout, "_init_distributed_env", return_value=None), patch.object( + with patch.object( + SGLangRollout, "_init_distributed_env", return_value=None + ), patch.object( SGLangRollout, "_init_inference_engine", return_value=None - ), patch.object(SGLangRollout, "_init_sampling_params", return_value=None): + ), patch.object( + SGLangRollout, "_init_sampling_params", return_value=None + ): rollout = SGLangRollout( actor_module="", config=sandbox_fusion_rollout_config, @@ -248,12 +271,16 @@ def test_tools_registration(self, mock_rollout): assert "code_interpreter" in mock_rollout._tool_map.keys() from verl.tools.sandbox_fusion_tools import SandboxFusionTool - assert isinstance(mock_rollout._tool_map["code_interpreter"], SandboxFusionTool) + assert isinstance( + mock_rollout._tool_map["code_interpreter"], + SandboxFusionTool) assert mock_rollout._tool_call_parser_type == "qwen25" def test_rollout_req_creation(self, mock_rollout, sandbox_data_proto): """Test request creation functionality""" - req_list = mock_rollout._preprocess_prompt_to_async_rollout_requests(sandbox_data_proto, n=1) + req_list = mock_rollout._preprocess_prompt_to_async_rollout_requests( + sandbox_data_proto, n=1 + ) assert len(req_list) == 1 assert req_list[0].state == AsyncRolloutRequestStateEnum.PENDING assert len(req_list[0].tool_schemas) == 1 @@ -278,16 +305,21 @@ def test_rollout_req_creation(self, mock_rollout, sandbox_data_proto): ), ) - def test_over_size_case(self, mock_rollout, sandbox_data_proto, sandbox_fusion_data): + def test_over_size_case( + self, mock_rollout, sandbox_data_proto, sandbox_fusion_data + ): """Test over-size response truncation case""" mock_rollout.config.multi_turn.max_assistant_turns = 1 - req = mock_rollout._preprocess_prompt_to_async_rollout_requests(sandbox_data_proto, n=1)[0] + req = mock_rollout._preprocess_prompt_to_async_rollout_requests( + sandbox_data_proto, n=1 + )[0] req = MagicMock(wraps=req, spec=AsyncRolloutRequest) req.finalize = MagicMock() req_list = [req] _, expect_turn_array, tool_return_array = sandbox_fusion_data - # here we mock a meta info with 'length'. indicate the response is truncate + # here we mock a meta info with 'length'. indicate the response is + # truncate mock_rollout._handle_engine_call = MagicMock() future = asyncio.Future() future.set_result( @@ -308,7 +340,10 @@ def test_over_size_case(self, mock_rollout, sandbox_data_proto, sandbox_fusion_d loop = asyncio.get_event_loop() output_req_list = loop.run_until_complete( asyncio.gather( - *[mock_rollout._async_rollout_a_request(req, True, False) for req in req_list], + *[ + mock_rollout._async_rollout_a_request(req, True, False) + for req in req_list + ], ) ) assert len(output_req_list) == 1 @@ -324,16 +359,21 @@ def test_over_size_case(self, mock_rollout, sandbox_data_proto, sandbox_fusion_d ) @skip_if_valid_sandbox(sandbox_url) - def test_tool_call_basic_case(self, mock_rollout, sandbox_data_proto, sandbox_fusion_data): + def test_tool_call_basic_case( + self, mock_rollout, sandbox_data_proto, sandbox_fusion_data + ): """Test basic tool call case""" mock_rollout.config.multi_turn.max_assistant_turns = 10 mock_rollout._tool_map["code_interpreter"].sandbox_fusion_url = sandbox_url - req = mock_rollout._preprocess_prompt_to_async_rollout_requests(sandbox_data_proto, n=1)[0] + req = mock_rollout._preprocess_prompt_to_async_rollout_requests( + sandbox_data_proto, n=1 + )[0] req = MagicMock(wraps=req, spec=AsyncRolloutRequest) req.finalize = MagicMock() req_list = [req] _, expect_turn_array, tool_return_array = sandbox_fusion_data - # here we mock a meta info with 'length'. indicate the response is truncate + # here we mock a meta info with 'length'. indicate the response is + # truncate mock_rollout._handle_engine_call = MagicMock() futures = [asyncio.Future() for i in expect_turn_array] for idx, (i, turn) in enumerate(zip(futures, expect_turn_array)): @@ -342,7 +382,13 @@ def test_tool_call_basic_case(self, mock_rollout, sandbox_data_proto, sandbox_fu "text": turn, "meta_info": { "id": "d1188d81cba840359df5b352b344bc8e", - "finish_reason": {"type": "tool_calls" if idx < len(expect_turn_array) - 1 else "stop"}, + "finish_reason": { + "type": ( + "tool_calls" + if idx < len(expect_turn_array) - 1 + else "stop" + ) + }, "prompt_tokens": len(turn), "completion_tokens": 100, "cached_tokens": 0, @@ -352,14 +398,18 @@ def test_tool_call_basic_case(self, mock_rollout, sandbox_data_proto, sandbox_fu ) if idx < len(expect_turn_array) - 1: assert mock_rollout._function_call_parser.has_tool_call(turn) - assert mock_rollout._function_call_parser.parse_non_stream(turn) + assert mock_rollout._function_call_parser.parse_non_stream( + turn) mock_rollout._handle_engine_call.side_effect = futures mock_rollout._tp_rank = 0 loop = asyncio.get_event_loop() output_req_list = loop.run_until_complete( asyncio.gather( - *[mock_rollout._async_rollout_a_request(req, True, False) for req in req_list], + *[ + mock_rollout._async_rollout_a_request(req, True, False) + for req in req_list + ], ) ) assert len(output_req_list) == 1 @@ -368,7 +418,8 @@ def test_tool_call_basic_case(self, mock_rollout, sandbox_data_proto, sandbox_fu # here we verify whether the code sandbox is executed correctly assert output_req.metrics == {"code_interpreter": ["3", "149"]} assert mock_rollout._handle_engine_call.call_count == 3 - assert len(output_req.messages) == 6 # user + 3*assistant + 2*tool_call + # user + 3*assistant + 2*tool_call + assert len(output_req.messages) == 6 code_counter = 0 for msg in output_req.messages: if msg.role == "tool": @@ -377,11 +428,15 @@ def test_tool_call_basic_case(self, mock_rollout, sandbox_data_proto, sandbox_fu assert code_counter == 2 @skip_if_valid_sandbox(sandbox_url) - def test_tool_call_batch_case(self, mock_rollout, sandbox_data_proto, sandbox_fusion_data): + def test_tool_call_batch_case( + self, mock_rollout, sandbox_data_proto, sandbox_fusion_data + ): """Test batch tool call case""" mock_rollout.config.multi_turn.max_assistant_turns = 10 mock_rollout._tool_map["code_interpreter"].sandbox_fusion_url = sandbox_url - req = mock_rollout._preprocess_prompt_to_async_rollout_requests(sandbox_data_proto, n=1)[0] + req = mock_rollout._preprocess_prompt_to_async_rollout_requests( + sandbox_data_proto, n=1 + )[0] req_nums = 100 req_list = [] req_turns_counter = {} @@ -392,7 +447,10 @@ def test_tool_call_batch_case(self, mock_rollout, sandbox_data_proto, sandbox_fu _temp_req = deepcopy(req) _temp_req.batch_data_id = i _temp_req.request_id = i - req_list.append(MagicMock(wraps=_temp_req, spec=AsyncRolloutRequest)) + req_list.append( + MagicMock( + wraps=_temp_req, + spec=AsyncRolloutRequest)) futures = [asyncio.Future() for i in expect_turn_array] for idx, (i, turn) in enumerate(zip(futures, expect_turn_array)): i.set_result( @@ -400,7 +458,13 @@ def test_tool_call_batch_case(self, mock_rollout, sandbox_data_proto, sandbox_fu "text": turn, "meta_info": { "id": "d1188d81cba840359df5b352b344bc8e", - "finish_reason": {"type": "tool_calls" if idx < len(expect_turn_array) - 1 else "stop"}, + "finish_reason": { + "type": ( + "tool_calls" + if idx < len(expect_turn_array) - 1 + else "stop" + ) + }, "prompt_tokens": len(turn), "completion_tokens": 100, "cached_tokens": 0, @@ -409,25 +473,38 @@ def test_tool_call_batch_case(self, mock_rollout, sandbox_data_proto, sandbox_fu } ) if idx < len(expect_turn_array) - 1: - assert mock_rollout._function_call_parser.has_tool_call(turn) - assert mock_rollout._function_call_parser.parse_non_stream(turn) + assert mock_rollout._function_call_parser.has_tool_call( + turn) + assert mock_rollout._function_call_parser.parse_non_stream( + turn) req_turns_map[_temp_req.batch_data_id] = futures req_turns_counter[_temp_req.batch_data_id] = 0 async def hacked_handle_engine_call( - self, _req: AsyncRolloutRequest, do_sample: bool, is_validate: bool, **kwargs + self, + _req: AsyncRolloutRequest, + do_sample: bool, + is_validate: bool, + **kwargs, ): - result = req_turns_map[_req.batch_data_id][req_turns_counter[_req.batch_data_id]] + result = req_turns_map[_req.batch_data_id][ + req_turns_counter[_req.batch_data_id] + ] req_turns_counter[_req.batch_data_id] += 1 re = await result return re - with patch.object(SGLangRollout, "_handle_engine_call", new=hacked_handle_engine_call): + with patch.object( + SGLangRollout, "_handle_engine_call", new=hacked_handle_engine_call + ): mock_rollout._tp_rank = 0 loop = asyncio.get_event_loop() output_req_list = loop.run_until_complete( asyncio.gather( - *[mock_rollout._async_rollout_a_request(req, True, False) for req in req_list], + *[ + mock_rollout._async_rollout_a_request(req, True, False) + for req in req_list + ], ) ) assert len(output_req_list) == req_nums @@ -437,7 +514,8 @@ async def hacked_handle_engine_call( assert output_req.state == AsyncRolloutRequestStateEnum.COMPLETED # here we verify whether the code sandbox is executed correctly assert output_req.metrics == {"code_interpreter": ["3", "149"]} - assert len(output_req.messages) == 6 # user + 3*assistant + 2*tool_call + # user + 3*assistant + 2*tool_call + assert len(output_req.messages) == 6 code_counter = 0 for msg in output_req.messages: if msg.role == "tool": @@ -562,9 +640,14 @@ def test_rate_limiter(self): # exec_worker = ExecutionWorker.options(max_concurrency=10).remote(enable_global_rate_limit=True, rate_limit=3) exec_worker = init_execution_pool( - num_workers=10, enable_global_rate_limit=True, rate_limit=3, mode=PoolMode.ThreadMode + num_workers=10, + enable_global_rate_limit=True, + rate_limit=3, + mode=PoolMode.ThreadMode, + ) + center = TestActor.options(get_if_exists=True, name="test-actor").remote( + self.rank, self.world_size ) - center = TestActor.options(get_if_exists=True, name="test-actor").remote(self.rank, self.world_size) ray.get(exec_worker.ping.remote()) def fn(i): @@ -584,7 +667,8 @@ def fn(i): print(f"Total time: {duration:.2f} seconds for rank: {self.rank}") assert results == list(range(6)) - # we have 6 task with rate limit of 3, therefore we need at least 2 round: 3*2=6 seconds + # we have 6 task with rate limit of 3, therefore we need at least 2 + # round: 3*2=6 seconds assert duration > 6 assert duration < 10 @@ -594,7 +678,10 @@ def test_rotten_execution(self): # exec_worker = ExecutionWorker.options(max_concurrency=10).remote(enable_global_rate_limit=True, rate_limit=6) exec_worker = init_execution_pool( - num_workers=10, enable_global_rate_limit=True, rate_limit=6, mode=PoolMode.ThreadMode + num_workers=10, + enable_global_rate_limit=True, + rate_limit=6, + mode=PoolMode.ThreadMode, ) ray.get(exec_worker.ping.remote()) @@ -609,8 +696,12 @@ def fn(i): results = loop.run_until_complete(asyncio.gather(*tasks)) expect_result = [None] + list(range(10)) + list(range(11, 20)) sorted_data = sorted(results, key=lambda x: (x is not None, x)) - assert sorted_data == expect_result, f"results: {results}, expect_result: {expect_result}" - rate_limiter = TokenBucketWorker.options(name="rate-limiter", get_if_exists=True).remote() + assert ( + sorted_data == expect_result + ), f"results: {results}, expect_result: {expect_result}" + rate_limiter = TokenBucketWorker.options( + name="rate-limiter", get_if_exists=True + ).remote() rate = ray.get(rate_limiter.get_current_count.remote()) assert rate == 0, f"rate: {rate}" @@ -626,9 +717,14 @@ def test_rate_limiter(self): # exec_worker = ExecutionWorker.options(max_concurrency=10).remote(enable_global_rate_limit=True, rate_limit=6) exec_worker = init_execution_pool( - num_workers=10, enable_global_rate_limit=True, rate_limit=6, mode=PoolMode.ThreadMode + num_workers=10, + enable_global_rate_limit=True, + rate_limit=6, + mode=PoolMode.ThreadMode, + ) + center = TestActor.options(get_if_exists=True, name="test-actor").remote( + self.rank, self.world_size ) - center = TestActor.options(get_if_exists=True, name="test-actor").remote(self.rank, self.world_size) ray.get(exec_worker.ping.remote()) def fn(i): diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_w_interaction.py b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_w_interaction.py index 3ccde18..3d37b89 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_w_interaction.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_w_interaction.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -61,21 +61,43 @@ def test_async_sglang_rollout_w_interaction(): ] ] interaction_kwargs = [ - {"name": "gsm8k", "query": "Who won the Champions League in 2019?", "ground_truth": "Real Madrid"}, - {"name": "gsm8k", "query": "The founder of Apple is", "ground_truth": "Steve Jobs"}, - {"name": "gsm8k", "query": "What's the best way to learn python?", "ground_truth": "Learn python from scratch"}, + { + "name": "gsm8k", + "query": "Who won the Champions League in 2019?", + "ground_truth": "Real Madrid", + }, + { + "name": "gsm8k", + "query": "The founder of Apple is", + "ground_truth": "Steve Jobs", + }, + { + "name": "gsm8k", + "query": "What's the best way to learn python?", + "ground_truth": "Learn python from scratch", + }, ] prompts = [ - tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True) + tokenizer.apply_chat_template( + message, tokenize=False, add_generation_prompt=True + ) for message in preencode_prompts ] - input_ids, attention_mask, position_ids = prepare_inputs(tokenizer, prompts, max_prompt_length) + input_ids, attention_mask, position_ids = prepare_inputs( + tokenizer, prompts, max_prompt_length + ) - hf_response_tokens = generate_hf_output(actor_model, input_ids, attention_mask, tokenizer, max_response_length) + hf_response_tokens = generate_hf_output( + actor_model, input_ids, attention_mask, tokenizer, max_response_length + ) - fsdp_device_mesh = init_device_mesh("cuda", mesh_shape=(tensor_parallel_size,), mesh_dim_names=("fsdp",)) + fsdp_device_mesh = init_device_mesh( + "cuda", mesh_shape=(tensor_parallel_size,), mesh_dim_names=("fsdp",) + ) inference_device_mesh_cpu = init_device_mesh( - "cpu", mesh_shape=(1, tensor_parallel_size, 1), mesh_dim_names=("dp", "infer_tp", "pp") + "cpu", + mesh_shape=(1, tensor_parallel_size, 1), + mesh_dim_names=("dp", "infer_tp", "pp"), ) fsdp_model = FSDP( @@ -94,16 +116,23 @@ def test_async_sglang_rollout_w_interaction(): interaction_config = { "interaction": [ - {"name": "gsm8k", "class_name": "verl.interactions.gsm8k_interaction.Gsm8kInteraction", "config": {}} - ] - } + { + "name": "gsm8k", + "class_name": "verl.interactions.gsm8k_interaction.Gsm8kInteraction", + "config": {}, + }]} with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: OmegaConf.save(interaction_config, f.name) interaction_config_path = f.name rollout_config = get_rollout_config( - max_response_length, max_prompt_length, dtype, tensor_parallel_size, None, interaction_config_path + max_response_length, + max_prompt_length, + dtype, + tensor_parallel_size, + None, + interaction_config_path, ) rollout = SGLangRollout( actor_module=local_model_path, @@ -135,7 +164,10 @@ def test_async_sglang_rollout_w_interaction(): messages = np.asarray(preencode_prompts) prompts = DataProto( batch=prompt_dict, - non_tensor_batch={"raw_prompt": messages, "interaction_kwargs": np.asarray(interaction_kwargs)}, + non_tensor_batch={ + "raw_prompt": messages, + "interaction_kwargs": np.asarray(interaction_kwargs), + }, ) prompts.meta_info.update( @@ -154,7 +186,8 @@ def test_async_sglang_rollout_w_interaction(): print(f"postprocessed {output.batch['responses'].shape=}") sglang_output = output.to("cpu") - sglang_response_tokens = tokenizer.batch_decode(sglang_output.batch["responses"]) + sglang_response_tokens = tokenizer.batch_decode( + sglang_output.batch["responses"]) print(f"hf response: {hf_response_tokens}") print(f"sglang response: {sglang_response_tokens}") diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_w_tools.py b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_w_tools.py index 20faab8..968ab94 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_w_tools.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_async_rollout_w_tools.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -61,16 +61,26 @@ def test_async_sglang_rollout_w_tool(): ] ] prompts = [ - tokenizer.apply_chat_template(message, tokenize=False, add_generation_prompt=True) + tokenizer.apply_chat_template( + message, tokenize=False, add_generation_prompt=True + ) for message in preencode_prompts ] - input_ids, attention_mask, position_ids = prepare_inputs(tokenizer, prompts, max_prompt_length) + input_ids, attention_mask, position_ids = prepare_inputs( + tokenizer, prompts, max_prompt_length + ) - hf_response_tokens = generate_hf_output(actor_model, input_ids, attention_mask, tokenizer, max_response_length) + hf_response_tokens = generate_hf_output( + actor_model, input_ids, attention_mask, tokenizer, max_response_length + ) - fsdp_device_mesh = init_device_mesh("cuda", mesh_shape=(tensor_parallel_size,), mesh_dim_names=("fsdp",)) + fsdp_device_mesh = init_device_mesh( + "cuda", mesh_shape=(tensor_parallel_size,), mesh_dim_names=("fsdp",) + ) inference_device_mesh_cpu = init_device_mesh( - "cpu", mesh_shape=(1, tensor_parallel_size, 1), mesh_dim_names=("dp", "infer_tp", "pp") + "cpu", + mesh_shape=(1, tensor_parallel_size, 1), + mesh_dim_names=("dp", "infer_tp", "pp"), ) fsdp_model = FSDP( @@ -141,7 +151,8 @@ def test_async_sglang_rollout_w_tool(): print(f"postprocessed {output.batch['responses'].shape=}") sglang_output = output.to("cpu") - sglang_response_tokens = tokenizer.batch_decode(sglang_output.batch["responses"]) + sglang_response_tokens = tokenizer.batch_decode( + sglang_output.batch["responses"]) print(f"hf response: {hf_response_tokens}") print(f"sglang response: {sglang_response_tokens}") diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_multi_interaction.py b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_multi_interaction.py index 465470f..f18e876 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_multi_interaction.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_multi_interaction.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -106,7 +106,8 @@ def create_mock_config_with_multi_interactions(): def setup_distributed(): """Initialize distributed environment if not already initialized.""" if not dist.is_initialized(): - dist.init_process_group(backend="nccl" if torch.cuda.is_available() else "gloo") + dist.init_process_group( + backend="nccl" if torch.cuda.is_available() else "gloo") class TestSGLangMultiInteraction: @@ -116,14 +117,19 @@ def test_initialize_multiple_interactions(self): config, temp_config_path = create_mock_config_with_multi_interactions() try: - # Mock SGLang engine and initialization methods like the reference test + # Mock SGLang engine and initialization methods like the reference + # test with ( patch.object(SGLangRollout, "_init_distributed_env", return_value=None), - patch.object(SGLangRollout, "_init_inference_engine", return_value=None), + patch.object( + SGLangRollout, "_init_inference_engine", return_value=None + ), patch.object(SGLangRollout, "_init_sampling_params", return_value=None), ): # Create a real tokenizer like the reference test - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B", padding_side="left") + tokenizer = AutoTokenizer.from_pretrained( + "Qwen/Qwen2.5-0.5B", padding_side="left" + ) tokenizer.pad_token = tokenizer.eos_token # Mock model config @@ -153,13 +159,25 @@ def test_initialize_multiple_interactions(self): assert "mock_agent1" in rollout.interaction_map assert "mock_agent2" in rollout.interaction_map - # Use class name comparison instead of isinstance for multi-process compatibility - assert rollout.interaction_map["mock_agent1"].__class__.__name__ == "MockInteraction" - assert rollout.interaction_map["mock_agent2"].__class__.__name__ == "MockInteraction" + # Use class name comparison instead of isinstance for + # multi-process compatibility + assert ( + rollout.interaction_map["mock_agent1"].__class__.__name__ + == "MockInteraction" + ) + assert ( + rollout.interaction_map["mock_agent2"].__class__.__name__ + == "MockInteraction" + ) - # Also check that they are instances of BaseInteraction (which should work across processes) - assert isinstance(rollout.interaction_map["mock_agent1"], BaseInteraction) - assert isinstance(rollout.interaction_map["mock_agent2"], BaseInteraction) + # Also check that they are instances of BaseInteraction (which + # should work across processes) + assert isinstance( + rollout.interaction_map["mock_agent1"], BaseInteraction + ) + assert isinstance( + rollout.interaction_map["mock_agent2"], BaseInteraction + ) # Check that names were set correctly assert rollout.interaction_map["mock_agent1"].name == "mock_agent1" @@ -176,10 +194,14 @@ def test_interaction_selection_by_name(self): try: with ( patch.object(SGLangRollout, "_init_distributed_env", return_value=None), - patch.object(SGLangRollout, "_init_inference_engine", return_value=None), + patch.object( + SGLangRollout, "_init_inference_engine", return_value=None + ), patch.object(SGLangRollout, "_init_sampling_params", return_value=None), ): - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B", padding_side="left") + tokenizer = AutoTokenizer.from_pretrained( + "Qwen/Qwen2.5-0.5B", padding_side="left" + ) tokenizer.pad_token = tokenizer.eos_token mock_model_config = MagicMock() @@ -201,14 +223,23 @@ def test_interaction_selection_by_name(self): ) # Test interaction selection logic - from verl.workers.rollout.schemas import AsyncRolloutRequest, AsyncRolloutRequestStateEnum, Message + from verl.workers.rollout.schemas import ( + AsyncRolloutRequest, + AsyncRolloutRequestStateEnum, + Message, + ) # Create a mock request with specific interaction name req = AsyncRolloutRequest( request_id="test_req", state=AsyncRolloutRequestStateEnum.INTERACTING, - messages=[Message(role="user", content="test message")], - interaction_kwargs={"name": "mock_agent2", "test_param": "value"}, + messages=[ + Message( + role="user", + content="test message")], + interaction_kwargs={ + "name": "mock_agent2", + "test_param": "value"}, input_ids=None, prompt_ids=None, response_ids=None, @@ -251,9 +282,7 @@ def test_fallback_to_default_interaction(self): "name": "gsm8k", "class_name": "tests.workers.rollout.test_sglang_multi_interaction.MockInteraction", "config": {}, - } - ] - } + }]} with tempfile.NamedTemporaryFile(mode="w", suffix=".yaml", delete=False) as f: OmegaConf.save(interaction_config, f.name) @@ -288,10 +317,14 @@ def test_fallback_to_default_interaction(self): try: with ( patch.object(SGLangRollout, "_init_distributed_env", return_value=None), - patch.object(SGLangRollout, "_init_inference_engine", return_value=None), + patch.object( + SGLangRollout, "_init_inference_engine", return_value=None + ), patch.object(SGLangRollout, "_init_sampling_params", return_value=None), ): - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B", padding_side="left") + tokenizer = AutoTokenizer.from_pretrained( + "Qwen/Qwen2.5-0.5B", padding_side="left" + ) tokenizer.pad_token = tokenizer.eos_token mock_model_config = MagicMock() @@ -314,7 +347,8 @@ def test_fallback_to_default_interaction(self): # Test that default interaction name works interaction_kwargs_without_name = {"test_param": "value"} - default_name = interaction_kwargs_without_name.get("name", "gsm8k") + default_name = interaction_kwargs_without_name.get( + "name", "gsm8k") assert default_name == "gsm8k" assert default_name in rollout.interaction_map @@ -329,10 +363,14 @@ def test_error_on_missing_interaction(self): try: with ( patch.object(SGLangRollout, "_init_distributed_env", return_value=None), - patch.object(SGLangRollout, "_init_inference_engine", return_value=None), + patch.object( + SGLangRollout, "_init_inference_engine", return_value=None + ), patch.object(SGLangRollout, "_init_sampling_params", return_value=None), ): - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B", padding_side="left") + tokenizer = AutoTokenizer.from_pretrained( + "Qwen/Qwen2.5-0.5B", padding_side="left" + ) tokenizer.pad_token = tokenizer.eos_token mock_model_config = MagicMock() @@ -401,7 +439,9 @@ def test_backward_compatibility_no_interaction_config(self): patch.object(SGLangRollout, "_init_inference_engine", return_value=None), patch.object(SGLangRollout, "_init_sampling_params", return_value=None), ): - tokenizer = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B", padding_side="left") + tokenizer = AutoTokenizer.from_pretrained( + "Qwen/Qwen2.5-0.5B", padding_side="left" + ) tokenizer.pad_token = tokenizer.eos_token mock_model_config = MagicMock() diff --git a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_spmd.py b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_spmd.py index e6b7256..3f59b72 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_spmd.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/test_sglang_spmd.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -35,7 +35,9 @@ def _pre_process_inputs(pad_token_id, prompt_token_ids: torch.Tensor): - non_pad_index = torch.nonzero(prompt_token_ids != pad_token_id, as_tuple=False)[0][0] + 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 @@ -51,14 +53,24 @@ def test_sglang_spmd(): local_model_path = "Qwen/Qwen2.5-0.5B" tokenizer, actor_model = load_tokenizer_and_model(local_model_path) - preencode_prompts = ["Who won the Champions League in 2019?", "The founder of Apple is", "What's your name?"] - input_ids, attention_mask, _ = prepare_inputs(tokenizer, preencode_prompts, max_prompt_length) + preencode_prompts = [ + "Who won the Champions League in 2019?", + "The founder of Apple is", + "What's your name?", + ] + input_ids, attention_mask, _ = prepare_inputs( + tokenizer, preencode_prompts, max_prompt_length + ) - hf_response_tokens = generate_hf_output(actor_model, input_ids, attention_mask, tokenizer, max_response_length) + hf_response_tokens = generate_hf_output( + actor_model, input_ids, attention_mask, tokenizer, max_response_length + ) tensor_parallel_size = 2 inference_device_mesh_cpu = init_device_mesh( - "cpu", mesh_shape=(1, tensor_parallel_size, 1), mesh_dim_names=["dp", "tp", "pp"] + "cpu", + mesh_shape=(1, tensor_parallel_size, 1), + mesh_dim_names=["dp", "tp", "pp"], ) tp_rank = inference_device_mesh_cpu["tp"].get_local_rank() @@ -74,7 +86,11 @@ def test_sglang_spmd(): input_ids = input_ids.cuda() idx_list = [] - pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + pad_token_id = ( + tokenizer.pad_token_id + if tokenizer.pad_token_id is not None + else tokenizer.eos_token_id + ) for i in range(input_ids.shape[0]): idx_list.append(_pre_process_inputs(pad_token_id, input_ids[i])) @@ -93,7 +109,10 @@ def test_sglang_spmd(): ) loop = asyncio.get_event_loop() - outputs = loop.run_until_complete(llm.async_generate(input_ids=idx_list, sampling_params=sampling_params)) + outputs = loop.run_until_complete( + llm.async_generate( + input_ids=idx_list, + sampling_params=sampling_params)) else: outputs = None @@ -108,7 +127,9 @@ def test_sglang_spmd(): sglang_response_tokens = [output["text"] for output in outputs] print(f"sglang response: {sglang_response_tokens}") - assert are_lists_similar(hf_response_tokens, sglang_response_tokens), "Strings differ more than 10%:\n" + assert are_lists_similar( + hf_response_tokens, sglang_response_tokens + ), "Strings differ more than 10%:\n" print("SPMD Test Passed!") torch.distributed.barrier() diff --git a/Agent0/executor_train/verl/tests/workers/rollout/utils_sglang.py b/Agent0/executor_train/verl/tests/workers/rollout/utils_sglang.py index 2e22e47..a4e328b 100644 --- a/Agent0/executor_train/verl/tests/workers/rollout/utils_sglang.py +++ b/Agent0/executor_train/verl/tests/workers/rollout/utils_sglang.py @@ -1,4 +1,4 @@ -# Copyright 2023-2024 SGLang Team +# Copyright 2023-2026 SGLang Team # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,7 +33,8 @@ def levenshtein(s1, s2): for i in range(1, m + 1): for j in range(1, n + 1): cost = 0 if s1[i - 1] == s2[j - 1] else 1 - dp[i][j] = min(dp[i - 1][j] + 1, dp[i][j - 1] + 1, dp[i - 1][j - 1] + cost) + dp[i][j] = min(dp[i - 1][j] + 1, dp[i][j - 1] + + 1, dp[i - 1][j - 1] + cost) return dp[m][n] @@ -57,7 +58,8 @@ def initialize_global_process_group(timeout_second=36000, spmd=False): if not torch.distributed.is_initialized(): # Check if already initialized print("Initializing process group...") - torch.distributed.init_process_group(timeout=timedelta(seconds=timeout_second)) + torch.distributed.init_process_group( + timeout=timedelta(seconds=timeout_second)) else: print("Process group already initialized.") @@ -74,7 +76,8 @@ def initialize_global_process_group(timeout_second=36000, spmd=False): else: CUDA_VISIBLE_DEVICES = str(local_rank) os.environ["CUDA_VISIBLE_DEVICES"] = CUDA_VISIBLE_DEVICES - print(f"CUDA_VISIBLE_DEVICES is not set, set to {CUDA_VISIBLE_DEVICES}") + print( + f"CUDA_VISIBLE_DEVICES is not set, set to {CUDA_VISIBLE_DEVICES}") return local_rank, rank, world_size @@ -86,25 +89,40 @@ def clean_torchelastic_env(): def load_tokenizer_and_model(local_model_path, dtype="bfloat16"): - tokenizer = AutoTokenizer.from_pretrained(local_model_path, padding_side="left") + tokenizer = AutoTokenizer.from_pretrained( + local_model_path, padding_side="left") tokenizer.pad_token = tokenizer.eos_token - model = AutoModelForCausalLM.from_pretrained(local_model_path, torch_dtype=getattr(torch, dtype), device_map="cuda") + model = AutoModelForCausalLM.from_pretrained( + local_model_path, torch_dtype=getattr(torch, dtype), device_map="cuda" + ) return tokenizer, model def prepare_inputs(tokenizer, prompts, max_prompt_length): - pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id + pad_token_id = ( + tokenizer.pad_token_id + if tokenizer.pad_token_id is not None + else tokenizer.eos_token_id + ) tokenized = tokenizer(prompts, return_tensors="pt", padding=True) - input_ids = pad_sequence_to_length(tokenized["input_ids"], max_prompt_length, pad_token_id, left_pad=True) - attention_mask = pad_sequence_to_length( - tokenized["attention_mask"], max_prompt_length, pad_token_id=0, left_pad=True + input_ids = pad_sequence_to_length( + tokenized["input_ids"], max_prompt_length, pad_token_id, left_pad=True ) + attention_mask = pad_sequence_to_length( + tokenized["attention_mask"], + max_prompt_length, + pad_token_id=0, + left_pad=True) position_ids = compute_position_id_with_mask(attention_mask) - position_ids = pad_sequence_to_length(position_ids, max_prompt_length, pad_token_id=0, left_pad=True) + position_ids = pad_sequence_to_length( + position_ids, max_prompt_length, pad_token_id=0, left_pad=True + ) return input_ids, attention_mask, position_ids -def generate_hf_output(model, input_ids, attention_mask, tokenizer, max_response_length): +def generate_hf_output( + model, input_ids, attention_mask, tokenizer, max_response_length +): generation_config = GenerationConfig(do_sample=False) output = model.generate( input_ids=input_ids.cuda(), @@ -118,7 +136,7 @@ def generate_hf_output(model, input_ids, attention_mask, tokenizer, max_response use_cache=False, ) seq = output.sequences - response = seq[:, input_ids.shape[1] :] + response = seq[:, input_ids.shape[1]:] return tokenizer.batch_decode(response) diff --git a/Agent0/executor_train/verl/verl/__init__.py b/Agent0/executor_train/verl/verl/__init__.py index 593f3dc..b14513d 100644 --- a/Agent0/executor_train/verl/verl/__init__.py +++ b/Agent0/executor_train/verl/verl/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -37,7 +37,9 @@ if os.getenv("VERL_USE_MODELSCOPE", "False").lower() == "true": if importlib.util.find_spec("modelscope") is None: - raise ImportError("You are using the modelscope hub, please install modelscope by `pip install modelscope -U`") + raise ImportError( + "You are using the modelscope hub, please install modelscope by `pip install modelscope -U`" + ) # Patch hub to download models from modelscope to speed up. from modelscope.utils.hf_util import patch_hub diff --git a/Agent0/executor_train/verl/verl/base_config.py b/Agent0/executor_train/verl/verl/base_config.py index d413160..868bc6d 100644 --- a/Agent0/executor_train/verl/verl/base_config.py +++ b/Agent0/executor_train/verl/verl/base_config.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,11 +13,13 @@ # limitations under the License. import collections -from dataclasses import fields # Import the fields function to inspect dataclass fields +# Import the fields function to inspect dataclass fields +from dataclasses import fields from typing import Any -# BaseConfig class inherits from collections.abc.Mapping, which means it can act like a dictionary +# BaseConfig class inherits from collections.abc.Mapping, which means it +# can act like a dictionary class BaseConfig(collections.abc.Mapping): """The BaseConfig provides omegaconf DictConfig-like interface for a dataclass config. diff --git a/Agent0/executor_train/verl/verl/experimental/__init__.py b/Agent0/executor_train/verl/verl/experimental/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/experimental/__init__.py +++ b/Agent0/executor_train/verl/verl/experimental/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/experimental/agent_loop/__init__.py b/Agent0/executor_train/verl/verl/experimental/agent_loop/__init__.py index c417811..159dc4b 100644 --- a/Agent0/executor_train/verl/verl/experimental/agent_loop/__init__.py +++ b/Agent0/executor_train/verl/verl/experimental/agent_loop/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/experimental/agent_loop/agent_loop.py b/Agent0/executor_train/verl/verl/experimental/agent_loop/agent_loop.py index e16f1a8..bf9da67 100644 --- a/Agent0/executor_train/verl/verl/experimental/agent_loop/agent_loop.py +++ b/Agent0/executor_train/verl/verl/experimental/agent_loop/agent_loop.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,7 +32,11 @@ from verl.single_controller.ray.base import RayWorkerGroup from verl.utils import hf_tokenizer from verl.utils.fs import copy_to_local -from verl.utils.rollout_trace import RolloutTraceConfig, rollout_trace_attr, rollout_trace_op +from verl.utils.rollout_trace import ( + RolloutTraceConfig, + rollout_trace_attr, + rollout_trace_op, +) from verl.workers.rollout.async_server import async_server_class logger = logging.getLogger(__file__) @@ -46,7 +50,12 @@ class AsyncLLMServerManager: - Sticky session: send multi-turn chat completions to same server for automatic prefix caching """ - def __init__(self, config: DictConfig, server_handles: list[ray.actor.ActorHandle], max_cache_size: int = 10000): + def __init__( + self, + config: DictConfig, + server_handles: list[ray.actor.ActorHandle], + max_cache_size: int = 10000, + ): """Initialize the AsyncLLMServerManager. Args: @@ -59,7 +68,9 @@ def __init__(self, config: DictConfig, server_handles: list[ray.actor.ActorHandl random.shuffle(self.server_handles) # Least requests load balancing - self.weighted_serveres = [[0, (hash(server), server)] for server in server_handles] + self.weighted_serveres = [ + [0, (hash(server), server)] for server in server_handles + ] heapq.heapify(self.weighted_serveres) # LRU cache to map request_id to server @@ -126,7 +137,12 @@ class AgentLoopBase(ABC): _class_initialized = False - def __init__(self, config: DictConfig, server_manager: AsyncLLMServerManager, tokenizer: AutoTokenizer): + def __init__( + self, + config: DictConfig, + server_manager: AsyncLLMServerManager, + tokenizer: AutoTokenizer, + ): """Initialize agent loop. Args: @@ -148,7 +164,9 @@ def init_class(cls, config: DictConfig, tokenizer: AutoTokenizer): cls._class_initialized = True @abstractmethod - async def run(self, messages: list[dict[str, Any]], sampling_params: dict[str, Any]) -> AgentLoopOutput: + async def run( + self, messages: list[dict[str, Any]], sampling_params: dict[str, Any] + ) -> AgentLoopOutput: """Run agent loop to interact with LLM server and environment. Args: @@ -165,7 +183,8 @@ async def run(self, messages: list[dict[str, Any]], sampling_params: dict[str, A class AgentLoopWorker: """Agent loop worker takes a batch of messages and run each message in an agent loop.""" - def __init__(self, config: DictConfig, server_handles: list[ray.actor.ActorHandle]): + def __init__(self, config: DictConfig, + server_handles: list[ray.actor.ActorHandle]): """Initialize agent loop manager. Args: @@ -224,7 +243,9 @@ async def generate_sequences(self, batch: DataProto) -> DataProto: # by default, we assume it's a single turn agent if "agent_name" not in batch.non_tensor_batch: - batch.non_tensor_batch["agent_name"] = np.array(["single_turn_agent"] * len(batch), dtype=object) + batch.non_tensor_batch["agent_name"] = np.array( + ["single_turn_agent"] * len(batch), dtype=object + ) tasks = [] agent_names = batch.non_tensor_batch["agent_name"] @@ -234,12 +255,20 @@ async def generate_sequences(self, batch: DataProto) -> DataProto: else: index = np.arange(len(raw_prompts)) - trajectory_info = await get_trajectory_info(batch.meta_info.get("global_steps", -1), index) + trajectory_info = await get_trajectory_info( + batch.meta_info.get("global_steps", -1), index + ) - for agent_name, messages, trajectory in zip(agent_names, raw_prompts, trajectory_info, strict=True): + for agent_name, messages, trajectory in zip( + agent_names, raw_prompts, trajectory_info, strict=True + ): tasks.append( - asyncio.create_task(self._run_agent_loop(agent_name, messages.tolist(), sampling_params, trajectory)) - ) + asyncio.create_task( + self._run_agent_loop( + agent_name, + messages.tolist(), + sampling_params, + trajectory))) outputs = await asyncio.gather(*tasks) output = self._postprocess(outputs) @@ -253,10 +282,14 @@ async def _run_agent_loop( trajectory: dict[str, Any], ) -> AgentLoopOutput: with rollout_trace_attr( - step=trajectory["step"], sample_index=trajectory["sample_index"], rollout_n=trajectory["rollout_n"] + step=trajectory["step"], + sample_index=trajectory["sample_index"], + rollout_n=trajectory["rollout_n"], ): agent_loop_class = self.get_agent_loop_class(agent_name) - agent_loop = agent_loop_class(self.config, self.server_manager, self.tokenizer) + agent_loop = agent_loop_class( + self.config, self.server_manager, self.tokenizer + ) output = await agent_loop.run(messages, sampling_params) return output @@ -276,7 +309,9 @@ def get_agent_loop_class(self, agent_name: str) -> type[AgentLoopBase]: ValueError: If the agent_name is not recognized. """ # TODO: add tool agent registrary - from verl.experimental.agent_loop.single_turn_agent_loop import SingleTurnAgentLoop + from verl.experimental.agent_loop.single_turn_agent_loop import ( + SingleTurnAgentLoop, + ) from verl.experimental.agent_loop.tool_agent_loop import ToolAgentLoop if agent_name == "single_turn_agent": @@ -302,7 +337,10 @@ def _postprocess(self, inputs: list[AgentLoopOutput]) -> DataProto: return_tensors="pt", return_attention_mask=True, ) - prompt_ids, prompt_attention_mask = outputs["input_ids"], outputs["attention_mask"] + prompt_ids, prompt_attention_mask = ( + outputs["input_ids"], + outputs["attention_mask"], + ) # responses self.tokenizer.padding_side = "right" @@ -313,7 +351,10 @@ def _postprocess(self, inputs: list[AgentLoopOutput]) -> DataProto: return_tensors="pt", return_attention_mask=True, ) - response_ids, response_attention_mask = outputs["input_ids"], outputs["attention_mask"] + response_ids, response_attention_mask = ( + outputs["input_ids"], + outputs["attention_mask"], + ) # response_mask outputs = self.tokenizer.pad( @@ -324,13 +365,16 @@ def _postprocess(self, inputs: list[AgentLoopOutput]) -> DataProto: return_attention_mask=False, ) response_mask = outputs["input_ids"] - assert response_ids.shape == response_mask.shape, ( - f"mismatch in response_ids and response_mask shape: {response_ids.shape} vs {response_mask.shape}" - ) + assert ( + response_ids.shape == response_mask.shape), f"mismatch in response_ids and response_mask shape: { + response_ids.shape} vs { + response_mask.shape}" response_mask = response_mask * response_attention_mask input_ids = torch.cat([prompt_ids, response_ids], dim=1) - attention_mask = torch.cat([prompt_attention_mask, response_attention_mask], dim=1) + attention_mask = torch.cat( + [prompt_attention_mask, response_attention_mask], dim=1 + ) position_ids = (attention_mask.cumsum(dim=1) - 1) * attention_mask batch = TensorDict( @@ -338,16 +382,24 @@ def _postprocess(self, inputs: list[AgentLoopOutput]) -> DataProto: "prompts": prompt_ids, # [bsz, prompt_length] "responses": response_ids, # [bsz, response_length] "response_mask": response_mask, # [bsz, response_length] - "input_ids": input_ids, # [bsz, prompt_length + response_length] - "attention_mask": attention_mask, # [bsz, prompt_length + response_length] - "position_ids": position_ids, # [bsz, prompt_length + response_length] + # [bsz, prompt_length + response_length] + "input_ids": input_ids, + # [bsz, prompt_length + response_length] + "attention_mask": attention_mask, + # [bsz, prompt_length + response_length] + "position_ids": position_ids, }, batch_size=len(input_ids), ) - num_turns = np.array([input.num_turns for input in inputs], dtype=np.int32) + num_turns = np.array( + [input.num_turns for input in inputs], dtype=np.int32) metrics = [input.metrics.model_dump() for input in inputs] - return DataProto(batch=batch, non_tensor_batch={"__num_turns__": num_turns}, meta_info={"metrics": metrics}) + return DataProto( + batch=batch, + non_tensor_batch={"__num_turns__": num_turns}, + meta_info={"metrics": metrics}, + ) async def get_trajectory_info(step, index): @@ -359,7 +411,9 @@ async def get_trajectory_info(step, index): rollout_n += 1 else: rollout_n = 0 - trajectory_info.append({"step": step, "sample_index": index[i], "rollout_n": rollout_n}) + trajectory_info.append( + {"step": step, "sample_index": index[i], "rollout_n": rollout_n} + ) return trajectory_info @@ -383,10 +437,14 @@ def __init__(self, config: DictConfig, worker_group: RayWorkerGroup): self.sleep() def _initialize_llm_servers(self): - self.rollout_tp_size = self.config.actor_rollout_ref.rollout.tensor_model_parallel_size + self.rollout_tp_size = ( + self.config.actor_rollout_ref.rollout.tensor_model_parallel_size + ) self.rollout_dp_size = self.worker_group.world_size // self.rollout_tp_size - register_center = ray.get_actor(f"{self.worker_group.name_prefix}_register_center") + register_center = ray.get_actor( + f"{self.worker_group.name_prefix}_register_center" + ) workers_info = ray.get(register_center.get_worker_info.remote()) assert len(workers_info) == self.worker_group.world_size @@ -400,20 +458,29 @@ def _initialize_llm_servers(self): rollout_backend_class=self.config.actor_rollout_ref.rollout.agent.custom_async_server.name, ) else: - server_class = async_server_class(rollout_backend=self.config.actor_rollout_ref.rollout.name) + server_class = async_server_class( + rollout_backend=self.config.actor_rollout_ref.rollout.name + ) # Start all server instances, restart if address already in use. unready_dp_ranks = set(range(self.rollout_dp_size)) while len(unready_dp_ranks) > 0: servers = { rollout_dp_rank: server_class.options( - # make sure AsyncvLLMServer colocates with its corresponding workers + # make sure AsyncvLLMServer colocates with its + # corresponding workers scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( - node_id=workers_info[rollout_dp_rank * self.rollout_tp_size], + node_id=workers_info[rollout_dp_rank * + self.rollout_tp_size], soft=False, ), name=f"async_llm_server_{rollout_dp_rank}", - ).remote(self.config, self.rollout_dp_size, rollout_dp_rank, self.worker_group.name_prefix) + ).remote( + self.config, + self.rollout_dp_size, + rollout_dp_rank, + self.worker_group.name_prefix, + ) for rollout_dp_rank in unready_dp_ranks } @@ -425,14 +492,17 @@ def _initialize_llm_servers(self): unready_dp_ranks.remove(rollout_dp_rank) except Exception: ray.kill(server) - print(f"rollout server {rollout_dp_rank} failed, maybe address already in use, restarting...") + print( + f"rollout server {rollout_dp_rank} failed, maybe address already in use, restarting...") # All server instances are ready, init AsyncLLM engine. - ray.get([server.init_engine.remote() for server in self.async_llm_servers]) + ray.get([server.init_engine.remote() + for server in self.async_llm_servers]) def _init_agent_loop_workers(self): self.agent_loop_workers = [] - for i in range(self.config.actor_rollout_ref.rollout.agent.num_workers): + for i in range( + self.config.actor_rollout_ref.rollout.agent.num_workers): self.agent_loop_workers.append( AgentLoopWorker.options( name=f"agent_loop_worker_{i}", @@ -462,16 +532,24 @@ def generate_sequences(self, prompts: DataProto) -> DataProto: self.sleep() # calculate performance metrics - metrics = [output.meta_info["metrics"] for output in outputs] # List[List[Dict[str, str]]] + metrics = [ + output.meta_info["metrics"] for output in outputs + ] # List[List[Dict[str, str]]] timing = self._performance_metrics(metrics, output) output.meta_info = {"timing": timing} return output - def _performance_metrics(self, metrics: list[list[dict[str, str]]], output: DataProto) -> dict[str, float]: + def _performance_metrics( + self, metrics: list[list[dict[str, str]]], output: DataProto + ) -> dict[str, float]: timing = {} - t_generate_sequences = np.array([metric["generate_sequences"] for chunk in metrics for metric in chunk]) - t_tool_calls = np.array([metric["tool_calls"] for chunk in metrics for metric in chunk]) + t_generate_sequences = np.array( + [metric["generate_sequences"] for chunk in metrics for metric in chunk] + ) + t_tool_calls = np.array( + [metric["tool_calls"] for chunk in metrics for metric in chunk] + ) timing["agent_loop/generate_sequences/min"] = t_generate_sequences.min() timing["agent_loop/generate_sequences/max"] = t_generate_sequences.max() timing["agent_loop/generate_sequences/mean"] = t_generate_sequences.mean() @@ -485,8 +563,12 @@ def _performance_metrics(self, metrics: list[list[dict[str, str]]], output: Data prompt_length = output.batch["prompts"].shape[1] timing["agent_loop/slowest/generate_sequences"] = t_generate_sequences[slowest] timing["agent_loop/slowest/tool_calls"] = t_tool_calls[slowest] - timing["agent_loop/slowest/prompt_length"] = attention_mask[:prompt_length].sum().item() - timing["agent_loop/slowest/response_length"] = attention_mask[prompt_length:].sum().item() + timing["agent_loop/slowest/prompt_length"] = ( + attention_mask[:prompt_length].sum().item() + ) + timing["agent_loop/slowest/response_length"] = ( + attention_mask[prompt_length:].sum().item() + ) return timing diff --git a/Agent0/executor_train/verl/verl/experimental/agent_loop/single_turn_agent_loop.py b/Agent0/executor_train/verl/verl/experimental/agent_loop/single_turn_agent_loop.py index e4021ef..45579ee 100644 --- a/Agent0/executor_train/verl/verl/experimental/agent_loop/single_turn_agent_loop.py +++ b/Agent0/executor_train/verl/verl/experimental/agent_loop/single_turn_agent_loop.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,16 +31,23 @@ def __init__(self, config, server_manager, tokenizer): self.prompt_length = config.actor_rollout_ref.rollout.prompt_length self.response_length = config.actor_rollout_ref.rollout.response_length - async def run(self, messages: list[dict[str, Any]], sampling_params: dict[str, Any]) -> AgentLoopOutput: + async def run( + self, messages: list[dict[str, Any]], sampling_params: dict[str, Any] + ) -> AgentLoopOutput: metrics = {} request_id = uuid4().hex prompt_ids = await self.loop.run_in_executor( - None, lambda: self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True) + None, + lambda: self.tokenizer.apply_chat_template( + messages, add_generation_prompt=True, tokenize=True + ), ) with simple_timer("generate_sequences", metrics): response_ids = await self.server_manager.generate( - request_id=request_id, prompt_ids=prompt_ids, sampling_params=sampling_params + request_id=request_id, + prompt_ids=prompt_ids, + sampling_params=sampling_params, ) response_mask = [1] * len(response_ids) diff --git a/Agent0/executor_train/verl/verl/experimental/agent_loop/tool_agent_loop.py b/Agent0/executor_train/verl/verl/experimental/agent_loop/tool_agent_loop.py index 2756668..2fec4c0 100644 --- a/Agent0/executor_train/verl/verl/experimental/agent_loop/tool_agent_loop.py +++ b/Agent0/executor_train/verl/verl/experimental/agent_loop/tool_agent_loop.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -46,7 +46,8 @@ class FunctionCall(BaseModel): class ToolParser(ABC): @abstractmethod - async def extract_tool_calls(self, responses_ids: list[int]) -> list[FunctionCall]: + async def extract_tool_calls( + self, responses_ids: list[int]) -> list[FunctionCall]: """Extract tool calls from the responses. Args: @@ -66,13 +67,18 @@ def __init__(self, tokenizer) -> None: self.tool_call_start_token: str = "" self.tool_call_end_token: str = "" - self.tool_call_regex = re.compile(r"(.*?)", re.DOTALL) + self.tool_call_regex = re.compile( + r"(.*?)", re.DOTALL) @rollout_trace_op - async def extract_tool_calls(self, responses_ids: list[int]) -> list[FunctionCall]: + async def extract_tool_calls( + self, responses_ids: list[int]) -> list[FunctionCall]: loop = asyncio.get_running_loop() text = await loop.run_in_executor(None, self.tokenizer.decode, responses_ids) - if self.tool_call_start_token not in text or self.tool_call_end_token not in text: + if ( + self.tool_call_start_token not in text + or self.tool_call_end_token not in text + ): return [] matches = self.tool_call_regex.findall(text) @@ -81,7 +87,12 @@ async def extract_tool_calls(self, responses_ids: list[int]) -> list[FunctionCal try: function_call = json.loads(match) name, arguments = function_call["name"], function_call["arguments"] - function_calls.append(FunctionCall(name=name, arguments=json.dumps(arguments, ensure_ascii=False))) + function_calls.append( + FunctionCall( + name=name, + arguments=json.dumps( + arguments, + ensure_ascii=False))) except Exception as e: logger.error(f"Failed to decode tool call: {e}") return function_calls @@ -101,29 +112,48 @@ def init_class(cls, config, tokenizer): # Initialize tools from config file cls.tokenizer = tokenizer cls.max_user_turns = config.actor_rollout_ref.rollout.multi_turn.max_user_turns - cls.max_assistant_turns = config.actor_rollout_ref.rollout.multi_turn.max_assistant_turns - cls.max_parallel_calls = config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls - cls.max_tool_response_length = config.actor_rollout_ref.rollout.multi_turn.max_tool_response_length - cls.tool_response_truncate_side = config.actor_rollout_ref.rollout.multi_turn.tool_response_truncate_side + cls.max_assistant_turns = ( + config.actor_rollout_ref.rollout.multi_turn.max_assistant_turns + ) + cls.max_parallel_calls = ( + config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls + ) + cls.max_tool_response_length = ( + config.actor_rollout_ref.rollout.multi_turn.max_tool_response_length) + cls.tool_response_truncate_side = ( + config.actor_rollout_ref.rollout.multi_turn.tool_response_truncate_side) tool_config_path = config.actor_rollout_ref.rollout.multi_turn.tool_config_path - tool_list = initialize_tools_from_config(tool_config_path) if tool_config_path else [] + tool_list = (initialize_tools_from_config( + tool_config_path) if tool_config_path else []) cls.tools = {tool.name: tool for tool in tool_list} - cls.tool_schemas = [tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) for tool in tool_list] - cls.tool_parser = cls.get_tool_parser(config.actor_rollout_ref.rollout.multi_turn.format) + cls.tool_schemas = [ + tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) + for tool in tool_list + ] + cls.tool_parser = cls.get_tool_parser( + config.actor_rollout_ref.rollout.multi_turn.format + ) print(f"Initialized tools: {cls.tools}") cls.prompt_length = config.actor_rollout_ref.rollout.prompt_length cls.response_length = config.actor_rollout_ref.rollout.response_length - cls.system_prompt = tokenizer.apply_chat_template([{}], add_generation_prompt=False, tokenize=True) + cls.system_prompt = tokenizer.apply_chat_template( + [{}], add_generation_prompt=False, tokenize=True + ) @rollout_trace_op - async def run(self, messages: list[dict[str, Any]], sampling_params: dict[str, Any]) -> AgentLoopOutput: + async def run( + self, messages: list[dict[str, Any]], sampling_params: dict[str, Any] + ) -> AgentLoopOutput: metrics = {} request_id = uuid4().hex prompt_ids = await self.loop.run_in_executor( None, lambda: self.tokenizer.apply_chat_template( - messages, tools=self.tool_schemas, add_generation_prompt=True, tokenize=True + messages, + tools=self.tool_schemas, + add_generation_prompt=True, + tokenize=True, ), ) response_mask = [] @@ -132,7 +162,9 @@ async def run(self, messages: list[dict[str, Any]], sampling_params: dict[str, A while True: with simple_timer("generate_sequences", metrics): response_ids = await self.server_manager.generate( - request_id=request_id, prompt_ids=prompt_ids, sampling_params=sampling_params + request_id=request_id, + prompt_ids=prompt_ids, + sampling_params=sampling_params, ) prompt_ids += response_ids response_mask += [1] * len(response_ids) @@ -171,18 +203,19 @@ async def run(self, messages: list[dict[str, Any]], sampling_params: dict[str, A messages, add_generation_prompt=True, tokenize=True ), ) - tool_response_ids = tool_response_ids[len(self.system_prompt) :] + tool_response_ids = tool_response_ids[len(self.system_prompt):] # NOTE: last turn should not be user turn, or the EOS token reward # can't be propagated to previous token in GAE. - if len(response_mask) + len(tool_response_ids) >= self.response_length: + if len(response_mask) + \ + len(tool_response_ids) >= self.response_length: break prompt_ids += tool_response_ids response_mask += [0] * len(tool_response_ids) user_turns += 1 - response_ids = prompt_ids[-len(response_mask) :] + response_ids = prompt_ids[-len(response_mask):] prompt_ids = prompt_ids[: len(prompt_ids) - len(response_mask)] output = AgentLoopOutput( @@ -198,7 +231,8 @@ async def _call_tool(self, tool_call: FunctionCall) -> dict[str, str]: """Call tool and return tool response.""" tool, instance_id = None, None try: - # TODO: append malformed tool_call to the prompt: invalid function name or arguments + # TODO: append malformed tool_call to the prompt: invalid function + # name or arguments tool_name = tool_call.name tool_args = json.loads(tool_call.arguments) tool = self.tools[tool_name] @@ -214,12 +248,20 @@ async def _call_tool(self, tool_call: FunctionCall) -> dict[str, str]: if len(tool_response) > self.max_tool_response_length: if self.tool_response_truncate_side == "left": - tool_response = tool_response[: self.max_tool_response_length] + "...(truncated)" + tool_response = ( + tool_response[: self.max_tool_response_length] + "...(truncated)" + ) elif self.tool_response_truncate_side == "right": - tool_response = "(truncated)..." + tool_response[-self.max_tool_response_length :] + tool_response = ( + "(truncated)..." + tool_response[-self.max_tool_response_length:] + ) else: length = self.max_tool_response_length // 2 - tool_response = tool_response[:length] + "...(truncated)..." + tool_response[-length:] + tool_response = ( + tool_response[:length] + + "...(truncated)..." + + tool_response[-length:] + ) return { "role": "tool", diff --git a/Agent0/executor_train/verl/verl/experimental/dynamic_dataset/dynamicgen_dataset.py b/Agent0/executor_train/verl/verl/experimental/dynamic_dataset/dynamicgen_dataset.py index a9532aa..e7b3708 100644 --- a/Agent0/executor_train/verl/verl/experimental/dynamic_dataset/dynamicgen_dataset.py +++ b/Agent0/executor_train/verl/verl/experimental/dynamic_dataset/dynamicgen_dataset.py @@ -80,26 +80,29 @@ def __init__( ): super().__init__(data_files, tokenizer, config, processor) self.datagen: AbstractDataGenerator = config.datagen - assert "datagen" in config and config.datagen.get("path", None) is not None, ( - f"datagen path is not set in config: {config}" - ) + assert ("datagen" in config and config.datagen.get("path", None) + is not None), f"datagen path is not set in config: {config}" # Dynamically load the custom datagen class - datagen_cls = load_extern_type(config.datagen.path, config.datagen.name) + datagen_cls = load_extern_type( + config.datagen.path, config.datagen.name) - # Verify that the custom datagen class inherits from AbstractDataGenerator + # Verify that the custom datagen class inherits from + # AbstractDataGenerator abs_cls = AbstractDataGenerator if not issubclass(datagen_cls, abs_cls): raise TypeError( - f"The custom datagen class '{config.datagen.name}' from '{config.datagen.path}'" - + " must inherit from {abs_cls}" - ) + f"The custom datagen class '{ + config.datagen.name}' from '{ + config.datagen.path}'" + + " must inherit from {abs_cls}") self.data_generator = datagen_cls(config.datagen) self.on_batch_end() def append_dataframe(self, new_dataframe: datasets.Dataset): new_dataframe = self.maybe_filter_out_long_prompts(new_dataframe) - self.dataframe = datasets.concatenate_datasets([self.dataframe, new_dataframe]) + self.dataframe = datasets.concatenate_datasets( + [self.dataframe, new_dataframe]) logger.info(f"new dataset len: {len(self.dataframe)}") diff --git a/Agent0/executor_train/verl/verl/interactions/__init__.py b/Agent0/executor_train/verl/verl/interactions/__init__.py index b6db0fc..084c798 100644 --- a/Agent0/executor_train/verl/verl/interactions/__init__.py +++ b/Agent0/executor_train/verl/verl/interactions/__init__.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/interactions/base.py b/Agent0/executor_train/verl/verl/interactions/base.py index 7c5d200..05f40fb 100644 --- a/Agent0/executor_train/verl/verl/interactions/base.py +++ b/Agent0/executor_train/verl/verl/interactions/base.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,9 +20,13 @@ class BaseInteraction: def __init__(self, config: dict[str, Any]): self.config = config - self.name: str = config.get("name", "interaction_agent") # More general agent default role name + self.name: str = config.get( + "name", "interaction_agent" + ) # More general agent default role name - async def start_interaction(self, instance_id: Optional[str] = None, **kwargs) -> str: + async def start_interaction( + self, instance_id: Optional[str] = None, **kwargs + ) -> str: """Create a tool instance. Args: @@ -38,7 +42,9 @@ async def start_interaction(self, instance_id: Optional[str] = None, **kwargs) - async def generate_response( self, instance_id: str, messages: list[dict[str, Any]], **kwargs - ) -> tuple[bool, str, float, dict[str, Any]]: # More clear response generation method + ) -> tuple[ + bool, str, float, dict[str, Any] + ]: # More clear response generation method """ Generates a response for the current turn of interaction. Returns a tuple containing: @@ -51,9 +57,15 @@ async def generate_response( response_content: str = "Your current result seems acceptable." current_turn_score: float = 0.8 additional_data: dict[str, Any] = {} - return should_terminate_sequence, response_content, current_turn_score, additional_data + return ( + should_terminate_sequence, + response_content, + current_turn_score, + additional_data, + ) - async def calculate_score(self) -> float: # More clear score calculation method + # More clear score calculation method + async def calculate_score(self) -> float: """ Calculates a score for the interaction, potentially considering aspects like partial exposure & in-context task switching. @@ -63,7 +75,9 @@ async def calculate_score(self) -> float: # More clear score calculation method score = 0.0 return score - async def finalize_interaction(self) -> None: # More clear interaction end and resource release method + async def finalize_interaction( + self, + ) -> None: # More clear interaction end and resource release method """ Finalizes the interaction session and releases any associated state or resources. Simulates: release state diff --git a/Agent0/executor_train/verl/verl/interactions/gsm8k_interaction.py b/Agent0/executor_train/verl/verl/interactions/gsm8k_interaction.py index 365cbb9..d74dc68 100644 --- a/Agent0/executor_train/verl/verl/interactions/gsm8k_interaction.py +++ b/Agent0/executor_train/verl/verl/interactions/gsm8k_interaction.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -41,7 +41,10 @@ def __init__(self, config: dict): self._instance_dict = {} async def start_interaction( - self, instance_id: Optional[str] = None, ground_truth: Optional[str] = None, **kwargs + self, + instance_id: Optional[str] = None, + ground_truth: Optional[str] = None, + **kwargs ) -> str: if instance_id is None: instance_id = str(uuid4()) @@ -65,7 +68,8 @@ async def generate_response( if content and content.startswith("#### "): self._instance_dict[instance_id]["response"] = content else: - self._instance_dict[instance_id]["response"] = "#### " + (content or "") + self._instance_dict[instance_id]["response"] = "#### " + \ + (content or "") reward = await self.calculate_score(instance_id) if reward == 1.0: diff --git a/Agent0/executor_train/verl/verl/interactions/utils/__init__.py b/Agent0/executor_train/verl/verl/interactions/utils/__init__.py index c4b932b..72375fe 100644 --- a/Agent0/executor_train/verl/verl/interactions/utils/__init__.py +++ b/Agent0/executor_train/verl/verl/interactions/utils/__init__.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/interactions/utils/interaction_registry.py b/Agent0/executor_train/verl/verl/interactions/utils/interaction_registry.py index df747af..8a7dfc0 100644 --- a/Agent0/executor_train/verl/verl/interactions/utils/interaction_registry.py +++ b/Agent0/executor_train/verl/verl/interactions/utils/interaction_registry.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -58,20 +58,25 @@ def initialize_interactions_from_config(interaction_config_file): # Extract config and name config = OmegaConf.to_container(interaction_item.config, resolve=True) - # Get the interaction name - either from config or derive from class name + # Get the interaction name - either from config or derive from class + # name name = interaction_item.get("name", None) if name is None: # If no name is specified, use the class name as default class_simple_name = cls_name.split(".")[-1] - # Remove "Interaction" suffix if present, otherwise use full class name + # Remove "Interaction" suffix if present, otherwise use full class + # name if class_simple_name.endswith("Interaction"): - name = class_simple_name[:-11].lower() # Remove "Interaction" (11 chars) + name = class_simple_name[ + :-11 + ].lower() # Remove "Interaction" (11 chars) else: name = class_simple_name.lower() # Check for duplicate names if name in interaction_map: - raise ValueError(f"Duplicate interaction name '{name}' found. Each interaction must have a unique name.") + raise ValueError( + f"Duplicate interaction name '{name}' found. Each interaction must have a unique name.") # Inject the name into the config config["name"] = name @@ -80,6 +85,7 @@ def initialize_interactions_from_config(interaction_config_file): interaction = interaction_cls(config=config) interaction_map[name] = interaction - logger.info(f"Initialized interaction '{name}' with class '{cls_name}'") + logger.info( + f"Initialized interaction '{name}' with class '{cls_name}'") return interaction_map diff --git a/Agent0/executor_train/verl/verl/model_merger/__init__.py b/Agent0/executor_train/verl/verl/model_merger/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/model_merger/__init__.py +++ b/Agent0/executor_train/verl/verl/model_merger/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/model_merger/__main__.py b/Agent0/executor_train/verl/verl/model_merger/__main__.py index 9d6a4e3..1714fca 100644 --- a/Agent0/executor_train/verl/verl/model_merger/__main__.py +++ b/Agent0/executor_train/verl/verl/model_merger/__main__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/model_merger/base_model_merger.py b/Agent0/executor_train/verl/verl/model_merger/base_model_merger.py index f13f5fb..d76cd4f 100644 --- a/Agent0/executor_train/verl/verl/model_merger/base_model_merger.py +++ b/Agent0/executor_train/verl/verl/model_merger/base_model_merger.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,13 +33,25 @@ def parse_args(): parser = argparse.ArgumentParser(description="verl model merger") - subparsers = parser.add_subparsers(dest="operation", required=True, help="Specify 'merge' or 'test' operation.") + subparsers = parser.add_subparsers( + dest="operation", + required=True, + help="Specify 'merge' or 'test' operation.") base_op_parser = argparse.ArgumentParser(add_help=False) base_op_parser.add_argument( - "--backend", type=str, required=True, choices=["fsdp", "megatron"], help="The backend of the model" + "--backend", + type=str, + required=True, + choices=["fsdp", "megatron"], + help="The backend of the model", + ) + base_op_parser.add_argument( + "--local_dir", + type=str, + default=None, + help="Path to the saved model checkpoints.", ) - base_op_parser.add_argument("--local_dir", type=str, default=None, help="Path to the saved model checkpoints.") base_op_parser.add_argument( "--tie-word-embedding", action="store_true", @@ -57,22 +69,38 @@ def parse_args(): "fit into GPU memory during initialization.", ) - merge_parser = subparsers.add_parser("merge", parents=[base_op_parser], help="Merge model checkpoints and save.") + merge_parser = subparsers.add_parser( + "merge", + parents=[base_op_parser], + help="Merge model checkpoints and save.") merge_parser.add_argument( - "--target_dir", default="tmp", type=str, help="Directory to save the merged huggingface model" + "--target_dir", + default="tmp", + type=str, + help="Directory to save the merged huggingface model", ) merge_parser.add_argument( - "--hf_upload_path", default=None, type=str, help="Hugging Face repository ID to upload the model" + "--hf_upload_path", + default=None, + type=str, + help="Hugging Face repository ID to upload the model", ) merge_parser.add_argument( - "--private", action="store_true", help="Whether to upload the model to a private Hugging Face repository" + "--private", + action="store_true", + help="Whether to upload the model to a private Hugging Face repository", ) test_parser = subparsers.add_parser( - "test", parents=[base_op_parser], help="Test merged model against a reference Hugging Face model" + "test", + parents=[base_op_parser], + help="Test merged model against a reference Hugging Face model", ) test_parser.add_argument( - "--test_hf_dir", type=str, required=True, help="Path to the reference Hugging Face model directory for testing" + "--test_hf_dir", + type=str, + required=True, + help="Path to the reference Hugging Face model directory for testing", ) args = parser.parse_args() @@ -95,7 +123,8 @@ class ModelMergerConfig: use_cpu_initialization: bool = False def __post_init__(self): - self.hf_upload = self.operation == "merge" and bool(self.hf_upload_path) + self.hf_upload = self.operation == "merge" and bool( + self.hf_upload_path) if self.operation == "test": self.target_dir = None self.hf_upload_path = None @@ -161,7 +190,8 @@ class BaseModelMerger(ABC): def __init__(self, config: ModelMergerConfig): self.config = config self.hf_model_config_path = config.hf_model_config_path - self.model_config = AutoConfig.from_pretrained(self.hf_model_config_path) + self.model_config = AutoConfig.from_pretrained( + self.hf_model_config_path) def get_transformers_auto_model_class(self): if "ForTokenClassification" in self.model_config.architectures[0]: @@ -171,7 +201,9 @@ def get_transformers_auto_model_class(self): elif "ForConditionalGeneration" in self.model_config.architectures[0]: return AutoModelForVision2Seq - raise NotImplementedError(f"Unknown architecture {self.model_config.architectures}") + raise NotImplementedError( + f"Unknown architecture {self.model_config.architectures}" + ) def patch_model_generation_config(self, model): """ @@ -182,7 +214,9 @@ def patch_model_generation_config(self, model): """ if model.can_generate(): try: - model.generation_config = GenerationConfig.from_pretrained(self.hf_model_config_path) + model.generation_config = GenerationConfig.from_pretrained( + self.hf_model_config_path + ) except OSError: print( f"Warning: Generation config file not found in {self.hf_model_config_path}, using a " @@ -200,7 +234,8 @@ def save_lora_adapter(self, state_dict: dict[str, torch.Tensor]): Note: This function change the 'state_dict' in place. """ - lora_params_names = [name for name in state_dict.keys() if "lora_" in name] + lora_params_names = [ + name for name in state_dict.keys() if "lora_" in name] if len(lora_params_names) == 0: return None @@ -220,22 +255,34 @@ def save_lora_adapter(self, state_dict: dict[str, torch.Tensor]): target_modules.add(lora_key.split(".")[-3]) lora_params[lora_key] = state_dict.pop(name) - lora_rank = min(lora_params[lora_key].shape[0], lora_params[lora_key].shape[1]) + lora_rank = min( + lora_params[lora_key].shape[0], + lora_params[lora_key].shape[1]) peft_dict = { "r": lora_rank, - "lora_alpha": 0, # lora_alpha is not set. An error should be raised to inform the user to set it manually. + # lora_alpha is not set. An error should be raised to inform the + # user to set it manually. + "lora_alpha": 0, "target_modules": list(target_modules), } peft_config = peft.LoraConfig(**peft_dict).to_dict() - peft_config["task_type"] = peft_config["task_type"].value if peft_config["task_type"] else None - peft_config["peft_type"] = peft_config["peft_type"].value if peft_config["peft_type"] else None + peft_config["task_type"] = ( + peft_config["task_type"].value if peft_config["task_type"] else None) + peft_config["peft_type"] = ( + peft_config["peft_type"].value if peft_config["peft_type"] else None) peft_config["target_modules"] = list(peft_config["target_modules"]) lora_path = os.path.join(self.config.target_dir, "lora_adapter") os.makedirs(lora_path, exist_ok=True) - with open(os.path.join(lora_path, "adapter_config.json"), "w", encoding="utf-8") as f: + with open( + os.path.join(lora_path, "adapter_config.json"), "w", encoding="utf-8" + ) as f: json.dump(peft_config, f, ensure_ascii=False, indent=4) - save_file(lora_params, os.path.join(lora_path, "adapter_model.safetensors")) + save_file( + lora_params, + os.path.join( + lora_path, + "adapter_model.safetensors")) for name in list(state_dict.keys()): key = ( @@ -250,7 +297,9 @@ def save_lora_adapter(self, state_dict: dict[str, torch.Tensor]): def save_hf_model_and_tokenizer(self, state_dict: dict[str, torch.Tensor]): auto_model_class = self.get_transformers_auto_model_class() with init_empty_weights(): - model = auto_model_class.from_config(self.model_config, torch_dtype=torch.bfloat16) + model = auto_model_class.from_config( + self.model_config, torch_dtype=torch.bfloat16 + ) model.to_empty(device="cpu") model = self.patch_model_generation_config(model) @@ -280,7 +329,11 @@ def upload_to_huggingface(self): api = HfApi() try: # Attempt to create repository - api.create_repo(repo_id=self.config.hf_upload_path, private=self.config.private, exist_ok=True) + api.create_repo( + repo_id=self.config.hf_upload_path, + private=self.config.private, + exist_ok=True, + ) except HfHubHTTPError as e: # Handle authentication/API errors if e.response.status_code == 401: @@ -288,26 +341,46 @@ def upload_to_huggingface(self): "Hugging Face authentication failed. Verify your token is valid and has write permissions." ) from e elif e.response.status_code == 404: - raise RepositoryNotFoundError(f"Repository path not found: {self.config.hf_upload_path}") from e + raise RepositoryNotFoundError( + f"Repository path not found: {self.config.hf_upload_path}" + ) from e else: - raise ConnectionError(f"Failed to create repository ({e.response.status_code}): {e}") from e + raise ConnectionError( + f"Failed to create repository ({ + e.response.status_code}): {e}") from e except requests.exceptions.ConnectionError as e: - raise ConnectionError("Network connection failed. Check your internet connection.") from e + raise ConnectionError( + "Network connection failed. Check your internet connection." + ) from e try: # Attempt folder upload - api.upload_folder(folder_path=self.config.target_dir, repo_id=self.config.hf_upload_path, repo_type="model") + api.upload_folder( + folder_path=self.config.target_dir, + repo_id=self.config.hf_upload_path, + repo_type="model", + ) except HfHubHTTPError as e: if e.response.status_code == 401: - raise PermissionError("Authentication failed during upload. Token may have expired.") from e + raise PermissionError( + "Authentication failed during upload. Token may have expired." + ) from e else: - raise RuntimeError(f"Upload failed ({e.response.status_code}): {e}") from e + raise RuntimeError( + f"Upload failed ({e.response.status_code}): {e}" + ) from e except requests.exceptions.ConnectionError as e: - raise ConnectionError("Network interruption during upload. Try again with stable connection.") from e + raise ConnectionError( + "Network interruption during upload. Try again with stable connection." + ) from e except OSError as e: - raise FileNotFoundError(f"Local folder error: {self.config.target_dir} - {str(e)}") from e + raise FileNotFoundError( + f"Local folder error: {self.config.target_dir} - {str(e)}" + ) from e except Exception as e: - raise RuntimeError(f"Unexpected error during upload: {str(e)}") from e + raise RuntimeError( + f"Unexpected error during upload: { + str(e)}") from e @abstractmethod def merge_and_save(self): @@ -315,4 +388,6 @@ def merge_and_save(self): @abstractmethod def cleanup(self): - raise NotImplementedError("Subclasses should implement this method to clean up resources if needed") + raise NotImplementedError( + "Subclasses should implement this method to clean up resources if needed" + ) diff --git a/Agent0/executor_train/verl/verl/model_merger/fsdp_model_merger.py b/Agent0/executor_train/verl/verl/model_merger/fsdp_model_merger.py index 7853b2b..d68d3f5 100644 --- a/Agent0/executor_train/verl/verl/model_merger/fsdp_model_merger.py +++ b/Agent0/executor_train/verl/verl/model_merger/fsdp_model_merger.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -74,7 +74,8 @@ def _get_world_size(self) -> int: """ config_path = Path(self.config.local_dir) / "fsdp_config.json" if not config_path.exists(): - raise FileNotFoundError(f"Config file {config_path} does not exist.") + raise FileNotFoundError( + f"Config file {config_path} does not exist.") with open(config_path) as f: config = json.load(f) @@ -88,12 +89,16 @@ def _get_world_size(self) -> int: def _load_rank_zero_state_dict(self, world_size: int) -> dict: return torch.load( - Path(self.config.local_dir) / f"model_world_size_{world_size}_rank_0.pt", + Path( + self.config.local_dir) / + f"model_world_size_{world_size}_rank_0.pt", map_location="cpu", weights_only=False, ) - def _extract_device_mesh_info(self, state_dict: dict, world_size: int) -> tuple[np.ndarray, tuple[str, ...]]: + def _extract_device_mesh_info( + self, state_dict: dict, world_size: int + ) -> tuple[np.ndarray, tuple[str, ...]]: """ Retrieves sharding information (device_mesh, mesh_dim_names) from a DTensor in the state_dict. If no DTensor is found, infers a simple FSDP mesh based on world_size. @@ -117,7 +122,10 @@ def _calculate_shard_configuration( self, mesh: np.ndarray, mesh_dim_names: tuple[str, ...] ) -> tuple[int, tuple[int, ...]]: """Calculates the total number of shards and the shape of the device mesh.""" - assert mesh_dim_names in (("fsdp",), ("ddp", "fsdp")), f"Unsupported mesh_dim_names {mesh_dim_names}" + assert mesh_dim_names in ( + ("fsdp",), + ("ddp", "fsdp"), + ), f"Unsupported mesh_dim_names {mesh_dim_names}" if "tp" in mesh_dim_names: # TODO: "tp" is not supported yet due to the above assert @@ -129,7 +137,9 @@ def _calculate_shard_configuration( return total_shards, mesh_shape - def _merge_by_placement(self, tensors: list[torch.Tensor], placement: Placement) -> torch.Tensor: + def _merge_by_placement( + self, tensors: list[torch.Tensor], placement: Placement + ) -> torch.Tensor: """Merges a list of tensors based on their DTensor placement""" if placement.is_replicate(): return tensors[0] @@ -141,19 +151,35 @@ def _merge_by_placement(self, tensors: list[torch.Tensor], placement: Placement) raise NotImplementedError(f"Unsupported placement: {placement}") def _load_and_merge_state_dicts( - self, world_size: int, total_shards: int, mesh_shape: tuple[int, ...], mesh_dim_names: tuple[str, ...] + self, + world_size: int, + total_shards: int, + mesh_shape: tuple[int, ...], + mesh_dim_names: tuple[str, ...], ) -> dict[str, torch.Tensor]: model_state_dict_lst = [None] * total_shards def process_one_shard(rank: int, model_state_dict_lst: list): - model_path = Path(self.config.local_dir) / f"model_world_size_{world_size}_rank_{rank}.pt" - state_dict = torch.load(model_path, map_location="cpu", weights_only=False) + model_path = ( + Path(self.config.local_dir) + / f"model_world_size_{world_size}_rank_{rank}.pt" + ) + state_dict = torch.load( + model_path, + map_location="cpu", + weights_only=False) model_state_dict_lst[rank] = state_dict return state_dict with ThreadPoolExecutor(max_workers=min(32, os.cpu_count())) as executor: - futures = [executor.submit(process_one_shard, rank, model_state_dict_lst) for rank in range(total_shards)] - for future in tqdm(futures, desc=f"Loading {total_shards} FSDP shards", total=total_shards): + futures = [ + executor.submit(process_one_shard, rank, model_state_dict_lst) + for rank in range(total_shards) + ] + for future in tqdm( + futures, + desc=f"Loading {total_shards} FSDP shards", + total=total_shards): future.result() # Merge state dicts from all shards @@ -194,7 +220,8 @@ def process_one_shard(rank: int, model_state_dict_lst: list): # 1-D list, FSDP without TP assert len(placements) == 1 shards = state_dict[key] - state_dict[key] = self._merge_by_placement(shards, placements[0]) + state_dict[key] = self._merge_by_placement( + shards, placements[0]) else: # 2-D list, FSDP + TP raise NotImplementedError("FSDP + TP is not supported yet") @@ -207,17 +234,25 @@ def merge_and_save(self): world_size = self._get_world_size() rank_zero_state_dict = self._load_rank_zero_state_dict(world_size) - mesh, mesh_dim_names = self._extract_device_mesh_info(rank_zero_state_dict, world_size) + mesh, mesh_dim_names = self._extract_device_mesh_info( + rank_zero_state_dict, world_size + ) print(f"Got device mesh {mesh}, mesh_dim_names {mesh_dim_names}") - total_shards, mesh_shape = self._calculate_shard_configuration(mesh, mesh_dim_names) - print(f"Processing model shards with {total_shards} {mesh_shape} in total") + total_shards, mesh_shape = self._calculate_shard_configuration( + mesh, mesh_dim_names + ) + print( + f"Processing model shards with {total_shards} {mesh_shape} in total") - merged_state_dict = self._load_and_merge_state_dicts(world_size, total_shards, mesh_shape, mesh_dim_names) + merged_state_dict = self._load_and_merge_state_dicts( + world_size, total_shards, mesh_shape, mesh_dim_names + ) if self.config.operation == "test": if not self.config.test_hf_dir: - raise ValueError("test_hf_dir must be provided for test operation") + raise ValueError( + "test_hf_dir must be provided for test operation") self._validate_state_dict(merged_state_dict) elif self.config.operation == "merge": self.save_hf_model_and_tokenizer(merged_state_dict) @@ -229,7 +264,9 @@ def merge_and_save(self): def _validate_state_dict(self, state_dict: dict[str, torch.Tensor]): auto_model_class = self.get_transformers_auto_model_class() - hf_model = auto_model_class.from_pretrained(self.config.test_hf_dir, torch_dtype=torch.bfloat16) + hf_model = auto_model_class.from_pretrained( + self.config.test_hf_dir, torch_dtype=torch.bfloat16 + ) hf_state_dict = hf_model.state_dict() del hf_model @@ -237,27 +274,36 @@ def _validate_state_dict(self, state_dict: dict[str, torch.Tensor]): collected_keys = set(state_dict.keys()) missing_keys = hf_model_keys - collected_keys - assert len(missing_keys) == 0, f"Missing keys in collected state dict: {list(sorted(missing_keys))}" + assert ( + len(missing_keys) == 0), f"Missing keys in collected state dict: { + list( + sorted(missing_keys))}" extra_keys = collected_keys - hf_model_keys - assert len(extra_keys) == 0, f"Extra keys in collected state dict: {list(sorted(extra_keys))}" + assert ( + len(extra_keys) == 0 + ), f"Extra keys in collected state dict: {list(sorted(extra_keys))}" for key in hf_model_keys: hf_shape = hf_state_dict[key].shape collected_shape = state_dict[key].shape - assert hf_shape == collected_shape, ( - f"Shape mismatch for key '{key}': original {hf_shape} vs collected {collected_shape}" - ) + assert ( + hf_shape == collected_shape + ), f"Shape mismatch for key '{key}': original {hf_shape} vs collected {collected_shape}" hf_dtype = hf_state_dict[key].dtype collected_dtype = state_dict[key].dtype - assert hf_dtype == collected_dtype, ( - f"Dtype mismatch for key '{key}': original {hf_dtype} vs collected {collected_dtype}" - ) + assert ( + hf_dtype == collected_dtype + ), f"Dtype mismatch for key '{key}': original {hf_dtype} vs collected {collected_dtype}" - torch.testing.assert_close(hf_state_dict[key], state_dict[key], atol=1e-6, rtol=1e-6) + torch.testing.assert_close( + hf_state_dict[key], state_dict[key], atol=1e-6, rtol=1e-6 + ) - print("FSDP checks passed: The merged state_dict matches the hf model saved by FSDPCheckpointManager.") + print( + "FSDP checks passed: The merged state_dict matches the hf model saved by FSDPCheckpointManager." + ) def cleanup(self): """Cleanup temporary files if needed.""" diff --git a/Agent0/executor_train/verl/verl/model_merger/megatron_model_merger.py b/Agent0/executor_train/verl/verl/model_merger/megatron_model_merger.py index c40bdf7..ea0122a 100644 --- a/Agent0/executor_train/verl/verl/model_merger/megatron_model_merger.py +++ b/Agent0/executor_train/verl/verl/model_merger/megatron_model_merger.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -86,7 +86,8 @@ class MegatronModelMerger(BaseModelMerger): def __init__(self, config: ModelMergerConfig): super().__init__(config) - # Currently we use only 1 rank to merge the dist_ckpt, we will move to multi-process save shortly afterwards + # Currently we use only 1 rank to merge the dist_ckpt, we will move to + # multi-process save shortly afterwards os.environ["RANK"] = "0" os.environ["WORLD_SIZE"] = "1" os.environ["MASTER_ADDR"] = "localhost" @@ -99,7 +100,8 @@ def __init__(self, config: ModelMergerConfig): expert_model_parallel_size=1, ) model_parallel_cuda_manual_seed(0) - self.hf_config = AutoConfig.from_pretrained(self.config.hf_model_config_path) + self.hf_config = AutoConfig.from_pretrained( + self.config.hf_model_config_path) print(self.hf_config, flush=True) self.params_mapping = { @@ -154,7 +156,8 @@ def _load_state_dicts(self, model_ckpt_path: str) -> dict[str, Any]: # init hf config tf_config = hf_to_mcore_config(self.hf_config, torch.bfloat16) tf_config.use_cpu_initialization = self.config.use_cpu_initialization - tie_word_embeddings = getattr(self.hf_config, "tie_word_embeddings", False) + tie_word_embeddings = getattr( + self.hf_config, "tie_word_embeddings", False) # init megatron model def megatron_model_provider(pre_process, post_process): @@ -182,8 +185,10 @@ def megatron_model_provider(pre_process, post_process): ) if self.config.use_cpu_initialization: - # convert meta device to empty tensor so it can use `copy_` function - whole_model[0].module = whole_model[0].module.to_empty(device="cpu") + # convert meta device to empty tensor so it can use `copy_` + # function + whole_model[0].module = whole_model[0].module.to_empty( + device="cpu") # load state dicts sharded_state_dict = {} @@ -191,7 +196,8 @@ def megatron_model_provider(pre_process, post_process): key = f"model{vpp_rank}" if len(whole_model) > 1 else "model" mpu.set_virtual_pipeline_model_parallel_rank(vpp_rank) sharded_state_dict[key] = model.sharded_state_dict() - model_state_dict = load_dist_checkpointing(sharded_state_dict, model_ckpt_path) + model_state_dict = load_dist_checkpointing( + sharded_state_dict, model_ckpt_path) model_state_dict_list = [] for vpp_rank, model in enumerate(whole_model): key = f"model{vpp_rank}" if len(whole_model) > 1 else "model" @@ -222,11 +228,14 @@ def _check_megatron_state_key(self, key: str) -> bool: # Exclude extra state keys if not key.startswith("decoder"): raise ValueError( - f"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder' in TransformerLayer." - ) + f"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder' in TransformerLayer.") def _split_tensors( - self, key: str, tensor: torch.Tensor, config: PretrainedConfig, is_value_model: bool = False + self, + key: str, + tensor: torch.Tensor, + config: PretrainedConfig, + is_value_model: bool = False, ) -> list[torch.Tensor]: """ Splits a tensor into multiple tensors based on the name. @@ -248,9 +257,9 @@ def _split_tensors( q_lst, k_lst, v_lst = [], [], [] assert config.num_attention_heads % config.num_key_value_heads == 0 num_q_per_kv = config.num_attention_heads // config.num_key_value_heads - assert tensor.shape[0] % (num_q_per_kv + 2) == 0, ( - f"Tensor shape {tensor.shape} is not divisible by {num_q_per_kv + 2}" - ) + assert ( + tensor.shape[0] % (num_q_per_kv + 2) == 0 + ), f"Tensor shape {tensor.shape} is not divisible by {num_q_per_kv + 2}" kv_size = tensor.shape[0] // (num_q_per_kv + 2) split_size = [kv_size * num_q_per_kv, kv_size, kv_size] @@ -266,11 +275,17 @@ def _split_tensors( k_lst.append(k) v_lst.append(v) - return [torch.cat(q_lst, dim=0), torch.cat(k_lst, dim=0), torch.cat(v_lst, dim=0)] + return [ + torch.cat(q_lst, dim=0), + torch.cat(k_lst, dim=0), + torch.cat(v_lst, dim=0), + ] else: return [tensor] - def _merge_state_dicts(self, model_state_dict_list: list[dict[str, Any]]) -> dict[str, torch.Tensor]: + def _merge_state_dicts( + self, model_state_dict_list: list[dict[str, Any]] + ) -> dict[str, torch.Tensor]: state_dict = {} layers_cum = 0 @@ -281,12 +296,16 @@ def _merge_state_dicts(self, model_state_dict_list: list[dict[str, Any]]) -> dic if "extra_state" in key: continue if self.config.tie_word_embedding and ("output_layer" in key): - print("skip lm_head and reward_head loading because of tie_word_embeddings") + print( + "skip lm_head and reward_head loading because of tie_word_embeddings" + ) continue self._check_megatron_state_key(key) hf_name = self._replace_name(key, self.params_mapping) - assert hf_name is not None, f"Failed to convert layer name [{key}] from megatron to huggingface." + assert ( + hf_name is not None + ), f"Failed to convert layer name [{key}] from megatron to huggingface." if "model.layers." in hf_name: local_layer_no = int(hf_name.split(".")[2]) layers_handled = max(local_layer_no, layers_handled) @@ -295,25 +314,36 @@ def _merge_state_dicts(self, model_state_dict_list: list[dict[str, Any]]) -> dic new_key_list[2] = str(global_layer_no) hf_name = ".".join(new_key_list) else: - warnings.warn(f"hf_name {hf_name} will not be fixed with layer number", stacklevel=2) + warnings.warn( + f"hf_name {hf_name} will not be fixed with layer number", + stacklevel=2, + ) tensor = model_state_dict[key] split_tensor = self._split_tensors( - key, tensor, self.hf_config, is_value_model=self.config.is_value_model + key, + tensor, + self.hf_config, + is_value_model=self.config.is_value_model, ) if len(split_tensor) == 1: state_dict[hf_name] = split_tensor[0] elif len(split_tensor) == 3: # split qkv - for n, d in zip(["q", "k", "v"], split_tensor, strict=True): + for n, d in zip(["q", "k", "v"], + split_tensor, strict=True): state_dict[hf_name.replace("qkv", n)] = d elif len(split_tensor) == 2: # split gate up - state_dict[hf_name.replace("gate_up", "gate")] = split_tensor[0] - state_dict[hf_name.replace("gate_up", "up")] = split_tensor[1] + state_dict[hf_name.replace( + "gate_up", "gate")] = split_tensor[0] + state_dict[hf_name.replace( + "gate_up", "up")] = split_tensor[1] shape_info = ( - split_tensor.shape if isinstance(split_tensor, torch.Tensor) else [t.shape for t in split_tensor] + split_tensor.shape + if isinstance(split_tensor, torch.Tensor) + else [t.shape for t in split_tensor] ) print(f"converted {key} to {hf_name} with shape {shape_info}") @@ -332,7 +362,8 @@ def merge_and_save(self): if self.config.operation == "test": if not self.config.test_hf_dir: - raise ValueError("test_hf_dir must be provided for test operation") + raise ValueError( + "test_hf_dir must be provided for test operation") self._validate_state_dict(merged_state_dict) elif self.config.operation == "merge": self.save_hf_model_and_tokenizer(merged_state_dict) @@ -346,11 +377,15 @@ def _validate_state_dict(self, state_dict: dict[str, torch.Tensor]): Compares the merged Megatron state_dict against a reference safetensors model. Applies necessary name mappings from Megatron to Hugging Face conventions using _replace_name. """ - ref_state_dict = load_file(Path(self.config.test_hf_dir) / "model.safetensors") + ref_state_dict = load_file( + Path( + self.config.test_hf_dir) / + "model.safetensors") for name, loaded_weight in state_dict.items(): # name = self._replace_name(original_name, self.params_mapping) - if not name or name.endswith(".bias") and name not in ref_state_dict: + if not name or name.endswith( + ".bias") and name not in ref_state_dict: continue if "rotary_emb.inv_freq" in name: continue @@ -361,9 +396,12 @@ def _validate_state_dict(self, state_dict: dict[str, torch.Tensor]): raise RuntimeError(f"key: {name} not exist in state_dict") param = ref_state_dict[name] assert loaded_weight.dtype == param.dtype - torch.testing.assert_close(loaded_weight.to("cpu"), param, atol=1e-2, rtol=5e-2) + torch.testing.assert_close( + loaded_weight.to("cpu"), param, atol=1e-2, rtol=5e-2 + ) - def _replace_name(self, megatron_name: str, name_mapping: dict[str, str]) -> str: + def _replace_name(self, megatron_name: str, + name_mapping: dict[str, str]) -> str: for m_name, v_name in name_mapping.items(): if m_name not in megatron_name: continue diff --git a/Agent0/executor_train/verl/verl/models/__init__.py b/Agent0/executor_train/verl/verl/models/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/models/__init__.py +++ b/Agent0/executor_train/verl/verl/models/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/models/llama/__init__.py b/Agent0/executor_train/verl/verl/models/llama/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/models/llama/__init__.py +++ b/Agent0/executor_train/verl/verl/models/llama/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/__init__.py b/Agent0/executor_train/verl/verl/models/llama/megatron/__init__.py index fc851ea..b9d86c3 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/__init__.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/__init__.py b/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/__init__.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_loader.py b/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_loader.py index dafecfd..d00488c 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_loader.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_loader.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -29,20 +29,25 @@ def _megatron_calc_layer_map(config): """ from megatron.core import mpu - print(f"get megatron data parallel size: {mpu.get_data_parallel_world_size()}") + print( + f"get megatron data parallel size: { + mpu.get_data_parallel_world_size()}") 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 + 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 - ) + 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, @@ -53,7 +58,12 @@ def _megatron_calc_layer_map(config): def load_state_dict_to_megatron_llama( - state_dict, wrapped_models, config, params_dtype, is_value_model=False, tie_word_embeddings=False + state_dict, + wrapped_models, + config, + params_dtype, + is_value_model=False, + tie_word_embeddings=False, ): """Load merged state_dict to sharded Megatron module in training.""" from megatron.core import DistributedDataParallel as LocalDDP @@ -72,7 +82,9 @@ def _get_gpt_model(model): def fetch_params(module): for param in module.parameters(): torch.distributed.fetch( - param.data, src=mpu.get_data_parallel_src_rank(), group=mpu.get_data_parallel_group() + param.data, + src=mpu.get_data_parallel_src_rank(), + group=mpu.get_data_parallel_group(), ) dp_rank = mpu.get_data_parallel_rank() @@ -82,7 +94,8 @@ def fetch_params(module): 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 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" @@ -91,7 +104,9 @@ def fetch_params(module): 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, ( + assert ( + num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers + ), ( f"num_layers_per_model: {num_layers_per_model} * pp_size: {pp_size} * virtual_pp_size " f"{virtual_pp_size} != config.num_hidden_layers: {config.num_hidden_layers}" ) @@ -99,7 +114,8 @@ def fetch_params(module): models = [None] * len(wrapped_models) for i, wrapped_model in enumerate(wrapped_models): - models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module)) + 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 @@ -109,7 +125,9 @@ def _fetch_tensor(tensor, name) -> torch.Tensor: if tensor is not None: tensor.data.copy_(state_dict[name]) - def _fetch_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor: + def _fetch_tp_shard_tensor_vocab( + tensor, name, chunk_dim=0, mutate_func=None + ) -> torch.Tensor: """fetch tensor in tp shards""" nonlocal state_dict tp_rank = mpu.get_tensor_model_parallel_rank() @@ -125,7 +143,9 @@ def _fetch_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> else: print(f"tp_shard tensor:[{name}] not in state_dict, skip loading") - def _fetch_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor: + def _fetch_tp_shard_tensor( + tensor, name, chunk_dim=0, mutate_func=None + ) -> torch.Tensor: """fetch tensor in tp shards""" nonlocal state_dict tp_rank = mpu.get_tensor_model_parallel_rank() @@ -141,7 +161,8 @@ def _fetch_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch else: print(f"tp_shard tensor:[{name}] not in state_dict, skip loading") - def _fetch_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor: + def _fetch_tp_shard_tensor_gate_up( + tensor, gate_name, up_name) -> torch.Tensor: """fetch gate_up tensor in tp shards""" nonlocal state_dict nonlocal mp_group @@ -151,23 +172,35 @@ def _fetch_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor: 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=get_device_id() + config.intermediate_size * 2, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) 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) - ) + 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) if tensor is not None: tensor.data.copy_(tensor_chunk[tp_rank]) else: - print(f"tp_shard tensor:[{gate_name}, {up_name}] not in state_dict, skip loading") - - def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor: + print( + f"tp_shard tensor:[{gate_name}, {up_name}] not in state_dict, skip loading") + + def _fetch_tp_shard_tensor_qkv( + tensor, + q_name, + k_name, + v_name) -> torch.Tensor: """fetch tensor in tp shards across mp_group""" nonlocal state_dict nonlocal mp_group @@ -185,28 +218,41 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor: 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=get_device_id() + total_size * tp_size, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) 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)) + 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=get_device_id() + total_size * tp_size, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) 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 + 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)) + 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) if tensor is not None: @@ -219,7 +265,9 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor: embed_tokens_weight = None if pp_rank == 0: embed_tokens_weight = gpt_model_module.model.embed_tokens.weight - _fetch_tp_shard_tensor_vocab(embed_tokens_weight, "model.embed_tokens.weight") + _fetch_tp_shard_tensor_vocab( + embed_tokens_weight, + "model.embed_tokens.weight") # Transformer layers # ------------------- @@ -235,10 +283,16 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor: for vpp_rank in range(vpp_size): num_layer_vpp_chunk = num_layer_per_pp // vpp_size num_layer_this_model = num_layer_vpp_chunk - offset = vpp_rank * (config.num_hidden_layers // mpu.get_virtual_pipeline_model_parallel_world_size()) + ( - mpu.get_pipeline_model_parallel_rank() * num_layer_vpp_chunk - ) - layer_list.extend(list(range(offset, offset + num_layer_this_model))) + offset = vpp_rank * ( + config.num_hidden_layers + // mpu.get_virtual_pipeline_model_parallel_world_size() + ) + (mpu.get_pipeline_model_parallel_rank() * num_layer_vpp_chunk) + layer_list.extend( + list( + range( + offset, + offset + + num_layer_this_model))) else: num_layer_this_model = num_layer_per_pp offset = pp_rank * num_layer_per_pp @@ -271,7 +325,11 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor: ) _fetch_tensor( - sync_layer.post_attention_layernorm.weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.post_attention_layernorm.weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.post_attention_layernorm.weight", ) @@ -300,10 +358,16 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor: lm_head_weight = gpt_model_module.lm_head.weight if is_value_model: - if "lm_head.weight" in state_dict and state_dict["lm_head.weight"].shape[0] == 1: + if ( + "lm_head.weight" in state_dict + and state_dict["lm_head.weight"].shape[0] == 1 + ): _fetch_tensor(lm_head_weight, "lm_head.weight") print_rank_0("load lm_head weight") - elif "reward_head.weight" in state_dict and state_dict["reward_head.weight"].shape[0] == 1: + elif ( + "reward_head.weight" in state_dict + and state_dict["reward_head.weight"].shape[0] == 1 + ): _fetch_tensor(lm_head_weight, "reward_head.weight") print_rank_0("load lm_head from value_head weight") else: @@ -314,4 +378,6 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor: dist.barrier() get_torch_device().empty_cache() - print_rank_0(f"loading megatron ckpt done, time elapsed {time.time() - start_time}s") + print_rank_0( + f"loading megatron ckpt done, time elapsed {time.time() - start_time}s" + ) diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_loader_depracated.py b/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_loader_depracated.py index 2f65bc6..42010f5 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_loader_depracated.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_loader_depracated.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -29,20 +29,25 @@ def _megatron_calc_layer_map(config): """ from megatron.core import mpu - print(f"get megatron data parallel size: {mpu.get_data_parallel_world_size()}") + print( + f"get megatron data parallel size: { + mpu.get_data_parallel_world_size()}") 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 + 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 - ) + 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, @@ -53,7 +58,12 @@ def _megatron_calc_layer_map(config): def load_state_dict_to_megatron_llama( - state_dict, wrapped_models, config, params_dtype, is_value_model=False, tie_word_embeddings=False + state_dict, + wrapped_models, + config, + params_dtype, + is_value_model=False, + tie_word_embeddings=False, ): """Load merged state_dict to sharded Megatron module in training.""" from megatron.core import DistributedDataParallel as LocalDDP @@ -72,7 +82,9 @@ def _get_gpt_model(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() + param.data, + src=mpu.get_data_parallel_src_rank(), + group=mpu.get_data_parallel_group(), ) dp_rank = mpu.get_data_parallel_rank() @@ -82,7 +94,8 @@ def broadcast_params(module): 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 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" @@ -91,7 +104,9 @@ def broadcast_params(module): 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, ( + assert ( + num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers + ), ( f"num_layers_per_model: {num_layers_per_model} * pp_size: {pp_size} * virtual_pp_size " f"{virtual_pp_size} != config.num_hidden_layers: {config.num_hidden_layers}" ) @@ -99,7 +114,8 @@ def broadcast_params(module): models = [None] * len(wrapped_models) for i, wrapped_model in enumerate(wrapped_models): - models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module)) + 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 @@ -137,7 +153,9 @@ def _broadcast_tensor(tensor, name) -> torch.Tensor: 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: + 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 @@ -162,7 +180,8 @@ def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None 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") + print_rank_0( + f"tp_shard tensor:[{name}] not in state_dict, skip loading") return if tensor is None: @@ -173,10 +192,13 @@ def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}" + 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=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: @@ -185,7 +207,9 @@ def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None 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: + 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 @@ -209,7 +233,8 @@ def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> t 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") + print_rank_0( + f"tp_shard tensor:[{name}] not in state_dict, skip loading") return if tensor is None: @@ -220,10 +245,13 @@ def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> t requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}" + 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=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: @@ -232,7 +260,8 @@ def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> t if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor) - def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor: + 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 @@ -243,15 +272,22 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens 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=get_device_id() + config.intermediate_size * 2, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) 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) - ) + 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 @@ -263,7 +299,10 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens 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") + print_rank_0( + f"tp_shard tensor:[{ + gate_name, + up_name}] not in state_dict, skip loading") return if tensor is None: @@ -274,11 +313,13 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank() == 0:} tensor {gate_name, up_name} shape " - f"{tensor.shape} != {chunk_shape}" + assert tensor.shape == chunk_shape, (f"rank #{ + torch.distributed.get_rank() == 0:} tensor { + gate_name, up_name} shape " f"{ + tensor.shape} != {chunk_shape}") + sync_tensor = torch.empty_like( + tensor, device=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: @@ -287,7 +328,8 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor) - def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor: + 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 @@ -295,7 +337,8 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tens 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 + 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] @@ -304,16 +347,24 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tens 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 + 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=get_device_id() + total_size * tp_size, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) 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_( + 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) ) @@ -322,15 +373,24 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tens 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=get_device_id() + total_size * tp_size, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) 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 + 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_( + new_weight_qkv[i * total_size: (i + 1) * total_size].copy_( torch.cat([q_part, k_part, v_part], dim=0) ) @@ -344,7 +404,11 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tens 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, k_name, v_name}] not in state_dict, skip loading") + print_rank_0( + f"tp_shard tensor:[{ + q_name, + k_name, + v_name}] not in state_dict, skip loading") return if tensor is None: @@ -355,10 +419,13 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tens requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank()} tensor {q_name} shape {tensor.shape} != {chunk_shape}" + 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=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: @@ -375,7 +442,9 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tens embed_tokens_weight = None if pp_rank == 0: embed_tokens_weight = gpt_model_module.model.embed_tokens.weight - _broadcast_tp_shard_tensor_vocab(embed_tokens_weight, "model.embed_tokens.weight") + _broadcast_tp_shard_tensor_vocab( + embed_tokens_weight, "model.embed_tokens.weight" + ) # Transformer layers # ------------------- @@ -395,7 +464,11 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tens ) _broadcast_tp_shard_tensor_qkv( - sync_layer.self_attn.qkv_proj.weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.self_attn.qkv_proj.weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.self_attn.q_proj.weight", f"{layer_name}.self_attn.k_proj.weight", f"{layer_name}.self_attn.v_proj.weight", @@ -408,7 +481,11 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tens ) _broadcast_tensor( - sync_layer.post_attention_layernorm.weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.post_attention_layernorm.weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.post_attention_layernorm.weight", ) @@ -438,10 +515,16 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tens lm_head_weight = gpt_model_module.lm_head.weight if is_value_model: - if "lm_head.weight" in state_dict and state_dict["lm_head.weight"].shape[0] == 1: + if ( + "lm_head.weight" in state_dict + and state_dict["lm_head.weight"].shape[0] == 1 + ): _broadcast_tensor(lm_head_weight, "lm_head.weight") print_rank_0("load lm_head weight") - elif "reward_head.weight" in state_dict and state_dict["reward_head.weight"].shape[0] == 1: + elif ( + "reward_head.weight" in state_dict + and state_dict["reward_head.weight"].shape[0] == 1 + ): _broadcast_tensor(lm_head_weight, "reward_head.weight") print_rank_0("load lm_head from value_head weight") else: @@ -455,4 +538,6 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tens broadcast_params(wrapped_model) get_torch_device().empty_cache() - print_rank_0(f"loading megatron ckpt done, time elapsed {time.time() - start_time}s") + print_rank_0( + f"loading megatron ckpt done, time elapsed {time.time() - start_time}s" + ) diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_saver.py b/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_saver.py index 595efcd..c06da31 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_saver.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/checkpoint_utils/llama_saver.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,15 +26,18 @@ from verl.utils.megatron_utils import unwrap_model -def _megatron_calc_global_rank(tp_rank: int = 0, dp_rank: int = 0, pp_rank: int = 0): +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.""" 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()}" - ) + 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()}" # We only support TP-DP-PP grouping, for correctness when resharding return (pp_rank * dp_size + dp_rank) * tp_size + tp_rank @@ -53,13 +56,16 @@ def _megatron_calc_layer_map(config): 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 + 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 - ) + 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, @@ -69,7 +75,12 @@ def _megatron_calc_layer_map(config): return layer_map -def merge_megatron_ckpt_llama(wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False): +def merge_megatron_ckpt_llama( + wrapped_models, + config, + dtype, + is_value_model=False, + tie_word_embeddings=False): """Merge sharded parameters of a Megatron module into a merged checkpoint. Args: @@ -96,7 +107,8 @@ def _get_gpt_model(model): 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 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" @@ -105,16 +117,18 @@ def _get_gpt_model(model): 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 + 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 - ) + 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() @@ -130,7 +144,8 @@ 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) + 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: @@ -165,12 +180,15 @@ def _broadcast_tensor(tensor, name, src_pp_rank) -> torch.Tensor: 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: + 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_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) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None @@ -179,7 +197,8 @@ def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_f 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") + print_rank_0( + f"tp_shard tensor:[{name}] not exist, skip collecting") return buffer_tensor = torch.empty( @@ -192,8 +211,14 @@ def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_f 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 + 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: @@ -205,12 +230,15 @@ def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_f 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: + 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_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) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None @@ -219,7 +247,10 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) 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") + print_rank_0( + f"tp_shard tensor:[{ + gate_name, + up_name}] not exist, skip collecting") return buffer_tensor = torch.empty( @@ -232,8 +263,14 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) 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 + 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: @@ -245,7 +282,8 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) 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_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) @@ -254,12 +292,14 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) state_dict[gate_name] = torch.cat(gate_weight_list, dim=0) state_dict[up_name] = torch.cat(up_weight_list, dim=0) - def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): + 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_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) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None @@ -268,7 +308,8 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): 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") + print_rank_0( + f"tp_shard tensor:[{q_name}] not exist, skip collecting") return buffer_tensor = torch.empty( @@ -281,8 +322,14 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): 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 + 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: @@ -297,13 +344,17 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): 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 + 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] + 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] + 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) @@ -312,10 +363,11 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): 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] + 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] + 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) @@ -406,23 +458,35 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): if is_value_model: if pp_rank == pp_size - 1: - print(f"gpt_model_module.lm_head.weight: {gpt_model_module.lm_head.weight.shape}") + print( + f"gpt_model_module.lm_head.weight: { + gpt_model_module.lm_head.weight.shape}") _broadcast_tensor( - gpt_model_module.lm_head.weight if pp_rank == pp_size - 1 else None, + gpt_model_module.lm_head.weight if pp_rank == pp_size - + 1 else None, "lm_head.weight", - src_pp_rank=pp_size - 1, + src_pp_rank=pp_size - + 1, ) _broadcast_tensor( - gpt_model_module.reward_head.weight - if pp_rank == pp_size - 1 and getattr(gpt_model_module, "reward_weight", None) is not None - else None, + ( + gpt_model_module.reward_head.weight if pp_rank == pp_size - + 1 and getattr( + gpt_model_module, + "reward_weight", + None) is not None else None), "reward_head.weight", - src_pp_rank=pp_size - 1, + src_pp_rank=pp_size - + 1, ) else: _broadcast_tp_shard_tensor( - getattr(gpt_model_module.lm_head, "weight", None) if pp_rank == pp_size - 1 else None, + ( + getattr(gpt_model_module.lm_head, "weight", None) + if pp_rank == pp_size - 1 + else None + ), "lm_head.weight", src_pp_rank=pp_size - 1, ) @@ -438,5 +502,8 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): if dtype != v.dtype: state_dict[k] = v.to(dtype) - print_rank_0(f"merge megatron ckpt done, time elapsed {time.time() - start_time}s") + print_rank_0( + f"merge megatron ckpt done, time elapsed { + time.time() - + start_time}s") return state_dict diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/__init__.py b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/__init__.py index 352bc56..99ecdeb 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/__init__.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_attention.py b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_attention.py index e8aacbd..31b31cd 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_attention.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_attention.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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 @@ -36,34 +36,53 @@ class LlamaRotaryEmbedding(nn.Module): - def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + 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)) + 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() + 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) + 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 + # 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) + 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) + 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), @@ -74,26 +93,49 @@ def forward(self, x, seq_len=None): 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): + 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 = 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 + # 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) + self.register_buffer( + "cos_cached", + emb.cos().to(dtype), + persistent=False) + self.register_buffer( + "sin_cached", + emb.sin().to(dtype), + persistent=False) 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): + 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) @@ -102,27 +144,51 @@ def _set_cos_sin_cache(self, seq_len, device, dtype): if seq_len > self.max_position_embeddings: base = self.base * ( - (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1) + (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)) + 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) + 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 + # 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) + self.register_buffer( + "cos_cached", + emb.cos().to(dtype), + persistent=False) + self.register_buffer( + "sin_cached", + emb.sin().to(dtype), + persistent=False) class LlamaLlama3ScalingRotaryEmbedding(LlamaRotaryEmbedding): - def __init__(self, dim, config, max_position_embeddings=2048, base=10000, device=None): + def __init__( + self, + dim, + config, + max_position_embeddings=2048, + base=10000, + device=None): super().__init__(dim, max_position_embeddings, base, device) - self.factor = config.rope_scaling["factor"] # `8` in the original implementation - self.high_freq_factor = config.rope_scaling["high_freq_factor"] # `1` in the original implementation - self.low_freq_factor = config.rope_scaling["low_freq_factor"] # `4` in the original implementation + self.factor = config.rope_scaling[ + "factor" + ] # `8` in the original implementation + self.high_freq_factor = config.rope_scaling[ + "high_freq_factor" + ] # `1` in the original implementation + self.low_freq_factor = config.rope_scaling[ + "low_freq_factor" + ] # `4` in the original implementation self.old_context_len = config.rope_scaling[ "original_max_position_embeddings" ] # `8192` in the original implementation @@ -131,28 +197,39 @@ def __init__(self, dim, config, max_position_embeddings=2048, base=10000, device high_freq_wavelen = self.old_context_len / self.high_freq_factor wavelen = 2 * math.pi / self.inv_freq - # wavelen < high_freq_wavelen: do nothing; wavelen > low_freq_wavelen: divide by factor - inv_freq_llama = torch.where(wavelen > low_freq_wavelen, self.inv_freq / self.factor, self.inv_freq) + # wavelen < high_freq_wavelen: do nothing; wavelen > low_freq_wavelen: + # divide by factor + inv_freq_llama = torch.where( + wavelen > low_freq_wavelen, + self.inv_freq / self.factor, + self.inv_freq) # otherwise: interpolate between the two, using a smooth factor - smooth_factor = (self.old_context_len / wavelen - self.low_freq_factor) / ( - self.high_freq_factor - self.low_freq_factor - ) - smoothed_inv_freq = (1 - smooth_factor) * inv_freq_llama / self.factor + smooth_factor * inv_freq_llama - is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen) - inv_freq = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama) + smooth_factor = (self.old_context_len / wavelen - self.low_freq_factor) / \ + (self.high_freq_factor - self.low_freq_factor) + smoothed_inv_freq = ( + 1 - smooth_factor + ) * inv_freq_llama / self.factor + smooth_factor * inv_freq_llama + is_medium_freq = ~(wavelen < high_freq_wavelen) * \ + ~(wavelen > low_freq_wavelen) + inv_freq = torch.where( + is_medium_freq, + smoothed_inv_freq, + inv_freq_llama) 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() + seq_len=max_position_embeddings, + device=self.inv_freq.device, + dtype=torch.get_default_dtype(), ) def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] + x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) @@ -172,14 +249,20 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: 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) + 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) class ParallelLlamaAttention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" - def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): + def __init__( + self, + config: LlamaConfig, + megatron_config: ModelParallelConfig): super().__init__() self.config = config self.megatron_config = megatron_config @@ -193,9 +276,10 @@ def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): # 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_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=" f"{self.num_key_value_heads}, tp_size={tp_size}" @@ -207,16 +291,18 @@ def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): 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} and " - f"`num_heads`: {self.num_heads})." - ) + f"hidden_size must be divisible by num_heads (got `hidden_size`: { + self.hidden_size} and " f"`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" + 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) @@ -255,7 +341,9 @@ def _init_rope(self): base=self.rope_theta, ) else: - rope_type_key = "type" if "type" in self.config.rope_scaling else "rope_type" + rope_type_key = ( + "type" if "type" in self.config.rope_scaling else "rope_type" + ) scaling_type = self.config.rope_scaling[rope_type_key] scaling_factor = self.config.rope_scaling["factor"] if scaling_type == "linear": @@ -283,53 +371,85 @@ def _init_rope(self): 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() + 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]]]: + 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, 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) + 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) + 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) + 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)}, " - f"but is {attn_weights.size()}" - ) + f"Attention weights should be of size { + ( + bsz, + self.num_heads_per_tp, + q_len, + kv_seq_len)}, " f"but is { + 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()}" - ) + 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_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)}, " - f"but is {attn_output.size()}" - ) + f"`attn_output` should be of size { + ( + bsz, + self.num_heads_per_tp, + q_len, + self.head_dim)}, " f"but is { + attn_output.size()}") attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, self.hidden_size_per_tp) @@ -349,18 +469,35 @@ def forward( from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa -def apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, sequence_length): +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) + 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) + 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 @@ -369,10 +506,22 @@ def apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, sequence_l # cos/sin shoudl be: (seq_length, rotary_dim / 2) 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 + 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 + k, + cos, + sin, + interleaved=False, + inplace=False, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, ) return q_embed, k_embed @@ -387,7 +536,9 @@ def forward( 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 + 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() @@ -407,14 +558,28 @@ def forward( # 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) + 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 + 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, + 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, @@ -449,12 +614,16 @@ def forward( ) 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() + 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 = 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 diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_decoder.py b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_decoder.py index f46e945..9c1e996 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_decoder.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_decoder.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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 @@ -33,23 +33,32 @@ class ParallelLlamaDecoderLayer(nn.Module): - def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, layer_idx: int): + def __init__( + self, + config: LlamaConfig, + megatron_config: ModelParallelConfig, + layer_idx: int): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) self.layer_idx = layer_idx self.hidden_size = config.hidden_size - self.self_attn = ParallelLlamaAttention(config=config, megatron_config=megatron_config) + 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]]]: + 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)` @@ -100,26 +109,35 @@ def forward( class ParallelLlamaDecoderLayerRmPad(nn.Module): - def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, layer_idx: int): + def __init__( + self, + config: LlamaConfig, + megatron_config: ModelParallelConfig, + layer_idx: int): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) self.layer_idx = layer_idx self.hidden_size = config.hidden_size - self.self_attn = ParallelLlamaAttentionRmPad(config=config, megatron_config=megatron_config) + 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]]]: + 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) diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_linear.py b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_linear.py index 043726c..6c5f59f 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_linear.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_linear.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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. @@ -11,7 +11,8 @@ # WITHOUT 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 +# Adapted from +# https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/linear.py import torch from megatron.core import tensor_parallel @@ -102,5 +103,7 @@ def forward( logits = super().forward(input_) logits = logits.float() if self.sequence_parallel: - logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False) + logits = tensor_parallel.gather_from_sequence_parallel_region( + logits, tensor_parallel_output_grad=False + ) return logits, None diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_mlp.py b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_mlp.py index 583a317..22943d2 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_mlp.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_mlp.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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 @@ -28,19 +28,25 @@ class ParallelLlamaMLP(nn.Module): - def __init__(self, config, megatron_config: ModelParallelConfig = None) -> None: + 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] + # 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" + 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) diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_rmsnorm.py b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_rmsnorm.py index bc2e9ae..a06db13 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_rmsnorm.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/layers/parallel_rmsnorm.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,7 +24,10 @@ class ParallelLlamaRMSNorm(nn.Module): - def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): + def __init__( + self, + config: LlamaConfig, + megatron_config: ModelParallelConfig): """ LlamaRMSNorm is equivalent to T5LayerNorm """ diff --git a/Agent0/executor_train/verl/verl/models/llama/megatron/modeling_llama_megatron.py b/Agent0/executor_train/verl/verl/models/llama/megatron/modeling_llama_megatron.py index ed5022e..90546f9 100644 --- a/Agent0/executor_train/verl/verl/models/llama/megatron/modeling_llama_megatron.py +++ b/Agent0/executor_train/verl/verl/models/llama/megatron/modeling_llama_megatron.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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 @@ -33,10 +33,14 @@ from verl.utils.megatron import tensor_parallel as tp_utils from verl.utils.megatron_utils import TransformerConfig, convert_config -from .layers import ParallelLlamaDecoderLayer, ParallelLlamaDecoderLayerRmPad, ParallelLlamaRMSNorm +from .layers import ( + ParallelLlamaDecoderLayer, + ParallelLlamaDecoderLayerRmPad, + ParallelLlamaRMSNorm, +) """ -TODO: +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 @@ -44,12 +48,18 @@ # Copied from transformers.models.bart.modeling_bart._make_causal_mask -def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device): +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 = 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) @@ -57,18 +67,24 @@ def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: t # Copied from transformers.models.bart.modeling_bart._expand_mask -def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): +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) + 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) + return inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(dtype).min + ) class ParallelLlamaModel(nn.Module): @@ -79,26 +95,41 @@ class ParallelLlamaModel(nn.Module): config: LlamaConfig """ - def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): + def __init__( + self, + config: LlamaConfig, + megatron_config: ModelParallelConfig): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) 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) + 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 + 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)] + [ + 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): + # 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 @@ -111,11 +142,13 @@ def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_em 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 - ) + 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 + expanded_attn_mask + if combined_attention_mask is None + else expanded_attn_mask + combined_attention_mask ) return combined_attention_mask @@ -140,7 +173,9 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) # embed positions - attention_mask = self._prepare_decoder_attention_mask(attention_mask, (batch_size, seq_length), inputs_embeds) + attention_mask = self._prepare_decoder_attention_mask( + attention_mask, (batch_size, seq_length), inputs_embeds + ) hidden_states = inputs_embeds @@ -159,16 +194,23 @@ def forward( class ParallelLlamaForCausalLM(nn.Module): - def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): + def __init__( + self, + config: LlamaConfig, + megatron_config: ModelParallelConfig): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) - self.model = ParallelLlamaModel(config, megatron_config=megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) + 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) + 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, @@ -195,7 +237,8 @@ def forward( Returns: ```""" - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + # decoder outputs consists of (dec_features, layer_state, dec_hidden, + # dec_attn) outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -205,7 +248,8 @@ def forward( hidden_states = outputs logits = self.lm_head(hidden_states)[0] - logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits) + logits = tensor_parallel.gather_from_tensor_model_parallel_region( + logits) logits = logits.float() return CausalLMOutputWithPast( @@ -228,22 +272,34 @@ class ParallelLlamaModelRmPad(nn.Module): config: LlamaConfig """ - def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): + def __init__( + self, + config: LlamaConfig, + megatron_config: ModelParallelConfig): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) 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) + 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 + 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)] + [ + ParallelLlamaDecoderLayerRmPad(config, megatron_config) + for _ in range(config.num_hidden_layers) + ] ) self.norm = ParallelLlamaRMSNorm(config, megatron_config) @@ -265,12 +321,15 @@ def forward( Returns: """ - inputs_embeds = self.embed_tokens(input_ids) # (1, total_nnz) -> (1, total_nnz, hidden_size) + 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) + inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region( + inputs_embeds) hidden_states = inputs_embeds for idx, decoder_layer in enumerate(self.layers): @@ -291,19 +350,26 @@ def forward( class ParallelLlamaForCausalLMRmPad(nn.Module): - def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig): + def __init__( + self, + config: LlamaConfig, + megatron_config: ModelParallelConfig): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) self.megatron_config = megatron_config - self.model = ParallelLlamaModelRmPad(config, megatron_config=megatron_config) + self.model = ParallelLlamaModelRmPad( + config, megatron_config=megatron_config) self.vocab_size = config.vocab_size self._init_head(config) def _init_head(self, config): 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) + 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, @@ -317,7 +383,9 @@ 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) + logits = tensor_parallel.gather_from_tensor_model_parallel_region( + logits + ) # (total_nnz_padded, 1, vocab_size) return logits def forward( @@ -343,7 +411,8 @@ def forward( ) # (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 + # 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) @@ -367,7 +436,8 @@ def forward( totol_nnz = cu_seqlens[-1] logits = logits[:totol_nnz] # (total_nnz_padded) - logits = torch.squeeze(logits, dim=1) # remove the artificial batch dimension + # remove the artificial batch dimension + logits = torch.squeeze(logits, dim=1) # add removed padding back logits = pad_input( logits, indices, batch_size, seqlen=sequence_length @@ -386,9 +456,13 @@ class ParallelLlamaForValueRmPad(ParallelLlamaForCausalLMRmPad): def _init_head(self, config): 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=config.hidden_size, out_features=1, bias=False) + 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=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) @@ -396,7 +470,9 @@ 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) + logits = tensor_parallel.gather_from_sequence_parallel_region( + logits, tensor_parallel_output_grad=False + ) return logits def forward( @@ -425,9 +501,16 @@ class ParallelLlamaModelRmPadPP(nn.Module): config: LlamaConfig """ - def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pre_process, post_process): + def __init__( + self, + config: LlamaConfig, + megatron_config: ModelParallelConfig, + pre_process, + post_process, + ): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.pre_process = pre_process @@ -435,11 +518,16 @@ def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pr 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) + 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 + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + **embedding_kwargs, ) else: self.embed_tokens = None @@ -454,14 +542,18 @@ def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pr self.layers = nn.ModuleList() self.num_layer_vpp_chunk = self.num_layer_per_pp // vpp_size self.num_layer_this_model = self.num_layer_vpp_chunk - offset = vpp_rank * (config.num_hidden_layers // vpp_size) + (pp_rank * self.num_layer_vpp_chunk) + offset = vpp_rank * (config.num_hidden_layers // vpp_size) + ( + pp_rank * self.num_layer_vpp_chunk + ) else: self.num_layer_this_model = self.num_layer_per_pp offset = pp_rank * self.num_layer_per_pp self.layers = nn.ModuleList() for i in range(self.num_layer_this_model): - layer = ParallelLlamaDecoderLayerRmPad(config, megatron_config, layer_idx=offset + i) + layer = ParallelLlamaDecoderLayerRmPad( + config, megatron_config, layer_idx=offset + i + ) self.layers.add_module(f"{i}", layer) if post_process: @@ -498,14 +590,17 @@ def forward( """ if self.pre_process: - inputs_embeds = self.embed_tokens(input_ids) # (1, total_nnz) -> (1, total_nnz, hidden_size) + 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) + inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region( + inputs_embeds) hidden_states = inputs_embeds else: @@ -540,14 +635,18 @@ def __init__( share_embeddings_and_output_weights=False, ): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) self.megatron_config = megatron_config self.model = ParallelLlamaModelRmPadPP( - config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process - ) - assert share_embeddings_and_output_weights is False, ( - "Llama Model not supports sharing embedding and output weights" + config, + megatron_config=megatron_config, + pre_process=pre_process, + post_process=post_process, ) + assert ( + share_embeddings_and_output_weights is False + ), "Llama Model not supports sharing embedding and output weights" self.share_embeddings_and_output_weights = share_embeddings_and_output_weights self.vocab_size = config.vocab_size self.pre_process = pre_process @@ -569,8 +668,10 @@ def set_input_tensor(self, input_tensor): def _init_head(self, config): 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) + 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, @@ -607,7 +708,8 @@ def forward( ```""" # 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 + # 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( @@ -615,9 +717,11 @@ def forward( ) # (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 + # 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 = sp_utils.pad_to_sequence_parallel( + input_ids_rmpad) input_ids_rmpad = input_ids_rmpad.transpose(0, 1) # (1, total_nnz+pad) @@ -632,15 +736,18 @@ def forward( if self.post_process: hidden_states = outputs - # print(f'hidden_states.shape = {hidden_states.shape}') # torch.Size([4, 32, 4096]) + # 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 the artificial batch dimension # torch.Size([8, 32, 16]) + logits = torch.squeeze(logits, dim=1) # 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 + # 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) @@ -660,9 +767,13 @@ class ParallelLlamaForValueRmPadPP(ParallelLlamaForCausalLMRmPadPP): def _init_head(self, config): 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=config.hidden_size, out_features=1, bias=False) + 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=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) @@ -670,7 +781,9 @@ 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) + logits = tensor_parallel.gather_from_sequence_parallel_region( + logits, tensor_parallel_output_grad=False + ) return logits def forward( @@ -680,7 +793,11 @@ def forward( attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> tuple | CausalLMOutputWithPast: - output = super().forward(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) + 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 diff --git a/Agent0/executor_train/verl/verl/models/mcore/__init__.py b/Agent0/executor_train/verl/verl/models/mcore/__init__.py index 29d0531..ed3c9ca 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/__init__.py +++ b/Agent0/executor_train/verl/verl/models/mcore/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); diff --git a/Agent0/executor_train/verl/verl/models/mcore/config_converter.py b/Agent0/executor_train/verl/verl/models/mcore/config_converter.py index 597afcd..1b4c28d 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/config_converter.py +++ b/Agent0/executor_train/verl/verl/models/mcore/config_converter.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # @@ -25,7 +25,9 @@ def _get_base_transformer_config( - hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs + hf_config: PretrainedConfig, + dtype: torch.dtype, + **override_transformer_config_kwargs, ) -> dict: """ Create a base TransformerConfig with common parameters across different model architectures. @@ -92,7 +94,10 @@ def _get_base_transformer_config( def _get_mla_transformer_config( - hf_config: PretrainedConfig, mla_rope_config: dict, dtype: torch.dtype, **override_transformer_config_kwargs + hf_config: PretrainedConfig, + mla_rope_config: dict, + dtype: torch.dtype, + **override_transformer_config_kwargs, ) -> dict: """ Create a MLATransformerConfig with common parameters across different model architectures. @@ -107,7 +112,9 @@ def _get_mla_transformer_config( Returns: MLATransformerConfig with common parameters """ - base_config = _get_base_transformer_config(hf_config=hf_config, dtype=dtype, **override_transformer_config_kwargs) + base_config = _get_base_transformer_config( + hf_config=hf_config, dtype=dtype, **override_transformer_config_kwargs + ) mla_config = { # MLA specific parameters "q_lora_rank": hf_config.q_lora_rank, @@ -130,10 +137,16 @@ def _get_mla_transformer_config( def hf_to_mcore_config_dense( - hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs + hf_config: PretrainedConfig, + dtype: torch.dtype, + **override_transformer_config_kwargs, ) -> TransformerConfig: # for LlamaForCausalLM or Qwen2ForCausalLM - qkv_bias = True if "Qwen2ForCausalLM" in hf_config.architectures else getattr(hf_config, "attention_bias", False) + qkv_bias = ( + True + if "Qwen2ForCausalLM" in hf_config.architectures + else getattr(hf_config, "attention_bias", False) + ) qk_layernorm = True if "Qwen3ForCausalLM" in hf_config.architectures else False args: dict = _get_base_transformer_config( @@ -151,7 +164,9 @@ def hf_to_mcore_config_dense( def hf_to_mcore_config_qwen2moe( - hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs + hf_config: PretrainedConfig, + dtype: torch.dtype, + **override_transformer_config_kwargs, ) -> TransformerConfig: args: dict = _get_base_transformer_config( hf_config=hf_config, @@ -167,7 +182,8 @@ def hf_to_mcore_config_qwen2moe( moe_shared_expert_intermediate_size=hf_config.shared_expert_intermediate_size, moe_aux_loss_coeff=hf_config.router_aux_loss_coef, # moe_aux_loss_coeff=0.0, - moe_router_load_balancing_type="none", # turn off aux_loss as it hurts perf in RL + moe_router_load_balancing_type="none", + # turn off aux_loss as it hurts perf in RL moe_shared_expert_overlap=True, moe_grouped_gemm=True, moe_router_score_function="softmax", @@ -186,7 +202,9 @@ def hf_to_mcore_config_qwen2moe( def hf_to_mcore_config_mixtral( - hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs + hf_config: PretrainedConfig, + dtype: torch.dtype, + **override_transformer_config_kwargs, ) -> TransformerConfig: args: dict = _get_base_transformer_config( hf_config=hf_config, @@ -199,7 +217,8 @@ def hf_to_mcore_config_mixtral( moe_aux_loss_coeff=hf_config.router_aux_loss_coef, moe_router_topk=hf_config.num_experts_per_tok, moe_router_pre_softmax=True, - moe_router_load_balancing_type="none", # turn off aux_loss as it hurts perf in RL + moe_router_load_balancing_type="none", + # turn off aux_loss as it hurts perf in RL moe_router_score_function="softmax", moe_shared_expert_intermediate_size=None, # mixtral has no shared expert moe_shared_expert_overlap=False, # mixtral has no shared expert @@ -220,7 +239,9 @@ def hf_to_mcore_config_mixtral( def hf_to_mcore_config_qwen3moe( - hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs + hf_config: PretrainedConfig, + dtype: torch.dtype, + **override_transformer_config_kwargs, ) -> TransformerConfig: args: dict = _get_base_transformer_config( hf_config=hf_config, @@ -235,7 +256,8 @@ def hf_to_mcore_config_qwen3moe( num_moe_experts=hf_config.num_experts, moe_aux_loss_coeff=hf_config.router_aux_loss_coef, # moe_aux_loss_coeff=0.0, - moe_router_load_balancing_type="none", # turn off aux_loss as it hurts perf in RL + moe_router_load_balancing_type="none", + # turn off aux_loss as it hurts perf in RL moe_grouped_gemm=True, moe_router_score_function="softmax", # Other optimizations @@ -253,7 +275,9 @@ def hf_to_mcore_config_qwen3moe( def hf_to_mcore_config_dpskv3( - hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs + hf_config: PretrainedConfig, + dtype: torch.dtype, + **override_transformer_config_kwargs, ) -> MLATransformerConfig: # DeepseekV3ForCausalLM from megatron.core.transformer.enums import AttnBackend @@ -274,17 +298,18 @@ def hf_to_mcore_config_dpskv3( if "rope_scaling" in hf_config and hf_config.rope_scaling is not None: mla_rope_config.update(hf_config.rope_scaling) moe_layer_freq = [1] * hf_config.num_hidden_layers - for i in range(min(hf_config.first_k_dense_replace, hf_config.num_hidden_layers)): + for i in range(min(hf_config.first_k_dense_replace, + hf_config.num_hidden_layers)): moe_layer_freq[i] = 0 # disable MTP and quantization for now if "num_nextn_predict_layers" in hf_config: - assert hf_config.num_nextn_predict_layers == 0, ( - "MTP is not supported for now, please modify the config.json to set num_nextn_predict_layers to 0" - ) - assert "quantization_config" not in hf_config or not hf_config.quantization_config, ( - "quantization is not supported for now, please modify the config.json to remove quantization_config" - ) + assert ( + hf_config.num_nextn_predict_layers == 0 + ), "MTP is not supported for now, please modify the config.json to set num_nextn_predict_layers to 0" + assert ( + "quantization_config" not in hf_config or not hf_config.quantization_config + ), "quantization is not supported for now, please modify the config.json to remove quantization_config" args: dict = _get_mla_transformer_config( hf_config=hf_config, @@ -302,7 +327,8 @@ def hf_to_mcore_config_dpskv3( moe_router_enable_expert_bias=True, moe_router_topk=hf_config.num_experts_per_tok, num_moe_experts=hf_config.n_routed_experts, - moe_shared_expert_intermediate_size=hf_config.moe_intermediate_size * hf_config.n_shared_experts, + moe_shared_expert_intermediate_size=hf_config.moe_intermediate_size + * hf_config.n_shared_experts, moe_aux_loss_coeff=getattr(hf_config, "aux_loss_alpha", 0.001), moe_router_load_balancing_type="seq_aux_loss", moe_shared_expert_overlap=True, @@ -335,7 +361,9 @@ def hf_to_mcore_config_dpskv3( def hf_to_mcore_config_qwen2_5_vl( - hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs + hf_config: PretrainedConfig, + dtype: torch.dtype, + **override_transformer_config_kwargs, ) -> TransformerConfig: # Qwen2_5_VLForConditionalGeneration @@ -354,7 +382,10 @@ def hf_to_mcore_config_qwen2_5_vl( def hf_to_mcore_config_llama4( - hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs + hf_config: PretrainedConfig, + dtype: torch.dtype, + **override_transformer_config_kwargs, ) -> TransformerConfig: # Llama4ForConditionalGeneration - raise NotImplementedError("Llama4ForConditionalGeneration is not supported yet") + raise NotImplementedError( + "Llama4ForConditionalGeneration is not supported yet") diff --git a/Agent0/executor_train/verl/verl/models/mcore/loader.py b/Agent0/executor_train/verl/verl/models/mcore/loader.py index 659b4ba..b280ae5 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/loader.py +++ b/Agent0/executor_train/verl/verl/models/mcore/loader.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -37,13 +37,16 @@ def _megatron_calc_layer_map(config): 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 + 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 - ) + 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, @@ -53,7 +56,9 @@ def _megatron_calc_layer_map(config): return layer_map -def load_state_dict_to_megatron_gptmodel(state_dict, wrapped_models, config, params_dtype, is_value_model=False): +def load_state_dict_to_megatron_gptmodel( + state_dict, wrapped_models, config, params_dtype, is_value_model=False +): """Load merged state_dict to sharded Megatron module in training.""" from megatron.core import DistributedDataParallel as LocalDDP from megatron.core import mpu @@ -71,19 +76,24 @@ def _get_gpt_model(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() + 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() cp_rank = mpu.get_context_parallel_rank() - src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=0, cp_rank=cp_rank) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=0, cp_rank=cp_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() == src_rank: - assert mp_group.rank() == 0, f"mp_rank:[{mp_group.rank}] != 0 on 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" @@ -92,12 +102,14 @@ def broadcast_params(module): 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 + 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)) + models[i] = unwrap_model( + wrapped_model, (torchDDP, LocalDDP, Float16Module)) gpt_model_module = _get_gpt_model(models[i]) assert len(gpt_model_module.decoder.layers) == num_layers_per_model @@ -135,7 +147,9 @@ def _broadcast_tensor(tensor, name) -> torch.Tensor: tensor.data.copy_(weight) dist.broadcast(tensor, src=src_rank, group=mp_group) - def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor: + 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 @@ -160,7 +174,8 @@ def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None 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") + print_rank_0( + f"tp_shard tensor:[{name}] not in state_dict, skip loading") return if tensor is None: @@ -171,10 +186,13 @@ def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}" + 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=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == src_rank: @@ -183,7 +201,9 @@ def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None 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: + 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 @@ -207,7 +227,8 @@ def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> t 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") + print_rank_0( + f"tp_shard tensor:[{name}] not in state_dict, skip loading") return if tensor is None: @@ -218,10 +239,13 @@ def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> t requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}" + 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=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == src_rank: @@ -230,7 +254,8 @@ def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> t if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor) - def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor: + 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 @@ -241,15 +266,22 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens 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=get_device_id() + config.intermediate_size * 2, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) 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) - ) + 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 @@ -261,7 +293,10 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens 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") + print_rank_0( + f"tp_shard tensor:[{ + gate_name, + up_name}] not in state_dict, skip loading") return if tensor is None: @@ -272,11 +307,13 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank() == src_rank:} tensor {gate_name, up_name} shape " - f"{tensor.shape} != {chunk_shape}" + assert tensor.shape == chunk_shape, (f"rank #{ + torch.distributed.get_rank() == src_rank:} tensor { + gate_name, up_name} shape " f"{ + tensor.shape} != {chunk_shape}") + sync_tensor = torch.empty_like( + tensor, device=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == src_rank: @@ -285,7 +322,9 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor) - def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> torch.Tensor: + def _broadcast_tp_shard_tensor_qkv( + tensor, q_name, k_name, v_name, bias=False + ) -> torch.Tensor: """broadcast tensor in tp shards across mp_group""" nonlocal state_dict nonlocal mp_group @@ -293,34 +332,65 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - tp_size = mpu.get_tensor_model_parallel_world_size() if torch.distributed.get_rank() == src_rank: - assert q_name in state_dict and k_name in state_dict and v_name in state_dict + 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 = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + hidden_size_per_head = getattr( + config, + "head_dim", + config.hidden_size // + config.num_attention_heads) if config.num_key_value_heads >= tp_size: q_size_tp = hidden_size_per_head * config.num_attention_heads // tp_size - kv_size_tp = hidden_size_per_head * config.num_key_value_heads // 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 sizes = [total_size * tp_size] if not bias: sizes.append(config.hidden_size) - new_weight_qkv = torch.empty(*sizes, dtype=params_dtype, device=get_device_id()) + new_weight_qkv = torch.empty( + *sizes, dtype=params_dtype, device=get_device_id() + ) 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] - num_query_groups_per_partition = models[0].config.num_query_groups // tp_size - new_weight_qkv_this_tp = new_weight_qkv[i * total_size : (i + 1) * total_size] - q_part_per_head = torch.chunk(q_part, num_query_groups_per_partition, dim=0) - k_part_per_head = torch.chunk(k_part, num_query_groups_per_partition, dim=0) - v_part_per_head = torch.chunk(v_part, num_query_groups_per_partition, dim=0) + 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] + num_query_groups_per_partition = ( + models[0].config.num_query_groups // tp_size + ) + new_weight_qkv_this_tp = new_weight_qkv[ + i * total_size: (i + 1) * total_size + ] + q_part_per_head = torch.chunk( + q_part, num_query_groups_per_partition, dim=0 + ) + k_part_per_head = torch.chunk( + k_part, num_query_groups_per_partition, dim=0 + ) + v_part_per_head = torch.chunk( + v_part, num_query_groups_per_partition, dim=0 + ) total_size_per_head = total_size // num_query_groups_per_partition for j in range(num_query_groups_per_partition): - new_weight_qkv_this_tp[j * total_size_per_head : (j + 1) * total_size_per_head].copy_( - torch.cat([q_part_per_head[j], k_part_per_head[j], v_part_per_head[j]], dim=0) + new_weight_qkv_this_tp[ + j * total_size_per_head: (j + 1) * total_size_per_head + ].copy_( + torch.cat( + [ + q_part_per_head[j], + k_part_per_head[j], + v_part_per_head[j], + ], + dim=0, + ) ) else: @@ -330,21 +400,46 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - sizes = [total_size * tp_size] if not bias: sizes.append(config.hidden_size) - new_weight_qkv = torch.empty(*sizes, dtype=params_dtype, device=get_device_id()) + new_weight_qkv = torch.empty( + *sizes, dtype=params_dtype, device=get_device_id() + ) 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 + 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_this_tp = new_weight_qkv[i * total_size : (i + 1) * total_size] - q_part_per_head = torch.chunk(q_part, config.num_attention_heads, dim=0) - k_part_per_head = torch.chunk(k_part, config.num_attention_heads, dim=0) - v_part_per_head = torch.chunk(v_part, config.num_attention_heads, dim=0) + new_weight_qkv_this_tp = new_weight_qkv[ + i * total_size: (i + 1) * total_size + ] + q_part_per_head = torch.chunk( + q_part, config.num_attention_heads, dim=0 + ) + k_part_per_head = torch.chunk( + k_part, config.num_attention_heads, dim=0 + ) + v_part_per_head = torch.chunk( + v_part, config.num_attention_heads, dim=0 + ) total_size_per_head = total_size // config.num_attention_heads for j in range(config.num_attention_heads): - new_weight_qkv_this_tp[j * total_size_per_head : (j + 1) * total_size_per_head].copy_( - torch.cat([q_part_per_head[j], k_part_per_head[j], v_part_per_head[j]], dim=0) + new_weight_qkv_this_tp[ + j * total_size_per_head: (j + 1) * total_size_per_head + ].copy_( + torch.cat( + [ + q_part_per_head[j], + k_part_per_head[j], + v_part_per_head[j], + ], + dim=0, + ) ) tensor_chunk = torch.chunk(new_weight_qkv, tp_size, dim=0) @@ -357,7 +452,11 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - 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, k_name, v_name}] not in state_dict, skip loading") + print_rank_0( + f"tp_shard tensor:[{ + q_name, + k_name, + v_name}] not in state_dict, skip loading") return if tensor is None: @@ -368,10 +467,13 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank()} tensor {q_name} shape {tensor.shape} != {chunk_shape}" + 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=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == src_rank: @@ -388,7 +490,9 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - embed_tokens_weight = None if pp_rank == 0: embed_tokens_weight = gpt_model_module.embedding.word_embeddings.weight - _broadcast_tp_shard_tensor_vocab(embed_tokens_weight, "model.embed_tokens.weight") + _broadcast_tp_shard_tensor_vocab( + embed_tokens_weight, "model.embed_tokens.weight" + ) # Transformer layers # ------------------- @@ -396,36 +500,57 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - for layer in range(config.num_hidden_layers): layer_name = f"model.layers.{layer}" - print_rank_0(f"loading layer #{layer}, with layer_name model.layers.{layer}...") + print_rank_0( + f"loading layer #{layer}, with layer_name model.layers.{layer}...") dst_pp_rank, dst_virtual_pp_rank, dst_layer_idx = layer_map[layer] gpt_model_module = _get_gpt_model(models[dst_virtual_pp_rank]) sync_layer = gpt_model_module.decoder.layers[dst_layer_idx] _broadcast_tensor( - sync_layer.self_attention.linear_qkv.layer_norm_weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.self_attention.linear_qkv.layer_norm_weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.input_layernorm.weight", ) if f"{layer_name}.self_attn.q_norm.weight" in state_dict: _broadcast_tensor( - sync_layer.self_attention.q_layernorm.weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.self_attention.q_layernorm.weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.self_attn.q_norm.weight", ) _broadcast_tensor( - sync_layer.self_attention.k_layernorm.weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.self_attention.k_layernorm.weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.self_attn.k_norm.weight", ) _broadcast_tp_shard_tensor_qkv( - sync_layer.self_attention.linear_qkv.weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.self_attention.linear_qkv.weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.self_attn.q_proj.weight", f"{layer_name}.self_attn.k_proj.weight", f"{layer_name}.self_attn.v_proj.weight", ) if f"{layer_name}.self_attn.q_proj.bias" in state_dict: _broadcast_tp_shard_tensor_qkv( - sync_layer.self_attention.linear_qkv.bias if dst_pp_rank == pp_rank else None, + ( + sync_layer.self_attention.linear_qkv.bias + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.self_attn.q_proj.bias", f"{layer_name}.self_attn.k_proj.bias", f"{layer_name}.self_attn.v_proj.bias", @@ -433,12 +558,20 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - ) _broadcast_tp_shard_tensor( - sync_layer.self_attention.linear_proj.weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.self_attention.linear_proj.weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.self_attn.o_proj.weight", chunk_dim=1, ) _broadcast_tensor( - sync_layer.mlp.linear_fc1.layer_norm_weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.mlp.linear_fc1.layer_norm_weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.post_attention_layernorm.weight", ) @@ -469,9 +602,15 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - if is_value_model: # if torch.distributed.get_rank() == src_rank: - if "lm_head.weight" in state_dict and state_dict["lm_head.weight"].shape[0] == 1: + if ( + "lm_head.weight" in state_dict + and state_dict["lm_head.weight"].shape[0] == 1 + ): _broadcast_tensor(lm_head_weight, "lm_head.weight") - elif "reward_head.weight" in state_dict and state_dict["reward_head.weight"].shape[0] == 1: + elif ( + "reward_head.weight" in state_dict + and state_dict["reward_head.weight"].shape[0] == 1 + ): _broadcast_tensor(lm_head_weight, "reward_head.weight") print_rank_0("load lm_head from value_head weight") else: @@ -489,4 +628,6 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - broadcast_params(wrapped_model) pass get_torch_device().empty_cache() - print_rank_0(f"loading megatron ckpt done, time elapsed {time.time() - start_time}s") + print_rank_0( + f"loading megatron ckpt done, time elapsed {time.time() - start_time}s" + ) diff --git a/Agent0/executor_train/verl/verl/models/mcore/mbridge.py b/Agent0/executor_train/verl/verl/models/mcore/mbridge.py index 35c32d6..ee9c69e 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/mbridge.py +++ b/Agent0/executor_train/verl/verl/models/mcore/mbridge.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -15,9 +15,14 @@ try: from mbridge import AutoBridge - from mbridge.utils.post_creation_callbacks import freeze_moe_router, make_value_model + from mbridge.utils.post_creation_callbacks import ( + freeze_moe_router, + make_value_model, + ) except ImportError: - print("mbridge package not found. Please install mbridge with `pip install verl[mcore]` or `pip install mbridge`") + print( + "mbridge package not found. Please install mbridge with `pip install verl[mcore]` or `pip install mbridge`" + ) raise __all__ = ["AutoBridge", "make_value_model", "freeze_moe_router"] diff --git a/Agent0/executor_train/verl/verl/models/mcore/model_forward.py b/Agent0/executor_train/verl/verl/models/mcore/model_forward.py index e70e11f..9fb7f70 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/model_forward.py +++ b/Agent0/executor_train/verl/verl/models/mcore/model_forward.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # @@ -16,7 +16,12 @@ from verl.utils.megatron_utils import unwrap_model -from .util import postprocess_packed_seqs, preprocess_packed_seqs, recover_left_padding, remove_left_padding +from .util import ( + postprocess_packed_seqs, + preprocess_packed_seqs, + recover_left_padding, + remove_left_padding, +) def gptmodel_forward( @@ -36,7 +41,9 @@ def gptmodel_forward( post_process = unwrap_model(model).post_process if pack_seqs: batch_size, seq_len = attention_mask.shape[:2] - input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=pre_process) + input_ids_rmpad, packed_seq_params = preprocess_packed_seqs( + input_ids, attention_mask, pre_process=pre_process + ) input_ids_rmpad = input_ids_rmpad.contiguous() output_orig = model( input_ids=input_ids_rmpad, @@ -52,23 +59,47 @@ def gptmodel_forward( output_dict = logits_processor(output_orig, **args) output = { k: postprocess_packed_seqs( - v, packed_seq_params, attention_mask, batch_size, seq_len, post_process=post_process + v, + packed_seq_params, + attention_mask, + batch_size, + seq_len, + post_process=post_process, ) for k, v in output_dict.items() } else: output = postprocess_packed_seqs( - output_orig, packed_seq_params, attention_mask, batch_size, seq_len, post_process=post_process + output_orig, + packed_seq_params, + attention_mask, + batch_size, + seq_len, + post_process=post_process, ) else: - assert logits_processor is None, "logits_processor is not supported for non-packed sequence" + assert ( + logits_processor is None + ), "logits_processor is not supported for non-packed sequence" batch_size, sequence_length = attention_mask.shape new_input_ids, new_attention_mask, new_position_ids = remove_left_padding( - input_ids, attention_mask, position_ids, sequence_parallel, pre_process=pre_process + input_ids, + attention_mask, + position_ids, + sequence_parallel, + pre_process=pre_process, + ) + output = model( + input_ids=new_input_ids, + attention_mask=new_attention_mask, + position_ids=new_position_ids, ) - output = model(input_ids=new_input_ids, attention_mask=new_attention_mask, position_ids=new_position_ids) output = recover_left_padding( - output, new_attention_mask, attention_mask, sequence_length, post_process=post_process + output, + new_attention_mask, + attention_mask, + sequence_length, + post_process=post_process, ) if value_model and post_process: output = output[..., 0] @@ -90,18 +121,26 @@ def gptmodel_forward_qwen2_5_vl( ): from megatron.core import parallel_state as mpu - assert mpu.get_context_parallel_world_size() == 1, "qwen2_5_vl's context parallel is not accurate yet" + assert ( + mpu.get_context_parallel_world_size() == 1 + ), "qwen2_5_vl's context parallel is not accurate yet" pre_process = unwrap_model(model).pre_process post_process = unwrap_model(model).post_process pixel_values = ( - multi_modal_inputs["pixel_values"].to(input_ids.device) if "pixel_values" in multi_modal_inputs else None + multi_modal_inputs["pixel_values"].to(input_ids.device) + if "pixel_values" in multi_modal_inputs + else None ) image_grid_thw = ( - multi_modal_inputs["image_grid_thw"].to(input_ids.device) if "image_grid_thw" in multi_modal_inputs else None + multi_modal_inputs["image_grid_thw"].to(input_ids.device) + if "image_grid_thw" in multi_modal_inputs + else None ) if pack_seqs: batch_size, seq_len = attention_mask.shape[:2] - input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=True) + input_ids_rmpad, packed_seq_params = preprocess_packed_seqs( + input_ids, attention_mask, pre_process=True + ) input_ids_rmpad = input_ids_rmpad.contiguous() output_orig = model( input_ids=input_ids_rmpad, @@ -120,18 +159,32 @@ def gptmodel_forward_qwen2_5_vl( output_dict = logits_processor(output_orig, **args) output = { k: postprocess_packed_seqs( - v, packed_seq_params, attention_mask, batch_size, seq_len, post_process=post_process + v, + packed_seq_params, + attention_mask, + batch_size, + seq_len, + post_process=post_process, ) for k, v in output_dict.items() } else: output = postprocess_packed_seqs( - output_orig, packed_seq_params, attention_mask, batch_size, seq_len, post_process=post_process + output_orig, + packed_seq_params, + attention_mask, + batch_size, + seq_len, + post_process=post_process, ) else: batch_size, sequence_length = attention_mask.shape new_input_ids, new_attention_mask, new_position_ids = remove_left_padding( - input_ids, attention_mask, position_ids, sequence_parallel, pre_process=pre_process + input_ids, + attention_mask, + position_ids, + sequence_parallel, + pre_process=pre_process, ) output = model( input_ids=new_input_ids, @@ -141,7 +194,11 @@ def gptmodel_forward_qwen2_5_vl( image_grid_thw=image_grid_thw, ) output = recover_left_padding( - output, new_attention_mask, attention_mask, sequence_length, post_process=post_process + output, + new_attention_mask, + attention_mask, + sequence_length, + post_process=post_process, ) if value_model and post_process: output = output[..., 0] diff --git a/Agent0/executor_train/verl/verl/models/mcore/model_forward_fused.py b/Agent0/executor_train/verl/verl/models/mcore/model_forward_fused.py index fc55ef1..8c1558e 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/model_forward_fused.py +++ b/Agent0/executor_train/verl/verl/models/mcore/model_forward_fused.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # @@ -76,10 +76,15 @@ def fused_forward_gptmodel( post_process: bool = unwrap_model(model).post_process batch_size, seq_len = attention_mask.shape[:2] - input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=pre_process) + input_ids_rmpad, packed_seq_params = preprocess_packed_seqs( + input_ids, attention_mask, pre_process=pre_process + ) input_ids_rmpad = input_ids_rmpad.contiguous() - labels_rmpad, _ = preprocess_packed_seqs(labels, attention_mask, pre_process=True) - labels_mask_rmpad, _ = preprocess_packed_seqs(labels_mask, attention_mask, pre_process=True) + labels_rmpad, _ = preprocess_packed_seqs( + labels, attention_mask, pre_process=True) + labels_mask_rmpad, _ = preprocess_packed_seqs( + labels_mask, attention_mask, pre_process=True + ) labels_rmpad = labels_rmpad.contiguous() labels_mask_rmpad = labels_mask_rmpad.contiguous() @@ -121,16 +126,25 @@ def fused_forward_qwen2_5_vl( post_process = unwrap_model(model).post_process pixel_values = ( - multi_modal_inputs["pixel_values"].to(input_ids.device) if "pixel_values" in multi_modal_inputs else None + multi_modal_inputs["pixel_values"].to(input_ids.device) + if "pixel_values" in multi_modal_inputs + else None ) image_grid_thw = ( - multi_modal_inputs["image_grid_thw"].to(input_ids.device) if "image_grid_thw" in multi_modal_inputs else None + multi_modal_inputs["image_grid_thw"].to(input_ids.device) + if "image_grid_thw" in multi_modal_inputs + else None ) batch_size, seq_len = attention_mask.shape[:2] - input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=True) - labels_rmpad, _ = preprocess_packed_seqs(labels, attention_mask, pre_process=True) - labels_mask_rmpad, _ = preprocess_packed_seqs(labels_mask, attention_mask, pre_process=True) + input_ids_rmpad, packed_seq_params = preprocess_packed_seqs( + input_ids, attention_mask, pre_process=True + ) + labels_rmpad, _ = preprocess_packed_seqs( + labels, attention_mask, pre_process=True) + labels_mask_rmpad, _ = preprocess_packed_seqs( + labels_mask, attention_mask, pre_process=True + ) labels_rmpad = labels_rmpad.contiguous() labels_mask_rmpad = labels_mask_rmpad.contiguous() input_ids_rmpad = input_ids_rmpad.contiguous() @@ -182,41 +196,57 @@ def _fused_GPTModel_forward( """ # If decoder_input is provided (not None), then input_ids and position_ids are ignored. - # Otherwise, apply embedding layer on input_ids and position_ids to get decoder_input. + # Otherwise, apply embedding layer on input_ids and position_ids to get + # decoder_input. # Decoder embedding. if decoder_input is not None: pass elif self.pre_process: - decoder_input = self.embedding(input_ids=input_ids, position_ids=position_ids) + decoder_input = self.embedding( + input_ids=input_ids, position_ids=position_ids) else: # intermediate stage of pipeline # decoder will get hidden_states from encoder.input_tensor decoder_input = None - # Rotary positional embeddings (embedding is None for PP intermediate devices) + # Rotary positional embeddings (embedding is None for PP intermediate + # devices) rotary_pos_emb = None rotary_pos_cos = None rotary_pos_sin = None - if self.position_embedding_type == "rope" and not self.config.multi_latent_attention: + if ( + self.position_embedding_type == "rope" + and not self.config.multi_latent_attention + ): if not self.training and self.config.flash_decode and inference_context: - assert inference_context.is_static_batching(), "GPTModel currently only supports static inference batching." + assert ( + inference_context.is_static_batching() + ), "GPTModel currently only supports static inference batching." # Flash decoding uses precomputed cos and sin for RoPE rotary_pos_cos, rotary_pos_sin = self.rotary_pos_emb_cache.setdefault( - inference_context.max_sequence_length, - self.rotary_pos_emb.get_cos_sin(inference_context.max_sequence_length), - ) + inference_context.max_sequence_length, self.rotary_pos_emb.get_cos_sin( + inference_context.max_sequence_length), ) else: rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( - inference_context, self.decoder, decoder_input, self.config, packed_seq_params + inference_context, + self.decoder, + decoder_input, + self.config, + packed_seq_params, ) rotary_pos_emb = self.rotary_pos_emb( rotary_seq_len, - packed_seq=packed_seq_params is not None and packed_seq_params.qkv_format == "thd", + packed_seq=packed_seq_params is not None + and packed_seq_params.qkv_format == "thd", ) - elif self.position_embedding_type == "mrope" and not self.config.multi_latent_attention: + elif ( + self.position_embedding_type == "mrope" + and not self.config.multi_latent_attention + ): if self.training or not self.config.flash_decode: - rotary_pos_emb = self.rotary_pos_emb(position_ids, self.mrope_section) + rotary_pos_emb = self.rotary_pos_emb( + position_ids, self.mrope_section) else: # Flash decoding uses precomputed cos and sin for RoPE raise NotImplementedError( @@ -231,7 +261,8 @@ def _fused_GPTModel_forward( and not self.training ): sequence_len_offset = torch.tensor( - [inference_context.sequence_len_offset] * inference_context.current_batch_size, + [inference_context.sequence_len_offset] + * inference_context.current_batch_size, dtype=torch.int32, device=rotary_pos_cos.device, # Co-locate this with the rotary tensors ) @@ -257,7 +288,9 @@ def _fused_GPTModel_forward( # Process inference output. if inference_context and not inference_context.is_static_batching(): - hidden_states = inference_context.last_token_logits(hidden_states.squeeze(1).unsqueeze(0)).unsqueeze(1) + hidden_states = inference_context.last_token_logits( + hidden_states.squeeze(1).unsqueeze(0) + ).unsqueeze(1) # logits and loss output_weight = None diff --git a/Agent0/executor_train/verl/verl/models/mcore/model_initializer.py b/Agent0/executor_train/verl/verl/models/mcore/model_initializer.py index 4c01b12..405646b 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/model_initializer.py +++ b/Agent0/executor_train/verl/verl/models/mcore/model_initializer.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # @@ -17,7 +17,10 @@ # use mcore transformer config to initialize the model from abc import ABC, abstractmethod -from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec, get_gpt_mtp_block_spec +from megatron.core.models.gpt.gpt_layer_specs import ( + get_gpt_decoder_block_spec, + get_gpt_mtp_block_spec, +) from megatron.core.models.gpt.gpt_model import GPTModel from .config_converter import PretrainedConfig, TransformerConfig @@ -26,14 +29,18 @@ class BaseModelInitializer(ABC): """Base class for model initializers.""" - def __init__(self, tfconfig: TransformerConfig, hf_config: PretrainedConfig): + def __init__( + self, + tfconfig: TransformerConfig, + hf_config: PretrainedConfig): self.tfconfig = tfconfig self.hf_config = hf_config @abstractmethod def get_transformer_layer_spec(self): """Get the transformer layer specification. - https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/gpt/gpt_layer_specs.py""" + https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/gpt/gpt_layer_specs.py + """ pass def get_rope_scaling_args(self) -> dict: @@ -42,7 +49,9 @@ def get_rope_scaling_args(self) -> dict: if "rope_scaling" in self.hf_config: if self.hf_config.rope_scaling is not None: # assert self.hf_config.rope_scaling["type"] == "linear", "only linear scaling is supported for now" - rope_scaling_args["seq_len_interpolation_factor"] = self.hf_config.rope_scaling["factor"] + rope_scaling_args["seq_len_interpolation_factor"] = ( + self.hf_config.rope_scaling["factor"] + ) return rope_scaling_args def initialize( @@ -83,10 +92,14 @@ def initialize( ) if post_process and value: - from verl.models.llama.megatron.layers.parallel_linear import LinearForLastLayer + from verl.models.llama.megatron.layers.parallel_linear import ( + LinearForLastLayer, + ) model.output_layer = LinearForLastLayer( - input_size=self.tfconfig.hidden_size, output_size=1, config=self.tfconfig + input_size=self.tfconfig.hidden_size, + output_size=1, + config=self.tfconfig, ) return model @@ -96,20 +109,29 @@ class DenseModel(BaseModelInitializer): """Initializer for dense models like Llama and Qwen2.""" def get_transformer_layer_spec(self): - assert self.tfconfig.normalization == "RMSNorm", "only RMSNorm is supported for now" - return get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True) + assert ( + self.tfconfig.normalization == "RMSNorm" + ), "only RMSNorm is supported for now" + return get_gpt_decoder_block_spec( + self.tfconfig, use_transformer_engine=True) class Qwen2MoEModel(BaseModelInitializer): """Initializer for Qwen2 MoE models.""" def get_transformer_layer_spec(self): - assert self.tfconfig.normalization == "RMSNorm", "only RMSNorm is supported for now" - transformer_layer_spec = get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True) + assert ( + self.tfconfig.normalization == "RMSNorm" + ), "only RMSNorm is supported for now" + transformer_layer_spec = get_gpt_decoder_block_spec( + self.tfconfig, use_transformer_engine=True + ) # Patch layer spec for shared experts for i in range(len(transformer_layer_spec.layer_specs)): - transformer_layer_spec.layer_specs[i].submodules.mlp.submodules.shared_experts.params["gate"] = True + transformer_layer_spec.layer_specs[ + i + ].submodules.mlp.submodules.shared_experts.params["gate"] = True return transformer_layer_spec @@ -127,8 +149,12 @@ class MixtralModel(BaseModelInitializer): """Initializer for Mixtral models.""" def get_transformer_layer_spec(self): - assert self.tfconfig.normalization == "RMSNorm", "only RMSNorm is supported for now" - transformer_layer_spec = get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True) + assert ( + self.tfconfig.normalization == "RMSNorm" + ), "only RMSNorm is supported for now" + transformer_layer_spec = get_gpt_decoder_block_spec( + self.tfconfig, use_transformer_engine=True + ) return transformer_layer_spec def initialize(self, **kwargs): @@ -144,8 +170,12 @@ class Qwen3MoEModel(BaseModelInitializer): """Initializer for Qwen3 MoE models.""" def get_transformer_layer_spec(self): - assert self.tfconfig.normalization == "RMSNorm", "only RMSNorm is supported for now" - transformer_layer_spec = get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True) + assert ( + self.tfconfig.normalization == "RMSNorm" + ), "only RMSNorm is supported for now" + transformer_layer_spec = get_gpt_decoder_block_spec( + self.tfconfig, use_transformer_engine=True + ) return transformer_layer_spec def initialize(self, **kwargs): @@ -162,7 +192,9 @@ class DeepseekV3Model(BaseModelInitializer): """Initializer for DeepseekV3 models.""" def get_transformer_layer_spec(self): - transformer_layer_spec = get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True) + transformer_layer_spec = get_gpt_decoder_block_spec( + self.tfconfig, use_transformer_engine=True + ) return transformer_layer_spec def get_rope_scaling_args(self) -> dict: @@ -180,7 +212,8 @@ def initialize( # MTP if self.tfconfig.mtp_num_layers is not None: transformer_layer_spec = self.get_transformer_layer_spec() - mtp_block_spec = get_gpt_mtp_block_spec(self.tfconfig, transformer_layer_spec, use_transformer_engine=True) + mtp_block_spec = get_gpt_mtp_block_spec( + self.tfconfig, transformer_layer_spec, use_transformer_engine=True) kwargs["mtp_block_spec"] = mtp_block_spec model = super().initialize(**kwargs) @@ -195,7 +228,9 @@ class Qwen25VLModel(BaseModelInitializer): """Initializer for Qwen2.5 VL models.""" def get_transformer_layer_spec(self): - transformer_layer_spec = get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True) + transformer_layer_spec = get_gpt_decoder_block_spec( + self.tfconfig, use_transformer_engine=True + ) return transformer_layer_spec def initialize( @@ -213,11 +248,20 @@ def initialize( transformer_layer_spec = self.get_transformer_layer_spec() - from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TERowParallelLinear + from megatron.core.extensions.transformer_engine import ( + TEColumnParallelLinear, + TERowParallelLinear, + ) from megatron.core.models.gpt.moe_module_specs import MLPSubmodules - from megatron.core.models.vision.vit_layer_specs import get_vit_layer_with_transformer_engine_spec + from megatron.core.models.vision.vit_layer_specs import ( + get_vit_layer_with_transformer_engine_spec, + ) - from .qwen2_5_vl import Qwen2_5VLModel, get_vision_model_config, get_vision_projection_config + from .qwen2_5_vl import ( + Qwen2_5VLModel, + get_vision_model_config, + get_vision_projection_config, + ) vision_transformer_config = get_vision_model_config(deepcopy(tfconfig)) vision_transformer_config.pipeline_model_parallel_size = 1 @@ -254,7 +298,9 @@ def initialize( ) if post_process and value: - from verl.models.llama.megatron.layers.parallel_linear import LinearForLastLayer + from verl.models.llama.megatron.layers.parallel_linear import ( + LinearForLastLayer, + ) qwen25_vl_model.language_model.output_layer = LinearForLastLayer( input_size=tfconfig.hidden_size, output_size=1, config=tfconfig diff --git a/Agent0/executor_train/verl/verl/models/mcore/patch_v012.py b/Agent0/executor_train/verl/verl/models/mcore/patch_v012.py index d54a3eb..be87654 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/patch_v012.py +++ b/Agent0/executor_train/verl/verl/models/mcore/patch_v012.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,7 +14,8 @@ # limitations under the License. # there is some bug in mcore 0.12, so we need to patch it -# 1. `get_query_key_value_tensors` in `multi_latent_attention.py` works wrong when packed_seq_params is not None +# 1. `get_query_key_value_tensors` in `multi_latent_attention.py` works +# wrong when packed_seq_params is not None def apply_patch(): @@ -44,22 +45,27 @@ def patch_get_query_key_value_tensors( """ # s = sequence length, b = batch size, h = hidden size, n = num attention heads # Attention heads [s, b, n*h] - assert hidden_states.ndim == 3, f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" + assert ( + hidden_states.ndim == 3 + ), f"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D" - inference_context = deprecate_inference_params(inference_context, inference_params) + inference_context = deprecate_inference_params( + inference_context, inference_params + ) # ========================================= # Prepare RoPE and seqlen related params # ========================================= rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len( - inference_context, None, hidden_states, self.config, packed_seq_params - ) + inference_context, None, hidden_states, self.config, packed_seq_params) # rotary_pos_emb:[s, b, 1, 64] mscale = 1.0 if self.config.rope_type == "rope": - packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == "thd" - rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq) + packed_seq = ( + packed_seq_params is not None and packed_seq_params.qkv_format == "thd") + rotary_pos_emb = self.rotary_pos_emb( + rotary_seq_len, packed_seq=packed_seq) else: rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len) @@ -79,9 +85,11 @@ def patch_get_query_key_value_tensors( # 2. Scatter sequence back to s / TP if sequence-parallel since it was # gathered by ColumnParallelLinear. if q_compressed.size(-1) != self.config.q_lora_rank: - q_compressed = gather_from_tensor_model_parallel_region(q_compressed) + q_compressed = gather_from_tensor_model_parallel_region( + q_compressed) if self.config.sequence_parallel: - q_compressed = scatter_to_sequence_parallel_region(q_compressed) + q_compressed = scatter_to_sequence_parallel_region( + q_compressed) q_compressed = self.q_layernorm(q_compressed) else: @@ -92,20 +100,30 @@ def patch_get_query_key_value_tensors( # elif linear_kv_down_proj is Linear: # kv_combined: [s / TP, b, (kv_lora_rank + qk_pos_emb_head_dim)] kv_combined, _ = self.linear_kv_down_proj(hidden_states) - if kv_combined.size(-1) != self.config.kv_lora_rank + self.config.qk_pos_emb_head_dim: + if ( + kv_combined.size(-1) + != self.config.kv_lora_rank + self.config.qk_pos_emb_head_dim + ): # kv_combined: [s, b, (kv_lora_rank + qk_pos_emb_head_dim)] kv_combined = gather_from_tensor_model_parallel_region(kv_combined) - # kv_compressed:[s, b, kv_lora_rank], k_pos_emb: [s, b, qk_pos_emb_head_dim] + # kv_compressed:[s, b, kv_lora_rank], k_pos_emb: [s, b, + # qk_pos_emb_head_dim] kv_compressed, k_pos_emb = torch.split( - kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1 + kv_combined, + [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], + dim=-1, ) if self.config.sequence_parallel: # kv_compressed:[s / TP, b, kv_lora_rank] - kv_compressed = scatter_to_sequence_parallel_region(kv_compressed) + kv_compressed = scatter_to_sequence_parallel_region( + kv_compressed) else: - # kv_compressed:[s / TP, b, kv_lora_rank], k_pos_emb: [s / TP, b, qk_pos_emb_head_dim] + # kv_compressed:[s / TP, b, kv_lora_rank], k_pos_emb: [s / TP, b, + # qk_pos_emb_head_dim] kv_compressed, k_pos_emb = torch.split( - kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1 + kv_combined, + [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], + dim=-1, ) if parallel_state.get_tensor_model_parallel_world_size() > 1: # k_pos_emb: [s, b, qk_pos_emb_head_dim] @@ -116,7 +134,9 @@ def patch_get_query_key_value_tensors( # ========================================= # QKV up projection and RoPE apply # ========================================= - def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb): + def qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ): if self.config.q_lora_rank is not None: q, _ = self.linear_q_up_proj(q_compressed) else: @@ -126,7 +146,11 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po q_len, bsz, _ = q.size() # q: [s, b, n, 192] - q = q.view(q_len, bsz, self.num_attention_heads_per_partition, self.q_head_dim) + q = q.view( + q_len, + bsz, + self.num_attention_heads_per_partition, + self.q_head_dim) # kv: [s, b, 2048] kv, _ = self.linear_kv_up_proj(kv_compressed) @@ -155,10 +179,14 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po k_pos_emb = torch.unsqueeze(k_pos_emb, 2) # q: [s, b, n, 128], q_pos_emb: [s, b, n, 64] - q_no_pe, q_pos_emb = torch.split(q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1) + q_no_pe, q_pos_emb = torch.split( + q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1 + ) # k_no_pe: [s, b, n, 128], value: [s, b, n, 128] - k_no_pe, value = torch.split(kv, [self.config.qk_head_dim, self.config.v_head_dim], dim=-1) + k_no_pe, value = torch.split( + kv, [self.config.qk_head_dim, self.config.v_head_dim], dim=-1 + ) if packed_seq_params is not None: cu_seqlens_q = packed_seq_params.cu_seqlens_q @@ -190,11 +218,15 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po # query: [s, b, n, 192] query = torch.cat([q_no_pe, q_pos_emb], dim=-1) if packed_seq_params is not None: - k_pos_emb = k_pos_emb.expand(-1, self.num_attention_heads_per_partition, -1) + k_pos_emb = k_pos_emb.expand( + -1, self.num_attention_heads_per_partition, -1 + ) key = torch.cat([k_no_pe, k_pos_emb], dim=-1) else: # key: [s, b, n, 192] - k_pos_emb = k_pos_emb.expand(-1, -1, self.num_attention_heads_per_partition, -1) + k_pos_emb = k_pos_emb.expand( + -1, -1, self.num_attention_heads_per_partition, -1 + ) key = torch.cat([k_no_pe, k_pos_emb], dim=-1) query = query.contiguous() @@ -205,10 +237,16 @@ def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_po if self.recompute_up_proj: self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput() query, key, value = self.qkv_up_checkpoint.checkpoint( - qkv_up_proj_and_rope_apply, q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + qkv_up_proj_and_rope_apply, + q_compressed, + kv_compressed, + k_pos_emb, + rotary_pos_emb, ) else: - query, key, value = qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb) + query, key, value = qkv_up_proj_and_rope_apply( + q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb + ) return query, key, value diff --git a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/__init__.py b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/__init__.py index 8842d02..f8b03e2 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/__init__.py +++ b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2024 Alibaba PAI Team. # @@ -18,4 +18,7 @@ from .model import Qwen2_5VLModel from .vision_config import get_vision_model_config, get_vision_projection_config -__all__ = ["Qwen2_5VLModel", "get_vision_model_config", "get_vision_projection_config"] +__all__ = [ + "Qwen2_5VLModel", + "get_vision_model_config", + "get_vision_projection_config"] diff --git a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/attention.py b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/attention.py index 91a27cc..f030823 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/attention.py +++ b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/attention.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2024 Alibaba PAI Team. # @@ -63,21 +63,29 @@ def forward( """ - inference_context = deprecate_inference_params(inference_context, inference_params) + inference_context = deprecate_inference_params( + inference_context, inference_params + ) if inference_context and inference_context.is_dynamic_batching(): - assert flash_decode_and_prefill_kernel is not None, ( - "Internal use only: install package `nvidia_chunked_flash_attn`." - ) + assert ( + flash_decode_and_prefill_kernel is not None + ), "Internal use only: install package `nvidia_chunked_flash_attn`." # hidden_states: [sq, b, h] - if self.config.flash_decode and not self.training and inference_context is not None: + if ( + self.config.flash_decode + and not self.training + and inference_context is not None + ): rotary_pos_emb = None else: assert rotary_pos_cos is None and rotary_pos_sin is None - # For self attention we just duplicate the rotary_pos_emb if it isn't already - if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple): + # For self attention we just duplicate the rotary_pos_emb if it isn't + # already + if rotary_pos_emb is not None and not isinstance( + rotary_pos_emb, tuple): rotary_pos_emb = (rotary_pos_emb,) * 2 # ===================== @@ -85,14 +93,17 @@ def forward( # ===================== # Get the query, key and value tensors based on the type of attention - # self or cross attn. - query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states) + query, key, value = self.get_query_key_value_tensors( + hidden_states, key_value_states + ) # =================================================== # Adjust key, value, and rotary_pos_emb for inference # =================================================== # This branch only runs in the decode phase of flash decoding and returns after the linear - # projection. This conditional is not used in the prefill phase or non-flash-decoding cases. + # projection. This conditional is not used in the prefill phase or + # non-flash-decoding cases. if ( self.config.flash_decode and inference_context is not None @@ -102,7 +113,9 @@ def forward( ): assert self.layer_number in inference_context.key_value_memory_dict assert inference_context.sequence_len_offset is not None - inference_key_memory, inference_value_memory = inference_context.key_value_memory_dict[self.layer_number] + inference_key_memory, inference_value_memory = ( + inference_context.key_value_memory_dict[self.layer_number] + ) output = self.flash_decode( sequence_len_offset=sequence_len_offset, query_layer=query, @@ -118,15 +131,17 @@ def forward( output, bias = self.linear_proj(context_layer) return output, bias - query, key, value, rotary_pos_emb, attn_mask_type = self._adjust_key_value_for_inference( - inference_context, - query, - key, - value, - rotary_pos_emb, - rotary_pos_cos, - rotary_pos_sin, - sequence_len_offset, + query, key, value, rotary_pos_emb, attn_mask_type = ( + self._adjust_key_value_for_inference( + inference_context, + query, + key, + value, + rotary_pos_emb, + rotary_pos_cos, + rotary_pos_sin, + sequence_len_offset, + ) ) if packed_seq_params is not None: @@ -155,11 +170,15 @@ def forward( if q_pos_emb is not None: # TODO VIJAY: simplify if inference_context is None or inference_context.is_static_batching(): - query = apply_rotary_pos_emb_absolute(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q) + query = apply_rotary_pos_emb_absolute( + query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q) else: - query = inference_context.apply_rotary_emb_query(query, q_pos_emb, self.config, cu_seqlens_q) + query = inference_context.apply_rotary_emb_query( + query, q_pos_emb, self.config, cu_seqlens_q + ) if k_pos_emb is not None: - key = apply_rotary_pos_emb_absolute(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv) + key = apply_rotary_pos_emb_absolute( + key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv) # TODO, can apply positional embedding to value_layer so it has # absolute positional embedding. @@ -200,10 +219,10 @@ def forward( cu_kv_lengths, max_seqlen_k = inference_context.cu_kv_lengths() core_attn_out = self.flash_decode_and_prefill( - q, k, v, max_seqlen_q, max_seqlen_k, cu_query_lengths, cu_kv_lengths - ) + q, k, v, max_seqlen_q, max_seqlen_k, cu_query_lengths, cu_kv_lengths) core_attn_out = core_attn_out.squeeze(0).unsqueeze(1) - core_attn_out = rearrange(core_attn_out, "s b h d -> s b (h d)") + core_attn_out = rearrange( + core_attn_out, "s b h d -> s b (h d)") if packed_seq_params is not None and packed_seq_params.qkv_format == "thd": # reshape to same output shape as unpacked case diff --git a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/model.py b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/model.py index 74e4406..640bfd7 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/model.py +++ b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/model.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2024 Alibaba PAI Team. # @@ -97,11 +97,15 @@ def __init__( super().__init__(config=language_transformer_config) # patch self_attention to use qwen2_5_vl attention - vision_transformer_layer_spec.submodules.self_attention.module = Qwen2_5VLSelfAttention + vision_transformer_layer_spec.submodules.self_attention.module = ( + Qwen2_5VLSelfAttention + ) for layer_spec in language_transformer_layer_spec.layer_specs: layer_spec.submodules.self_attention.module = Qwen2_5VLSelfAttention - logging.getLogger(__name__).warning("Qwen2VL model is under development and may be missing features.") + logging.getLogger(__name__).warning( + "Qwen2VL model is under development and may be missing features." + ) self.pre_process = pre_process self.post_process = post_process @@ -115,10 +119,14 @@ def __init__( self.image_token_id = image_token_id self.video_token_id = video_token_id - self.square_merge_size = vision_projection_config.ffn_hidden_size // vision_transformer_config.hidden_size + self.square_merge_size = ( + vision_projection_config.ffn_hidden_size + // vision_transformer_config.hidden_size + ) # This attribute is needed to check if an all-reduce is required - # on the word embeddings inside `finalize_model_grads._allreduce_word_embedding_grads`. + # on the word embeddings inside + # `finalize_model_grads._allreduce_word_embedding_grads`. self.share_embeddings_and_output_weights = False if self.pre_process: self.vision_model = Qwen2_5VisionModel( @@ -147,7 +155,9 @@ def __init__( scatter_embedding_sequence_parallel=False, ) - self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights + self.share_embeddings_and_output_weights = ( + self.language_model.share_embeddings_and_output_weights + ) def shared_embedding_or_output_weight(self): """This is a convenience method to surface the language model's word embeddings, which is @@ -161,14 +171,21 @@ def set_input_tensor(self, input_tensor) -> None: # gives us non-lists or None if not isinstance(input_tensor, list): input_tensor = [input_tensor] - assert len(input_tensor) == 1, "input_tensor should only be length 1 for Qwen2VL" + assert ( + len(input_tensor) == 1 + ), "input_tensor should only be length 1 for Qwen2VL" if self.pre_process: self.encoder_hidden_state = input_tensor[0] else: self.language_model.set_input_tensor(input_tensor[0]) - def freeze(self, freeze_language_model: bool, freeze_vision_model: bool, freeze_vision_projection: bool): + def freeze( + self, + freeze_language_model: bool, + freeze_vision_model: bool, + freeze_vision_projection: bool, + ): """Freeze model modules. Make specific modules non-trainable by setting requires_grad to False for the module's parameters. @@ -234,14 +251,17 @@ def forward( video_start_index = image_mask.sum().item() if video_grid_thw is not None: video_mask = input_ids == self.video_token_id - vision_grid_thw = torch.cat([vision_grid_thw, video_grid_thw], dim=0) + vision_grid_thw = torch.cat( + [vision_grid_thw, video_grid_thw], dim=0) vision_data = torch.cat([vision_data, pixel_values_videos], dim=0) video_start_index = image_mask.sum().item() + video_mask.sum().item() use_inference_kv_cache = ( - inference_params is not None and "image_tokens_count" in inference_params.key_value_memory_dict + inference_params is not None + and "image_tokens_count" in inference_params.key_value_memory_dict ) use_inference_kv_cache = ( - inference_params is not None and "image_tokens_count" in inference_params.key_value_memory_dict + inference_params is not None + and "image_tokens_count" in inference_params.key_value_memory_dict ) if use_inference_kv_cache: raise NotImplementedError() @@ -250,12 +270,15 @@ def forward( vision_embeds = None if vision_grid_thw is not None and vision_grid_thw.shape[0] > 0: vision_embeds = self.vision_model( - vision_data=vision_data, # If None, vision model should use intermediate outputs (EPP > 1) + vision_data=vision_data, + # If None, vision model should use intermediate outputs + # (EPP > 1) grid_thw=vision_grid_thw, # should provided in each EPP stage ) # If running inference, the language model KV cache will be updated for image token positions. - # Here we store the image tokens sequence length, which can be used as an offset to the KV cache later. + # Here we store the image tokens sequence length, which can be used + # as an offset to the KV cache later. if inference_params is not None: raise NotImplementedError() # inference_params.key_value_memory_dict["image_tokens_count"] = ( @@ -269,7 +292,8 @@ def forward( input_ids=input_ids, position_ids=None, # NOTE: disable ) # [text_seq_len, b, h_language] - # NOTE: why not cat here? is it the combined embeddings useless? + # NOTE: why not cat here? is it the combined embeddings + # useless? combined_embeddings = language_embeddings elif vision_embeds is not None: if video_start_index == 0: @@ -283,9 +307,8 @@ def forward( video_embeds = vision_embeds[video_start_index:] else: raise ValueError( - f"Expect video token start index in range [0, {vision_embeds.shape[0]}], but got " - f"{video_start_index}" - ) + f"Expect video token start index in range [0, { + vision_embeds.shape[0]}], but got " f"{video_start_index}") combined_embeddings = self.language_model.embedding( input_ids=input_ids, @@ -293,22 +316,30 @@ def forward( ) # [text_seq_len, b, h_language] if image_embeds is not None or video_embeds is not None: - combined_embeddings = combined_embeddings.transpose(0, 1).contiguous() + combined_embeddings = combined_embeddings.transpose( + 0, 1 + ).contiguous() if image_embeds is not None: - image_mask = (input_ids == self.image_token_id).contiguous() + image_mask = ( + input_ids == self.image_token_id).contiguous() if image_mask.sum() > 0: combined_embeddings = combined_embeddings.clone() combined_embeddings[image_mask] = image_embeds.to( - dtype=combined_embeddings.dtype, device=combined_embeddings.device + dtype=combined_embeddings.dtype, + device=combined_embeddings.device, ) if video_embeds is not None: - video_mask = (input_ids == self.video_token_id).contiguous() + video_mask = ( + input_ids == self.video_token_id).contiguous() if video_mask.sum() > 0: combined_embeddings = combined_embeddings.clone() combined_embeddings[video_mask] = video_embeds.to( - dtype=combined_embeddings.dtype, device=combined_embeddings.device + dtype=combined_embeddings.dtype, + device=combined_embeddings.device, ) - combined_embeddings = combined_embeddings.transpose(0, 1).contiguous() + combined_embeddings = combined_embeddings.transpose( + 0, 1 + ).contiguous() else: combined_embeddings = self.language_model.embedding( @@ -316,21 +347,29 @@ def forward( position_ids=None, # NOTE: disable ) # [text_seq_len, b, h_language] if self.config.sequence_parallel: - combined_embeddings = tensor_parallel.scatter_to_sequence_parallel_region(combined_embeddings) + combined_embeddings = ( + tensor_parallel.scatter_to_sequence_parallel_region( + combined_embeddings + ) + ) combined_embeddings = combined_embeddings.contiguous() else: combined_embeddings = None from .rope_utils import get_rope_index position_ids, _ = get_rope_index( - input_ids, image_grid_thw=image_grid_thw, video_grid_thw=video_grid_thw, attention_mask=attention_mask + input_ids, + image_grid_thw=image_grid_thw, + video_grid_thw=video_grid_thw, + attention_mask=attention_mask, ) output = self.language_model( input_ids=None, position_ids=position_ids, # None in encoder attention_mask=attention_mask, # None in encoder - decoder_input=combined_embeddings, # only not None in the first decoder PP stage + decoder_input=combined_embeddings, + # only not None in the first decoder PP stage labels=labels, # only not None in the last decoder PP stage # inference_params=inference_params, # currently always None packed_seq_params=packed_seq_params, # currently always None diff --git a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/rope_utils.py b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/rope_utils.py index fadc74d..b1a8fdb 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/rope_utils.py +++ b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/rope_utils.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2024 Alibaba PAI Team. # @@ -107,7 +107,9 @@ def get_rope_index( video_token_id = 151656 vision_start_token_id = 151652 mrope_position_deltas = [] - if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): total_input_ids = input_ids if attention_mask is None: attention_mask = torch.ones_like(total_input_ids) @@ -123,7 +125,9 @@ def get_rope_index( for i, input_ids in enumerate(total_input_ids): input_ids = input_ids[attention_mask[i] == 1] image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + vision_start_indices = torch.argwhere( + input_ids == vision_start_token_id + ).squeeze(1) vision_tokens = input_ids[vision_start_indices + 1] image_nums = (vision_tokens == image_token_id).sum() video_nums = (vision_tokens == video_token_id).sum() @@ -171,39 +175,67 @@ def get_rope_index( ) text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + st_idx = (llm_pos_ids_list[-1].max() + + 1 if len(llm_pos_ids_list) > 0 else 0) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) range_tensor = torch.arange(llm_grid_t).view(-1, 1) - expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w) + expanded_range = range_tensor.expand(-1, + llm_grid_h * llm_grid_w) time_tensor = expanded_range * second_per_grid_t * tokens_per_second time_tensor_long = time_tensor.long() t_index = time_tensor_long.flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append(torch.stack( + [t_index, h_index, w_index]) + text_len + st_idx) st = ed + llm_grid_t * llm_grid_h * llm_grid_w if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + st_idx = (llm_pos_ids_list[-1].max() + + 1 if len(llm_pos_ids_list) > 0 else 0) text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device) - mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i])) - mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1) + position_ids[..., i, attention_mask[i] == 1] = llm_positions.to( + position_ids.device + ) + mrope_position_deltas.append( + llm_positions.max() + 1 - len(total_input_ids[i]) + ) + mrope_position_deltas = torch.tensor( + mrope_position_deltas, device=input_ids.device + ).unsqueeze(1) return position_ids, mrope_position_deltas else: if attention_mask is not None: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) - position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device) - max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0] - mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1] + position_ids = (position_ids.unsqueeze(0).expand( + 3, -1, -1).to(attention_mask.device)) + max_position_ids = position_ids.max(0, keepdim=False)[0].max( + -1, keepdim=True + )[0] + mrope_position_deltas = max_position_ids + \ + 1 - attention_mask.shape[-1] else: position_ids = ( torch.arange(input_ids.shape[1], device=input_ids.device) @@ -220,8 +252,10 @@ def get_rope_index( def apply_rotary_pos_emb_thd_absolute( - t: Tensor, cu_seqlens: Tensor, freqs: Tensor, rotary_interleaved: bool = False -) -> Tensor: + t: Tensor, + cu_seqlens: Tensor, + freqs: Tensor, + rotary_interleaved: bool = False) -> Tensor: """A baseline implementation of applying RoPE for `thd` format. Args: @@ -233,7 +267,9 @@ def apply_rotary_pos_emb_thd_absolute( Returns: Tensor: Shape [t, h, d]. The input tensor after applying RoPE. """ - return _apply_rotary_pos_emb_bshd(t[:, None], freqs, rotary_interleaved=rotary_interleaved).squeeze(1) + return _apply_rotary_pos_emb_bshd( + t[:, None], freqs, rotary_interleaved=rotary_interleaved + ).squeeze(1) def apply_rotary_pos_emb_absolute( @@ -253,7 +289,9 @@ def apply_rotary_pos_emb_absolute( if cu_seqlens is None: # NOTE: TE backends do not support mRoPE in bshd format when bs > 1 if freqs.shape[1] > 1: - return _apply_rotary_pos_emb_bshd(t, freqs, rotary_interleaved=config.rotary_interleaved) + return _apply_rotary_pos_emb_bshd( + t, freqs, rotary_interleaved=config.rotary_interleaved + ) else: return fused_apply_rotary_pos_emb(t, freqs) else: @@ -261,6 +299,9 @@ def apply_rotary_pos_emb_absolute( return fused_apply_rotary_pos_emb(t[:, None], freqs).squeeze(1) else: if cu_seqlens is None: - return _apply_rotary_pos_emb_bshd(t, freqs, rotary_interleaved=config.rotary_interleaved) + return _apply_rotary_pos_emb_bshd( + t, freqs, rotary_interleaved=config.rotary_interleaved + ) else: - return apply_rotary_pos_emb_thd_absolute(t, cu_seqlens, freqs, rotary_interleaved=config.rotary_interleaved) + return apply_rotary_pos_emb_thd_absolute( + t, cu_seqlens, freqs, rotary_interleaved=config.rotary_interleaved) diff --git a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_config.py b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_config.py index 0631c90..1fc9cab 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_config.py +++ b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_config.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2024 Alibaba PAI Team. # @@ -24,14 +24,17 @@ def get_vision_model_config(config: TransformerConfig) -> TransformerConfig: # diff: out_hidden_size & intermediate_size # mlp: hidden_size -> intermediate_size -> embed_dim, silu - # NOTE: here we provide a workaround to solve the wrong layer amount when VPP of decoder is on + # NOTE: here we provide a workaround to solve the wrong layer amount when + # VPP of decoder is on if config.num_layers in [28, 36]: config.ffn_hidden_size = 3420 else: config.ffn_hidden_size = 3456 if parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None: - config.num_layers = 32 * parallel_state.get_virtual_pipeline_model_parallel_world_size() # depth + config.num_layers = ( + 32 * parallel_state.get_virtual_pipeline_model_parallel_world_size() + ) # depth else: config.num_layers = 32 # depth config.num_attention_heads = 16 # num_heads diff --git a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_model.py b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_model.py index 06b4fd3..3775907 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_model.py +++ b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_model.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2024 Alibaba PAI Team. # @@ -30,7 +30,8 @@ from .vision_transformer_block import Qwen2_5VisionTransformerBlock as TransformerBlock -# copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py +# copied from +# https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py class PatchEmbed(nn.Module): def __init__( self, @@ -46,26 +47,42 @@ def __init__( self.embed_dim = embed_dim kernel_size = [temporal_patch_size, patch_size, patch_size] - self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False) + self.proj = nn.Conv3d( + in_channels, + embed_dim, + kernel_size=kernel_size, + stride=kernel_size, + bias=False, + ) def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: target_dtype = self.proj.weight.dtype hidden_states = hidden_states.view( - -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size + -1, + self.in_channels, + self.temporal_patch_size, + self.patch_size, + self.patch_size, + ) + hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view( + -1, self.embed_dim ) - hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim) return hidden_states -# copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py +# copied from +# https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py class VisionRotaryEmbedding(nn.Module): def __init__(self, dim: int, theta: float = 10000.0) -> None: super().__init__() - inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) + inv_freq = 1.0 / \ + (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim)) self.register_buffer("inv_freq", inv_freq, persistent=False) def forward(self, seqlen: int) -> torch.Tensor: - seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype) + seq = torch.arange( + seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype + ) freqs = torch.outer(seq, self.inv_freq) return freqs.float() @@ -127,7 +144,8 @@ def __init__( # Transformer layers. # TODO: Follow-up changes will make pre and post_process configurable. They are needed for supporting # pipeline parallelism. - # NOTE: a final layer norm and/or linear layer present in some implementations are omitted here. + # NOTE: a final layer norm and/or linear layer present in some + # implementations are omitted here. self.decoder = TransformerBlock( config=transformer_config, spec=transformer_layer_spec, @@ -141,7 +159,10 @@ def __init__( if self.post_process: self.projection = MultimodalProjector( - projection_config, projection_layer_spec, projection_type, projection_config.ffn_hidden_size + projection_config, + projection_layer_spec, + projection_type, + projection_config.ffn_hidden_size, ) else: self.projection = None @@ -181,10 +202,12 @@ def rot_pos_emb(self, grid_thw): ) wpos_ids = wpos_ids.permute(0, 2, 1, 3) wpos_ids = wpos_ids.flatten() - pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids.append(torch.stack( + [hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) pos_ids = torch.cat(pos_ids, dim=0).to(grid_thw.device) max_grid_size = grid_thw[:, 1:].max() - rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size).to(grid_thw.device) + rotary_pos_emb_full = self.rotary_pos_emb( + max_grid_size).to(grid_thw.device) rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) return rotary_pos_emb @@ -192,14 +215,18 @@ def get_window_index(self, grid_thw): window_index: list = [] cu_window_seqlens: list = [0] window_index_id = 0 - vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size + vit_merger_window_size = ( + self.window_size // self.spatial_merge_size // self.patch_size + ) for grid_t, grid_h, grid_w in grid_thw: llm_grid_h, llm_grid_w = ( grid_h // self.spatial_merge_size, grid_w // self.spatial_merge_size, ) - index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w) + index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape( + grid_t, llm_grid_h, llm_grid_w + ) pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size @@ -222,7 +249,8 @@ def get_window_index(self, grid_thw): index_padded = index_padded.reshape(-1) index_new = index_padded[index_padded != -100] window_index.append(index_new + window_index_id) - cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1] + cu_seqlens_tmp = (seqlens.cumsum( + 0) * self.spatial_merge_unit + cu_window_seqlens[-1]) cu_window_seqlens.extend(cu_seqlens_tmp.tolist()) window_index_id += (grid_t * llm_grid_h * llm_grid_w).item() window_index = torch.cat(window_index, dim=0) @@ -251,7 +279,8 @@ def forward( assert self.input_tensor is None assert inference_params is None - # Rotary positional embeddings (embedding is None for PP intermediate devices) + # Rotary positional embeddings (embedding is None for PP intermediate + # devices) vision_data = self.patch_embed(vision_data) window_index, cu_window_seqlens = self.get_window_index(grid_thw) cu_window_seqlens = torch.tensor( @@ -262,14 +291,19 @@ def forward( cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens) seq_len, _ = vision_data.size() - vision_data = vision_data.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + vision_data = vision_data.reshape( + seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1 + ) vision_data = vision_data[window_index, :, :] vision_data = vision_data.reshape(seq_len, 1, -1) rotary_pos_emb = self.rot_pos_emb(grid_thw) - rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1) + rotary_pos_emb = rotary_pos_emb.reshape( + seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1 + ) rotary_pos_emb = rotary_pos_emb[window_index, :, :] - rotary_pos_emb = rotary_pos_emb.reshape(seq_len, 1, 1, -1).repeat(1, 1, 1, 2) + rotary_pos_emb = rotary_pos_emb.reshape( + seq_len, 1, 1, -1).repeat(1, 1, 1, 2) hidden_states = self.decoder( hidden_states=vision_data, @@ -282,7 +316,8 @@ def forward( **(extra_block_kwargs or {}), ) - hidden_states = self.projection(hidden_states.view(-1, self.merge_hidden_size)) + hidden_states = self.projection( + hidden_states.view(-1, self.merge_hidden_size)) reverse_indices = torch.argsort(window_index) return hidden_states[reverse_indices, :] @@ -293,7 +328,9 @@ def build_packed_seq_params( ) -> PackedSeqParams: # NOTE: each frame is a sequence (rather than each grid) if grid_thw is not None: - seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]) + seqlens = torch.repeat_interleave( + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0] + ) cu_seqlens = seqlens.cumsum(dim=0) cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).int() else: diff --git a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_transformer_block.py b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_transformer_block.py index 8f765a0..2bdb886 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_transformer_block.py +++ b/Agent0/executor_train/verl/verl/models/mcore/qwen2_5_vl/vision_transformer_block.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright (c) 2024 Alibaba PAI Team. # @@ -34,7 +34,12 @@ def _checkpointed_forward( """Forward method with activation checkpointing.""" def custom(start: int, end: int): - def custom_forward(hidden_states, attention_mask, context, context_mask, rotary_pos_emb): + def custom_forward( + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb): for index in range(start, end): if index in fullatt_block_indexes: packed_seq_params_now = packed_seq_params_full @@ -95,7 +100,8 @@ def checkpoint_handler(forward_func): elif self.config.recompute_method == "block": # Checkpoint the input activation of only a set number of individual # Transformer layers and skip the rest. - # A method fully use the device memory removing redundant re-computation. + # A method fully use the device memory removing redundant + # re-computation. recompute_skip_num_layers = 0 for layer_idx in range(self.num_layers_per_pipeline_rank): # Skip recomputation when input grad computation is not needed. @@ -105,12 +111,19 @@ def checkpoint_handler(forward_func): recompute_skip_num_layers += 1 if ( layer_idx >= recompute_skip_num_layers - and layer_idx < self.config.recompute_num_layers + recompute_skip_num_layers + and layer_idx + < self.config.recompute_num_layers + recompute_skip_num_layers ): - hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1)) + hidden_states, context = checkpoint_handler( + custom(layer_idx, layer_idx + 1) + ) else: hidden_states, context = custom(layer_idx, layer_idx + 1)( - hidden_states, attention_mask, context, context_mask, rotary_pos_emb + hidden_states, + attention_mask, + context, + context_mask, + rotary_pos_emb, ) else: raise ValueError("Invalid activation recompute method.") @@ -164,9 +177,12 @@ def forward( [s, b, h], and optionally the updated context tensor if cross-attention is used. """ - inference_context = deprecate_inference_params(inference_context, inference_params) + inference_context = deprecate_inference_params( + inference_context, inference_params + ) - # Delete the obsolete reference to the initial input tensor if necessary + # Delete the obsolete reference to the initial input tensor if + # necessary if isinstance(hidden_states, WrappedTensor): hidden_states = hidden_states.unwrap() @@ -174,7 +190,8 @@ def forward( # See set_input_tensor() hidden_states = self.input_tensor - # Update the inference parameters with the current batch size in case it is variable + # Update the inference parameters with the current batch size in case + # it is variable if inference_context and not self.training: inference_context.current_batch_size = hidden_states.size(1) @@ -193,7 +210,9 @@ def forward( # likely redundant, since p2p_communication.py (likely originator) # already creates viewless tensors. That said, make_viewless_tensor() # is called here to be future-proof and corner-case-proof. - hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=True, keep_graph=True + ) if self.config.sequence_parallel: rng_context = tensor_parallel.get_cuda_rng_tracker().fork() @@ -205,9 +224,14 @@ def forward( # if we are using other fp8 recipes, then the context manager enter&exit are free # we can wrap fp8_context within the for loop over layers, so that we can fine-grained # control which layer will be fp8 or bf16 - use_outer_fp8_context = self.config.fp8 and self.config.fp8_recipe == Fp8Recipe.delayed - use_inner_fp8_context = self.config.fp8 and self.config.fp8_recipe != Fp8Recipe.delayed - outer_fp8_context = get_fp8_context(self.config) if use_outer_fp8_context else nullcontext() + use_outer_fp8_context = ( + self.config.fp8 and self.config.fp8_recipe == Fp8Recipe.delayed + ) + use_inner_fp8_context = ( + self.config.fp8 and self.config.fp8_recipe != Fp8Recipe.delayed + ) + outer_fp8_context = (get_fp8_context(self.config) + if use_outer_fp8_context else nullcontext()) with rng_context, outer_fp8_context: # Forward pass. @@ -226,7 +250,9 @@ def forward( else: for l_no, layer in enumerate(self.layers): inner_fp8_context = ( - get_fp8_context(self.config, layer.layer_number - 1) if use_inner_fp8_context else nullcontext() + get_fp8_context(self.config, layer.layer_number - 1) + if use_inner_fp8_context + else nullcontext() ) if l_no in fullatt_block_indexes: packed_seq_params_now = packed_seq_params_full @@ -252,7 +278,8 @@ def forward( and self.config.cpu_offloading and self.group_prefetch_offload_commit_async is not None ): - hidden_states = self.group_prefetch_offload_commit_async(hidden_states) + hidden_states = self.group_prefetch_offload_commit_async( + hidden_states) # Final layer norm. if self.final_layernorm is not None: @@ -260,6 +287,8 @@ def forward( # TENorm produces a "viewed" tensor. This will result in schedule.py's # deallocate_output_tensor() throwing an error, so a viewless tensor is # created to prevent this. - hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True) + hidden_states = make_viewless_tensor( + inp=hidden_states, requires_grad=True, keep_graph=True + ) return hidden_states diff --git a/Agent0/executor_train/verl/verl/models/mcore/registry.py b/Agent0/executor_train/verl/verl/models/mcore/registry.py index 23f01e8..039a8f2 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/registry.py +++ b/Agent0/executor_train/verl/verl/models/mcore/registry.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -73,18 +73,20 @@ class SupportedModel(Enum): # Registry for model configuration converters -MODEL_CONFIG_CONVERTER_REGISTRY: dict[SupportedModel, Callable[[PretrainedConfig, torch.dtype], TransformerConfig]] = { - SupportedModel.LLAMA: hf_to_mcore_config_dense, - SupportedModel.QWEN2: hf_to_mcore_config_dense, - SupportedModel.QWEN2_MOE: hf_to_mcore_config_qwen2moe, - SupportedModel.DEEPSEEK_V3: hf_to_mcore_config_dpskv3, - SupportedModel.MIXTRAL: hf_to_mcore_config_mixtral, - SupportedModel.QWEN2_5_VL: hf_to_mcore_config_qwen2_5_vl, - SupportedModel.LLAMA4: hf_to_mcore_config_llama4, - SupportedModel.QWEN3: hf_to_mcore_config_dense, - SupportedModel.QWEN3_MOE: hf_to_mcore_config_qwen3moe, - SupportedModel.QWEN2_5_VL: hf_to_mcore_config_qwen2_5_vl, -} +MODEL_CONFIG_CONVERTER_REGISTRY: dict[SupportedModel, + Callable[[PretrainedConfig, + torch.dtype], + TransformerConfig]] = {SupportedModel.LLAMA: hf_to_mcore_config_dense, + SupportedModel.QWEN2: hf_to_mcore_config_dense, + SupportedModel.QWEN2_MOE: hf_to_mcore_config_qwen2moe, + SupportedModel.DEEPSEEK_V3: hf_to_mcore_config_dpskv3, + SupportedModel.MIXTRAL: hf_to_mcore_config_mixtral, + SupportedModel.QWEN2_5_VL: hf_to_mcore_config_qwen2_5_vl, + SupportedModel.LLAMA4: hf_to_mcore_config_llama4, + SupportedModel.QWEN3: hf_to_mcore_config_dense, + SupportedModel.QWEN3_MOE: hf_to_mcore_config_qwen3moe, + SupportedModel.QWEN2_5_VL: hf_to_mcore_config_qwen2_5_vl, + } # Registry for model initializers MODEL_INITIALIZER_REGISTRY: dict[SupportedModel, type[BaseModelInitializer]] = { @@ -154,7 +156,9 @@ def get_supported_model(model_type: str) -> SupportedModel: def hf_to_mcore_config( - hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs + hf_config: PretrainedConfig, + dtype: torch.dtype, + **override_transformer_config_kwargs, ) -> TransformerConfig: """Convert huggingface PretrainedConfig to mcore TransformerConfig. @@ -166,9 +170,13 @@ def hf_to_mcore_config( Returns: The mcore TransformerConfig. """ - assert len(hf_config.architectures) == 1, "Only one architecture is supported for now" + assert ( + len(hf_config.architectures) == 1 + ), "Only one architecture is supported for now" model = get_supported_model(hf_config.architectures[0]) - return MODEL_CONFIG_CONVERTER_REGISTRY[model](hf_config, dtype, **override_transformer_config_kwargs) + return MODEL_CONFIG_CONVERTER_REGISTRY[model]( + hf_config, dtype, **override_transformer_config_kwargs + ) def init_mcore_model( @@ -196,7 +204,9 @@ def init_mcore_model( Returns: The initialized model. """ - assert len(hf_config.architectures) == 1, "Only one architecture is supported for now" + assert ( + len(hf_config.architectures) == 1 + ), "Only one architecture is supported for now" model = get_supported_model(hf_config.architectures[0]) initializer_cls = MODEL_INITIALIZER_REGISTRY[model] initializer = initializer_cls(tfconfig, hf_config) @@ -213,7 +223,9 @@ def get_mcore_forward_fn(hf_config: PretrainedConfig) -> Callable: """ Get the forward function for given model architecture. """ - assert len(hf_config.architectures) == 1, "Only one architecture is supported for now" + assert ( + len(hf_config.architectures) == 1 + ), "Only one architecture is supported for now" model = get_supported_model(hf_config.architectures[0]) return MODEL_FORWARD_REGISTRY[model] @@ -222,16 +234,22 @@ def get_mcore_forward_fused_fn(hf_config: PretrainedConfig) -> Callable: """ Get the forward function for given model architecture. """ - assert len(hf_config.architectures) == 1, "Only one architecture is supported for now" + assert ( + len(hf_config.architectures) == 1 + ), "Only one architecture is supported for now" model = get_supported_model(hf_config.architectures[0]) return MODEL_FORWARD_FUSED_REGISTRY[model] -def get_mcore_weight_converter(hf_config: PretrainedConfig, dtype: torch.dtype) -> Callable: +def get_mcore_weight_converter( + hf_config: PretrainedConfig, dtype: torch.dtype +) -> Callable: """ Get the weight converter for given model architecture. """ - assert len(hf_config.architectures) == 1, "Only one architecture is supported for now" + assert ( + len(hf_config.architectures) == 1 + ), "Only one architecture is supported for now" model = get_supported_model(hf_config.architectures[0]) tfconfig = hf_to_mcore_config(hf_config, dtype) return MODEL_WEIGHT_CONVERTER_REGISTRY[model](hf_config, tfconfig) diff --git a/Agent0/executor_train/verl/verl/models/mcore/saver.py b/Agent0/executor_train/verl/verl/models/mcore/saver.py index 2a954b2..e9e4fd5 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/saver.py +++ b/Agent0/executor_train/verl/verl/models/mcore/saver.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -28,7 +28,11 @@ def _megatron_calc_global_rank( - tp_rank: int = 0, dp_rank: int = 0, pp_rank: int = 0, cp_rank: int = 0, ep_rank: int = 0 + tp_rank: int = 0, + dp_rank: int = 0, + pp_rank: int = 0, + cp_rank: int = 0, + ep_rank: int = 0, ): """Calculate global rank with support for CP/EP parallelism""" @@ -39,15 +43,17 @@ def _megatron_calc_global_rank( cp_size = mpu.get_context_parallel_world_size() # ep_size = mpu.get_expert_model_parallel_world_size() - # Verify total GPU count matches (must be consistent with parallel_state.py) + # Verify total GPU count matches (must be consistent with + # parallel_state.py) total_size = tp_size * dp_size * pp_size * cp_size - assert total_size == torch.distributed.get_world_size(), ( - f"{tp_size}x{dp_size}x{pp_size}x{cp_size} != {torch.distributed.get_world_size()}" - ) + assert ( + total_size == torch.distributed.get_world_size() + ), f"{tp_size}x{dp_size}x{pp_size}x{cp_size} != {torch.distributed.get_world_size()}" # Core calculation logic (corresponds to RankGenerator order parameter) # Assumes default order is "tp-cp-ep-dp-pp" - return ((pp_rank * dp_size + dp_rank) * cp_size + cp_rank) * tp_size + tp_rank + return ((pp_rank * dp_size + dp_rank) * + cp_size + cp_rank) * tp_size + tp_rank def _megatron_calc_layer_map(config): @@ -64,13 +70,16 @@ def _megatron_calc_layer_map(config): 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 + 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 - ) + 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, @@ -80,7 +89,12 @@ def _megatron_calc_layer_map(config): return layer_map -def merge_megatron_ckpt_gptmodel(wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False): +def merge_megatron_ckpt_gptmodel( + wrapped_models, + config, + dtype, + is_value_model=False, + tie_word_embeddings=False): """Merge sharded parameters of a Megatron module into a merged checkpoint. Args: @@ -108,7 +122,8 @@ def _get_gpt_model(model): 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 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" @@ -117,16 +132,18 @@ def _get_gpt_model(model): 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 + 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].decoder.layers) == num_layers_per_model, ( - "len model layers {} not equal to num_layers_per_model {}".format( - len(models[i].decoder.layers), num_layers_per_model - ) + models[i] = unwrap_model( + wrapped_model, (torchDDP, LocalDDP, Float16Module)) + assert ( + len(models[i].decoder.layers) == num_layers_per_model + ), "len model layers {} not equal to num_layers_per_model {}".format( + len(models[i].decoder.layers), num_layers_per_model ) state_dict = dict() @@ -142,7 +159,9 @@ 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, cp_rank=cp_rank) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank + ) if torch.distributed.get_rank() == src_rank: if tensor is None: @@ -177,13 +196,17 @@ def _broadcast_tensor(tensor, name, src_pp_rank) -> torch.Tensor: 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: + 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, cp_rank=cp_rank) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank + ) chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None @@ -192,7 +215,8 @@ def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_f 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") + print_rank_0( + f"tp_shard tensor:[{name}] not exist, skip collecting") return buffer_tensor = torch.empty( @@ -205,8 +229,14 @@ def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_f 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, cp_rank=cp_rank) - sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor + cur_src_rank = _megatron_calc_global_rank( + tp_rank=i, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_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: @@ -218,13 +248,17 @@ def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_f 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: + 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, cp_rank=cp_rank) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank + ) chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None @@ -233,7 +267,10 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) 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") + print_rank_0( + f"tp_shard tensor:[{ + gate_name, + up_name}] not exist, skip collecting") return buffer_tensor = torch.empty( @@ -246,8 +283,14 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) 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, cp_rank=cp_rank) - sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor + cur_src_rank = _megatron_calc_global_rank( + tp_rank=i, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_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: @@ -259,7 +302,8 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) 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_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) @@ -268,13 +312,16 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) state_dict[gate_name] = torch.cat(gate_weight_list, dim=0) state_dict[up_name] = torch.cat(up_weight_list, dim=0) - def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): + 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, cp_rank=cp_rank) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank + ) chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None @@ -283,7 +330,8 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): 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") + print_rank_0( + f"tp_shard tensor:[{q_name}] not exist, skip collecting") return buffer_tensor = torch.empty( @@ -296,8 +344,14 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): 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, cp_rank=cp_rank) - sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor + cur_src_rank = _megatron_calc_global_rank( + tp_rank=i, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_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: @@ -308,21 +362,35 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): q_weight_list = [] k_weight_list = [] v_weight_list = [] - hidden_size_per_head = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + hidden_size_per_head = getattr( + config, + "head_dim", + config.hidden_size // + config.num_attention_heads) if config.num_key_value_heads >= tp_size: q_size_tp = hidden_size_per_head * config.num_attention_heads // tp_size - kv_size_tp = hidden_size_per_head * config.num_key_value_heads // 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): - num_query_groups_per_partition = wrapped_models[0].config.num_query_groups // tp_size - qkv_part = full_tensor[i * total_size : (i + 1) * total_size] + num_query_groups_per_partition = ( + wrapped_models[0].config.num_query_groups // tp_size + ) + qkv_part = full_tensor[i * + total_size: (i + 1) * total_size] q_size_chunk = q_size_tp // num_query_groups_per_partition kv_size_chunk = kv_size_tp // num_query_groups_per_partition - for qkv_part_chunk in qkv_part.chunk(num_query_groups_per_partition): + for qkv_part_chunk in qkv_part.chunk( + num_query_groups_per_partition + ): q_part = qkv_part_chunk[:q_size_chunk] - k_part = qkv_part_chunk[q_size_chunk : q_size_chunk + kv_size_chunk] - v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk :] + k_part = qkv_part_chunk[ + q_size_chunk: q_size_chunk + kv_size_chunk + ] + v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk:] q_weight_list.append(q_part) k_weight_list.append(k_part) v_weight_list.append(v_part) @@ -331,14 +399,21 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): kv_size_tp = hidden_size_per_head total_size = q_size_tp + 2 * kv_size_tp for i in range(tp_size): - num_query_groups_per_partition = wrapped_models[0].config.num_query_groups // tp_size - qkv_part = full_tensor[i * total_size : (i + 1) * total_size] + num_query_groups_per_partition = ( + wrapped_models[0].config.num_query_groups // tp_size + ) + qkv_part = full_tensor[i * + total_size: (i + 1) * total_size] q_size_chunk = q_size_tp // num_query_groups_per_partition kv_size_chunk = kv_size_tp // num_query_groups_per_partition - for qkv_part_chunk in qkv_part.chunk(num_query_groups_per_partition): + for qkv_part_chunk in qkv_part.chunk( + num_query_groups_per_partition + ): q_part = qkv_part_chunk[:q_size_chunk] - k_part = qkv_part_chunk[q_size_chunk : q_size_chunk + kv_size_chunk] - v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk :] + k_part = qkv_part_chunk[ + q_size_chunk: q_size_chunk + kv_size_chunk + ] + v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk:] q_weight_list.append(q_part) if i * config.num_key_value_heads % tp_size == 0: k_weight_list.append(k_part) @@ -454,12 +529,20 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): if is_value_model: lm_head_weight = None if pp_rank == pp_size - 1: - lm_head_weight = getattr(gpt_model_module.output_layer, "weight", None) - _broadcast_tensor(lm_head_weight, "lm_head.weight", src_pp_rank=pp_size - 1) + lm_head_weight = getattr( + gpt_model_module.output_layer, "weight", None + ) + _broadcast_tensor( + lm_head_weight, "lm_head.weight", src_pp_rank=pp_size - 1 + ) else: _broadcast_tp_shard_tensor( - getattr(gpt_model_module.output_layer, "weight", None) if pp_rank == pp_size - 1 else None, + ( + getattr(gpt_model_module.output_layer, "weight", None) + if pp_rank == pp_size - 1 + else None + ), "lm_head.weight", src_pp_rank=pp_size - 1, ) @@ -471,27 +554,50 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): if dtype != v.dtype: state_dict[k] = v.to(dtype) - print_rank_0(f"merge megatron ckpt done, time elapsed {time.time() - start_time}s") + print_rank_0( + f"merge megatron ckpt done, time elapsed { + time.time() - + start_time}s") return state_dict def merge_megatron_ckpt_gptmodel_qwen_moe( - wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False -): - raise NotImplementedError("merge_megatron_ckpt_gptmodel_qwen_moe is not implemented") + wrapped_models, + config, + dtype, + is_value_model=False, + tie_word_embeddings=False): + raise NotImplementedError( + "merge_megatron_ckpt_gptmodel_qwen_moe is not implemented" + ) def merge_megatron_ckpt_gptmodel_qwen2_5_vl( - wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False -): - raise NotImplementedError("merge_megatron_ckpt_gptmodel_qwen2_5_vl is not implemented") + wrapped_models, + config, + dtype, + is_value_model=False, + tie_word_embeddings=False): + raise NotImplementedError( + "merge_megatron_ckpt_gptmodel_qwen2_5_vl is not implemented" + ) -def merge_megatron_ckpt_gptmodel_dpskv3(wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False): - raise NotImplementedError("merge_megatron_ckpt_gptmodel_dpskv3 is not implemented") +def merge_megatron_ckpt_gptmodel_dpskv3( + wrapped_models, + config, + dtype, + is_value_model=False, + tie_word_embeddings=False): + raise NotImplementedError( + "merge_megatron_ckpt_gptmodel_dpskv3 is not implemented") def merge_megatron_ckpt_gptmodel_mixtral( - wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False -): - raise NotImplementedError("merge_megatron_ckpt_gptmodel_mixtral is not implemented") + wrapped_models, + config, + dtype, + is_value_model=False, + tie_word_embeddings=False): + raise NotImplementedError( + "merge_megatron_ckpt_gptmodel_mixtral is not implemented") diff --git a/Agent0/executor_train/verl/verl/models/mcore/util.py b/Agent0/executor_train/verl/verl/models/mcore/util.py index c1ef7a2..8d17a80 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/util.py +++ b/Agent0/executor_train/verl/verl/models/mcore/util.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -39,38 +39,49 @@ def preprocess_packed_seqs( pad_size = (align_size - seqlens_in_batch % align_size) % align_size seqlens_in_batch_padded = seqlens_in_batch + pad_size - cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device) + cu_seqlens = torch.zeros( + batch_size + 1, + dtype=torch.int32, + device=input_ids.device) cu_seqlens[1:] = torch.cumsum(seqlens_in_batch, dim=0) - cu_seqlens_padded = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device) + cu_seqlens_padded = torch.zeros( + batch_size + 1, dtype=torch.int32, device=input_ids.device + ) cu_seqlens_padded[1:] = torch.cumsum(seqlens_in_batch_padded, dim=0) max_seqlen_in_batch = seqlens_in_batch_padded.max().item() shape = list(input_ids.shape[1:]) shape[0] = seqlens_in_batch_padded.sum().item() // cp_size if pre_process: - input_ids_rmpad = torch.zeros(shape, dtype=input_ids.dtype, device=input_ids.device) + input_ids_rmpad = torch.zeros( + shape, dtype=input_ids.dtype, device=input_ids.device + ) for i in range(batch_size): if cp_size <= 1: seqlen = seqlens_in_batch[i] - input_ids_rmpad[cu_seqlens_padded[i] : cu_seqlens_padded[i] + seqlen] = input_ids[i, attention_mask[i]] + input_ids_rmpad[ + cu_seqlens_padded[i]: cu_seqlens_padded[i] + seqlen + ] = input_ids[i, attention_mask[i]] continue seqlen = seqlens_in_batch_padded[i] // cp_size half_seqlen = seqlen // 2 start_idx = cu_seqlens_padded[i] // cp_size # split to 2 chunks d = input_ids[i, attention_mask[i]] - input_ids_rmpad[start_idx : start_idx + half_seqlen] = d[ - half_seqlen * cp_rank : half_seqlen * (cp_rank + 1) + input_ids_rmpad[start_idx: start_idx + half_seqlen] = d[ + half_seqlen * cp_rank: half_seqlen * (cp_rank + 1) ] - remain_start = seqlens_in_batch_padded[i] - half_seqlen * (cp_rank + 1) + remain_start = seqlens_in_batch_padded[i] - \ + half_seqlen * (cp_rank + 1) remain_end = seqlens_in_batch_padded[i] - half_seqlen * cp_rank remain_end = min(remain_end, d.shape[0]) remain_len = remain_end - remain_start if remain_len > 0: - input_ids_rmpad[start_idx + half_seqlen : start_idx + half_seqlen + remain_len] = d[ - remain_start:remain_end - ] + input_ids_rmpad[start_idx + + half_seqlen: start_idx + + half_seqlen + + remain_len] = d[remain_start:remain_end] packed_seq_params = PackedSeqParams( qkv_format="thd", @@ -100,7 +111,9 @@ def postprocess_packed_seqs( """ if not post_process: return output - shape = [batch_size, seq_len] + list(output.shape[2:]) # 1,packed, dim -> batch_size, seq_len, dim + shape = [batch_size, seq_len] + list( + output.shape[2:] + ) # 1,packed, dim -> batch_size, seq_len, dim output_new = torch.zeros(shape, dtype=output.dtype, device=output.device) cp_size = mpu.get_context_parallel_world_size() @@ -109,7 +122,10 @@ def postprocess_packed_seqs( # output shape: [1, packed_len, hidden_dim] # need to gather across cp group and concatenate in sequence dimension output_list = [torch.empty_like(output) for _ in range(cp_size)] - torch.distributed.all_gather(output_list, output.detach(), group=mpu.get_context_parallel_group()) + torch.distributed.all_gather( + output_list, + output.detach(), + group=mpu.get_context_parallel_group()) output_list[mpu.get_context_parallel_rank()] = output else: output_list = [output] @@ -117,26 +133,36 @@ def postprocess_packed_seqs( if cp_size <= 1: s = attention_mask[i].sum().item() output_new[i, attention_mask[i]] = output[0][ - packed_seq_params.cu_seqlens_q_padded[i] : packed_seq_params.cu_seqlens_q_padded[i] + s + packed_seq_params.cu_seqlens_q_padded[ + i + ]: packed_seq_params.cu_seqlens_q_padded[i] + + s ] continue s_len_padded_chunk = ( - packed_seq_params.cu_seqlens_q_padded[i + 1] - packed_seq_params.cu_seqlens_q_padded[i] + packed_seq_params.cu_seqlens_q_padded[i + 1] + - packed_seq_params.cu_seqlens_q_padded[i] ) // cp_size half_seqlen = s_len_padded_chunk // 2 s_len = attention_mask[i].sum().item() s_len_padded = s_len_padded_chunk * cp_size - tmp = torch.empty(s_len_padded, *output.shape[2:], device=output.device) + tmp = torch.empty(s_len_padded, * + output.shape[2:], device=output.device) for j in range(cp_size): o = output_list[j][0] # split to 2 chunks packed_start_idx = packed_seq_params.cu_seqlens_q_padded[i] // cp_size o0, o1 = ( - o[packed_start_idx : packed_start_idx + half_seqlen], - o[packed_start_idx + half_seqlen : packed_start_idx + s_len_padded_chunk], + o[packed_start_idx: packed_start_idx + half_seqlen], + o[ + packed_start_idx + + half_seqlen: packed_start_idx + + s_len_padded_chunk + ], ) - tmp[j * half_seqlen : (j + 1) * half_seqlen] = o0 - tmp[s_len_padded - (j + 1) * half_seqlen : s_len_padded - j * half_seqlen] = o1 + tmp[j * half_seqlen: (j + 1) * half_seqlen] = o0 + tmp[s_len_padded - + (j + 1) * half_seqlen: s_len_padded - j * half_seqlen] = o1 output_new[i, attention_mask[i]] = tmp[:s_len] return output_new @@ -167,15 +193,25 @@ def remove_left_padding( seq_len = seq_len + pad_size shape[1] = seq_len if pre_process: - new_input_ids = torch.zeros(dtype=input_ids.dtype, device=input_ids.device, size=shape) + new_input_ids = torch.zeros( + dtype=input_ids.dtype, device=input_ids.device, size=shape + ) new_attention_mask = torch.zeros( - dtype=attention_mask.dtype, device=attention_mask.device, size=(batch_size, seq_len) + dtype=attention_mask.dtype, + device=attention_mask.device, + size=(batch_size, seq_len), ) - new_position_ids = torch.zeros(dtype=position_ids.dtype, device=position_ids.device, size=(batch_size, seq_len)) + new_position_ids = torch.zeros( + dtype=position_ids.dtype, + device=position_ids.device, + size=( + batch_size, + seq_len)) for i in range(batch_size): if pre_process: new_input_ids[i, : seq_lens[i]] = input_ids[i, attention_mask[i]] - new_attention_mask[i, : seq_lens[i]] = attention_mask[i, attention_mask[i]] + new_attention_mask[i, : seq_lens[i] + ] = attention_mask[i, attention_mask[i]] new_position_ids[i, : seq_lens[i]] = position_ids[i, attention_mask[i]] if pre_process: return new_input_ids, new_attention_mask, new_position_ids @@ -199,9 +235,13 @@ def recover_left_padding( shape = list(result.shape) batch_size = shape[0] shape[1] = origin_seqlen - new_result = torch.zeros(dtype=result.dtype, device=result.device, size=shape) + new_result = torch.zeros( + dtype=result.dtype, + device=result.device, + size=shape) for i in range(batch_size): - new_result[i, original_attention_mask[i]] = result[i, attention_mask[i]] + new_result[i, original_attention_mask[i] + ] = result[i, attention_mask[i]] return new_result @@ -232,9 +272,19 @@ def postprocess_packed_seqs_for_dict_output( output.log_probs = output.log_probs.view(1, -1) output.log_probs = output.log_probs.masked_fill(~labels_mask, 0.0) ret["entropy"] = postprocess_packed_seqs( - output.entropy, packed_seq_params, attention_mask, batch_size, seq_len, post_process=post_process + output.entropy, + packed_seq_params, + attention_mask, + batch_size, + seq_len, + post_process=post_process, ) ret["log_probs"] = postprocess_packed_seqs( - output.log_probs, packed_seq_params, attention_mask, batch_size, seq_len, post_process=post_process + output.log_probs, + packed_seq_params, + attention_mask, + batch_size, + seq_len, + post_process=post_process, ) return ret diff --git a/Agent0/executor_train/verl/verl/models/mcore/weight_converter.py b/Agent0/executor_train/verl/verl/models/mcore/weight_converter.py index 791513f..d825663 100644 --- a/Agent0/executor_train/verl/verl/models/mcore/weight_converter.py +++ b/Agent0/executor_train/verl/verl/models/mcore/weight_converter.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. # Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. # @@ -23,46 +23,66 @@ class McoreToHFWeightConverterBase: - def __init__(self, hf_config: PretrainedConfig, mcore_config: TransformerConfig): + def __init__(self, hf_config: PretrainedConfig, + mcore_config: TransformerConfig): self.hf_config = hf_config self.mcore_config = mcore_config - def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> torch.Tensor: + def convert_param( + self, name: str, params_one_group: list[torch.Tensor] + ) -> torch.Tensor: raise NotImplementedError class McoreToHFWeightConverterDense(McoreToHFWeightConverterBase): - def _convert_attention_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def _convert_attention_param( + self, name: str, params: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: # 'decoder.layers.0.self_attention.linear_proj.weight' # 'decoder.layers.0.self_attention.linear_qkv.layer_norm_weight' # 'decoder.layers.0.self_attention.linear_qkv.weight' # 'decoder.layers.0.self_attention.linear_qkv.bias' layer_number = name.split(".")[2] convert_names = [] - if "self_attention.linear_qkv.bias" in name or "self_attention.linear_qkv.weight" in name: + if ( + "self_attention.linear_qkv.bias" in name + or "self_attention.linear_qkv.weight" in name + ): param_type = name.split(".")[-1] assert param_type == "bias" or param_type == "weight" - convert_names.append(f"model.layers.{layer_number}.self_attn.q_proj.{param_type}") - convert_names.append(f"model.layers.{layer_number}.self_attn.k_proj.{param_type}") - convert_names.append(f"model.layers.{layer_number}.self_attn.v_proj.{param_type}") + convert_names.append( + f"model.layers.{layer_number}.self_attn.q_proj.{param_type}" + ) + convert_names.append( + f"model.layers.{layer_number}.self_attn.k_proj.{param_type}" + ) + convert_names.append( + f"model.layers.{layer_number}.self_attn.v_proj.{param_type}" + ) assert len(params) == 3 elif "self_attention.linear_proj.weight" in name: - convert_names.append(f"model.layers.{layer_number}.self_attn.o_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.self_attn.o_proj.weight") assert len(params) == 1 elif "self_attention.linear_qkv.layer_norm_weight" in name: - convert_names.append(f"model.layers.{layer_number}.input_layernorm.weight") + convert_names.append( + f"model.layers.{layer_number}.input_layernorm.weight") assert len(params) == 1 elif "self_attention.q_layernorm.weight" in name: - convert_names.append(f"model.layers.{layer_number}.self_attn.q_norm.weight") + convert_names.append( + f"model.layers.{layer_number}.self_attn.q_norm.weight") assert len(params) == 1 elif "self_attention.k_layernorm.weight" in name: - convert_names.append(f"model.layers.{layer_number}.self_attn.k_norm.weight") + convert_names.append( + f"model.layers.{layer_number}.self_attn.k_norm.weight") assert len(params) == 1 else: raise NotImplementedError(f"Unsupported parameter name: {name}") return convert_names, params - def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def _convert_mlp_param( + self, name: str, params: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: # 'decoder.layers.0.mlp.linear_fc1.layer_norm_weight' # 'decoder.layers.0.mlp.linear_fc1.weight' # 'decoder.layers.0.mlp.linear_fc2.weight' @@ -70,20 +90,27 @@ def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[lis convert_names = [] if "mlp.linear_fc1.weight" in name: # split gate_proj and up_proj - convert_names.append(f"model.layers.{layer_number}.mlp.gate_proj.weight") - convert_names.append(f"model.layers.{layer_number}.mlp.up_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.gate_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.up_proj.weight") assert len(params) == 2 elif "mlp.linear_fc1.layer_norm_weight" in name: - convert_names.append(f"model.layers.{layer_number}.post_attention_layernorm.weight") + convert_names.append( + f"model.layers.{layer_number}.post_attention_layernorm.weight" + ) assert len(params) == 1 elif "mlp.linear_fc2.weight" in name: - convert_names.append(f"model.layers.{layer_number}.mlp.down_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.down_proj.weight") assert len(params) == 1 else: raise NotImplementedError(f"Unsupported parameter name: {name}") return convert_names, params - def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def convert_param( + self, name: str, params_one_group: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: direct_name_mapping = { "embedding.word_embeddings.weight": "model.embed_tokens.weight", "decoder.final_layernorm.weight": "model.norm.weight", @@ -101,7 +128,9 @@ def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> tupl class McoreToHFWeightConverterQwen2Moe(McoreToHFWeightConverterDense): - def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def _convert_mlp_param( + self, name: str, params: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: # 'decoder.layers.0.pre_mlp_layernorm.weight', # 'decoder.layers.0.mlp.router.weight', # 'decoder.layers.0.mlp.shared_experts.gate_weight', @@ -118,29 +147,41 @@ def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[lis layer_number = name.split(".")[2] convert_names = [] if "pre_mlp_layernorm" in name: - convert_names.append(f"model.layers.{layer_number}.post_attention_layernorm.weight") + convert_names.append( + f"model.layers.{layer_number}.post_attention_layernorm.weight" + ) assert len(params) == 1 elif "mlp.router.weight" in name: - convert_names.append(f"model.layers.{layer_number}.mlp.gate.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.gate.weight") assert len(params) == 1 elif "shared_experts.gate_weight" in name: - convert_names.append(f"model.layers.{layer_number}.mlp.shared_expert_gate.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.shared_expert_gate.weight" + ) assert len(params) == 1 elif "shared_experts.linear_fc1.weight" in name: # split gate_proj and up_proj - convert_names.append(f"model.layers.{layer_number}.mlp.shared_expert.gate_proj.weight") - convert_names.append(f"model.layers.{layer_number}.mlp.shared_expert.up_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.shared_expert.gate_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.shared_expert.up_proj.weight" + ) assert len(params) == 2 elif "shared_experts.linear_fc2.weight" in name: - convert_names.append(f"model.layers.{layer_number}.mlp.shared_expert.down_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.shared_expert.down_proj.weight") assert len(params) == 1 elif "mlp.experts.linear_fc1" in name: # split gate_proj and up_proj expert_id = name.split("weight")[-1] - convert_names.append(f"model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight") - convert_names.append(f"model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight") assert len(params) == 2 elif "mlp.experts.linear_fc2" in name: expert_id = name.split("weight")[-1] - convert_names.append(f"model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight") assert len(params) == 1 else: raise NotImplementedError(f"Unsupported parameter name: {name}") @@ -148,7 +189,9 @@ def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[lis class McoreToHFWeightConverterQwen2_5_VL(McoreToHFWeightConverterDense): - def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def convert_param( + self, name: str, params_one_group: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: direct_name_mapping = { "language_model.embedding.word_embeddings.weight": "model.embed_tokens.weight", "language_model.decoder.final_layernorm.weight": "model.norm.weight", @@ -170,7 +213,9 @@ def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> tupl else: raise NotImplementedError(f"Unsupported parameter name: {name}") - def _convert_attention_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def _convert_attention_param( + self, name: str, params: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: model_type, _, _, layer_number = name.split(".")[:4] convert_names = [] @@ -197,7 +242,8 @@ def _convert_attention_param(self, name: str, params: list[torch.Tensor]) -> tup convert_names.append(f"model.layers.{layer_number}.{one}") else: assert len(params) == 1 - convert_names.append(f"model.layers.{layer_number}.{mapped_name}") + convert_names.append( + f"model.layers.{layer_number}.{mapped_name}") elif model_type == "vision_model": name_map_after_layer = { "self_attention.linear_proj.weight": "attn.proj.weight", @@ -212,24 +258,34 @@ def _convert_attention_param(self, name: str, params: list[torch.Tensor]) -> tup new_param = torch.cat(params, dim=0) params = [new_param] if "bias" in name_after_layer: - convert_names.append(f"visual.blocks.{layer_number}.attn.qkv.bias") + convert_names.append( + f"visual.blocks.{layer_number}.attn.qkv.bias") else: - convert_names.append(f"visual.blocks.{layer_number}.attn.qkv.weight") + convert_names.append( + f"visual.blocks.{layer_number}.attn.qkv.weight" + ) else: assert len(params) == 1 - convert_names.append(f"visual.blocks.{layer_number}.{mapped_name}") + convert_names.append( + f"visual.blocks.{layer_number}.{mapped_name}") else: raise NotImplementedError(f"Unsupported model type: {model_type}") return convert_names, params - def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def _convert_mlp_param( + self, name: str, params: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: model_type, _, _, layer_number = name.split(".")[:4] convert_names = [] if model_type == "language_model": name_map_after_layer = { - "mlp.linear_fc1.weight": ["mlp.gate_proj.weight", "mlp.up_proj.weight"], - "mlp.linear_fc1.bias": ["mlp.gate_proj.bias", "mlp.up_proj.bias"], + "mlp.linear_fc1.weight": [ + "mlp.gate_proj.weight", + "mlp.up_proj.weight"], + "mlp.linear_fc1.bias": [ + "mlp.gate_proj.bias", + "mlp.up_proj.bias"], "mlp.linear_fc2.weight": "mlp.down_proj.weight", "mlp.linear_fc2.bias": "mlp.down_proj.bias", "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight", @@ -242,12 +298,17 @@ def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[lis convert_names.append(f"model.layers.{layer_number}.{one}") else: assert len(params) == 1 - convert_names.append(f"model.layers.{layer_number}.{mapped_name}") + convert_names.append( + f"model.layers.{layer_number}.{mapped_name}") elif model_type == "vision_model": name_map_after_layer = { - "mlp.linear_fc1.weight": ["mlp.gate_proj.weight", "mlp.up_proj.weight"], - "mlp.linear_fc1.bias": ["mlp.gate_proj.bias", "mlp.up_proj.bias"], + "mlp.linear_fc1.weight": [ + "mlp.gate_proj.weight", + "mlp.up_proj.weight"], + "mlp.linear_fc1.bias": [ + "mlp.gate_proj.bias", + "mlp.up_proj.bias"], "mlp.linear_fc2.weight": "mlp.down_proj.weight", "mlp.linear_fc2.bias": "mlp.down_proj.bias", "mlp.linear_fc1.layer_norm_weight": "norm2.weight", @@ -260,14 +321,17 @@ def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[lis convert_names.append(f"visual.blocks.{layer_number}.{one}") else: assert len(params) == 1 - convert_names.append(f"visual.blocks.{layer_number}.{mapped_name}") + convert_names.append( + f"visual.blocks.{layer_number}.{mapped_name}") else: raise NotImplementedError(f"Unsupported model type: {model_type}") return convert_names, params class McoreToHFWeightConverterDpskv3(McoreToHFWeightConverterBase): - def _convert_attention_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def _convert_attention_param( + self, name: str, params: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: # mcore # 'decoder.layers.0.input_layernorm.weight' # 'decoder.layers.0.self_attention.linear_proj.weight' @@ -303,10 +367,14 @@ def _convert_attention_param(self, name: str, params: list[torch.Tensor]) -> tup convert_names = [] layer_number = name.split(".")[2] name_after_layer = name.split(f".{layer_number}.")[1] - convert_names.append(f"model.layers.{layer_number}.{name_map_after_layer[name_after_layer]}") + convert_names.append( + f"model.layers.{layer_number}.{ + name_map_after_layer[name_after_layer]}") return convert_names, params - def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def _convert_mlp_param( + self, name: str, params: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: # mcore dense # 'decoder.layers.0.mlp.linear_fc1.layer_norm_weight' # 'decoder.layers.0.mlp.linear_fc2.weight' @@ -343,7 +411,9 @@ def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[lis "mlp.linear_fc1.layer_norm_weight": "post_attention_layernorm.weight", "mlp.linear_fc2.weight": "mlp.down_proj.weight", "mlp.shared_experts.linear_fc2.weight": "mlp.shared_experts.down_proj.weight", - "mlp.linear_fc1.weight": ["mlp.gate_proj.weight", "mlp.up_proj.weight"], + "mlp.linear_fc1.weight": [ + "mlp.gate_proj.weight", + "mlp.up_proj.weight"], "mlp.shared_experts.linear_fc1.weight": [ "mlp.shared_experts.gate_proj.weight", "mlp.shared_experts.up_proj.weight", @@ -363,24 +433,33 @@ def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[lis convert_names.append(f"model.layers.{layer_number}.{one}") else: assert len(params) == 1 - convert_names.append(f"model.layers.{layer_number}.{mapped_name}") + convert_names.append( + f"model.layers.{layer_number}.{mapped_name}") else: if "mlp.experts.linear_fc1.weight" in name: expert_id = name.split("weight")[-1] - convert_names.append(f"model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight") - convert_names.append(f"model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight") assert len(params) == 2 elif "mlp.experts.linear_fc2.weight" in name: expert_id = name.split("weight")[-1] - convert_names.append(f"model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight") assert len(params) == 1 else: - raise NotImplementedError(f"Unsupported parameter name: {name}") + raise NotImplementedError( + f"Unsupported parameter name: {name}") return convert_names, params - def _convert_mtp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: - assert self.mcore_config.mtp_num_layers == 1, "only support one mtp layer for now" + def _convert_mtp_param( + self, name: str, params: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: + assert ( + self.mcore_config.mtp_num_layers == 1 + ), "only support one mtp layer for now" assert self.mcore_config.num_layers == 61, "only support 61 layers for now" direct_name_mapping = { "mtp.layers.0.enorm.weight": "model.layers.61.enorm.weight", @@ -390,18 +469,25 @@ def _convert_mtp_param(self, name: str, params: list[torch.Tensor]) -> tuple[lis } if name in direct_name_mapping: return [direct_name_mapping[name]], [params[0]] - assert "mtp.layers.0.transformer_layer" in name, "only support transformer layer for now" + assert ( + "mtp.layers.0.transformer_layer" in name + ), "only support transformer layer for now" # use proxy name to convert - proxy_name = name.replace("mtp.layers.0.transformer_layer", "decoder.layers.61") + proxy_name = name.replace( + "mtp.layers.0.transformer_layer", + "decoder.layers.61") if "self_attention" in proxy_name or "input_layernorm.weight" in proxy_name: - convert_names, params = self._convert_attention_param(proxy_name, params) + convert_names, params = self._convert_attention_param( + proxy_name, params) elif "mlp" in proxy_name: convert_names, params = self._convert_mlp_param(proxy_name, params) else: raise NotImplementedError(f"Unsupported parameter name: {name}") return convert_names, params - def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def convert_param( + self, name: str, params_one_group: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: direct_name_mapping = { "embedding.word_embeddings.weight": "model.embed_tokens.weight", "decoder.final_layernorm.weight": "model.norm.weight", @@ -420,7 +506,9 @@ def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> tupl class McoreToHFWeightConverterMixtral(McoreToHFWeightConverterDense): - def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def _convert_mlp_param( + self, name: str, params: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: # decoder.layers.0.mlp.router.weight # decoder.layers.0.mlp.experts.linear_fc1.weight0 - weight7 # decoder.layers.0.mlp.experts.linear_fc2.weight0 - weight7 @@ -428,23 +516,32 @@ def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[lis layer_number = name.split(".")[2] convert_names = [] if "pre_mlp_layernorm" in name: - convert_names.append(f"model.layers.{layer_number}.post_attention_layernorm.weight") + convert_names.append( + f"model.layers.{layer_number}.post_attention_layernorm.weight" + ) elif "mlp.router.weight" in name: - convert_names.append(f"model.layers.{layer_number}.block_sparse_moe.gate.weight") + convert_names.append( + f"model.layers.{layer_number}.block_sparse_moe.gate.weight" + ) elif "mlp.experts.linear_fc1.weight" in name: expert_id = name.split("weight")[-1] - convert_names.append(f"model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w1.weight") - convert_names.append(f"model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w3.weight") + convert_names.append( + f"model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w1.weight") + convert_names.append( + f"model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w3.weight") elif "mlp.experts.linear_fc2.weight" in name: expert_id = name.split("weight")[-1] - convert_names.append(f"model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w2.weight") + convert_names.append( + f"model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w2.weight") else: raise NotImplementedError(f"Unsupported parameter name: {name}") return convert_names, params class McoreToHFWeightConverterQwen3Moe(McoreToHFWeightConverterDense): - def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]: + def _convert_mlp_param( + self, name: str, params: list[torch.Tensor] + ) -> tuple[list[str], list[torch.Tensor]]: # qwen3 moe no share expert # 'decoder.layers.0.pre_mlp_layernorm.weight', @@ -460,19 +557,25 @@ def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[lis layer_number = name.split(".")[2] convert_names = [] if "pre_mlp_layernorm" in name: - convert_names.append(f"model.layers.{layer_number}.post_attention_layernorm.weight") + convert_names.append( + f"model.layers.{layer_number}.post_attention_layernorm.weight" + ) assert len(params) == 1 elif "mlp.router.weight" in name: - convert_names.append(f"model.layers.{layer_number}.mlp.gate.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.gate.weight") assert len(params) == 1 elif "mlp.experts.linear_fc1" in name: # split gate_proj and up_proj expert_id = name.split("weight")[-1] - convert_names.append(f"model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight") - convert_names.append(f"model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight") assert len(params) == 2 elif "mlp.experts.linear_fc2" in name: expert_id = name.split("weight")[-1] - convert_names.append(f"model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight") + convert_names.append( + f"model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight") assert len(params) == 1 else: raise NotImplementedError(f"Unsupported parameter name: {name}") diff --git a/Agent0/executor_train/verl/verl/models/qwen2/__init__.py b/Agent0/executor_train/verl/verl/models/qwen2/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/__init__.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/__init__.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/__init__.py index 57e33ee..0af23fe 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/__init__.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/__init__.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/__init__.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_loader.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_loader.py index 3168635..a56fb7c 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_loader.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_loader.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -34,13 +34,16 @@ def _megatron_calc_layer_map(config): 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 + 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 - ) + 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, @@ -51,7 +54,12 @@ def _megatron_calc_layer_map(config): def load_state_dict_to_megatron_qwen2( - state_dict, wrapped_models, config, params_dtype, is_value_model=False, tie_word_embeddings=False + state_dict, + wrapped_models, + config, + params_dtype, + is_value_model=False, + tie_word_embeddings=False, ): """Load merged state_dict to sharded Megatron module in training.""" from megatron.core import DistributedDataParallel as LocalDDP @@ -70,7 +78,9 @@ def _get_gpt_model(model): def fetch_params(module): for param in module.parameters(): torch.distributed.fetch( - param.data, src=mpu.get_data_parallel_src_rank(), group=mpu.get_data_parallel_group() + param.data, + src=mpu.get_data_parallel_src_rank(), + group=mpu.get_data_parallel_group(), ) dp_rank = mpu.get_data_parallel_rank() @@ -80,7 +90,8 @@ def fetch_params(module): 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 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" @@ -89,7 +100,9 @@ def fetch_params(module): 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, ( + assert ( + num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers + ), ( f"num_layers_per_model: {num_layers_per_model} * pp_size: {pp_size} * virtual_pp_size: " f"{virtual_pp_size} != config.num_hidden_layers: {config.num_hidden_layers}" ) @@ -97,7 +110,8 @@ def fetch_params(module): models = [None] * len(wrapped_models) for i, wrapped_model in enumerate(wrapped_models): - models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module)) + 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 @@ -107,7 +121,9 @@ def _fetch_tensor(tensor, name) -> torch.Tensor: if tensor is not None: tensor = tensor.data.copy_(state_dict[name], non_blocking=True) - def _fetch_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor: + def _fetch_tp_shard_tensor_vocab( + tensor, name, chunk_dim=0, mutate_func=None + ) -> torch.Tensor: """fetch tensor in tp shards""" nonlocal state_dict tp_rank = mpu.get_tensor_model_parallel_rank() @@ -119,11 +135,14 @@ def _fetch_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> full_weight = mutate_func(full_weight) tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim) if tensor is not None: - tensor = tensor.data.copy_(tensor_chunk[tp_rank], non_blocking=True) + tensor = tensor.data.copy_( + tensor_chunk[tp_rank], non_blocking=True) else: print(f"tp_shard tensor:[{name}] not in state_dict, skip loading") - def _fetch_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor: + def _fetch_tp_shard_tensor( + tensor, name, chunk_dim=0, mutate_func=None + ) -> torch.Tensor: """fetch tensor in tp shards""" nonlocal state_dict tp_rank = mpu.get_tensor_model_parallel_rank() @@ -135,11 +154,13 @@ def _fetch_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch full_weight = mutate_func(full_weight) tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim) if tensor is not None: - tensor = tensor.data.copy_(tensor_chunk[tp_rank], non_blocking=True) + tensor = tensor.data.copy_( + tensor_chunk[tp_rank], non_blocking=True) else: print(f"tp_shard tensor:[{name}] not in state_dict, skip loading") - def _fetch_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor: + def _fetch_tp_shard_tensor_gate_up( + tensor, gate_name, up_name) -> torch.Tensor: """fetch gate_up tensor in tp shards""" nonlocal state_dict nonlocal mp_group @@ -149,23 +170,34 @@ def _fetch_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor: 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=get_device_id() + config.intermediate_size * 2, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) 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) - ) + 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) if tensor is not None: - tensor = tensor.data.copy_(tensor_chunk[tp_rank], non_blocking=True) + tensor = tensor.data.copy_( + tensor_chunk[tp_rank], non_blocking=True) else: - print(f"tp_shard tensor:[{gate_name}, {up_name}] not in state_dict, skip loading") + print( + f"tp_shard tensor:[{gate_name}, {up_name}] not in state_dict, skip loading") - def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> torch.Tensor: + def _fetch_tp_shard_tensor_qkv( + tensor, q_name, k_name, v_name, bias=False + ) -> torch.Tensor: """fetch tensor in tp shards across mp_group""" nonlocal state_dict nonlocal mp_group @@ -184,15 +216,23 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> to total_size = q_size_tp + 2 * kv_size_tp if not bias: new_weight_qkv = torch.empty( - total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id() + total_size * tp_size, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) else: - new_weight_qkv = torch.empty(total_size * tp_size, dtype=params_dtype, device=get_device_id()) + new_weight_qkv = torch.empty( + total_size * tp_size, + dtype=params_dtype, + device=get_device_id()) 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)) + 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 @@ -200,21 +240,33 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> to total_size = q_size_tp + 2 * kv_size_tp if not bias: new_weight_qkv = torch.empty( - total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id() + total_size * tp_size, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) else: - new_weight_qkv = torch.empty(total_size * tp_size, dtype=params_dtype, device=get_device_id()) + new_weight_qkv = torch.empty( + total_size * tp_size, + dtype=params_dtype, + device=get_device_id()) 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 + 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)) + 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) if tensor is not None: - tensor = tensor.data.copy_(tensor_chunk[tp_rank], non_blocking=True) + tensor = tensor.data.copy_( + tensor_chunk[tp_rank], non_blocking=True) # Embeddings # ------------------- @@ -222,7 +274,9 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> to gpt_model_module = _get_gpt_model(models[0]) if pp_rank == 0: embed_tokens_weight = gpt_model_module.model.embed_tokens.weight - _fetch_tp_shard_tensor_vocab(embed_tokens_weight, "model.embed_tokens.weight") + _fetch_tp_shard_tensor_vocab( + embed_tokens_weight, + "model.embed_tokens.weight") # Transformer layers # ------------------- @@ -238,10 +292,16 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> to for vpp_rank in range(vpp_size): num_layer_vpp_chunk = num_layer_per_pp // vpp_size num_layer_this_model = num_layer_vpp_chunk - offset = vpp_rank * (config.num_hidden_layers // mpu.get_virtual_pipeline_model_parallel_world_size()) + ( - mpu.get_pipeline_model_parallel_rank() * num_layer_vpp_chunk - ) - layer_list.extend(list(range(offset, offset + num_layer_this_model))) + offset = vpp_rank * ( + config.num_hidden_layers + // mpu.get_virtual_pipeline_model_parallel_world_size() + ) + (mpu.get_pipeline_model_parallel_rank() * num_layer_vpp_chunk) + layer_list.extend( + list( + range( + offset, + offset + + num_layer_this_model))) else: num_layer_this_model = num_layer_per_pp offset = pp_rank * num_layer_per_pp @@ -287,7 +347,11 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> to ) _fetch_tensor( - sync_layer.post_attention_layernorm.weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.post_attention_layernorm.weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.post_attention_layernorm.weight", ) @@ -319,10 +383,16 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> to lm_head_weight = gpt_model_module.lm_head.weight if is_value_model: - if "lm_head.weight" in state_dict and state_dict["lm_head.weight"].shape[0] == 1: + if ( + "lm_head.weight" in state_dict + and state_dict["lm_head.weight"].shape[0] == 1 + ): _fetch_tensor(lm_head_weight, "lm_head.weight") print_rank_0("load lm_head from value_head weight") - elif "reward_head.weight" in state_dict and state_dict["reward_head.weight"].shape[0] == 1: + elif ( + "reward_head.weight" in state_dict + and state_dict["reward_head.weight"].shape[0] == 1 + ): _fetch_tensor(lm_head_weight, "reward_head.weight") print_rank_0("load lm_head from value_head weight") else: @@ -334,4 +404,6 @@ def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> to dist.barrier() get_torch_device().empty_cache() - print_rank_0(f"loading megatron ckpt done, time elapsed {time.time() - start_time}s") + print_rank_0( + f"loading megatron ckpt done, time elapsed {time.time() - start_time}s" + ) diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_loader_depracated.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_loader_depracated.py index 770e365..ebdfd6b 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_loader_depracated.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_loader_depracated.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -34,13 +34,16 @@ def _megatron_calc_layer_map(config): 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 + 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 - ) + 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, @@ -51,7 +54,12 @@ def _megatron_calc_layer_map(config): def load_state_dict_to_megatron_qwen2( - state_dict, wrapped_models, config, params_dtype, is_value_model=False, tie_word_embeddings=False + state_dict, + wrapped_models, + config, + params_dtype, + is_value_model=False, + tie_word_embeddings=False, ): """Load merged state_dict to sharded Megatron module in training.""" from megatron.core import DistributedDataParallel as LocalDDP @@ -70,7 +78,9 @@ def _get_gpt_model(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() + param.data, + src=mpu.get_data_parallel_src_rank(), + group=mpu.get_data_parallel_group(), ) dp_rank = mpu.get_data_parallel_rank() @@ -80,7 +90,8 @@ def broadcast_params(module): 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 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" @@ -89,7 +100,9 @@ def broadcast_params(module): 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, ( + assert ( + num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers + ), ( f"num_layers_per_model: {num_layers_per_model} * pp_size: {pp_size} * virtual_pp_size: " f"{virtual_pp_size} != config.num_hidden_layers: {config.num_hidden_layers}" ) @@ -97,7 +110,8 @@ def broadcast_params(module): models = [None] * len(wrapped_models) for i, wrapped_model in enumerate(wrapped_models): - models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module)) + 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 @@ -135,7 +149,9 @@ def _broadcast_tensor(tensor, name) -> torch.Tensor: 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: + 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 @@ -160,7 +176,8 @@ def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None 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") + print_rank_0( + f"tp_shard tensor:[{name}] not in state_dict, skip loading") return if tensor is None: @@ -171,10 +188,13 @@ def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}" + 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=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: @@ -183,7 +203,9 @@ def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None 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: + 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 @@ -207,7 +229,8 @@ def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> t 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") + print_rank_0( + f"tp_shard tensor:[{name}] not in state_dict, skip loading") return if tensor is None: @@ -218,10 +241,13 @@ def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> t requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}" + 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=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: @@ -230,7 +256,8 @@ def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> t if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor) - def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor: + 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 @@ -241,15 +268,22 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens 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=get_device_id() + config.intermediate_size * 2, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) 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) - ) + 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 @@ -261,7 +295,10 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens 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") + print_rank_0( + f"tp_shard tensor:[{ + gate_name, + up_name}] not in state_dict, skip loading") return if tensor is None: @@ -272,11 +309,13 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank() == 0:} tensor {gate_name, up_name} shape " - f"{tensor.shape} != {chunk_shape}" + assert tensor.shape == chunk_shape, (f"rank #{ + torch.distributed.get_rank() == 0:} tensor { + gate_name, up_name} shape " f"{ + tensor.shape} != {chunk_shape}") + sync_tensor = torch.empty_like( + tensor, device=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: @@ -285,7 +324,9 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tens if (i == tp_rank) and (tensor is not None): tensor.data.copy_(sync_tensor) - def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> torch.Tensor: + def _broadcast_tp_shard_tensor_qkv( + tensor, q_name, k_name, v_name, bias=False + ) -> torch.Tensor: """broadcast tensor in tp shards across mp_group""" nonlocal state_dict nonlocal mp_group @@ -293,7 +334,8 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - 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 + 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] @@ -302,19 +344,28 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - 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 + kv_size_tp = ( + hidden_size_per_head * + config.num_key_value_heads // + tp_size) total_size = q_size_tp + 2 * kv_size_tp if not bias: new_weight_qkv = torch.empty( - total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id() + total_size * tp_size, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) else: - new_weight_qkv = torch.empty(total_size * tp_size, dtype=params_dtype, device=get_device_id()) + new_weight_qkv = torch.empty( + total_size * tp_size, dtype=params_dtype, device=get_device_id()) 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_( + 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) ) @@ -324,17 +375,27 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - total_size = q_size_tp + 2 * kv_size_tp if not bias: new_weight_qkv = torch.empty( - total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id() + total_size * tp_size, + config.hidden_size, + dtype=params_dtype, + device=get_device_id(), ) else: - new_weight_qkv = torch.empty(total_size * tp_size, dtype=params_dtype, device=get_device_id()) + new_weight_qkv = torch.empty( + total_size * tp_size, dtype=params_dtype, device=get_device_id()) 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 + 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_( + new_weight_qkv[i * total_size: (i + 1) * total_size].copy_( torch.cat([q_part, k_part, v_part], dim=0) ) @@ -348,7 +409,11 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - 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, k_name, v_name}] not in state_dict, skip loading") + print_rank_0( + f"tp_shard tensor:[{ + q_name, + k_name, + v_name}] not in state_dict, skip loading") return if tensor is None: @@ -359,10 +424,13 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - requires_grad=False, ) else: - assert tensor.shape == chunk_shape, ( - f"rank #{torch.distributed.get_rank()} tensor {q_name} shape {tensor.shape} != {chunk_shape}" + 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=get_device_id(), requires_grad=False ) - sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False) for i in range(tp_size): if torch.distributed.get_rank() == 0: @@ -379,7 +447,9 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - embed_tokens_weight = None if pp_rank == 0: embed_tokens_weight = gpt_model_module.model.embed_tokens.weight - _broadcast_tp_shard_tensor_vocab(embed_tokens_weight, "model.embed_tokens.weight") + _broadcast_tp_shard_tensor_vocab( + embed_tokens_weight, "model.embed_tokens.weight" + ) # Transformer layers # ------------------- @@ -399,7 +469,11 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - ) _broadcast_tp_shard_tensor_qkv( - sync_layer.self_attn.qkv_proj.weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.self_attn.qkv_proj.weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.self_attn.q_proj.weight", f"{layer_name}.self_attn.k_proj.weight", f"{layer_name}.self_attn.v_proj.weight", @@ -420,7 +494,11 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - ) _broadcast_tensor( - sync_layer.post_attention_layernorm.weight if dst_pp_rank == pp_rank else None, + ( + sync_layer.post_attention_layernorm.weight + if dst_pp_rank == pp_rank + else None + ), f"{layer_name}.post_attention_layernorm.weight", ) @@ -453,10 +531,16 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - lm_head_weight = gpt_model_module.lm_head.weight if is_value_model: - if "lm_head.weight" in state_dict and state_dict["lm_head.weight"].shape[0] == 1: + if ( + "lm_head.weight" in state_dict + and state_dict["lm_head.weight"].shape[0] == 1 + ): _broadcast_tensor(lm_head_weight, "lm_head.weight") print_rank_0("load lm_head from value_head weight") - elif "reward_head.weight" in state_dict and state_dict["reward_head.weight"].shape[0] == 1: + elif ( + "reward_head.weight" in state_dict + and state_dict["reward_head.weight"].shape[0] == 1 + ): _broadcast_tensor(lm_head_weight, "reward_head.weight") print_rank_0("load lm_head from value_head weight") else: @@ -472,4 +556,6 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) - broadcast_params(wrapped_model) get_torch_device().empty_cache() - print_rank_0(f"loading megatron ckpt done, time elapsed {time.time() - start_time}s") + print_rank_0( + f"loading megatron ckpt done, time elapsed {time.time() - start_time}s" + ) diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_saver.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_saver.py index 737f73b..a3d91eb 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_saver.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/checkpoint_utils/qwen2_saver.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,15 +26,18 @@ from verl.utils.megatron_utils import unwrap_model -def _megatron_calc_global_rank(tp_rank: int = 0, dp_rank: int = 0, pp_rank: int = 0): +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.""" 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()}" - ) + 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()}" # We only support TP-DP-PP grouping, for correctness when resharding return (pp_rank * dp_size + dp_rank) * tp_size + tp_rank @@ -53,13 +56,16 @@ def _megatron_calc_layer_map(config): 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 + 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 - ) + 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, @@ -69,7 +75,12 @@ def _megatron_calc_layer_map(config): return layer_map -def merge_megatron_ckpt_qwen2(wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False): +def merge_megatron_ckpt_qwen2( + wrapped_models, + config, + dtype, + is_value_model=False, + tie_word_embeddings=False): """Merge sharded parameters of a Megatron module into a merged checkpoint. Args: @@ -96,7 +107,8 @@ def _get_gpt_model(model): 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 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" @@ -105,16 +117,18 @@ def _get_gpt_model(model): 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 + 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 - ) + 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() @@ -130,7 +144,8 @@ 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) + 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: @@ -165,12 +180,15 @@ def _broadcast_tensor(tensor, name, src_pp_rank) -> torch.Tensor: 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: + 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_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) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None @@ -179,7 +197,8 @@ def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_f 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") + print_rank_0( + f"tp_shard tensor:[{name}] not exist, skip collecting") return buffer_tensor = torch.empty( @@ -192,8 +211,14 @@ def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_f 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 + 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: @@ -205,12 +230,15 @@ def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_f 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: + 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_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) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None @@ -219,7 +247,10 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) 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") + print_rank_0( + f"tp_shard tensor:[{ + gate_name, + up_name}] not exist, skip collecting") return buffer_tensor = torch.empty( @@ -232,8 +263,14 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) 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 + 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: @@ -245,7 +282,8 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) 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_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) @@ -254,12 +292,14 @@ def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) state_dict[gate_name] = torch.cat(gate_weight_list, dim=0) state_dict[up_name] = torch.cat(up_weight_list, dim=0) - def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): + 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_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) + src_rank = _megatron_calc_global_rank( + tp_rank=0, dp_rank=0, pp_rank=src_pp_rank) chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None @@ -268,7 +308,8 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): 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") + print_rank_0( + f"tp_shard tensor:[{q_name}] not exist, skip collecting") return buffer_tensor = torch.empty( @@ -281,8 +322,14 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): 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 + 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: @@ -297,13 +344,17 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): 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 + 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] + 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] + 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) @@ -312,10 +363,11 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): 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] + 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] + 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) @@ -417,21 +469,31 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): if is_value_model: _broadcast_tensor( - gpt_model_module.lm_head.weight if pp_rank == pp_size - 1 else None, + gpt_model_module.lm_head.weight if pp_rank == pp_size - + 1 else None, "lm_head.weight", - src_pp_rank=pp_size - 1, + src_pp_rank=pp_size - + 1, ) _broadcast_tensor( - gpt_model_module.reward_head.weight - if pp_rank == pp_size - 1 and getattr(gpt_model_module, "reward_weight", None) is not None - else None, + ( + gpt_model_module.reward_head.weight if pp_rank == pp_size - + 1 and getattr( + gpt_model_module, + "reward_weight", + None) is not None else None), "reward_head.weight", - src_pp_rank=pp_size - 1, + src_pp_rank=pp_size - + 1, ) else: _broadcast_tp_shard_tensor( - getattr(gpt_model_module.lm_head, "weight", None) if pp_rank == pp_size - 1 else None, + ( + getattr(gpt_model_module.lm_head, "weight", None) + if pp_rank == pp_size - 1 + else None + ), "lm_head.weight", src_pp_rank=pp_size - 1, ) @@ -444,5 +506,8 @@ def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank): if dtype != v.dtype: state_dict[k] = v.to(dtype) - print_rank_0(f"merge megatron ckpt done, time elapsed {time.time() - start_time}s") + print_rank_0( + f"merge megatron ckpt done, time elapsed { + time.time() - + start_time}s") return state_dict diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/__init__.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/__init__.py index 263ea59..0ae513b 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/__init__.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_attention.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_attention.py index 702c429..c7396cf 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_attention.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_attention.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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 @@ -40,34 +40,53 @@ class Qwen2RotaryEmbedding(nn.Module): - def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None): + 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)) + 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() + 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) + 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 + # 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) + 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) + 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), @@ -78,26 +97,49 @@ def forward(self, x, seq_len=None): class Qwen2LinearScalingRotaryEmbedding(Qwen2RotaryEmbedding): """Qwen2RotaryEmbedding 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): + 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 = 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 + # 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) + self.register_buffer( + "cos_cached", + emb.cos().to(dtype), + persistent=False) + self.register_buffer( + "sin_cached", + emb.sin().to(dtype), + persistent=False) class Qwen2DynamicNTKScalingRotaryEmbedding(Qwen2RotaryEmbedding): """Qwen2RotaryEmbedding 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): + 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) @@ -106,24 +148,36 @@ def _set_cos_sin_cache(self, seq_len, device, dtype): if seq_len > self.max_position_embeddings: base = self.base * ( - (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1) + (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)) + 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) + 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 + # 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) + self.register_buffer( + "cos_cached", + emb.cos().to(dtype), + persistent=False) + self.register_buffer( + "sin_cached", + emb.sin().to(dtype), + persistent=False) def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] + x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) @@ -143,14 +197,20 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: 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) + 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) class ParallelQwen2Attention(nn.Module): """Multi-headed attention from 'Attention Is All You Need' paper""" - def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig): + def __init__( + self, + config: Qwen2Config, + megatron_config: ModelParallelConfig): super().__init__() self.config = config self.megatron_config = megatron_config @@ -164,9 +224,10 @@ def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig): # 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_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=" f"{self.num_key_value_heads}, tp_size={tp_size}" @@ -178,16 +239,18 @@ def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig): 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} and " - f"`num_heads`: {self.num_heads})." - ) + f"hidden_size must be divisible by num_heads (got `hidden_size`: { + self.hidden_size} and " f"`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" + 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) @@ -228,53 +291,85 @@ def _init_rope(self): ) 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() + 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]]]: + 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, 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) + 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) + 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) + 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)}, " - f"but is {attn_weights.size()}" - ) + f"Attention weights should be of size { + ( + bsz, + self.num_heads_per_tp, + q_len, + kv_seq_len)}, " f"but is { + 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()}" - ) + 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_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)}, " - f"but is {attn_output.size()}" - ) + f"`attn_output` should be of size { + ( + bsz, + self.num_heads_per_tp, + q_len, + self.head_dim)}, " f"but is { + attn_output.size()}") attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.reshape(bsz, q_len, self.hidden_size_per_tp) @@ -289,18 +384,35 @@ def forward( """ -def apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, sequence_length): +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) + 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) + 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 @@ -309,10 +421,22 @@ def apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, sequence_l # cos/sin shoudl be: (seq_length, rotary_dim / 2) 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 + 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 + k, + cos, + sin, + interleaved=False, + inplace=False, + cu_seqlens=cu_seqlens, + max_seqlen=max_seqlen, ) return q_embed, k_embed @@ -327,7 +451,9 @@ def forward( 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 + 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() @@ -347,14 +473,28 @@ def forward( # 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) + 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 + 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, + 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, @@ -388,12 +528,16 @@ def forward( ) 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() + 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 = 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 diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_decoder.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_decoder.py index 3c8a2a6..1355f81 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_decoder.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_decoder.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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 @@ -33,23 +33,32 @@ class ParallelQwen2DecoderLayer(nn.Module): - def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig, layer_idx: int): + def __init__( + self, + config: Qwen2Config, + megatron_config: ModelParallelConfig, + layer_idx: int): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) self.layer_idx = layer_idx self.hidden_size = config.hidden_size - self.self_attn = ParallelQwen2Attention(config=config, megatron_config=megatron_config) + self.self_attn = ParallelQwen2Attention( + config=config, megatron_config=megatron_config + ) self.mlp = ParallelQwen2MLP(config, megatron_config=megatron_config) self.input_layernorm = ParallelQwen2RMSNorm(config, megatron_config) - self.post_attention_layernorm = ParallelQwen2RMSNorm(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]]]: + self.post_attention_layernorm = ParallelQwen2RMSNorm( + 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)` @@ -100,26 +109,35 @@ def forward( class ParallelQwen2DecoderLayerRmPad(nn.Module): - def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig, layer_idx: int): + def __init__( + self, + config: Qwen2Config, + megatron_config: ModelParallelConfig, + layer_idx: int): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) self.hidden_size = config.hidden_size self.layer_idx = layer_idx - self.self_attn = ParallelQwen2AttentionRmPad(config=config, megatron_config=megatron_config) + self.self_attn = ParallelQwen2AttentionRmPad( + config=config, megatron_config=megatron_config + ) self.mlp = ParallelQwen2MLP(config, megatron_config=megatron_config) self.input_layernorm = ParallelQwen2RMSNorm(config, megatron_config) - self.post_attention_layernorm = ParallelQwen2RMSNorm(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]]]: + self.post_attention_layernorm = ParallelQwen2RMSNorm( + 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) diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_linear.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_linear.py index e6d4a09..de90d28 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_linear.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_linear.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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. @@ -11,7 +11,8 @@ # WITHOUT 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 +# Adapted from +# https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/linear.py from megatron.core import tensor_parallel diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_mlp.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_mlp.py index 672908a..096c561 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_mlp.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_mlp.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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 @@ -28,19 +28,25 @@ class ParallelQwen2MLP(nn.Module): - def __init__(self, config, megatron_config: ModelParallelConfig = None) -> None: + 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] + # 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" + 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) diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_rmsnorm.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_rmsnorm.py index 2f4c90d..e9c12f2 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_rmsnorm.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/layers/parallel_rmsnorm.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,7 +24,10 @@ class ParallelQwen2RMSNorm(nn.Module): - def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig): + def __init__( + self, + config: Qwen2Config, + megatron_config: ModelParallelConfig): """ Qwen2RMSNorm is equivalent to T5LayerNorm """ diff --git a/Agent0/executor_train/verl/verl/models/qwen2/megatron/modeling_qwen2_megatron.py b/Agent0/executor_train/verl/verl/models/qwen2/megatron/modeling_qwen2_megatron.py index 92e81be..3a72c4b 100644 --- a/Agent0/executor_train/verl/verl/models/qwen2/megatron/modeling_qwen2_megatron.py +++ b/Agent0/executor_train/verl/verl/models/qwen2/megatron/modeling_qwen2_megatron.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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 @@ -34,10 +34,14 @@ from verl.utils.megatron import tensor_parallel as tp_utils from verl.utils.megatron_utils import TransformerConfig, convert_config -from .layers import ParallelQwen2DecoderLayer, ParallelQwen2DecoderLayerRmPad, ParallelQwen2RMSNorm +from .layers import ( + ParallelQwen2DecoderLayer, + ParallelQwen2DecoderLayerRmPad, + ParallelQwen2RMSNorm, +) """ -TODO: +TODO: 1. Add weight initialization. Here we need to be careful on TP weight init. 2. Add sequence parallel 3. Load checkpoint from Qwen2 pretrained checkpoint @@ -45,12 +49,18 @@ # Copied from transformers.models.bart.modeling_bart._make_causal_mask -def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device): +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 = 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) @@ -58,18 +68,24 @@ def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: t # Copied from transformers.models.bart.modeling_bart._expand_mask -def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): +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) + 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) + return inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(dtype).min + ) class ParallelQwen2Model(nn.Module): @@ -80,26 +96,41 @@ class ParallelQwen2Model(nn.Module): config: Qwen2Config """ - def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig): + def __init__( + self, + config: Qwen2Config, + megatron_config: ModelParallelConfig): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) 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, megatron_config) + assert embedding_kwargs.get( + "config", False + ), "must have ModelParallelConfig" + tp_utils.update_kwargs_with_config( + embedding_kwargs, megatron_config) self.embed_tokens = tensor_parallel.VocabParallelEmbedding( - num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + **embedding_kwargs, ) self.layers = nn.ModuleList( - [ParallelQwen2DecoderLayer(config, megatron_config) for _ in range(config.num_hidden_layers)] + [ + ParallelQwen2DecoderLayer(config, megatron_config) + for _ in range(config.num_hidden_layers) + ] ) self.norm = ParallelQwen2RMSNorm(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): + # 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 @@ -112,11 +143,13 @@ def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_em 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 - ) + 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 + expanded_attn_mask + if combined_attention_mask is None + else expanded_attn_mask + combined_attention_mask ) return combined_attention_mask @@ -141,7 +174,9 @@ def forward( inputs_embeds = self.embed_tokens(input_ids) # embed positions - attention_mask = self._prepare_decoder_attention_mask(attention_mask, (batch_size, seq_length), inputs_embeds) + attention_mask = self._prepare_decoder_attention_mask( + attention_mask, (batch_size, seq_length), inputs_embeds + ) hidden_states = inputs_embeds @@ -160,16 +195,23 @@ def forward( class ParallelQwen2ForCausalLM(nn.Module): - def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig): + def __init__( + self, + config: Qwen2Config, + megatron_config: ModelParallelConfig): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) - self.model = ParallelQwen2Model(config, megatron_config=megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) + self.model = ParallelQwen2Model( + 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) + 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, @@ -196,7 +238,8 @@ def forward( Returns: ```""" - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + # decoder outputs consists of (dec_features, layer_state, dec_hidden, + # dec_attn) outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -206,7 +249,8 @@ def forward( hidden_states = outputs logits = self.lm_head(hidden_states)[0] - logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits) + logits = tensor_parallel.gather_from_tensor_model_parallel_region( + logits) logits = logits.float() return CausalLMOutputWithPast( @@ -229,22 +273,34 @@ class ParallelQwen2ModelRmPad(nn.Module): config: Qwen2Config """ - def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig): + def __init__( + self, + config: Qwen2Config, + megatron_config: ModelParallelConfig): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) 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) + 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 + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + **embedding_kwargs, ) self.layers = nn.ModuleList( - [ParallelQwen2DecoderLayerRmPad(config, megatron_config) for _ in range(config.num_hidden_layers)] + [ + ParallelQwen2DecoderLayerRmPad(config, megatron_config) + for _ in range(config.num_hidden_layers) + ] ) self.norm = ParallelQwen2RMSNorm(config, megatron_config) @@ -266,12 +322,15 @@ def forward( Returns: """ - inputs_embeds = self.embed_tokens(input_ids) # (1, total_nnz) -> (1, total_nnz, hidden_size) + 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) + inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region( + inputs_embeds) hidden_states = inputs_embeds for idx, decoder_layer in enumerate(self.layers): @@ -292,19 +351,26 @@ def forward( class ParallelQwen2ForCausalLMRmPad(nn.Module): - def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig): + def __init__( + self, + config: Qwen2Config, + megatron_config: ModelParallelConfig): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) self.megatron_config = megatron_config - self.model = ParallelQwen2ModelRmPad(config, megatron_config=megatron_config) + self.model = ParallelQwen2ModelRmPad( + config, megatron_config=megatron_config) self.vocab_size = config.vocab_size self._init_head(config) def _init_head(self, config: Qwen2Config): 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) + 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, @@ -318,7 +384,9 @@ 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) + logits = tensor_parallel.gather_from_tensor_model_parallel_region( + logits + ) # (total_nnz_padded, 1, vocab_size) return logits def forward( @@ -344,7 +412,8 @@ def forward( ) # (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 + # 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) @@ -368,7 +437,8 @@ def forward( totol_nnz = cu_seqlens[-1] logits = logits[:totol_nnz] # (total_nnz_padded) - logits = torch.squeeze(logits, dim=1) # remove the artificial batch dimension + # remove the artificial batch dimension + logits = torch.squeeze(logits, dim=1) # add removed padding back logits = pad_input( logits, indices, batch_size, seqlen=sequence_length @@ -387,9 +457,13 @@ class ParallelQwen2ForValueRmPad(ParallelQwen2ForCausalLMRmPad): def _init_head(self, config): 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=config.hidden_size, out_features=1, bias=False) + 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=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) @@ -397,7 +471,9 @@ 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) + logits = tensor_parallel.gather_from_sequence_parallel_region( + logits, tensor_parallel_output_grad=False + ) return logits def forward( @@ -426,9 +502,16 @@ class ParallelQwen2ModelRmPadPP(nn.Module): config: Qwen2Config """ - def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig, pre_process, post_process): + def __init__( + self, + config: Qwen2Config, + megatron_config: ModelParallelConfig, + pre_process, + post_process, + ): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) self.padding_idx = config.pad_token_id self.vocab_size = config.vocab_size self.pre_process = pre_process @@ -436,11 +519,16 @@ def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig, pr 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) + 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 + num_embeddings=config.vocab_size, + embedding_dim=config.hidden_size, + **embedding_kwargs, ) else: self.embed_tokens = None @@ -454,14 +542,18 @@ def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig, pr 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 - offset = vpp_rank * (config.num_hidden_layers // vpp_size) + (pp_rank * self.num_layer_vpp_chunk) + offset = vpp_rank * (config.num_hidden_layers // vpp_size) + ( + pp_rank * self.num_layer_vpp_chunk + ) else: self.num_layer_this_model = self.num_layer_per_pp offset = pp_rank * self.num_layer_per_pp self.layers = nn.ModuleList() for i in range(self.num_layer_this_model): - layer = ParallelQwen2DecoderLayerRmPad(config, megatron_config, layer_idx=i + offset) + layer = ParallelQwen2DecoderLayerRmPad( + config, megatron_config, layer_idx=i + offset + ) self.layers.add_module(f"{i}", layer) if post_process: @@ -498,14 +590,17 @@ def forward( """ if self.pre_process: - inputs_embeds = self.embed_tokens(input_ids) # (1, total_nnz) -> (1, total_nnz, hidden_size) + 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) + inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region( + inputs_embeds) hidden_states = inputs_embeds else: @@ -540,10 +635,14 @@ def __init__( share_embeddings_and_output_weights, ): super().__init__() - self.config: TransformerConfig = convert_config(config, megatron_config) + self.config: TransformerConfig = convert_config( + config, megatron_config) self.megatron_config = megatron_config self.model = ParallelQwen2ModelRmPadPP( - config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process + config, + megatron_config=megatron_config, + pre_process=pre_process, + post_process=post_process, ) self.share_embeddings_and_output_weights = share_embeddings_and_output_weights self.vocab_size = config.vocab_size @@ -568,15 +667,18 @@ def set_input_tensor(self, input_tensor): def _init_head(self, config): 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) + 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, - skip_weight_param_allocation=self.pre_process and self.share_embeddings_and_output_weights, + skip_weight_param_allocation=self.pre_process + and self.share_embeddings_and_output_weights, **column_kwargs, ) @@ -599,11 +701,16 @@ def setup_embeddings_and_output_layer(self) -> None: if parallel_state.get_pipeline_model_parallel_world_size() == 1: # Zero out wgrad if sharing embeddings between two layers on same # pipeline stage to make sure grad accumulation into main_grad is - # correct and does not include garbage values (e.g., from torch.empty). + # correct and does not include garbage values (e.g., from + # torch.empty). self.shared_embedding_or_output_weight().zero_out_wgrad = True return - if parallel_state.is_pipeline_first_stage() and self.pre_process and not self.post_process: + if ( + parallel_state.is_pipeline_first_stage() + and self.pre_process + and not self.post_process + ): self.shared_embedding_or_output_weight().shared_embedding = True if self.post_process and not self.pre_process: @@ -614,10 +721,15 @@ def setup_embeddings_and_output_layer(self) -> None: self.lm_head.weight.shared = True self.lm_head.weight.shared_embedding = True - if torch.distributed.is_initialized() and parallel_state.is_rank_in_embedding_group(): + if ( + torch.distributed.is_initialized() + and parallel_state.is_rank_in_embedding_group() + ): weight = self.shared_embedding_or_output_weight() weight.data = weight.data.to(get_device_name()) - torch.distributed.all_reduce(weight.data, group=parallel_state.get_embedding_group()) + torch.distributed.all_reduce( + weight.data, group=parallel_state.get_embedding_group() + ) def shared_embedding_or_output_weight(self) -> torch.Tensor: if self.pre_process: @@ -634,7 +746,8 @@ def _forward_head(self, hidden_states): if self.share_embeddings_and_output_weights: output_weight = self.shared_embedding_or_output_weight() logits = self.lm_head(hidden_states, weight=output_weight)[0] - # print(f'logits shape after forward_head: {logits.shape}') # [8, 32, 8] + # print(f'logits shape after forward_head: {logits.shape}') # [8, 32, + # 8] logits = logits.float() # (total_nnz_padded, 1, vocab_size // tp) return logits @@ -657,7 +770,8 @@ def forward( ```""" # 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 + # 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( @@ -665,9 +779,11 @@ def forward( ) # (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 + # 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 = sp_utils.pad_to_sequence_parallel( + input_ids_rmpad) input_ids_rmpad = input_ids_rmpad.transpose(0, 1) # (1, total_nnz+pad) @@ -683,13 +799,15 @@ def forward( if self.post_process: hidden_states = outputs logits = self._forward_head(hidden_states) - logits = torch.squeeze(logits, dim=1) # remove the artificial batch dimension # torch.Size([8, 32, 16]) + # remove the artificial batch dimension # torch.Size([8, 32, 16]) + logits = torch.squeeze(logits, dim=1) # 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 + # 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) @@ -709,9 +827,13 @@ class ParallelQwen2ForValueRmPadPP(ParallelQwen2ForCausalLMRmPadPP): def _init_head(self, config): 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=config.hidden_size, out_features=1, bias=False) + 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=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) @@ -719,7 +841,9 @@ 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) + logits = tensor_parallel.gather_from_sequence_parallel_region( + logits, tensor_parallel_output_grad=False + ) return logits def forward( @@ -729,7 +853,11 @@ def forward( attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, ) -> tuple | CausalLMOutputWithPast: - output = super().forward(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids) + 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 diff --git a/Agent0/executor_train/verl/verl/models/registry.py b/Agent0/executor_train/verl/verl/models/registry.py index 829b9e2..1c7ddc7 100644 --- a/Agent0/executor_train/verl/verl/models/registry.py +++ b/Agent0/executor_train/verl/verl/models/registry.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -22,15 +22,27 @@ _MODELS = { "LlamaForCausalLM": ( "llama", - ("ParallelLlamaForCausalLMRmPadPP", "ParallelLlamaForValueRmPadPP", "ParallelLlamaForCausalLMRmPad"), + ( + "ParallelLlamaForCausalLMRmPadPP", + "ParallelLlamaForValueRmPadPP", + "ParallelLlamaForCausalLMRmPad", + ), ), "Qwen2ForCausalLM": ( "qwen2", - ("ParallelQwen2ForCausalLMRmPadPP", "ParallelQwen2ForValueRmPadPP", "ParallelQwen2ForCausalLMRmPad"), + ( + "ParallelQwen2ForCausalLMRmPadPP", + "ParallelQwen2ForValueRmPadPP", + "ParallelQwen2ForCausalLMRmPad", + ), ), "MistralForCausalLM": ( "mistral", - ("ParallelMistralForCausalLMRmPadPP", "ParallelMistralForValueRmPadPP", "ParallelMistralForCausalLMRmPad"), + ( + "ParallelMistralForCausalLMRmPadPP", + "ParallelMistralForValueRmPadPP", + "ParallelMistralForCausalLMRmPad", + ), ), } @@ -38,7 +50,8 @@ # return model class class ModelRegistry: @staticmethod - def load_model_cls(model_arch: str, value=False) -> Optional[type[nn.Module]]: + def load_model_cls(model_arch: str, + value=False) -> Optional[type[nn.Module]]: if model_arch not in _MODELS: return None @@ -50,7 +63,8 @@ def load_model_cls(model_arch: str, value=False) -> Optional[type[nn.Module]]: 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") + module = importlib.import_module( + f"verl.models.{module_name}.{megatron}.modeling_{module_name}_megatron") return getattr(module, model_cls_name, None) @staticmethod diff --git a/Agent0/executor_train/verl/verl/models/transformers/__init__.py b/Agent0/executor_train/verl/verl/models/transformers/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/models/transformers/__init__.py +++ b/Agent0/executor_train/verl/verl/models/transformers/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/models/transformers/dense_common.py b/Agent0/executor_train/verl/verl/models/transformers/dense_common.py index 56fe293..6ea75a6 100644 --- a/Agent0/executor_train/verl/verl/models/transformers/dense_common.py +++ b/Agent0/executor_train/verl/verl/models/transformers/dense_common.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -46,12 +46,19 @@ def forward_base_model( This function should be generic enough for all pure text models. ```""" - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states ) - # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) + # decoder outputs consists of (dec_features, layer_state, dec_hidden, + # dec_attn) outputs = self.model( input_ids=input_ids, attention_mask=attention_mask, @@ -103,7 +110,8 @@ def forward_with_torch_backend( hidden_states = outputs[0] if not return_dict: - raise NotImplementedError("forward_with_torch_backend has to return_dict") + raise NotImplementedError( + "forward_with_torch_backend has to return_dict") # Loss calculations if labels is not None: @@ -111,7 +119,9 @@ def forward_with_torch_backend( elif input_ids is not None: rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1) else: - raise RuntimeError("To use forward_with_torch_backend, either labels or input_ids must be provided.") + raise RuntimeError( + "To use forward_with_torch_backend, either labels or input_ids must be provided." + ) fused_linear_for_ppo = FusedLinearForPPO() log_probs, entropy = fused_linear_for_ppo.forward( @@ -166,7 +176,8 @@ def forward_with_triton_backend( hidden_states = outputs[0] if not return_dict: - raise NotImplementedError("forward_with_triton_backend has to return_dict") + raise NotImplementedError( + "forward_with_triton_backend has to return_dict") # Loss calculations if labels is not None: @@ -174,7 +185,9 @@ def forward_with_triton_backend( elif input_ids is not None: rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1) else: - raise RuntimeError("To use forward_with_triton_backend, either labels or input_ids must be provided.") + raise RuntimeError( + "To use forward_with_triton_backend, either labels or input_ids must be provided." + ) log_probs, entropy = linear_cross_entropy( hidden_states, diff --git a/Agent0/executor_train/verl/verl/models/transformers/kimi_vl.py b/Agent0/executor_train/verl/verl/models/transformers/kimi_vl.py index edd7936..86b49d1 100644 --- a/Agent0/executor_train/verl/verl/models/transformers/kimi_vl.py +++ b/Agent0/executor_train/verl/verl/models/transformers/kimi_vl.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,7 +31,7 @@ def rotate_half(x): """Rotates half the hidden dims of the input.""" x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] + x2 = x[..., x.shape[-1] // 2:] return torch.cat((-x2, x1), dim=-1) @@ -80,8 +80,11 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: 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) + 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) def _ulysses_flash_attn_forward( @@ -106,49 +109,69 @@ def _ulysses_flash_attn_forward( # batch_size x seq_length x head_dim x hidden_dim # therefore we just need to keep the original shape compressed_kv = self.kv_a_proj_with_mqa(hidden_states) - compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) + compressed_kv, k_pe = torch.split( + compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2) kv = ( - self.kv_b_proj(self.kv_a_layernorm(compressed_kv)) - .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim) - .transpose(1, 2) + self.kv_b_proj( + self.kv_a_layernorm(compressed_kv)) .view( + bsz, + q_len, + self.num_heads, + self.qk_nope_head_dim + + self.v_head_dim) .transpose( + 1, + 2)) + + k_nope, value_states = torch.split( + kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 ) - k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) - # patch ulysses_sp_size = get_ulysses_sequence_parallel_world_size() if ulysses_sp_size > 1: validate_ulysses_config(self.num_heads, ulysses_sp_size) - num_key_value_groups = self.config.num_attention_heads // self.config.num_key_value_heads + num_key_value_groups = ( + self.config.num_attention_heads // self.config.num_key_value_heads + ) k_pe = repeat_kv(k_pe, ulysses_sp_size) # to keep heads=1 after a2a k_nope = repeat_kv(k_nope, num_key_value_groups) value_states = repeat_kv(value_states, num_key_value_groups) q = gather_seq_scatter_heads(q, seq_dim=2, head_dim=1) k_pe = gather_seq_scatter_heads(k_pe, seq_dim=2, head_dim=1) k_nope = gather_seq_scatter_heads(k_nope, seq_dim=2, head_dim=1) - value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1) + value_states = gather_seq_scatter_heads( + value_states, seq_dim=2, head_dim=1) # (batch_size, num_head / sp_size, seq_length, head_size) full_q_len = q.size(2) # full_q_len = seq_length else: full_q_len = q_len - q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) + q_nope, q_pe = torch.split( + q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) cos, sin = self.rotary_emb(value_states, seq_len=full_q_len) q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids) - query_states = k_pe.new_empty(bsz, self.num_heads // ulysses_sp_size, full_q_len, self.q_head_dim) + query_states = k_pe.new_empty( + bsz, self.num_heads // ulysses_sp_size, full_q_len, self.q_head_dim + ) query_states[:, :, :, : self.qk_nope_head_dim] = q_nope - query_states[:, :, :, self.qk_nope_head_dim :] = q_pe + query_states[:, :, :, self.qk_nope_head_dim:] = q_pe - key_states = k_pe.new_empty(bsz, self.num_heads // ulysses_sp_size, full_q_len, self.q_head_dim) + key_states = k_pe.new_empty( + bsz, self.num_heads // ulysses_sp_size, full_q_len, self.q_head_dim + ) key_states[:, :, :, : self.qk_nope_head_dim] = k_nope - key_states[:, :, :, self.qk_nope_head_dim :] = k_pe + key_states[:, :, :, self.qk_nope_head_dim:] = k_pe if self.q_head_dim != self.v_head_dim: - value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim]) + value_states = F.pad( + value_states, [ + 0, self.q_head_dim - self.v_head_dim]) # 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 @@ -174,12 +197,15 @@ def _ulysses_flash_attn_forward( ) if ulysses_sp_size > 1: - attn_output = gather_heads_scatter_seq(attn_output, head_dim=2, seq_dim=1) + attn_output = gather_heads_scatter_seq( + attn_output, head_dim=2, seq_dim=1) if self.q_head_dim != self.v_head_dim: attn_output = attn_output[:, :, :, : self.v_head_dim] - attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim).contiguous() + attn_output = attn_output.reshape( + bsz, q_len, self.num_heads * self.v_head_dim + ).contiguous() attn_output = self.o_proj(attn_output) return attn_output, None, None diff --git a/Agent0/executor_train/verl/verl/models/transformers/llama.py b/Agent0/executor_train/verl/verl/models/transformers/llama.py index 687ceab..581c936 100644 --- a/Agent0/executor_train/verl/verl/models/transformers/llama.py +++ b/Agent0/executor_train/verl/verl/models/transformers/llama.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -46,7 +46,9 @@ def llama_flash_attn_forward( 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 + 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]]]: """ @@ -65,9 +67,15 @@ def llama_flash_attn_forward( # 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) + 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) @@ -80,9 +88,12 @@ def llama_flash_attn_forward( validate_ulysses_config(self.num_heads, ulysses_sp_size) # (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) + 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 @@ -91,17 +102,23 @@ def llama_flash_attn_forward( "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." - ) + "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) + 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) + # 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 @@ -152,10 +169,12 @@ def llama_flash_attn_forward( **kwargs, ) - attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous() + 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 = 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) @@ -184,40 +203,58 @@ def llama_attn_forward( bsz, q_len, _ = hidden_states.shape - query_states = self.q_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - key_states = self.k_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2) - value_states = self.v_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2) + query_states = (self.q_proj(hidden_states).view( + bsz, q_len, -1, self.head_dim).transpose(1, 2)) + key_states = (self.k_proj(hidden_states).view( + bsz, q_len, -1, self.head_dim).transpose(1, 2)) + value_states = (self.v_proj(hidden_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: - validate_ulysses_config(self.config.num_attention_heads, ulysses_sp_size) + validate_ulysses_config( + self.config.num_attention_heads, + ulysses_sp_size) - 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) + 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) cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + 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) + # 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 + ) attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": - if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + if self.config._attn_implementation == "sdpa" and kwargs.get( + "output_attentions", False + ): logger.warning_once( "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. " "Falling back to eager attention. This warning can be removed using the argument " - '`attn_implementation="eager"` when loading the model.' - ) + '`attn_implementation="eager"` when loading the model.') else: - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attention_interface = ALL_ATTENTION_FUNCTIONS[ + self.config._attn_implementation + ] attn_output, attn_weights = attention_interface( self, @@ -230,10 +267,12 @@ def llama_attn_forward( **kwargs, ) - attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous() + 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 = 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) return attn_output, attn_weights diff --git a/Agent0/executor_train/verl/verl/models/transformers/monkey_patch.py b/Agent0/executor_train/verl/verl/models/transformers/monkey_patch.py index d6be65a..c59c129 100644 --- a/Agent0/executor_train/verl/verl/models/transformers/monkey_patch.py +++ b/Agent0/executor_train/verl/verl/models/transformers/monkey_patch.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -43,8 +43,11 @@ def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: batch, slen, num_key_value_heads, head_dim = hidden_states.shape if n_rep == 1: return hidden_states - hidden_states = hidden_states[:, :, :, None, :].expand(batch, slen, num_key_value_heads, n_rep, head_dim) - return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim) + hidden_states = hidden_states[:, :, :, None, :].expand( + batch, slen, num_key_value_heads, n_rep, head_dim + ) + return hidden_states.reshape( + batch, slen, num_key_value_heads * n_rep, head_dim) def _ulysses_flash_attention_forward( @@ -71,7 +74,9 @@ def _ulysses_flash_attention_forward( ########## AlltoAll for Ulysses ########## if ulysses_sp_size > 1: - assert position_ids is not None, "position_ids is required for Ulysses sequence parallelism" + assert ( + position_ids is not None + ), "position_ids is required for Ulysses sequence parallelism" # NOTE: repeat kv heads to be divided by sequence parallel. Instead of repeating nheads_q//nheads_k, # we choose to repeat sp_size//nheads_k, since flash_attention supports MQA/GQA. @@ -84,28 +89,42 @@ def _ulysses_flash_attention_forward( value_states = repeat_kv(value_states, repeats) # (bsz, seq_len/n, n_head, head_dim) -> (bsz, seq_len, n_head/n, head_dim) - query_states = gather_seq_scatter_heads(query_states, seq_dim=1, head_dim=2) - key_states = gather_seq_scatter_heads(key_states, seq_dim=1, head_dim=2) - value_states = gather_seq_scatter_heads(value_states, seq_dim=1, head_dim=2) + query_states = gather_seq_scatter_heads( + query_states, seq_dim=1, head_dim=2) + key_states = gather_seq_scatter_heads( + key_states, seq_dim=1, head_dim=2) + value_states = gather_seq_scatter_heads( + value_states, seq_dim=1, head_dim=2) # TODO: all_gather position_ids because `prepare_fa2_from_position_ids` needs it, we can eliminate # this all_gather by passing cu_seq_lens_q, cu_seq_lens_k, max_length_k, max_length_q explicitly. # https://github.com/huggingface/transformers/pull/33932 # (bsz, seq_len/n) -> (bsz, seq_len) - position_ids_list = [torch.empty_like(position_ids) for _ in range(ulysses_sp_size)] - torch.distributed.all_gather(position_ids_list, position_ids, group=get_ulysses_sequence_parallel_group()) + position_ids_list = [ + torch.empty_like(position_ids) for _ in range(ulysses_sp_size) + ] + torch.distributed.all_gather( + position_ids_list, + position_ids, + group=get_ulysses_sequence_parallel_group()) position_ids = torch.concat(position_ids_list, dim=-1) # (bsz, seq_len, n_head/n, head_dim) attn_output = _flash_attention_forward( - query_states, key_states, value_states, *args, position_ids=position_ids, **kwargs + query_states, + key_states, + value_states, + *args, + position_ids=position_ids, + **kwargs, ) ########## AlltoAll for Ulysses ########## if ulysses_sp_size > 1: # (bsz, seq_len, n_head/n, head_dim) -> (bsz, seq_len/n, n_head, head_dim) - attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2) + attn_output = gather_heads_scatter_seq( + attn_output, seq_dim=1, head_dim=2) return attn_output @@ -129,7 +148,9 @@ def ulysses_wrapped_decoder_forward(self, *args, **kwargs): and getattr(self, "_needs_initial_slice", True) ) if slice_now: - call_kwargs["inputs_embeds"] = slice_input_tensor(inputs_embeds, dim=1, padding=False) + call_kwargs["inputs_embeds"] = slice_input_tensor( + inputs_embeds, dim=1, padding=False + ) self._needs_initial_slice = False try: return original_forward(self, *args, **call_kwargs) @@ -142,7 +163,9 @@ def ulysses_wrapped_decoder_forward(self, *args, **kwargs): original_forward = model_class.forward wrapped_forward = _create_ulysses_wrapped_decoder_forward(original_forward) model_class.forward = wrapped_forward - print(f"Monkey patch {model_class.__name__}.forward for Ulysses SP input slicing.") + print( + f"Monkey patch { + model_class.__name__}.forward for Ulysses SP input slicing.") def patch_forward_with_backends( @@ -157,7 +180,8 @@ def patch_forward_with_backends( use_fused_kernels (bool): Whether to use fused kernels. fused_kernels_backend (str): The backend to use for fused kernels. """ - if not use_fused_kernels or fused_kernels_backend not in ["triton", "torch"]: + if not use_fused_kernels or fused_kernels_backend not in [ + "triton", "torch"]: print( f"Skipping monkey patch for {model.__class__.__name__} as use_fused_kernels is " f"{use_fused_kernels} or fused_kernels_backend is {fused_kernels_backend}" @@ -167,29 +191,43 @@ def patch_forward_with_backends( forward_with_torch_backend_function = model.__class__.forward forward_with_triton_backend_function = model.__class__.forward if model.config.model_type == "qwen2_5_vl": - from verl.models.transformers.qwen2_5_vl import forward_with_torch_backend, forward_with_triton_backend + from verl.models.transformers.qwen2_5_vl import ( + forward_with_torch_backend, + forward_with_triton_backend, + ) forward_with_torch_backend_function = forward_with_torch_backend forward_with_triton_backend_function = forward_with_triton_backend elif model.config.model_type == "qwen2_vl": - from verl.models.transformers.qwen2_vl import forward_with_torch_backend, forward_with_triton_backend + from verl.models.transformers.qwen2_vl import ( + forward_with_torch_backend, + forward_with_triton_backend, + ) forward_with_torch_backend_function = forward_with_torch_backend forward_with_triton_backend_function = forward_with_triton_backend else: - from verl.models.transformers.dense_common import forward_with_torch_backend, forward_with_triton_backend + from verl.models.transformers.dense_common import ( + forward_with_torch_backend, + forward_with_triton_backend, + ) forward_with_torch_backend_function = forward_with_torch_backend forward_with_triton_backend_function = forward_with_triton_backend if fused_kernels_backend == "triton": model.__class__.forward = forward_with_triton_backend_function - print(f"Using Triton backend for fused kernels in {model.__class__.__name__}") + print( + f"Using Triton backend for fused kernels in { + model.__class__.__name__}") elif fused_kernels_backend == "torch": model.__class__.forward = forward_with_torch_backend_function - print(f"Using Torch backend for fused kernels in {model.__class__.__name__}") + print( + f"Using Torch backend for fused kernels in { + model.__class__.__name__}") else: - raise ValueError(f"Unsupported fused_kernels_backend: {fused_kernels_backend}. Choose 'triton' or 'torch'.") + raise ValueError( + f"Unsupported fused_kernels_backend: {fused_kernels_backend}. Choose 'triton' or 'torch'.") def apply_monkey_patch( @@ -210,17 +248,23 @@ def apply_monkey_patch( module = sys.modules[model.__module__] try: - num_attention_heads, num_key_value_heads = model.config.num_attention_heads, model.config.num_key_value_heads + num_attention_heads, num_key_value_heads = ( + model.config.num_attention_heads, + model.config.num_key_value_heads, + ) except AttributeError: num_attention_heads, num_key_value_heads = ( model.config.text_config.num_attention_heads, model.config.text_config.num_key_value_heads, ) - assert num_attention_heads % ulysses_sp_size == 0, ( - f"num_attention_heads {num_attention_heads} must be divisible by ulysses_sp_size {ulysses_sp_size}" - ) - assert num_key_value_heads % ulysses_sp_size == 0 or ulysses_sp_size % num_key_value_heads == 0, ( + assert ( + num_attention_heads % ulysses_sp_size == 0 + ), f"num_attention_heads {num_attention_heads} must be divisible by ulysses_sp_size {ulysses_sp_size}" + assert ( + num_key_value_heads % ulysses_sp_size == 0 + or ulysses_sp_size % num_key_value_heads == 0 + ), ( f"num_key_value_heads {num_key_value_heads} must be divisible by ulysses_sp_size " f"{ulysses_sp_size}or vise versa. Upon ulysses_sp_size % num_key_value_heads == 0," f"kv heads are repeated to ensure correctness." @@ -238,7 +282,9 @@ def state_dict(self, *args, **kwargs): # TODO: VLM models only, unify monkey patch to LLM models. if model.config.model_type == "qwen2_5_vl": if is_transformers_version_in_range(min_version="4.53.0"): - from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLAttention + from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VLAttention, + ) # TODO: Support transformers 4.53 raise ValueError("Transformers 4.53 is not supported") @@ -255,11 +301,13 @@ def state_dict(self, *args, **kwargs): if ulysses_sp_size > 1: if is_transformers_version_in_range(min_version="4.52.0"): - from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLTextModel + from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VLTextModel, ) patch_vlm_for_ulysses_input_slicing(Qwen2_5_VLTextModel) else: - from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLModel + from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import ( + Qwen2_5_VLModel, ) patch_vlm_for_ulysses_input_slicing(Qwen2_5_VLModel) @@ -270,7 +318,9 @@ def state_dict(self, *args, **kwargs): # TODO: Support transformers 4.53 raise ValueError("Transformers 4.53 is not supported") else: - from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLFlashAttention2 as Qwen2VLAttention + from transformers.models.qwen2_vl.modeling_qwen2_vl import ( + Qwen2VLFlashAttention2 as Qwen2VLAttention, + ) if use_remove_padding or ulysses_sp_size > 1: from verl.models.transformers.qwen2_vl import ulysses_flash_attn_forward @@ -280,7 +330,9 @@ def state_dict(self, *args, **kwargs): if ulysses_sp_size > 1: if is_transformers_version_in_range(min_version="4.52.0"): - from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLTextModel + from transformers.models.qwen2_vl.modeling_qwen2_vl import ( + Qwen2VLTextModel, + ) patch_vlm_for_ulysses_input_slicing(Qwen2VLTextModel) else: @@ -308,24 +360,35 @@ def state_dict(self, *args, **kwargs): if use_remove_padding or ulysses_sp_size > 1: if hasattr(module, "_flash_attention_forward"): module._flash_attention_forward = _ulysses_flash_attention_forward - print(f"Monkey patch _flash_attention_forward in {model.__module__}") + print( + f"Monkey patch _flash_attention_forward in { + model.__module__}") else: # transformers>=4.48.0 from transformers.integrations import flash_attention flash_attention._flash_attention_forward = _ulysses_flash_attention_forward - print(f"Monkey patch _flash_attention_forward in {flash_attention.__name__}") - - patch_forward_with_backends(model, use_fused_kernels=use_fused_kernels, fused_kernels_backend=fused_kernels_backend) + print( + f"Monkey patch _flash_attention_forward in { + flash_attention.__name__}") + + patch_forward_with_backends( + model, + use_fused_kernels=use_fused_kernels, + fused_kernels_backend=fused_kernels_backend, + ) @lru_cache -def is_transformers_version_in_range(min_version: Optional[str] = None, max_version: Optional[str] = None) -> bool: +def is_transformers_version_in_range( + min_version: Optional[str] = None, max_version: Optional[str] = None +) -> bool: try: # Get the installed version of the transformers library transformers_version_str = importlib.metadata.version("transformers") except importlib.metadata.PackageNotFoundError as e: - raise ModuleNotFoundError("The `transformers` package is not installed.") from e + raise ModuleNotFoundError( + "The `transformers` package is not installed.") from e transformers_version = version.parse(transformers_version_str) diff --git a/Agent0/executor_train/verl/verl/models/transformers/npu_patch.py b/Agent0/executor_train/verl/verl/models/transformers/npu_patch.py index e6bb373..38d057e 100644 --- a/Agent0/executor_train/verl/verl/models/transformers/npu_patch.py +++ b/Agent0/executor_train/verl/verl/models/transformers/npu_patch.py @@ -1,50 +1,55 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# -# Copyright 2025 The Qwen Team and The HuggingFace Inc. team -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - - -import torch -import torch_npu -from torch_npu import npu_rotary_mul as apply_rotary_emb -from transformers.models.qwen2_5_vl import modeling_qwen2_5_vl -from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2RMSNorm - - -# This patch takes effect when using apply_rotary_pos_emb_flashatt on qwen2_5_vl and will be removed in -# subsequent versions -# https://github.com/huggingface/transformers/pull/38491 -def apply_rotary_pos_emb_flashatt_npu( - q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor -) -> tuple[torch.Tensor, torch.Tensor]: - cos = cos.chunk(2, dim=-1)[0].contiguous() - sin = sin.chunk(2, dim=-1)[0].contiguous() - cos = cos.repeat(1, 2) - sin = sin.repeat(1, 2) - q_embed = apply_rotary_emb( - q.float(), cos.unsqueeze(0).unsqueeze(2).float(), sin.unsqueeze(0).unsqueeze(2).float() - ).type_as(q) - k_embed = apply_rotary_emb( - k.float(), cos.unsqueeze(0).unsqueeze(2).float(), sin.unsqueeze(0).unsqueeze(2).float() - ).type_as(k) - return q_embed, k_embed - - -# This api can improve performance on ASCEND NPU -def rms_norm_forward(self, x): - return torch_npu.npu_rms_norm(x, self.weight, epsilon=self.variance_epsilon)[0] - - -Qwen2RMSNorm.forward = rms_norm_forward -modeling_qwen2_5_vl.apply_rotary_pos_emb_flashatt = apply_rotary_pos_emb_flashatt_npu +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates +# +# Copyright 2025 The Qwen Team and The HuggingFace Inc. team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import torch +import torch_npu +from torch_npu import npu_rotary_mul as apply_rotary_emb +from transformers.models.qwen2_5_vl import modeling_qwen2_5_vl +from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2RMSNorm + + +# This patch takes effect when using apply_rotary_pos_emb_flashatt on qwen2_5_vl and will be removed in +# subsequent versions +# https://github.com/huggingface/transformers/pull/38491 +def apply_rotary_pos_emb_flashatt_npu( + q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor +) -> tuple[torch.Tensor, torch.Tensor]: + cos = cos.chunk(2, dim=-1)[0].contiguous() + sin = sin.chunk(2, dim=-1)[0].contiguous() + cos = cos.repeat(1, 2) + sin = sin.repeat(1, 2) + q_embed = apply_rotary_emb( + q.float(), + cos.unsqueeze(0).unsqueeze(2).float(), + sin.unsqueeze(0).unsqueeze(2).float(), + ).type_as(q) + k_embed = apply_rotary_emb( + k.float(), + cos.unsqueeze(0).unsqueeze(2).float(), + sin.unsqueeze(0).unsqueeze(2).float(), + ).type_as(k) + return q_embed, k_embed + + +# This api can improve performance on ASCEND NPU +def rms_norm_forward(self, x): + return torch_npu.npu_rms_norm( + x, self.weight, epsilon=self.variance_epsilon)[0] + + +Qwen2RMSNorm.forward = rms_norm_forward +modeling_qwen2_5_vl.apply_rotary_pos_emb_flashatt = apply_rotary_pos_emb_flashatt_npu diff --git a/Agent0/executor_train/verl/verl/models/transformers/qwen2.py b/Agent0/executor_train/verl/verl/models/transformers/qwen2.py index e55fb26..3303f6d 100644 --- a/Agent0/executor_train/verl/verl/models/transformers/qwen2.py +++ b/Agent0/executor_train/verl/verl/models/transformers/qwen2.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -39,7 +39,9 @@ def qwen2_flash_attn_forward( 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 + position_embeddings: Optional[ + tuple[torch.Tensor, torch.Tensor] + ] = None, # will become mandatory in v4.46 ): """ Adapted from transformers 4.47.1 to support Ulysses sequence parallelism. @@ -52,9 +54,11 @@ def qwen2_flash_attn_forward( 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) + 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) + 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() @@ -63,9 +67,12 @@ def qwen2_flash_attn_forward( validate_ulysses_config(self.num_heads, ulysses_sp_size) # (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) + 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 @@ -74,16 +81,22 @@ def qwen2_flash_attn_forward( "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." - ) + "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) + 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) + 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) @@ -141,10 +154,12 @@ def qwen2_flash_attn_forward( ) # use full_q_len to reshape - attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous() + 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 = 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) @@ -173,30 +188,44 @@ def qwen2_attn_forward( bsz, q_len, _ = hidden_states.shape hidden_shape = (bsz, q_len, -1, self.head_dim) - query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2) + query_states = self.q_proj(hidden_states).view( + hidden_shape).transpose(1, 2) key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2) - value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + value_states = self.v_proj(hidden_states).view( + hidden_shape).transpose(1, 2) ########## AlltoAll for Ulysses ########## ulysses_sp_size = get_ulysses_sequence_parallel_world_size() if ulysses_sp_size > 1: - validate_ulysses_config(self.config.num_attention_heads, ulysses_sp_size) + validate_ulysses_config( + self.config.num_attention_heads, + ulysses_sp_size) # (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) + 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) cos, sin = position_embeddings - query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + 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) + # 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 + ) sliding_window = None if ( @@ -210,14 +239,17 @@ def qwen2_attn_forward( attention_interface: Callable = eager_attention_forward if self.config._attn_implementation != "eager": - if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False): + if self.config._attn_implementation == "sdpa" and kwargs.get( + "output_attentions", False + ): logger.warning_once( "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. " "Falling back to eager attention. This warning can be removed using the argument " - '`attn_implementation="eager"` when loading the model.' - ) + '`attn_implementation="eager"` when loading the model.') else: - attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + attention_interface = ALL_ATTENTION_FUNCTIONS[ + self.config._attn_implementation + ] attn_output, attn_weights = attention_interface( self, @@ -231,11 +263,13 @@ def qwen2_attn_forward( **kwargs, ) - attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous() + attn_output = attn_output.reshape( + bsz, full_q_len, -1, self.head_dim).contiguous() ########## AlltoAll for Ulysses ########## if ulysses_sp_size > 1: # (bsz, seq_len, n_head/n, head_dim) -> (bsz, seq_len/n, n_head, head_dim) - attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2) + 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) return attn_output, attn_weights diff --git a/Agent0/executor_train/verl/verl/models/transformers/qwen2_5_vl.py b/Agent0/executor_train/verl/verl/models/transformers/qwen2_5_vl.py index 51d9753..81d43bf 100644 --- a/Agent0/executor_train/verl/verl/models/transformers/qwen2_5_vl.py +++ b/Agent0/executor_train/verl/verl/models/transformers/qwen2_5_vl.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -51,18 +51,27 @@ def forward_base_model( Copy paste Qwen2_5_VL's forward https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/transformers/model/qwen2_5_vl.py ```""" - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) if pixel_values is not None: pixel_values = pixel_values.type(self.visual.dtype) image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) - n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_tokens = ( + input_ids == self.config.image_token_id).sum().item() n_image_features = image_embeds.shape[0] if n_image_tokens != n_image_features: raise ValueError( @@ -75,13 +84,17 @@ def forward_base_model( mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) image_mask = mask_expanded.to(inputs_embeds.device) - image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) - inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter( + image_mask, image_embeds) if pixel_values_videos is not None: pixel_values_videos = pixel_values_videos.type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + video_embeds = self.visual( + pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = ( + input_ids == self.config.video_token_id).sum().item() n_video_features = video_embeds.shape[0] if n_video_tokens != n_video_features: raise ValueError( @@ -94,16 +107,22 @@ def forward_base_model( mask_expanded = mask_unsqueezed.expand_as(inputs_embeds) video_mask = mask_expanded.to(inputs_embeds.device) - video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) - inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter( + video_mask, video_embeds) if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) - # if we get 4D attention mask we cannot calculate rope deltas anymore. TODO @raushan fixme - if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): + # if we get 4D attention mask we cannot calculate rope deltas anymore. + # TODO @raushan fixme + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2): # calculate RoPE index once per generation in the pre-fill stage only - if (cache_position is not None and cache_position[0] == 0) or self.rope_deltas is None: + if ( + cache_position is not None and cache_position[0] == 0 + ) or self.rope_deltas is None: position_ids, rope_deltas = self.get_rope_index( input_ids, image_grid_thw, @@ -112,14 +131,21 @@ def forward_base_model( attention_mask, ) self.rope_deltas = rope_deltas - # then use the prev pre-calculated rope-deltas to get the correct position ids + # then use the prev pre-calculated rope-deltas to get the correct + # position ids else: batch_size, seq_length, _ = inputs_embeds.shape - delta = (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) if cache_position is not None else 0 - position_ids = torch.arange(seq_length, device=inputs_embeds.device) + delta = ( + (cache_position[0] + self.rope_deltas).to(inputs_embeds.device) + if cache_position is not None + else 0 + ) + position_ids = torch.arange( + seq_length, device=inputs_embeds.device) position_ids = position_ids.view(1, -1).expand(batch_size, -1) if cache_position is not None: # otherwise `deltas` is an int `0` - delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + delta = delta.repeat_interleave( + batch_size // delta.shape[0], dim=0) position_ids = position_ids.add(delta) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) @@ -185,7 +211,8 @@ def forward_with_torch_backend( hidden_states = outputs[0] if not return_dict: - raise NotImplementedError("forward_with_torch_backend has to return_dict") + raise NotImplementedError( + "forward_with_torch_backend has to return_dict") # Loss calculations if labels is not None: @@ -193,7 +220,9 @@ def forward_with_torch_backend( elif input_ids is not None: rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1) else: - raise RuntimeError("To use forward_with_torch_backend, either labels or input_ids must be provided.") + raise RuntimeError( + "To use forward_with_torch_backend, either labels or input_ids must be provided." + ) fused_linear_for_ppo = FusedLinearForPPO() log_probs, entropy = fused_linear_for_ppo.forward( @@ -260,7 +289,8 @@ def forward_with_triton_backend( hidden_states = outputs[0] if not return_dict: - raise NotImplementedError("forward_with_triton_backend has to return_dict") + raise NotImplementedError( + "forward_with_triton_backend has to return_dict") # Loss calculations if labels is not None: @@ -268,7 +298,9 @@ def forward_with_triton_backend( elif input_ids is not None: rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1) else: - raise RuntimeError("To use forward_with_triton_backend, either labels or input_ids must be provided.") + raise RuntimeError( + "To use forward_with_triton_backend, either labels or input_ids must be provided." + ) log_probs, entropy = linear_cross_entropy( hidden_states, diff --git a/Agent0/executor_train/verl/verl/models/transformers/qwen2_vl.py b/Agent0/executor_train/verl/verl/models/transformers/qwen2_vl.py index 358b00b..c91b8c6 100644 --- a/Agent0/executor_train/verl/verl/models/transformers/qwen2_vl.py +++ b/Agent0/executor_train/verl/verl/models/transformers/qwen2_vl.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,9 +33,14 @@ ) try: - from transformers.modeling_flash_attention_utils import flash_attn_func, flash_attn_varlen_func + from transformers.modeling_flash_attention_utils import ( + flash_attn_func, + flash_attn_varlen_func, + ) - _flash_supports_window_size = "window_size" in list(inspect.signature(flash_attn_func).parameters) + _flash_supports_window_size = "window_size" in list( + inspect.signature(flash_attn_func).parameters + ) except ImportError: flash_attn_varlen_func = None @@ -57,16 +62,25 @@ def get_rope_index( tokens_per_second = 2 image_token_id = processor.tokenizer.convert_tokens_to_ids("<|image_pad|>") video_token_id = processor.tokenizer.convert_tokens_to_ids("<|video_pad|>") - vision_start_token_id = processor.tokenizer.convert_tokens_to_ids("<|vision_start|>") - if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None): + vision_start_token_id = processor.tokenizer.convert_tokens_to_ids( + "<|vision_start|>" + ) + if input_ids is not None and ( + image_grid_thw is not None or video_grid_thw is not None + ): if attention_mask is None: attention_mask = torch.ones_like(input_ids) - position_ids = torch.ones(3, input_ids.size(0), dtype=input_ids.dtype, device=input_ids.device) # (3, seqlen) + position_ids = torch.ones( + 3, + input_ids.size(0), + dtype=input_ids.dtype, + device=input_ids.device) # (3, seqlen) image_index, video_index = 0, 0 input_ids = input_ids[attention_mask == 1] image_nums, video_nums = 0, 0 - vision_start_indices = torch.argwhere(input_ids == vision_start_token_id) + vision_start_indices = torch.argwhere( + input_ids == vision_start_token_id) vision_tokens = input_ids[vision_start_indices + 1] image_nums = (vision_tokens == image_token_id).sum() video_nums = (vision_tokens == video_token_id).sum() @@ -99,7 +113,11 @@ def get_rope_index( video_grid_thw[video_index][1], video_grid_thw[video_index][2], ) - second_per_grid_t = second_per_grid_ts[video_index] if second_per_grid_ts is not None else 1.0 + second_per_grid_t = ( + second_per_grid_ts[video_index] + if second_per_grid_ts is not None + else 1.0 + ) video_index += 1 remain_videos -= 1 @@ -112,50 +130,88 @@ def get_rope_index( ) text_len = ed - st - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + st_idx = llm_pos_ids_list[-1].max() + \ + 1 if len(llm_pos_ids_list) > 0 else 0 + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) - t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w) - t_index = (t_index * second_per_grid_t * tokens_per_second).long().flatten() - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten() - llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx) + t_index = (torch.arange(llm_grid_t).view(-1, + 1).expand(-1, llm_grid_h * llm_grid_w)) + t_index = ( + t_index * + second_per_grid_t * + tokens_per_second).long().flatten() + h_index = ( + torch.arange(llm_grid_h) + .view(1, -1, 1) + .expand(llm_grid_t, -1, llm_grid_w) + .flatten() + ) + w_index = ( + torch.arange(llm_grid_w) + .view(1, 1, -1) + .expand(llm_grid_t, llm_grid_h, -1) + .flatten() + ) + llm_pos_ids_list.append( + torch.stack([t_index, h_index, w_index]) + text_len + st_idx + ) st = ed + llm_grid_t * llm_grid_h * llm_grid_w if st < len(input_tokens): - st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + st_idx = llm_pos_ids_list[-1].max() + \ + 1 if len(llm_pos_ids_list) > 0 else 0 text_len = len(input_tokens) - st - llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx) + llm_pos_ids_list.append( + torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx + ) llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) - position_ids[..., attention_mask == 1] = llm_positions.to(position_ids.device) + position_ids[..., attention_mask == + 1] = llm_positions.to(position_ids.device) else: if attention_mask is not None: position_ids = attention_mask.long().cumsum(-1) - 1 position_ids.masked_fill_(attention_mask == 0, 1) - position_ids = position_ids.unsqueeze(0).expand(3, -1).to(input_ids.device) + position_ids = position_ids.unsqueeze( + 0).expand(3, -1).to(input_ids.device) else: - position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).view(1, -1).expand(3, -1) + position_ids = ( + torch.arange(input_ids.shape[1], device=input_ids.device) + .view(1, -1) + .expand(3, -1) + ) return position_ids def prepare_fa2_from_position_ids( - query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, position_ids: torch.Tensor + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + position_ids: torch.Tensor, ): query = query.view(-1, query.size(-2), query.size(-1)) key = key.view(-1, key.size(-2), key.size(-1)) value = value.view(-1, value.size(-2), value.size(-1)) position_ids = position_ids.flatten() - indices_q = torch.arange(position_ids.size(0), device=position_ids.device, dtype=torch.int32) - cu_seqlens = torch.cat( - ( - indices_q[position_ids == 0], - torch.tensor(position_ids.size(), device=position_ids.device, dtype=torch.int32), - ) + indices_q = torch.arange( + position_ids.size(0), device=position_ids.device, dtype=torch.int32 + ) + cu_seqlens = torch.cat((indices_q[position_ids == 0], torch.tensor( + position_ids.size(), device=position_ids.device, dtype=torch.int32), )) + max_length = ( + cu_seqlens.diff().max() + ) # use cu_seqlens to infer max_length for qwen2vl mrope + return ( + query, + key, + value, + indices_q, + (cu_seqlens, cu_seqlens), + (max_length, max_length), ) - max_length = cu_seqlens.diff().max() # use cu_seqlens to infer max_length for qwen2vl mrope - return (query, key, value, indices_q, (cu_seqlens, cu_seqlens), (max_length, max_length)) def flash_attention_forward( @@ -176,21 +232,32 @@ def flash_attention_forward( """ causal = is_causal if not use_top_left_mask else is_causal and query_length != 1 - # Assuming 4D tensors, key_states.shape[1] is the key/value sequence length (source length). + # Assuming 4D tensors, key_states.shape[1] is the key/value sequence + # length (source length). use_sliding_windows = ( - _flash_supports_window_size and sliding_window is not None and key_states.shape[1] > sliding_window + _flash_supports_window_size + and sliding_window is not None + and key_states.shape[1] > sliding_window ) - flash_kwargs = {"window_size": (sliding_window, sliding_window)} if use_sliding_windows else {} + flash_kwargs = ({"window_size": (sliding_window, + sliding_window)} if use_sliding_windows else {}) if is_flash_attn_greater_or_equal("2.4.1"): if deterministic is None: - deterministic = os.environ.get("FLASH_ATTENTION_DETERMINISTIC", "0") == "1" + deterministic = os.environ.get( + "FLASH_ATTENTION_DETERMINISTIC", "0") == "1" flash_kwargs["deterministic"] = deterministic - if position_ids is not None and query_length != 1 and not (torch.diff(position_ids[0], dim=-1) >= 0).all(): + if ( + position_ids is not None + and query_length != 1 + and not (torch.diff(position_ids[0], dim=-1) >= 0).all() + ): batch_size = query_states.size(0) - query_states, key_states, value_states, _, cu_seq_lens, max_seq_lens = prepare_fa2_from_position_ids( - query_states, key_states, value_states, position_ids[0] + query_states, key_states, value_states, _, cu_seq_lens, max_seq_lens = ( + prepare_fa2_from_position_ids( + query_states, key_states, value_states, position_ids[0] + ) ) # remove channel dimension cu_seqlens_q, cu_seqlens_k = cu_seq_lens max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens @@ -207,7 +274,9 @@ def flash_attention_forward( causal=causal, **flash_kwargs, ) - attn_output = attn_output.view(batch_size, -1, attn_output.size(-2), attn_output.size(-1)) + attn_output = attn_output.view( + batch_size, -1, attn_output.size(-2), attn_output.size(-1) + ) else: attn_output = _flash_attention_forward( query_states, @@ -230,19 +299,32 @@ def ulysses_flash_attn_forward( hidden_states: torch.Tensor, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, - position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None, # will become mandatory in v4.46 + position_embeddings: Optional[ + tuple[torch.Tensor, torch.Tensor] + ] = None, # will become mandatory in v4.46 **kwargs, ) -> tuple[torch.Tensor, None, None]: - from transformers.models.qwen2_vl.modeling_qwen2_vl import apply_multimodal_rotary_pos_emb, repeat_kv + from transformers.models.qwen2_vl.modeling_qwen2_vl import ( + apply_multimodal_rotary_pos_emb, + repeat_kv, + ) bsz, q_len, _ = hidden_states.size() # q_len = seq_length / sp_size - query_states = self.q_proj(hidden_states) # (batch_size, seq_length / sp_size, num_heads * head_size) + query_states = self.q_proj( + hidden_states + ) # (batch_size, seq_length / sp_size, num_heads * head_size) key_states = self.k_proj(hidden_states) value_states = self.v_proj(hidden_states) - 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) + 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) ulysses_sp_size = get_ulysses_sequence_parallel_world_size() @@ -251,15 +333,19 @@ def ulysses_flash_attn_forward( key_states = repeat_kv(key_states, self.num_key_value_groups) value_states = repeat_kv(value_states, self.num_key_value_groups) - 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) + 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) # (batch_size, num_head / sp_size, seq_length, head_size) full_q_len = query_states.size(2) # full_q_len = seq_length else: full_q_len = q_len - # Because the input can be padded, the absolute sequence length depends on the max position id. + # Because the input can be padded, the absolute sequence length depends on + # the max position id. if position_embeddings is None: cos, sin = self.rotary_emb(value_states, position_ids) else: @@ -297,9 +383,11 @@ def ulysses_flash_attn_forward( position_ids=position_ids, # important: pass position ids ) # (batch_size, seq_length, num_head / sp_size, head_size) if ulysses_sp_size > 1: - attn_output = gather_heads_scatter_seq(attn_output, head_dim=2, seq_dim=1) + attn_output = gather_heads_scatter_seq( + attn_output, head_dim=2, seq_dim=1) - attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous() + attn_output = attn_output.reshape( + bsz, q_len, self.hidden_size).contiguous() attn_output = self.o_proj(attn_output) return attn_output, None, None @@ -332,18 +420,27 @@ def forward_base_model( Copy paste Qwen2VL's forward https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/transformers/model/qwen2_vl.py ```""" - output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) output_hidden_states = ( - output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict ) - return_dict = return_dict if return_dict is not None else self.config.use_return_dict if inputs_embeds is None: inputs_embeds = self.model.embed_tokens(input_ids) if pixel_values is not None: pixel_values = pixel_values.type(self.visual.get_dtype()) image_embeds = self.visual(pixel_values, grid_thw=image_grid_thw) - n_image_tokens = (input_ids == self.config.image_token_id).sum().item() + n_image_tokens = ( + input_ids == self.config.image_token_id).sum().item() n_image_features = image_embeds.shape[0] if n_image_tokens != n_image_features: raise ValueError( @@ -356,13 +453,18 @@ def forward_base_model( .expand_as(inputs_embeds) .to(inputs_embeds.device) ) - image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype) - inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds) + image_embeds = image_embeds.to( + inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter( + image_mask, image_embeds) if pixel_values_videos is not None: - pixel_values_videos = pixel_values_videos.type(self.visual.get_dtype()) - video_embeds = self.visual(pixel_values_videos, grid_thw=video_grid_thw) - n_video_tokens = (input_ids == self.config.video_token_id).sum().item() + pixel_values_videos = pixel_values_videos.type( + self.visual.get_dtype()) + video_embeds = self.visual( + pixel_values_videos, grid_thw=video_grid_thw) + n_video_tokens = ( + input_ids == self.config.video_token_id).sum().item() n_video_features = video_embeds.shape[0] if n_video_tokens != n_video_features: raise ValueError( @@ -375,25 +477,39 @@ def forward_base_model( .expand_as(inputs_embeds) .to(inputs_embeds.device) ) - video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype) - inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds) + video_embeds = video_embeds.to( + inputs_embeds.device, inputs_embeds.dtype) + inputs_embeds = inputs_embeds.masked_scatter( + video_mask, video_embeds) if attention_mask is not None: attention_mask = attention_mask.to(inputs_embeds.device) - if position_ids is None and (attention_mask is None or attention_mask.ndim == 2): + if position_ids is None and ( + attention_mask is None or attention_mask.ndim == 2): # calculate RoPE index once per generation in the pre-fill stage only - if (cache_position is not None and cache_position[0] == 0) or self.rope_deltas is None: - position_ids, rope_deltas = self.get_rope_index(input_ids, image_grid_thw, video_grid_thw, attention_mask) + if ( + cache_position is not None and cache_position[0] == 0 + ) or self.rope_deltas is None: + position_ids, rope_deltas = self.get_rope_index( + input_ids, image_grid_thw, video_grid_thw, attention_mask + ) self.rope_deltas = rope_deltas - # then use the prev pre-calculated rope-deltas to get the correct position ids + # then use the prev pre-calculated rope-deltas to get the correct + # position ids else: batch_size, seq_length, _ = inputs_embeds.shape - delta = cache_position[0] + self.rope_deltas if cache_position is not None else 0 - position_ids = torch.arange(seq_length, device=inputs_embeds.device) + delta = ( + cache_position[0] + self.rope_deltas + if cache_position is not None + else 0 + ) + position_ids = torch.arange( + seq_length, device=inputs_embeds.device) position_ids = position_ids.view(1, -1).expand(batch_size, -1) if cache_position is not None: # otherwise `deltas` is an int `0` - delta = delta.repeat_interleave(batch_size // delta.shape[0], dim=0) + delta = delta.repeat_interleave( + batch_size // delta.shape[0], dim=0) position_ids = position_ids.add(delta) position_ids = position_ids.unsqueeze(0).expand(3, -1, -1) @@ -458,7 +574,8 @@ def forward_with_torch_backend( hidden_states = outputs[0] if not return_dict: - raise NotImplementedError("forward_with_torch_backend has to return_dict") + raise NotImplementedError( + "forward_with_torch_backend has to return_dict") # Loss calculations if labels is not None: @@ -466,7 +583,9 @@ def forward_with_torch_backend( elif input_ids is not None: rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1) else: - raise RuntimeError("To use forward_with_torch_backend, either labels or input_ids must be provided.") + raise RuntimeError( + "To use forward_with_torch_backend, either labels or input_ids must be provided." + ) fused_linear_for_ppo = FusedLinearForPPO() log_probs, entropy = fused_linear_for_ppo.forward( @@ -531,7 +650,8 @@ def forward_with_triton_backend( hidden_states = outputs[0] if not return_dict: - raise NotImplementedError("forward_with_triton_backend has to return_dict") + raise NotImplementedError( + "forward_with_triton_backend has to return_dict") # Loss calculations if labels is not None: @@ -539,7 +659,9 @@ def forward_with_triton_backend( elif input_ids is not None: rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1) else: - raise RuntimeError("To use forward_with_triton_backend, either labels or input_ids must be provided.") + raise RuntimeError( + "To use forward_with_triton_backend, either labels or input_ids must be provided." + ) log_probs, entropy = linear_cross_entropy( hidden_states, diff --git a/Agent0/executor_train/verl/verl/models/weight_loader_registry.py b/Agent0/executor_train/verl/verl/models/weight_loader_registry.py index 8aa3bc7..5ffbca3 100644 --- a/Agent0/executor_train/verl/verl/models/weight_loader_registry.py +++ b/Agent0/executor_train/verl/verl/models/weight_loader_registry.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/protocol.py b/Agent0/executor_train/verl/verl/protocol.py index 0029913..eff7d68 100644 --- a/Agent0/executor_train/verl/verl/protocol.py +++ b/Agent0/executor_train/verl/verl/protocol.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -51,12 +51,17 @@ class _DataProtoConfigMeta(type): @property def auto_padding(cls): - enabled_by_env = os.getenv("VERL_AUTO_PADDING", "FALSE").upper() in ["TRUE", "1"] + enabled_by_env = os.getenv("VERL_AUTO_PADDING", "FALSE").upper() in [ + "TRUE", + "1", + ] return enabled_by_env or cls._config.get(cls.auto_padding_key, False) @auto_padding.setter def auto_padding(cls, enabled: bool): - assert isinstance(enabled, bool), f"enabled must be a boolean, got {enabled} as {type(enabled)}" + assert isinstance( + enabled, bool + ), f"enabled must be a boolean, got {enabled} as {type(enabled)}" cls._config[cls.auto_padding_key] = enabled @@ -89,7 +94,8 @@ def pad_dataproto_to_divisor(data: "DataProto", size_divisor: int): data_padded = DataProto.concat([data] + padding_protos) else: if len(data) == 0: - logging.warning("padding a DataProto with no item, no changed made") + logging.warning( + "padding a DataProto with no item, no changed made") pad_size = 0 data_padded = data return data_padded, pad_size @@ -102,31 +108,36 @@ def unpad_dataproto(data: "DataProto", pad_size): return data -def union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict: +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}" - ) + 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" - ) + assert tensor_dict1[key].equal( + tensor_dict2[key] + ), f"{key} in tensor_dict1 and tensor_dict2 are not the same object" return tensor_dict1 -def union_numpy_dict(tensor_dict1: dict[str, np.ndarray], tensor_dict2: dict[str, np.ndarray]) -> dict[str, np.ndarray]: +def union_numpy_dict( + tensor_dict1: dict[str, np.ndarray], tensor_dict2: dict[str, np.ndarray] +) -> dict[str, 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) # to properly deal with nan and object type - assert pd.DataFrame(tensor_dict2[key]).equals(pd.DataFrame(tensor_dict1[key])), ( - f"{key} in tensor_dict1 and tensor_dict2 are not the same object" - ) + assert pd.DataFrame(tensor_dict2[key]).equals( + pd.DataFrame(tensor_dict1[key]) + ), f"{key} in tensor_dict1 and tensor_dict2 are not the same object" tensor_dict1[key] = val return tensor_dict1 @@ -159,9 +170,12 @@ def fold_batch_dim(data: "DataProto", new_batch_size): 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:])) + non_tensor[key] = np.reshape(val, newshape=( + new_batch_size, -1, *val.shape[1:])) - return type(data)(batch=tensor, non_tensor_batch=non_tensor, meta_info=data.meta_info) + return type(data)( + batch=tensor, non_tensor_batch=non_tensor, meta_info=data.meta_info + ) def unfold_batch_dim(data: "DataProto", batch_dims=2): @@ -178,9 +192,13 @@ def unfold_batch_dim(data: "DataProto", batch_dims=2): non_tensor_new = {} for key, val in non_tensor.items(): - non_tensor_new[key] = np.reshape(val, newshape=(batch_size, *val.shape[batch_dims:])) + non_tensor_new[key] = np.reshape( + val, newshape=(batch_size, *val.shape[batch_dims:]) + ) - return type(data)(batch=tensor, non_tensor_batch=non_tensor_new, meta_info=data.meta_info) + return type(data)( + batch=tensor, non_tensor_batch=non_tensor_new, meta_info=data.meta_info + ) def collate_fn(x: list["DataProtoItem"]): @@ -254,11 +272,18 @@ def __getitem__(self, item): elif isinstance(item, list | np.ndarray | torch.Tensor): return self.select_idxs(item) - # Case 3: Single integer - return DataProtoItem for backward compatibility + # Case 3: Single integer - return DataProtoItem for backward + # compatibility elif isinstance(item, int | np.integer): tensor_data = self.batch[item] if self.batch is not None else None - 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) + 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, + ) # # Case 4: Unsupported type else: @@ -268,7 +293,10 @@ def __getstate__(self): import io buffer = io.BytesIO() - if version.parse(tensordict.__version__) >= version.parse("0.5.0") and self.batch is not None: + if ( + version.parse(tensordict.__version__) >= version.parse("0.5.0") + and self.batch is not None + ): self.batch = self.batch.contiguous() self.batch = self.batch.consolidate() torch.save(self.batch, buffer) @@ -322,15 +350,22 @@ def check_consistency(self): 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" + 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 self.non_tensor_batch is not None and len(self.non_tensor_batch) != 0: + if ( + self.batch is not None + and self.non_tensor_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." + 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(): @@ -338,12 +373,17 @@ def check_consistency(self): f"data in the non_tensor_batch must be a numpy.array with dtype=object, but for " f"{key=}, got {type(val)=}" ) - assert val.shape[0] == batch_size, ( - f"key {key} length {len(val)} is not equal to batch size {batch_size}" - ) + 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, torch.Tensor | np.ndarray], meta_info=None, auto_padding=False): + def from_single_dict( + cls, + data: dict[str, torch.Tensor | np.ndarray], + meta_info=None, + auto_padding=False, + ): """Create a DataProto from a dict of tensors and non_tensors""" tensors = {} non_tensors = {} @@ -356,7 +396,12 @@ def from_single_dict(cls, data: dict[str, torch.Tensor | np.ndarray], meta_info= else: raise ValueError(f"Unsupported type in data {type(val)}") - return cls.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info, auto_padding=auto_padding) + return cls.from_dict( + tensors=tensors, + non_tensors=non_tensors, + meta_info=meta_info, + auto_padding=auto_padding, + ) @classmethod def from_dict( @@ -374,7 +419,9 @@ def from_dict( 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." + assert ( + num_batch_dims == 1 + ), "only support num_batch_dims=1 when non_tensors is not None." if tensors is None: tensors = {} @@ -403,10 +450,16 @@ def from_dict( if not isinstance(val, np.ndarray): non_tensors[key] = np.array(val, dtype=object) - tensor_dict = TensorDict(source=tensors, batch_size=batch_size) if tensors else None + tensor_dict = ( + TensorDict( + source=tensors, + batch_size=batch_size) if tensors else None) if auto_padding: meta_info[DataProtoConfig.auto_padding_key] = True - return cls(batch=tensor_dict, non_tensor_batch=non_tensors, meta_info=meta_info) + return cls( + batch=tensor_dict, + non_tensor_batch=non_tensors, + meta_info=meta_info) def to(self, device) -> "DataProto": """move the batch to device @@ -422,7 +475,13 @@ def to(self, device) -> "DataProto": self.batch = self.batch.to(device) return self - def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None, deepcopy=False) -> "DataProto": + 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: @@ -440,7 +499,11 @@ def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=Non 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} + 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 @@ -448,14 +511,19 @@ def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=Non 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} + 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 type(self)(batch=sub_batch, non_tensor_batch=non_tensor_batch, meta_info=sub_meta_info) + return type(self)( + batch=sub_batch, + non_tensor_batch=non_tensor_batch, + meta_info=sub_meta_info) def select_idxs(self, idxs): """ @@ -479,13 +547,18 @@ def select_idxs(self, idxs): idxs_torch = idxs idxs_np = idxs.detach().cpu().numpy() - batch_size = int(idxs_np.sum()) if idxs_np.dtype == bool else idxs_np.shape[0] + batch_size = int( + idxs_np.sum()) if idxs_np.dtype == bool else idxs_np.shape[0] if self.batch is not None: # Use TensorDict's built-in indexing capabilities selected_batch = TensorDict( - source={key: tensor[idxs_torch] for key, tensor in self.batch.items()}, - batch_size=(batch_size,), + source={ + key: tensor[idxs_torch] for key, + tensor in self.batch.items()}, + batch_size=( + batch_size, + ), device=self.batch.device, ) else: @@ -495,7 +568,11 @@ def select_idxs(self, idxs): for key, val in self.non_tensor_batch.items(): selected_non_tensor[key] = val[idxs_np] - return type(self)(batch=selected_batch, non_tensor_batch=selected_non_tensor, meta_info=self.meta_info) + return type(self)( + batch=selected_batch, + non_tensor_batch=selected_non_tensor, + meta_info=self.meta_info, + ) def slice(self, start=None, end=None, step=None): """ @@ -541,9 +618,15 @@ def slice(self, start=None, end=None, step=None): sliced_non_tensor[key] = val[slice_obj] # Return a new DataProto object - return type(self)(batch=sliced_batch, non_tensor_batch=sliced_non_tensor, meta_info=self.meta_info) + return type(self)( + batch=sliced_batch, + non_tensor_batch=sliced_non_tensor, + meta_info=self.meta_info, + ) - def pop(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None) -> "DataProto": + 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: @@ -574,7 +657,9 @@ def pop(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None) 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) + return DataProto.from_dict( + tensors=tensors, non_tensors=non_tensors, meta_info=meta_info + ) def rename(self, old_keys=None, new_keys=None) -> "DataProto": """ @@ -588,7 +673,9 @@ def validate_input(keys): elif isinstance(keys, list): pass else: - raise TypeError(f"keys must be a list or a string, but got {type(keys)}") + raise TypeError( + f"keys must be a list or a string, but got { + type(keys)}") return keys old_keys = validate_input(old_keys) @@ -596,8 +683,9 @@ def validate_input(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)}" - ) + 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)) @@ -618,11 +706,18 @@ def union(self, other: "DataProto") -> "DataProto": 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.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 - def make_iterator(self, mini_batch_size, epochs, seed=None, dataloader_kwargs=None): + def make_iterator( + self, + mini_batch_size, + epochs, + seed=None, + dataloader_kwargs=None): r"""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. @@ -638,7 +733,9 @@ def make_iterator(self, mini_batch_size, epochs, seed=None, dataloader_kwargs=No 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" + 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 = {} @@ -651,7 +748,11 @@ def make_iterator(self, mini_batch_size, epochs, seed=None, dataloader_kwargs=No assert isinstance(dataloader_kwargs, dict) train_dataloader = DataLoader( - dataset=self, batch_size=mini_batch_size, collate_fn=collate_fn, generator=generator, **dataloader_kwargs + dataset=self, + batch_size=mini_batch_size, + collate_fn=collate_fn, + generator=generator, + **dataloader_kwargs, ) def get_data(): @@ -668,7 +769,9 @@ def is_padding_enabled(self): Returns: bool: True if padding is enabled, False otherwise. """ - dataproto_specific_padding = self.meta_info.get(DataProtoConfig.auto_padding_key, False) + dataproto_specific_padding = self.meta_info.get( + DataProtoConfig.auto_padding_key, False + ) return dataproto_specific_padding or DataProtoConfig.auto_padding def padding(self, padding_size, padding_candidate=""): @@ -680,7 +783,9 @@ def padding(self, padding_size, padding_candidate=""): """ if padding_size == 0: return - padding_candidate = self.select_idxs([0 if padding_candidate == "first" else len(self) - 1]) + padding_candidate = self.select_idxs( + [0 if padding_candidate == "first" else len(self) - 1] + ) padding_part = padding_candidate.repeat(padding_size) padded_dp = DataProto.concat([self, padding_part]) self.batch = padded_dp.batch @@ -696,14 +801,16 @@ def chunk(self, chunks: int) -> list["DataProto"]: List[DataProto]: a list of DataProto after splitting """ if not self.is_padding_enabled(): - assert len(self) % chunks == 0, ( - f"only support equal chunk. Got size of DataProto {len(self)} and chunk {chunks}." - ) + assert ( + len(self) % + chunks == 0), f"only support equal chunk. Got size of DataProto { + len(self)} and chunk {chunks}." bsz_in_batch = None if self.batch is not None: batch_lst = self.batch.chunk(chunks=chunks, dim=0) - bsz_in_batch = np.array([batch.batch_size[0] for batch in batch_lst]) + bsz_in_batch = np.array([batch.batch_size[0] + for batch in batch_lst]) chunk_indices = np.cumsum(bsz_in_batch)[:-1] else: batch_lst = [None for _ in range(chunks)] @@ -722,7 +829,11 @@ def chunk(self, chunks: int) -> list["DataProto"]: output = [] for i in range(chunks): output.append( - type(self)(batch=batch_lst[i], non_tensor_batch=non_tensor_batch_lst[i], meta_info=self.meta_info) + type(self)( + batch=batch_lst[i], + non_tensor_batch=non_tensor_batch_lst[i], + meta_info=self.meta_info, + ) ) return output @@ -741,14 +852,21 @@ def concat(data: list["DataProto"]) -> "DataProto": batch_lst = [] for batch in data: batch_lst.append(batch.batch) - new_batch = torch.cat(batch_lst, dim=0) if batch_lst[0] is not None else None + new_batch = torch.cat( + batch_lst, dim=0) if batch_lst[0] is not None else None - non_tensor_batch = list_of_dict_to_dict_of_list(list_of_dict=[d.non_tensor_batch for d in data]) + 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) cls = type(data[0]) if len(data) > 0 else DataProto - return cls(batch=new_batch, non_tensor_batch=non_tensor_batch, meta_info=data[0].meta_info) + return cls( + batch=new_batch, + non_tensor_batch=non_tensor_batch, + meta_info=data[0].meta_info, + ) def reorder(self, indices): """ @@ -756,7 +874,9 @@ def reorder(self, indices): """ 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()} + self.non_tensor_batch = { + key: val[indices_np] for key, val in self.non_tensor_batch.items() + } def repeat(self, repeat_times=2, interleave=True): """ @@ -773,12 +893,15 @@ def repeat(self, repeat_times=2, interleave=True): if interleave: # Interleave the data repeated_tensors = { - key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items() + 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:]) + key: tensor.unsqueeze(0) + .expand(repeat_times, *tensor.shape) + .reshape(-1, *tensor.shape[1:]) for key, tensor in self.batch.items() } @@ -792,9 +915,12 @@ def repeat(self, repeat_times=2, interleave=True): 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) + 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)) + repeated_non_tensor_batch[key] = np.tile( + val, (repeat_times,) + (1,) * (val.ndim - 1) + ) return type(self)( batch=repeated_batch, @@ -802,7 +928,9 @@ def repeat(self, repeat_times=2, interleave=True): meta_info=self.meta_info, ) - def unfold_column_chunks(self, n_split: int, split_keys: Optional[list[str]] = None): + def unfold_column_chunks( + self, n_split: int, split_keys: Optional[list[str]] = None + ): """Split along the second dim into `n_split`, unfold it to the first dim (batch dim) Useful in passing grouped tensors that doesn't want to be shuffled in dataset. keys not in split_keys are repeated to match the shape @@ -817,10 +945,15 @@ def unfold_column_chunks(self, n_split: int, split_keys: Optional[list[str]] = N shape[1] = self.batch[key].shape[1] // n_split unfolded_batch[key] = self.batch[key].reshape(*shape) else: - unfolded_batch[key] = torch.repeat_interleave(self.batch[key], n_split, dim=0) - # locate the `unfolded_batch` as a TensorDict on the same device as the original batch + unfolded_batch[key] = torch.repeat_interleave( + self.batch[key], n_split, dim=0 + ) + # locate the `unfolded_batch` as a TensorDict on the same device as + # the original batch unfolded_batch = TensorDict( - source=unfolded_batch, batch_size=(self.batch.batch_size[0] * n_split,), device=self.batch.device + source=unfolded_batch, + batch_size=(self.batch.batch_size[0] * n_split,), + device=self.batch.device, ) else: unfolded_batch = None @@ -833,7 +966,8 @@ def unfold_column_chunks(self, n_split: int, split_keys: Optional[list[str]] = N shape[1] = val.shape[1] // n_split repeated_non_tensor_batch[key] = val.reshape(*shape) else: - repeated_non_tensor_batch[key] = np.repeat(val, n_split, axis=0) + repeated_non_tensor_batch[key] = np.repeat( + val, n_split, axis=0) return type(self)( batch=unfolded_batch, @@ -860,15 +994,16 @@ def sample_level_repeat(self, repeat_times): assert len(repeat_times.shape) == 1 repeat_times = repeat_times.tolist() else: - assert isinstance(repeat_times, list), ( - f"repeat_times type must be in [list, torch.Tensor, np.ndarray, tuple], got {type(repeat_times)}" - ) + assert isinstance( + repeat_times, list), f"repeat_times type must be in [list, torch.Tensor, np.ndarray, tuple], got { + type(repeat_times)}" repeat_times = torch.tensor(repeat_times) if self.batch is not None: # Interleave the data repeated_tensors = { - key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items() + key: tensor.repeat_interleave(repeat_times, dim=0) + for key, tensor in self.batch.items() } repeated_batch = TensorDict( @@ -881,7 +1016,8 @@ def sample_level_repeat(self, repeat_times): repeated_non_tensor_batch = {} for key, val in self.non_tensor_batch.items(): - repeated_non_tensor_batch[key] = np.repeat(val, repeat_times, axis=0) + repeated_non_tensor_batch[key] = np.repeat( + val, repeat_times, axis=0) return type(self)( batch=repeated_batch, @@ -924,7 +1060,9 @@ 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 + 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 @@ -935,19 +1073,28 @@ def get(self): 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 + # split in batch dim, select using dp + output = self.dispatch_fn(output) return output def all_gather_data_proto(data: DataProto, process_group): - # Note that this is an inplace operator just like torch.distributed.all_gather + # Note that this is an inplace operator just like + # torch.distributed.all_gather group_size = torch.distributed.get_world_size(group=process_group) assert isinstance(data, DataProto) prev_device = data.batch.device data.batch = data.batch.to(get_device_id()) - data.batch = allgather_dict_tensors(data.batch.contiguous(), size=group_size, group=process_group, dim=0) + data.batch = allgather_dict_tensors( + data.batch.contiguous(), size=group_size, group=process_group, dim=0 + ) data.batch = data.batch.to(prev_device) # all gather non_tensor_batch all_non_tensor_batch = [None for _ in range(group_size)] - torch.distributed.all_gather_object(all_non_tensor_batch, data.non_tensor_batch, group=process_group) - data.non_tensor_batch = {k: np.concatenate([d[k] for d in all_non_tensor_batch]) for k in data.non_tensor_batch} + torch.distributed.all_gather_object( + all_non_tensor_batch, data.non_tensor_batch, group=process_group + ) + data.non_tensor_batch = { + k: np.concatenate([d[k] for d in all_non_tensor_batch]) + for k in data.non_tensor_batch + } diff --git a/Agent0/executor_train/verl/verl/single_controller/__init__.py b/Agent0/executor_train/verl/verl/single_controller/__init__.py index ad6c42a..238deec 100644 --- a/Agent0/executor_train/verl/verl/single_controller/__init__.py +++ b/Agent0/executor_train/verl/verl/single_controller/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,7 +19,9 @@ version_folder = os.path.dirname(os.path.join(os.path.abspath(__file__))) # Note(haibin.lin): single_controller.__version__ is deprecated -with open(os.path.join(os.path.join(version_folder, os.pardir), "version/version")) as f: +with open( + os.path.join(os.path.join(version_folder, os.pardir), "version/version") +) as f: __version__ = f.read().strip() diff --git a/Agent0/executor_train/verl/verl/single_controller/base/__init__.py b/Agent0/executor_train/verl/verl/single_controller/base/__init__.py index b24bd99..cea972b 100644 --- a/Agent0/executor_train/verl/verl/single_controller/base/__init__.py +++ b/Agent0/executor_train/verl/verl/single_controller/base/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/single_controller/base/decorator.py b/Agent0/executor_train/verl/verl/single_controller/base/decorator.py index 1008a79..71de637 100644 --- a/Agent0/executor_train/verl/verl/single_controller/base/decorator.py +++ b/Agent0/executor_train/verl/verl/single_controller/base/decorator.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,7 +19,8 @@ from verl.protocol import DataProtoFuture, _padding_size_key from verl.utils.py_functional import DynamicEnum -# here we add a magic number of avoid user-defined function already have this attribute +# here we add a magic number of avoid user-defined function already have +# this attribute MAGIC_ATTR = "attrs_3141562937" @@ -103,12 +104,16 @@ def _split_args_kwargs_data_proto_with_auto_padding(chunks, *args, **kwargs): # for padding, we only support DataProto with same length if data_proto_len is None: data_proto_len = len(arg) - padding_size = (chunks - (data_proto_len % chunks)) if (data_proto_len % chunks > 0) else 0 + padding_size = ( + (chunks - (data_proto_len % chunks)) + if (data_proto_len % chunks > 0) + else 0 + ) splitted_kwargs[_padding_size_key] = padding_size else: - assert data_proto_len == len(arg), ( - f"expecting all arg share same length of {data_proto_len}, but got {len(arg)}" - ) + assert data_proto_len == len( + arg + ), f"expecting all arg share same length of {data_proto_len}, but got {len(arg)}" data_proto_len = len(arg) arg.padding(padding_size=padding_size) @@ -123,9 +128,9 @@ def _split_args_kwargs_data_proto_with_auto_padding(chunks, *args, **kwargs): padding_size = chunks - (data_proto_len % chunks) splitted_kwargs[_padding_size_key] = padding_size else: - assert data_proto_len == len(val), ( - f"expecting all arg share same length of {data_proto_len}, but got {len(val)}" - ) + assert data_proto_len == len( + val + ), f"expecting all arg share same length of {data_proto_len}, but got {len(val)}" data_proto_len = len(val) splitted_kwargs[key] = val.chunk(chunks=chunks) @@ -156,9 +161,9 @@ def dispatch_megatron_compute(worker_group, *args, **kwargs): """ 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)}" - ) + assert isinstance( + worker_group, MegatronWorkerGroup + ), f"worker_group must be MegatronWorkerGroup, Got {type(worker_group)}" # ray put all the args in advance to avoid duplicate serialization cost import ray @@ -168,7 +173,8 @@ def dispatch_megatron_compute(worker_group, *args, **kwargs): all_args = [] for arg in args: - assert isinstance(arg, tuple | list) and len(arg) == worker_group.dp_size + 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 @@ -198,7 +204,11 @@ def collect_megatron_compute(worker_group, output): 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 and local_rank_info.cp_rank == 0: + if ( + local_rank_info.tp_rank == 0 + and local_rank_info.pp_rank == pp_size - 1 + and local_rank_info.cp_rank == 0 + ): output_in_dp.append(output[global_rank]) return output_in_dp @@ -211,8 +221,11 @@ def dispatch_megatron_compute_data_proto(worker_group, *args, **kwargs): 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) + 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) def _concat_data_proto_or_future(output: list): @@ -244,7 +257,9 @@ def collect_megatron_compute_data_proto(worker_group, output): 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)}" + assert isinstance( + o, DataProto | ray.ObjectRef + ), f"expecting {o} to be DataProto, but got {type(o)}" return _concat_data_proto_or_future(output) @@ -289,13 +304,16 @@ def dispatch_megatron_pp_as_dp(worker_group, *args, **kwargs): all_kwargs = {} for k, v in kwargs.items(): - assert isinstance(v, list | tuple) and len(v) == pp_dp_cp_size, f"expect len(v)=={pp_dp_cp_size}, got {len(v)}" + assert ( + isinstance(v, list | tuple) and len(v) == pp_dp_cp_size + ), f"expect len(v)=={pp_dp_cp_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 local_cp_rank = worker_group.get_megatron_rank_info(rank=i).cp_rank - # compute the rank in arg. Note that the order is dp then cp then pp + # compute the rank in arg. Note that the order is dp then cp then + # pp dp_cp_rank = local_cp_rank * dp_size + local_dp_rank arg_rank = dp_cp_rank * pp_size + local_pp_rank transformed_v.append(v[arg_rank]) @@ -338,9 +356,13 @@ def dispatch_megatron_pp_as_dp_data_proto(worker_group, *args, **kwargs): assert isinstance(worker_group, MegatronWorkerGroup) - pp_dp_cp_size = worker_group.dp_size * worker_group.pp_size * worker_group.cp_size - splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(pp_dp_cp_size, *args, **kwargs) - ret = dispatch_megatron_pp_as_dp(worker_group, *splitted_args, **splitted_kwargs) + pp_dp_cp_size = worker_group.dp_size * \ + worker_group.pp_size * worker_group.cp_size + splitted_args, splitted_kwargs = _split_args_kwargs_data_proto( + pp_dp_cp_size, *args, **kwargs + ) + ret = dispatch_megatron_pp_as_dp( + worker_group, *splitted_args, **splitted_kwargs) return ret @@ -358,9 +380,11 @@ def dispatch_dp_compute(worker_group, *args, **kwargs): assert isinstance(worker_group, WorkerGroup) for arg in args: - assert isinstance(arg, tuple | list) and len(arg) == worker_group.world_size + 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 + assert isinstance(v, tuple | list) and len( + v) == worker_group.world_size return args, kwargs @@ -378,10 +402,7 @@ def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs): assert isinstance(worker_group, WorkerGroup) # Note: enable auto padding for dp compute DatapProto splitted_args, splitted_kwargs = _split_args_kwargs_data_proto_with_auto_padding( - worker_group.world_size, - *args, - **kwargs, - ) + worker_group.world_size, *args, **kwargs, ) return splitted_args, splitted_kwargs @@ -389,10 +410,14 @@ 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 isinstance(args[0], FunctionType) # NOTE: The first one args is a function! + # NOTE: The first one args is a function! + assert isinstance(args[0], FunctionType) - 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 + 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 @@ -402,7 +427,9 @@ def collect_dp_compute_data_proto(worker_group, output): from verl.protocol import DataProto for o in output: - assert isinstance(o, DataProto | ray.ObjectRef), f"expecting {o} to be DataProto, but got {type(o)}" + 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) @@ -426,7 +453,10 @@ def collect_dp_compute_data_proto(worker_group, output): "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_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, @@ -435,7 +465,10 @@ def collect_dp_compute_data_proto(worker_group, output): "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: { + "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, @@ -444,7 +477,10 @@ def collect_dp_compute_data_proto(worker_group, output): "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}, + Dispatch.DP_COMPUTE_METRIC: { + "dispatch_fn": dispatch_dp_compute_data_proto, + "collect_fn": collect_dp_compute, + }, Dispatch.DIRECT_ROLLOUT_METHOD: { "dispatch_fn": dummy_direct_rollout_call, "collect_fn": dummy_direct_rollout_call, @@ -462,8 +498,13 @@ def register_dispatch_mode(dispatch_mode_name, dispatch_fn, collect_fn): """ dispatch_mode = Dispatch.register(dispatch_mode_name) _check_dispatch_mode(dispatch_mode) - assert dispatch_mode not in DISPATCH_MODE_FN_REGISTRY, f"dispatch_mode_name {dispatch_mode_name} already exists" - DISPATCH_MODE_FN_REGISTRY[dispatch_mode] = {"dispatch_fn": dispatch_fn, "collect_fn": collect_fn} + assert ( + dispatch_mode not in DISPATCH_MODE_FN_REGISTRY + ), f"dispatch_mode_name {dispatch_mode_name} already exists" + DISPATCH_MODE_FN_REGISTRY[dispatch_mode] = { + "dispatch_fn": dispatch_fn, + "collect_fn": collect_fn, + } def update_dispatch_mode(dispatch_mode, dispatch_fn, collect_fn): @@ -471,8 +512,13 @@ def update_dispatch_mode(dispatch_mode, dispatch_fn, collect_fn): Update the dispatch mode. """ _check_dispatch_mode(dispatch_mode) - assert dispatch_mode in DISPATCH_MODE_FN_REGISTRY, f"dispatch_mode {dispatch_mode} not found" - DISPATCH_MODE_FN_REGISTRY[dispatch_mode] = {"dispatch_fn": dispatch_fn, "collect_fn": collect_fn} + assert ( + dispatch_mode in DISPATCH_MODE_FN_REGISTRY + ), f"dispatch_mode {dispatch_mode} not found" + DISPATCH_MODE_FN_REGISTRY[dispatch_mode] = { + "dispatch_fn": dispatch_fn, + "collect_fn": collect_fn, + } def get_predefined_execute_fn(execute_mode): @@ -488,17 +534,21 @@ def get_predefined_execute_fn(execute_mode): 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}" - ) + 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" + assert ( + key in dispatch_mode + ), f"key {key} should be in dispatch_mode if it is a dictionary" def _check_execute_mode(execute_mode): - assert isinstance(execute_mode, Execute), f"execute_mode must be a Execute. Got {execute_mode}" + assert isinstance( + execute_mode, Execute + ), f"execute_mode must be a Execute. Got {execute_mode}" def _materialize_futures(*args, **kwargs): @@ -516,7 +566,12 @@ def _materialize_futures(*args, **kwargs): return new_args, kwargs -def register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, blocking=True, materialize_futures=True): +def register( + dispatch_mode=Dispatch.ALL_TO_ALL, + execute_mode=Execute.ALL, + blocking=True, + materialize_futures=True, +): """Register a function with distributed execution configuration. This decorator registers a function with specific dispatch and execution modes @@ -554,7 +609,11 @@ async def async_inner(*args, **kwargs): return await func(*args, **kwargs) wrapper = async_inner if inspect.iscoroutinefunction(func) else inner - attrs = {"dispatch_mode": dispatch_mode, "execute_mode": execute_mode, "blocking": blocking} + attrs = { + "dispatch_mode": dispatch_mode, + "execute_mode": execute_mode, + "blocking": blocking, + } setattr(wrapper, MAGIC_ATTR, attrs) return wrapper diff --git a/Agent0/executor_train/verl/verl/single_controller/base/megatron/__init__.py b/Agent0/executor_train/verl/verl/single_controller/base/megatron/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/single_controller/base/megatron/__init__.py +++ b/Agent0/executor_train/verl/verl/single_controller/base/megatron/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/single_controller/base/megatron/worker.py b/Agent0/executor_train/verl/verl/single_controller/base/megatron/worker.py index baf6eb8..4aa3bf3 100644 --- a/Agent0/executor_train/verl/verl/single_controller/base/megatron/worker.py +++ b/Agent0/executor_train/verl/verl/single_controller/base/megatron/worker.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,7 +26,9 @@ def get_megatron_global_info(self): dp_size = mpu.get_data_parallel_world_size() pp_size = mpu.get_pipeline_model_parallel_world_size() cp_size = mpu.get_context_parallel_world_size() - info = DistGlobalInfo(tp_size=tp_size, dp_size=dp_size, pp_size=pp_size, cp_size=cp_size) + info = DistGlobalInfo( + tp_size=tp_size, dp_size=dp_size, pp_size=pp_size, cp_size=cp_size + ) return info def get_megatron_rank_info(self): @@ -36,7 +38,9 @@ def get_megatron_rank_info(self): dp_rank = mpu.get_data_parallel_rank() pp_rank = mpu.get_pipeline_model_parallel_rank() cp_rank = mpu.get_context_parallel_rank() - info = DistRankInfo(tp_rank=tp_rank, dp_rank=dp_rank, pp_rank=pp_rank, cp_rank=cp_rank) + info = DistRankInfo( + tp_rank=tp_rank, dp_rank=dp_rank, pp_rank=pp_rank, cp_rank=cp_rank + ) return info def _init_hf_config_and_tf_config( @@ -59,11 +63,19 @@ def _init_hf_config_and_tf_config( # Step 1: initialize the tokenizer self.local_path = copy_to_local(model_path) if tokenizer_or_path is None: - self.tokenizer = hf_tokenizer(self.local_path, trust_remote_code=trust_remote_code) - self.processor = hf_processor(self.local_path, trust_remote_code=trust_remote_code) + self.tokenizer = hf_tokenizer( + self.local_path, trust_remote_code=trust_remote_code + ) + self.processor = hf_processor( + self.local_path, trust_remote_code=trust_remote_code + ) elif isinstance(tokenizer_or_path, str): - self.tokenizer = hf_tokenizer(copy_to_local(tokenizer_or_path), trust_remote_code=trust_remote_code) - self.processor = hf_processor(copy_to_local(tokenizer_or_path), trust_remote_code=trust_remote_code) + self.tokenizer = hf_tokenizer( + copy_to_local(tokenizer_or_path), + trust_remote_code=trust_remote_code) + self.processor = hf_processor( + copy_to_local(tokenizer_or_path), + trust_remote_code=trust_remote_code) else: self.tokenizer = tokenizer_or_path self.processor = tokenizer_or_path @@ -75,7 +87,9 @@ def _init_hf_config_and_tf_config( self.tokenizer.chat_template = self.config.model.custom_chat_template # Step 2: get the hf - hf_config = AutoConfig.from_pretrained(self.local_path, trust_remote_code=trust_remote_code) + hf_config = AutoConfig.from_pretrained( + self.local_path, trust_remote_code=trust_remote_code + ) # Step 3: override the hf config override_config_kwargs = { @@ -83,23 +97,34 @@ def _init_hf_config_and_tf_config( "eos_token_id": self.tokenizer.eos_token_id, "pad_token_id": self.tokenizer.pad_token_id, } - override_config_kwargs.update(override_model_config.get("model_config", {})) - self.share_embeddings_and_output_weights = getattr(hf_config, "tie_word_embeddings", False) - update_model_config(hf_config, override_config_kwargs=override_config_kwargs) + override_config_kwargs.update( + override_model_config.get( + "model_config", {})) + self.share_embeddings_and_output_weights = getattr( + hf_config, "tie_word_embeddings", False + ) + update_model_config( + hf_config, + override_config_kwargs=override_config_kwargs) self.architectures = getattr(hf_config, "architectures", None) if self.rank == 0: print(f"Model config after override: {hf_config}") - tf_config = hf_to_mcore_config(hf_config, dtype, **override_transformer_config) + tf_config = hf_to_mcore_config( + hf_config, dtype, **override_transformer_config) def add_optimization_config_to_tf_config(tf_config): # add optimization config to tf_config, e.g. checkpointing if self.config.model.get("enable_gradient_checkpointing", False): - gradient_checkpointing_cfg = dict(self.config.model.get("gradient_checkpointing_kwargs", dict())) - tf_config.recompute_method = gradient_checkpointing_cfg.get("activations_checkpoint_method", "full") - tf_config.recompute_granularity = gradient_checkpointing_cfg.get( - "activations_checkpoint_granularity", "full" + gradient_checkpointing_cfg = dict( + self.config.model.get( + "gradient_checkpointing_kwargs", dict())) + tf_config.recompute_method = gradient_checkpointing_cfg.get( + "activations_checkpoint_method", "full" ) - tf_config.recompute_num_layers = gradient_checkpointing_cfg.get("activations_checkpoint_num_layers", -1) + tf_config.recompute_granularity = gradient_checkpointing_cfg.get( + "activations_checkpoint_granularity", "full") + tf_config.recompute_num_layers = gradient_checkpointing_cfg.get( + "activations_checkpoint_num_layers", -1) if megatron_config := self.config.get("megatron", {}): if extra := megatron_config.get("extra", {}): for k, v in extra.items(): diff --git a/Agent0/executor_train/verl/verl/single_controller/base/megatron/worker_group.py b/Agent0/executor_train/verl/verl/single_controller/base/megatron/worker_group.py index b9beb84..77fb95d 100644 --- a/Agent0/executor_train/verl/verl/single_controller/base/megatron/worker_group.py +++ b/Agent0/executor_train/verl/verl/single_controller/base/megatron/worker_group.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -25,30 +25,42 @@ def __init__(self, resource_pool: ResourcePool, **kwargs): self._megatron_global_info: DistGlobalInfo = None def init_megatron(self, default_megatron_kwargs: dict = None): - raise NotImplementedError("MegatronWorkerGroup.init_megatron should be overwritten") + raise NotImplementedError( + "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}" + 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" + 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" + 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" + assert ( + self._megatron_global_info is not None + ), "MegatronWorkerGroup._megatron_global_info must be initialized" return self._megatron_global_info.pp_size @property def cp_size(self): - assert self._megatron_global_info is not None, "MegatronWorkerGroup._megatron_global_info must be initialized" + assert ( + self._megatron_global_info is not None + ), "MegatronWorkerGroup._megatron_global_info must be initialized" return self._megatron_global_info.cp_size def get_megatron_global_info(self): diff --git a/Agent0/executor_train/verl/verl/single_controller/base/register_center/__init__.py b/Agent0/executor_train/verl/verl/single_controller/base/register_center/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/single_controller/base/register_center/__init__.py +++ b/Agent0/executor_train/verl/verl/single_controller/base/register_center/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/single_controller/base/register_center/ray.py b/Agent0/executor_train/verl/verl/single_controller/base/register_center/ray.py index ac071cd..7663a9e 100644 --- a/Agent0/executor_train/verl/verl/single_controller/base/register_center/ray.py +++ b/Agent0/executor_train/verl/verl/single_controller/base/register_center/ray.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/single_controller/base/worker.py b/Agent0/executor_train/verl/verl/single_controller/base/worker.py index 561b9ba..0512749 100644 --- a/Agent0/executor_train/verl/verl/single_controller/base/worker.py +++ b/Agent0/executor_train/verl/verl/single_controller/base/worker.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -50,7 +50,8 @@ def get_node_ip_by_sdk(): return ray._private.services.get_node_ip_address() else: - raise NotImplementedError("WG_BACKEND now just support ray mode.") + raise NotImplementedError( + "WG_BACKEND now just support ray mode.") host_ipv4 = os.getenv("MY_HOST_IP", None) host_ipv6 = os.getenv("MY_HOST_IPV6", None) @@ -95,9 +96,15 @@ def __new__(cls, *args, **kwargs): 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)) + # 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 @@ -110,7 +117,9 @@ def _configure_before_init(self, register_center_name: str, rank: int): rank (int): Rank of the worker in the distributed setup """ - assert isinstance(rank, int), f"rank must be int, instead of {type(rank)}" + 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() @@ -120,7 +129,9 @@ def _configure_before_init(self, register_center_name: str, rank: int): } if os.getenv("WG_BACKEND", None) == "ray": - from verl.single_controller.base.register_center.ray import create_worker_group_register_center + 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 @@ -131,7 +142,11 @@ def _configure_before_init(self, register_center_name: str, rank: int): self.register_center = ray.get_actor(register_center_name) # set worker info for node affinity scheduling - ray.get(self.register_center.set_worker_info.remote(rank, ray.get_runtime_context().get_node_id())) + ray.get( + self.register_center.set_worker_info.remote( + rank, ray.get_runtime_context().get_node_id() + ) + ) @classmethod def env_keys(cls): @@ -179,7 +194,8 @@ def __init__(self, cuda_visible_devices=None) -> None: "_master_port": master_port, } if cuda_visible_devices is not None: - store[f"_{get_visible_devices_keyword()}".lower()] = cuda_visible_devices + store[f"_{get_visible_devices_keyword()}".lower() + ] = cuda_visible_devices self._configure_with_store(store=store) @@ -230,7 +246,9 @@ def _setup_env_cuda_visible_devices(self): # Otherwise, we will set ROCR_VISIBLE_DEVICES to CUDA_VISIBLE_DEVICES # and remove ROCR_VISIBLE_DEVICES. if cuda_val: - raise ValueError("Please don't set ROCR_VISIBLE_DEVICES when HIP/CUDA_VISIBLE_DEVICES is set.") + raise ValueError( + "Please don't set ROCR_VISIBLE_DEVICES when HIP/CUDA_VISIBLE_DEVICES is set." + ) cuda_val = os.environ.pop("ROCR_VISIBLE_DEVICES") os.environ["CUDA_VISIBLE_DEVICES"] = cuda_val @@ -249,7 +267,10 @@ def _configure_with_store(self, store: dict): """ This function should only be called inside by WorkerGroup """ - store_env_dict = {f"_{key.lower()}": store.get(f"_{key.lower()}", None) for key in type(self).env_keys()} + store_env_dict = { + f"_{key.lower()}": store.get(f"_{key.lower()}", None) + for key in type(self).env_keys() + } self.__dict__.update(store_env_dict) # this is hacky # print(f"__dict__: {self.__dict__}") for key in type(self).env_keys(): @@ -258,7 +279,9 @@ def _configure_with_store(self, store: dict): # 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 "" + str(self._master_addr).replace("[", "").replace("]", "") + if self._master_addr + else "" ) def get_master_addr_port(self): @@ -269,7 +292,9 @@ def get_cuda_visible_devices(self): """Get the CUDA visible devices configuration.""" import os - visible_devices = os.environ.get(get_visible_devices_keyword().upper(), "not set") + visible_devices = os.environ.get( + get_visible_devices_keyword().upper(), "not set" + ) return visible_devices @property @@ -297,7 +322,8 @@ 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) + @register(dispatch_mode=Dispatch.ALL_TO_ALL, + execute_mode=Execute.RANK_ZERO) def execute_func_rank_zero(self, func, *args, **kwargs): """Execute a function in rank zero execution mode. diff --git a/Agent0/executor_train/verl/verl/single_controller/base/worker_group.py b/Agent0/executor_train/verl/verl/single_controller/base/worker_group.py index cb86ab4..e0391ab 100644 --- a/Agent0/executor_train/verl/verl/single_controller/base/worker_group.py +++ b/Agent0/executor_train/verl/verl/single_controller/base/worker_group.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -21,7 +21,12 @@ import time from typing import Any, Callable -from .decorator import MAGIC_ATTR, Dispatch, get_predefined_dispatch_fn, get_predefined_execute_fn +from .decorator import ( + MAGIC_ATTR, + Dispatch, + get_predefined_dispatch_fn, + get_predefined_execute_fn, +) class ResourcePool: @@ -31,7 +36,11 @@ class ResourcePool: across all nodes in the pool. """ - def __init__(self, process_on_nodes=None, max_colocate_count: int = 10, n_gpus_per_node=8) -> None: + def __init__( + self, + process_on_nodes=None, + max_colocate_count: int = 10, + n_gpus_per_node=8) -> None: """Initialize the ResourcePool with node processes and GPU configuration. Args: @@ -43,7 +52,8 @@ def __init__(self, process_on_nodes=None, max_colocate_count: int = 10, n_gpus_p process_on_nodes = [] self._store = process_on_nodes self.max_colocate_count = max_colocate_count - self.n_gpus_per_node = n_gpus_per_node # this is left for future huawei GPU that contains 16 GPUs per node + # this is left for future huawei GPU that contains 16 GPUs per node + self.n_gpus_per_node = n_gpus_per_node def add_node(self, process_count): self._store.append(process_count) @@ -63,13 +73,16 @@ def store(self): def local_world_size_list(self) -> list[int]: """Returns a flat list where each process has its local world size.""" nested_local_world_size_list = [ - [local_world_size for _ in range(local_world_size)] for local_world_size in self._store + [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]: """Returns a flat list of local ranks for all processes across all nodes.""" - nested_local_rank_list = [[i for i in range(local_world_size)] for local_world_size in self._store] + 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] @@ -99,7 +112,10 @@ def __call__(self) -> Any: return self.cls(*self.args, **self.kwargs) -def check_workers_alive(workers: list, is_alive: Callable, gap_time: float = 1) -> None: +def check_workers_alive( + workers: list, + is_alive: Callable, + gap_time: float = 1) -> None: """Continuously monitors worker processes and raises SIGABRT if any worker dies. Args: @@ -115,7 +131,8 @@ def check_workers_alive(workers: list, is_alive: Callable, gap_time: float = 1) while True: for worker in workers: if not is_alive(worker): - logging.warning(f"worker {worker} is not alive sending signal to main thread") + logging.warning( + f"worker {worker} is not alive sending signal to main thread") signal.raise_signal(signal.SIGABRT) time.sleep(gap_time) @@ -149,12 +166,15 @@ def __init__(self, resource_pool: ResourcePool, **kwargs) -> None: def _is_worker_alive(self, worker): """Check if a worker is alive. Must be implemented by derived classes.""" - raise NotImplementedError("WorkerGroup._is_worker_alive called, should be implemented in derived class.") + raise NotImplementedError( + "WorkerGroup._is_worker_alive called, should be implemented in derived class." + ) def _block_until_all_workers_alive(self) -> None: """Blocks until all workers in the group are alive.""" while True: - all_state = [self._is_worker_alive(worker) for worker in self._workers] + all_state = [self._is_worker_alive( + worker) for worker in self._workers] if False in all_state: time.sleep(1) else: @@ -166,11 +186,13 @@ def start_worker_aliveness_check(self, every_n_seconds=1) -> None: Args: every_n_seconds (int): Interval between aliveness checks """ - # before starting checking worker aliveness, make sure all workers are already alive + # 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) + target=check_workers_alive, + args=(self._workers, self._is_worker_alive, every_n_seconds), ) self._checker_thread.start() @@ -193,16 +215,23 @@ def _bind_worker_method(self, user_defined_cls, func_generator): for method_name in dir(user_defined_cls): try: method = getattr(user_defined_cls, method_name) - assert callable(method), f"{method_name} in {user_defined_cls} is not callable" + assert callable( + method + ), f"{method_name} in {user_defined_cls} is not callable" except Exception: - # if it is a property, it will fail because Class doesn't have instance property + # 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, "attribute must contain dispatch_mode in its key" + assert isinstance( + attribute, dict + ), f"attribute must be a dictionary. Got {type(attribute)}" + assert ( + "dispatch_mode" in attribute + ), "attribute must contain dispatch_mode in its key" dispatch_mode = attribute["dispatch_mode"] execute_mode = attribute["execute_mode"] @@ -211,7 +240,8 @@ def _bind_worker_method(self, user_defined_cls, func_generator): # get dispatch fn if isinstance(dispatch_mode, Dispatch): # get default dispatch fn - fn = get_predefined_dispatch_fn(dispatch_mode=dispatch_mode) + fn = get_predefined_dispatch_fn( + dispatch_mode=dispatch_mode) dispatch_fn = fn["dispatch_fn"] collect_fn = fn["collect_fn"] else: @@ -222,7 +252,8 @@ def _bind_worker_method(self, user_defined_cls, func_generator): collect_fn = dispatch_mode["collect_fn"] # get execute_fn_name - execute_mode = get_predefined_execute_fn(execute_mode=execute_mode) + execute_mode = get_predefined_execute_fn( + execute_mode=execute_mode) wg_execute_fn_name = execute_mode["execute_fn_name"] # get execute_fn from string @@ -247,6 +278,7 @@ def _bind_worker_method(self, user_defined_cls, func_generator): setattr(self, method_name, func) method_names.append(method_name) except Exception as e: - raise ValueError(f"Fail to set method_name {method_name}") from e + raise ValueError( + f"Fail to set method_name {method_name}") from e return method_names diff --git a/Agent0/executor_train/verl/verl/single_controller/ray/__init__.py b/Agent0/executor_train/verl/verl/single_controller/ray/__init__.py index d2a5d6d..aff5b39 100644 --- a/Agent0/executor_train/verl/verl/single_controller/ray/__init__.py +++ b/Agent0/executor_train/verl/verl/single_controller/ray/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/single_controller/ray/base.py b/Agent0/executor_train/verl/verl/single_controller/ray/base.py index bfcf87b..723caa4 100644 --- a/Agent0/executor_train/verl/verl/single_controller/ray/base.py +++ b/Agent0/executor_train/verl/verl/single_controller/ray/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,10 +24,18 @@ from ray.experimental.state.api import get_actor from ray.util import list_named_actors from ray.util.placement_group import PlacementGroup, placement_group -from ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy, PlacementGroupSchedulingStrategy +from ray.util.scheduling_strategies import ( + NodeAffinitySchedulingStrategy, + PlacementGroupSchedulingStrategy, +) from verl.protocol import DataProto, _padding_size_key -from verl.single_controller.base import ClassWithInitArgs, ResourcePool, Worker, WorkerGroup +from verl.single_controller.base import ( + ClassWithInitArgs, + ResourcePool, + Worker, + WorkerGroup, +) from verl.single_controller.base.decorator import MAGIC_ATTR, Dispatch __all__ = ["Worker"] @@ -41,7 +49,13 @@ def get_random_string(length: int) -> str: return "".join(random.choice(letters_digits) for _ in range(length)) -def func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking): +def func_generator( + self, + method_name, + dispatch_fn, + collect_fn, + execute_fn, + blocking): class Functor: def __call__(this, *args, **kwargs): args, kwargs = dispatch_fn(self, *args, **kwargs) @@ -62,7 +76,8 @@ def __call__(this, *args, **kwargs): return type(method_name, (Functor,), {})() -def sort_placement_group_by_node_ip(pgs: list[PlacementGroup]) -> list[PlacementGroup]: +def sort_placement_group_by_node_ip( + pgs: list[PlacementGroup]) -> list[PlacementGroup]: """ Sort the placement groups by node ip, all bundles in a single placement group should be on the same node. @@ -72,7 +87,8 @@ def sort_placement_group_by_node_ip(pgs: list[PlacementGroup]) -> list[Placement With this function, if there's only one resource pool and there's no node change, RANK should be consistent across nodes in multiple ray jobs, even if the whole ray cluster is restarted. """ - node_ip = {node["NodeID"]: node["NodeManagerAddress"] for node in ray.nodes()} + node_ip = {node["NodeID"]: node["NodeManagerAddress"] + for node in ray.nodes()} pg_ip = {} for pg in pgs: specs = ray._private.state.state.placement_group_table(pg.id) @@ -95,17 +111,23 @@ def __init__( super().__init__(process_on_nodes, max_colocate_count) self.use_gpu = use_gpu # print(f"in RayProcessDispatchConfiguration: name_prefix = {name_prefix}") - self.name_prefix = get_random_string(length=6) if name_prefix is None else name_prefix + self.name_prefix = ( + get_random_string(length=6) if name_prefix is None else name_prefix + ) self.pgs = None self.detached = detached self.accelerator_type = accelerator_type - def get_placement_groups(self, strategy="STRICT_PACK", name=None, device_name="cuda"): + def get_placement_groups( + self, strategy="STRICT_PACK", name=None, device_name="cuda" + ): 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])}:" + 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}") if device_name == "npu": @@ -118,12 +140,20 @@ def get_placement_groups(self, strategy="STRICT_PACK", name=None, device_name="c bundle[device_name] = 1 if self.accelerator_type is not None: bundle[self.accelerator_type] = 1e-4 - pg_scheme = [[bundle.copy() for _ in range(process_count)] for process_count in self._store] + pg_scheme = [ + [bundle.copy() 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) + placement_group( + bundles=bundles, + strategy=strategy, + name=pg_name_prefix + str(idx), + lifetime=lifetime, + ) for idx, bundles in enumerate(pg_scheme) ] @@ -134,7 +164,9 @@ def get_placement_groups(self, strategy="STRICT_PACK", name=None, device_name="c def extract_pg_from_exist( - resource_pools: dict[str, RayResourcePool], src_role_names: list[str], resource_pool: RayResourcePool + resource_pools: dict[str, RayResourcePool], + src_role_names: list[str], + resource_pool: RayResourcePool, ) -> list: src_pgs = [ pg @@ -143,31 +175,48 @@ def extract_pg_from_exist( 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) + 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" - ) + 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)] -def merge_resource_pool(rp1: RayResourcePool, rp2: RayResourcePool) -> RayResourcePool: +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_colocate_count == rp2.max_colocate_count, "Both RayResourcePool must has the same max_colocate_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" + assert ( + rp1.max_colocate_count == rp2.max_colocate_count + ), "Both RayResourcePool must has the same max_colocate_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 = type(rp1)(new_store, rp1.use_gpu, f"{rp1.name_prefix}_{rp2.name_prefix}") + merged = type(rp1)( + new_store, rp1.use_gpu, f"{ + rp1.name_prefix}_{ + rp2.name_prefix}") merged.pgs = rp1.get_placement_groups() + rp2.get_placement_groups() return merged @@ -227,13 +276,21 @@ def __call__( """ if sharing_with is not None: target_node_id = ray.get(sharing_with.get_node_id.remote()) - 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=visible_devices, **self.kwargs) + 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=visible_devices, **self.kwargs + ) options = { "scheduling_strategy": PlacementGroupSchedulingStrategy( - placement_group=placement_group, placement_group_bundle_index=placement_group_bundle_idx + placement_group=placement_group, + placement_group_bundle_index=placement_group_bundle_idx, ) } options.update(self._options) @@ -288,9 +345,12 @@ def __init__( """ 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 + self.name_prefix = ( + get_random_string(length=6) if name_prefix is None else name_prefix + ) self._ray_wait_register_center_timeout = ray_wait_register_center_timeout - # Whether the WorkerGroup is a Colocate WorkerGroup created by FusedWorker. + # Whether the WorkerGroup is a Colocate WorkerGroup created by + # FusedWorker. self.fused_worker_used = ray_cls_with_init.fused_worker_used # if a WorkerGroup is spawned from Colocate WorkerGroup, this indicates which sub-class is binded to # this WorkerGroup. @@ -298,22 +358,33 @@ def __init__( self.device_name = device_name self.profile_steps = kwargs.get("profile_steps", None) self.worker_nsight_options = kwargs.get("worker_nsight_options", None) - if self.worker_nsight_options is not None and self.worker_nsight_options["capture-range-end"] is None: - self.worker_nsight_options["capture-range-end"] = f"repeat-shutdown:{6 * len(self.profile_steps)}" + if ( + self.worker_nsight_options is not None + and self.worker_nsight_options["capture-range-end"] is None + ): + self.worker_nsight_options["capture-range-end"] = ( + f"repeat-shutdown:{6 * len(self.profile_steps)}" + ) if worker_names is not None and (not self.fused_worker_used): 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, worker_handles=worker_handles) + self._init_with_detached_workers( + worker_names=worker_names, worker_handles=worker_handles + ) else: self._init_with_resource_pool( - resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, bin_pack=bin_pack, detached=detached + 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) + self._bind_worker_method( + self.ray_cls_with_init.cls, func_generator) self.wg_dict = None self.method_names = [] @@ -328,18 +399,28 @@ def _is_worker_alive(self, worker: ray.actor.ActorHandle): bool: True if the worker is alive, False otherwise """ 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 + 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, worker_handles): # ray.get_actor holds a weak reference to the actor, which causes actors garbage collected unexpectedly # if we only hold spawn RayWorkerGroup. By passing actor handle explicitly, spawn RayWorkerGroup have # strong reference to these actors. # https://github.com/ray-project/ray/pull/45699 - workers = worker_handles if worker_handles else [ray.get_actor(name=name) for name in worker_names] + workers = ( + worker_handles + if worker_handles + else [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): + def _init_with_resource_pool( + self, resource_pool, ray_cls_with_init, bin_pack, detached + ): """Initialize the worker group by creating new workers from a resource pool. Args: @@ -353,7 +434,9 @@ def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, d strategy = "PACK" if bin_pack: strategy = "STRICT_PACK" - pgs = resource_pool.get_placement_groups(strategy=strategy, device_name=self.device_name) + pgs = resource_pool.get_placement_groups( + strategy=strategy, device_name=self.device_name + ) world_size = resource_pool.world_size self._world_size = world_size # cia.add_kwarg("_world_size", world_size) @@ -362,11 +445,14 @@ def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, d rank = -1 local_world_size = resource_pool.store[0] for pg_idx, pg in enumerate(sort_placement_group_by_node_ip(pgs)): - assert local_world_size <= pg.bundle_count, f"when generating for {self.name_prefix}, for the " + 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 + # 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), @@ -382,9 +468,14 @@ def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, d 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 + 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" + # e.g. Worker_2:5 + name = f"{self.name_prefix}{cia_name}_{pg_idx}:{local_rank}" if self.profile_steps and self.device_name == "cuda": ray_cls_with_init.update_options( @@ -397,7 +488,9 @@ def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, d } ) else: - ray_cls_with_init.update_options({"runtime_env": {"env_vars": env_vars}, "name": name}) + ray_cls_with_init.update_options( + {"runtime_env": {"env_vars": env_vars}, "name": name} + ) if detached: ray_cls_with_init.update_options({"lifetime": "detached"}) @@ -418,7 +511,10 @@ def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, d actor_name = f"{self.name_prefix}_register_center" start_time = time.time() - while time.time() - start_time < self._ray_wait_register_center_timeout: + while ( + time.time() - start_time + < self._ray_wait_register_center_timeout + ): if actor_name in list_named_actors(): register_center_actor = ray.get_actor(actor_name) break @@ -427,11 +523,7 @@ def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, d if elapsed % 30 == 0: logging.warning( "Waiting for register center actor %s to be ready. Elapsed time: %s seconds out of " - "%s seconds.", - actor_name, - elapsed, - self._ray_wait_register_center_timeout, - ) + "%s seconds.", actor_name, elapsed, self._ray_wait_register_center_timeout, ) time.sleep(1) if register_center_actor is None: @@ -445,8 +537,13 @@ def _init_with_resource_pool(self, resource_pool, ray_cls_with_init, bin_pack, d "`trainer.ray_wait_register_center_timeout`." ) - 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"] + 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}") @@ -530,7 +627,9 @@ def spawn_fused(self, prefix_set): wg_dict = dict() for key in prefix_set: new_wg = deepcopy(self) - new_wg._bind_worker_method(self.ray_cls_with_init.cls.raw_cls_dict[key], func_generator) + new_wg._bind_worker_method( + self.ray_cls_with_init.cls.raw_cls_dict[key], func_generator + ) new_wg.sub_cls_name = key wg_dict[key] = new_wg return wg_dict @@ -545,9 +644,16 @@ def fuse(self, prefix_set): self.wg_dict = self.spawn(prefix_set) for role_name, role_wg in self.wg_dict.items(): setattr(self, role_name, role_wg) - self.method_names = self._bind_worker_method(self.ray_cls_with_init.cls, func_generator) + self.method_names = self._bind_worker_method( + self.ray_cls_with_init.cls, func_generator + ) - def _execute_remote_single_worker(self, worker, method_name: str, *args, **kwargs): + def _execute_remote_single_worker( + self, + worker, + method_name: str, + *args, + **kwargs): """Execute a method on a single worker remotely. Args: @@ -561,7 +667,9 @@ def _execute_remote_single_worker(self, worker, method_name: str, *args, **kwarg """ if self.fused_worker_used and method_name not in self.method_names: remote_call = getattr(worker, self.fused_worker_execute_fn_name) - return remote_call.remote(f"{self.sub_cls_name}_fwmn_{method_name}", *args, **kwargs) + return remote_call.remote( + f"{self.sub_cls_name}_fwmn_{method_name}", *args, **kwargs + ) # fused worker not used remote_call = getattr(worker, method_name) return remote_call.remote(*args, **kwargs) @@ -577,7 +685,9 @@ def execute_rank_zero_sync(self, method_name: str, *args, **kwargs): Returns: Result of the method execution """ - return ray.get(self.execute_rank_zero_async(method_name, *args, **kwargs)) + return ray.get( + self.execute_rank_zero_async( + method_name, *args, **kwargs)) def execute_rank_zero_async(self, method_name: str, *args, **kwargs): """Execute a method on rank zero worker asynchronously. @@ -590,7 +700,9 @@ def execute_rank_zero_async(self, method_name: str, *args, **kwargs): Returns: Remote object reference to the method execution """ - return self._execute_remote_single_worker(self._workers[0], method_name, *args, **kwargs) + return self._execute_remote_single_worker( + self._workers[0], method_name, *args, **kwargs + ) def execute_rank_zero(self, method_name: str, *args, **kwargs): """Alias for execute_rank_zero_async. @@ -647,19 +759,29 @@ def execute_all_async(self, method_name: str, *args, **kwargs): # element in these lists to the corresponding 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()): + 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()} result.append( - self._execute_remote_single_worker(self._workers[i], method_name, *sliced_args, **sliced_kwargs) - ) + self._execute_remote_single_worker( + self._workers[i], + method_name, + *sliced_args, + **sliced_kwargs)) return result - return [self._execute_remote_single_worker(worker, method_name, *args, **kwargs) for worker in self._workers] + return [ + self._execute_remote_single_worker(worker, method_name, *args, **kwargs) + for worker in self._workers + ] @property def master_address(self): @@ -694,9 +816,12 @@ def _bind_workers_method_to_parent(cls, key, user_defined_cls): 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" + assert callable( + method + ), f"{method_name} in {user_defined_cls} is not callable" except Exception: - # if it is a property, it will fail because Class doesn't have instance property + # if it is a property, it will fail because Class doesn't have + # instance property continue if hasattr(method, MAGIC_ATTR): @@ -704,13 +829,17 @@ def _bind_workers_method_to_parent(cls, key, user_defined_cls): def generate_function(name, key=key): def func(self, *args, **kwargs): # dispatch to the actual worker - return getattr(self.worker_dict[key], name)(*args, **kwargs) + return getattr( + self.worker_dict[key], name)( + *args, **kwargs) async def async_func(self, *args, **kwargs): # dispatch to the actual worker return await getattr(self.worker_dict[key], name)(*args, **kwargs) - wrapper = async_func if inspect.iscoroutinefunction(method) else func # noqa: B023 + wrapper = ( + async_func if inspect.iscoroutinefunction(method) else func + ) # noqa: B023 return wrapper @@ -720,17 +849,22 @@ async def async_func(self, *args, **kwargs): setattr(func, MAGIC_ATTR, attrs) try: # bind direct rollout method to class without prefix - if attrs["dispatch_mode"] == Dispatch.DIRECT_ROLLOUT_METHOD and "rollout" in key: - assert not hasattr(cls, method_name), ( - f"conflict direct rollout method {method_name} with role {key}" - ) + if ( + attrs["dispatch_mode"] == Dispatch.DIRECT_ROLLOUT_METHOD + and "rollout" in key + ): + assert not hasattr( + cls, method_name + ), f"conflict direct rollout method {method_name} with role {key}" setattr(cls, method_name, func) - print(f"bind role {key} method {method_name} to class {cls}") + print( + f"bind role {key} method {method_name} to class {cls}") else: method_name_with_prefix = key + "_" + method_name setattr(cls, method_name_with_prefix, func) except Exception as e: - raise ValueError(f"Fail to set method_name {method_name}") from e + raise ValueError( + f"Fail to set method_name {method_name}") from e def _unwrap_ray_remote(cls): @@ -763,7 +897,9 @@ def create_colocated_worker_cls(class_dict: dict[str, RayClassWithInitArgs]): worker_cls = _determine_fsdp_megatron_base_class( [cls.cls.__ray_actor_class__.__mro__ for cls in class_dict.values()] ) - assert issubclass(worker_cls, Worker), f"worker_cls {worker_cls} should be a subclass of Worker" + assert issubclass( + worker_cls, Worker + ), f"worker_cls {worker_cls} should be a subclass of Worker" print(f"colocated worker base class {worker_cls}") for key, cls in class_dict.items(): @@ -784,7 +920,8 @@ def __init__(self): # when DISABLE_WORKER_INIT == 1 it will return immediately 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", {}) + *init_args_dict[key].get("args", ()), + **init_args_dict[key].get("kwargs", {}), ) # now monkey-patch the methods from inner class to WorkerDict @@ -800,7 +937,8 @@ def __init__(self): FusedWorkerCLSName = "FusedWorker" -def create_colocated_worker_raw_cls(class_dict: dict[str, RayClassWithInitArgs]): +def create_colocated_worker_raw_cls( + class_dict: dict[str, RayClassWithInitArgs]): """ This function returns a FusedWorker class. @@ -818,9 +956,16 @@ def create_colocated_worker_raw_cls(class_dict: dict[str, RayClassWithInitArgs]) The same as `FusedWorker.fused_worker_dict`, enables underlying class to access other underlying classes. """ - raw_cls_dict = {cls_name: _unwrap_ray_remote(cia.cls) for cls_name, cia in class_dict.items()} - init_args_dict = {cls_name: cia.args for cls_name, cia in class_dict.items()} - init_kwargs_dict = {cls_name: cia.kwargs for cls_name, cia in class_dict.items()} + raw_cls_dict = { + cls_name: _unwrap_ray_remote( + cia.cls) for cls_name, + cia in class_dict.items()} + init_args_dict = { + cls_name: cia.args for cls_name, + cia in class_dict.items()} + init_kwargs_dict = { + cls_name: cia.kwargs for cls_name, + cia in class_dict.items()} cls_names = list(class_dict.keys()) # FusedWorker_Actor_Critic @@ -842,15 +987,25 @@ def __init__(self, *args, **kwargs): strict=True, ): with patch.dict(os.environ, {"DISABLE_WORKER_INIT": "1"}): - udc._get_ray_actor_cls_name = lambda x, name_renamed=class_name_renamed: name_renamed - udc._get_ray_method_prefix = lambda x, name_prefixed=cls_name: f"{name_prefixed}_" - # cls_name = "actor", "critic", udc = ActorWorker, CriticWorker - self.fused_worker_dict[cls_name] = udc(*ud_args, **ud_kwargs) + udc._get_ray_actor_cls_name = ( + lambda x, name_renamed=class_name_renamed: name_renamed + ) + udc._get_ray_method_prefix = ( + lambda x, name_prefixed=cls_name: f"{name_prefixed}_" + ) + # cls_name = "actor", "critic", udc = ActorWorker, + # CriticWorker + self.fused_worker_dict[cls_name] = udc( + *ud_args, **ud_kwargs) setattr(self, cls_name, self.fused_worker_dict[cls_name]) - # injecting fused_worker to each sub worker so they can be aware of existence of each other + # injecting fused_worker to each sub worker so they can be aware of + # existence of each other for _, worker in self.fused_worker_dict.items(): - setattr(worker, Worker.fused_worker_attr_name, self.fused_worker_dict) + setattr( + worker, + Worker.fused_worker_attr_name, + self.fused_worker_dict) def _fuw_execute(self, method_name: str, *args, **kwargs): # for fused_worker, method_name is in a form of "{cls_name}_fwmn_{method_name}" @@ -859,9 +1014,9 @@ def _fuw_execute(self, method_name: str, *args, **kwargs): cls_name = names[0] method_name = names[1] - assert cls_name in self.fused_worker_dict, ( - f"calling {cls_name}'s {method_name}, but {cls_name} not in fused_worker_dict" - ) + assert ( + cls_name in self.fused_worker_dict + ), f"calling {cls_name}'s {method_name}, but {cls_name} not in fused_worker_dict" udc_method = getattr(self.fused_worker_dict[cls_name], method_name) return udc_method(*args, **kwargs) @@ -872,7 +1027,8 @@ def _fuw_execute(self, method_name: str, *args, **kwargs): return renamed_fused_worker_cls -def create_colocated_worker_cls_fused(class_dict: dict[str, RayClassWithInitArgs]): +def create_colocated_worker_cls_fused( + class_dict: dict[str, RayClassWithInitArgs]): """ This function returns a RayClassWithInitArgs instance of FusedWorker, which is an replacement of `create_colocated_worker_cls`. WorkerGroup constructed using this class will be a colocated diff --git a/Agent0/executor_train/verl/verl/single_controller/ray/megatron.py b/Agent0/executor_train/verl/verl/single_controller/ray/megatron.py index b46fe44..f881455 100644 --- a/Agent0/executor_train/verl/verl/single_controller/ray/megatron.py +++ b/Agent0/executor_train/verl/verl/single_controller/ray/megatron.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -29,7 +29,12 @@ class NVMegatronRayWorkerGroup(RayWorkerGroup, MegatronWorkerGroup): so that the dispatcher can use it to dispatch data. """ - def __init__(self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWithInitArgs, **kwargs): + def __init__( + self, + resource_pool: RayResourcePool, + ray_cls_with_init: RayClassWithInitArgs, + **kwargs, + ): """ Initialize the NVMegatronRayWorkerGroup. @@ -38,8 +43,13 @@ def __init__(self, resource_pool: RayResourcePool, ray_cls_with_init: RayClassWi ray_cls_with_init (RayClassWithInitArgs): The Ray class with initialization arguments **kwargs: Additional keyword arguments to pass to the parent class """ - 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") + 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") ) @@ -65,7 +75,9 @@ def __init__( **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_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") ) @@ -74,4 +86,7 @@ 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) + self.execute_all_sync( + method_name="init_megatron", + default_megatron_kwargs=default_megatron_kwargs, + ) diff --git a/Agent0/executor_train/verl/verl/third_party/__init__.py b/Agent0/executor_train/verl/verl/third_party/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/third_party/__init__.py +++ b/Agent0/executor_train/verl/verl/third_party/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/third_party/sglang/__init__.py b/Agent0/executor_train/verl/verl/third_party/sglang/__init__.py index 15593ca..55c9b80 100644 --- a/Agent0/executor_train/verl/verl/third_party/sglang/__init__.py +++ b/Agent0/executor_train/verl/verl/third_party/sglang/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2023-2024 SGLang Team +# Copyright 2023-2026 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at @@ -11,7 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/third_party/sglang/parallel_state.py b/Agent0/executor_train/verl/verl/third_party/sglang/parallel_state.py index cdec743..71d99ca 100644 --- a/Agent0/executor_train/verl/verl/third_party/sglang/parallel_state.py +++ b/Agent0/executor_train/verl/verl/third_party/sglang/parallel_state.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Copyright 2023 The SGlang team. # Adapted from # https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/parallel_state.py @@ -57,7 +57,9 @@ def initialize_parallel_state( # 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) + init_distributed_environment( + world_size, rank, distributed_init_method, local_rank, backend + ) if torch.distributed.get_world_size() > 1: # NOTE: build a separate inference group with infer tp & micro dp initialize_model_parallel_for_sglang( @@ -65,7 +67,9 @@ def initialize_parallel_state( 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) + initialize_model_parallel( + tensor_model_parallel_size, pipeline_model_parallel_size, backend + ) # NOTE(linjunrong): After init SGLang rollout using class EngineFragment, user should always remember to call @@ -84,9 +88,12 @@ def ensure_model_parallel_initialized( 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) + 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) + initialize_model_parallel( + tensor_model_parallel_size, pipeline_model_parallel_size, backend + ) return assert get_tensor_model_parallel_world_size() == tensor_model_parallel_size, ( @@ -140,7 +147,9 @@ def initialize_model_parallel_for_sglang( 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) + 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, @@ -158,15 +167,24 @@ def initialize_model_parallel_for_sglang( # 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 + 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): + 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)) + 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) @@ -197,7 +215,11 @@ def initialize_model_parallel_for_sglang( 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) + _PP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_custom_allreduce=False) ps._PP = _PP # for verl @@ -234,7 +256,9 @@ def initialize_model_parallel( # 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) + 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 @@ -251,7 +275,10 @@ def initialize_model_parallel( 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)) + 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 @@ -280,7 +307,12 @@ def initialize_model_parallel( if ps._TP is not None: _PP = ps._TP else: - _PP = init_model_parallel_group(group_ranks, get_world_group().local_rank, backend, use_custom_allreduce=False) + _PP = init_model_parallel_group( + group_ranks, + get_world_group().local_rank, + backend, + use_custom_allreduce=False, + ) ps._PP = _PP @@ -302,7 +334,8 @@ def get_device_mesh(): # NOTE(linjunrong): In the vllm version parallel_state.py. verl created its own _TP and _PP as verl want to use # the process group for some extra purpose. Under the hood, there is no difference between them and the original # one in vllm.distributed.parallel_state. However, the implementation need to hack the init process of inference -# engine, as we do not maintain another SGLang here, I just use the original _TP and _PP directly. +# engine, as we do not maintain another SGLang here, I just use the +# original _TP and _PP directly. def get_tensor_model_parallel_group(): """Get the tensor model parallel group the caller rank belongs to.""" @@ -312,7 +345,8 @@ def get_tensor_model_parallel_group(): 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()) + return torch.distributed.get_world_size( + group=get_tensor_model_parallel_group()) def get_tensor_model_parallel_rank(): diff --git a/Agent0/executor_train/verl/verl/third_party/vllm/__init__.py b/Agent0/executor_train/verl/verl/third_party/vllm/__init__.py index 76fe51b..624b1c4 100644 --- a/Agent0/executor_train/verl/verl/third_party/vllm/__init__.py +++ b/Agent0/executor_train/verl/verl/third_party/vllm/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/tools/__init__.py b/Agent0/executor_train/verl/verl/tools/__init__.py index c4b932b..72375fe 100644 --- a/Agent0/executor_train/verl/verl/tools/__init__.py +++ b/Agent0/executor_train/verl/verl/tools/__init__.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/tools/base_tool.py b/Agent0/executor_train/verl/verl/tools/base_tool.py index 9a1189d..6cff2d3 100644 --- a/Agent0/executor_train/verl/verl/tools/base_tool.py +++ b/Agent0/executor_train/verl/verl/tools/base_tool.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -38,7 +38,13 @@ def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): self.tool_schema = tool_schema or self.get_openai_tool_schema() assert self.tool_schema is not None, "Tool schema is not set!" self.name = self.tool_schema.function.name - print(json.dumps(self.tool_schema.model_dump(exclude_unset=True, exclude_none=True), indent=2)) + print( + json.dumps( + self.tool_schema.model_dump( + exclude_unset=True, + exclude_none=True), + indent=2, + )) def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: return self.tool_schema @@ -58,7 +64,9 @@ async def create(self, instance_id: Optional[str] = None, **kwargs) -> str: return instance_id @rollout_trace_op - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: """Execute the tool. Args: diff --git a/Agent0/executor_train/verl/verl/tools/geo3k_tool.py b/Agent0/executor_train/verl/verl/tools/geo3k_tool.py index 6ffd6fb..63a6d30 100644 --- a/Agent0/executor_train/verl/verl/tools/geo3k_tool.py +++ b/Agent0/executor_train/verl/verl/tools/geo3k_tool.py @@ -1,6 +1,6 @@ -# Copyright 2023-2025 SGLang Team +# Copyright 2023-2026 SGLang Team # Copyright Amazon.com, Inc. or its affiliates. -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -64,7 +64,12 @@ def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: return self.tool_schema - async def create(self, instance_id: Optional[str] = None, ground_truth: Optional[str] = None, **kwargs) -> str: + async def create( + self, + instance_id: Optional[str] = None, + ground_truth: Optional[str] = None, + **kwargs, + ) -> str: if instance_id is None: instance_id = str(uuid4()) self._instance_dict[instance_id] = { @@ -75,14 +80,18 @@ async def create(self, instance_id: Optional[str] = None, ground_truth: Optional return instance_id, None @rollout_trace_op - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: answer = parameters.get("answer", "") if not isinstance(answer, str): answer = str(answer) self._instance_dict[instance_id]["response"] = answer reward = await self.calc_reward(instance_id) # penalty for non improved answer submission - tool_reward = 0.0 if reward > self._instance_dict[instance_id]["reward"] else -0.05 + tool_reward = ( + 0.0 if reward > self._instance_dict[instance_id]["reward"] else - + 0.05) # update the reward self._instance_dict[instance_id]["reward"] = reward return f"Current parsed {answer=} {reward=}", tool_reward, {} diff --git a/Agent0/executor_train/verl/verl/tools/gsm8k_tool.py b/Agent0/executor_train/verl/verl/tools/gsm8k_tool.py index f6d8913..c4a4d67 100644 --- a/Agent0/executor_train/verl/verl/tools/gsm8k_tool.py +++ b/Agent0/executor_train/verl/verl/tools/gsm8k_tool.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -64,7 +64,12 @@ def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: return self.tool_schema - async def create(self, instance_id: Optional[str] = None, ground_truth: Optional[str] = None, **kwargs) -> str: + async def create( + self, + instance_id: Optional[str] = None, + ground_truth: Optional[str] = None, + **kwargs, + ) -> str: if instance_id is None: instance_id = str(uuid4()) self._instance_dict[instance_id] = { @@ -75,7 +80,9 @@ async def create(self, instance_id: Optional[str] = None, ground_truth: Optional return instance_id @rollout_trace_op - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: answer = parameters.get("answer", "") if not isinstance(answer, str): answer = str(answer) @@ -87,7 +94,9 @@ async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) reward = await self.calc_reward(instance_id) # penalty for non improved answer submission - tool_reward = 0.0 if reward > self._instance_dict[instance_id]["reward"] else -0.05 + tool_reward = ( + 0.0 if reward > self._instance_dict[instance_id]["reward"] else - + 0.05) # update the reward self._instance_dict[instance_id]["reward"] = reward diff --git a/Agent0/executor_train/verl/verl/tools/mcp_base_tool.py b/Agent0/executor_train/verl/verl/tools/mcp_base_tool.py index dacd18e..c72724e 100644 --- a/Agent0/executor_train/verl/verl/tools/mcp_base_tool.py +++ b/Agent0/executor_train/verl/verl/tools/mcp_base_tool.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -36,7 +36,8 @@ def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): self._instance_dict = {} self.timeout = config.get("timeout", 30) - # TODO(hechanghao): create a global client manager to manage the rate limit, client and pool + # TODO(hechanghao): create a global client manager to manage the rate + # limit, client and pool logger.info(f"Initialized MCPBaseTool with config: {config}") def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: @@ -63,7 +64,9 @@ async def create(self, instance_id: Optional[str] = None, **kwargs) -> str: async def _call_tool(self, instance_id, parameters) -> tuple[str, dict]: err_msg = "" try: - call_tool_result = await ClientManager.call_tool(self.name, parameters, self.timeout) + call_tool_result = await ClientManager.call_tool( + self.name, parameters, self.timeout + ) except ClientError as e: err_msg = f"\n Tool call failed: {e}" except ConnectionError as e: @@ -71,23 +74,31 @@ async def _call_tool(self, instance_id, parameters) -> tuple[str, dict]: except Exception as e: err_msg = f"\n An unexpected error occurred: {e}" - logger.debug(f"Tool result for instance {instance_id} with tool {self.name}: {call_tool_result.content}") + logger.debug( + f"Tool result for instance {instance_id} with tool { + self.name}: { + call_tool_result.content}") result, metadata = self._parse_tool_result(call_tool_result.content) metadata["api_request_error"] += err_msg return result, metadata @rollout_trace_op - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: if self.name == "" or self.name is None or parameters is None: error_msg = "Error: 'parameters' is missing or empty." - logger.error(f"[MCPTool] {error_msg} Received tool name: {self.name}, parameters: {parameters}") + logger.error( + f"[MCPTool] {error_msg} Received tool name: { + self.name}, parameters: {parameters}") return json.dumps({"result": error_msg}), 0.0, {} try: result_text, metadata = await self._call_tool(instance_id, parameters) # Store results in instance dictionary - self._instance_dict[instance_id]["reward"].append(result_text.strip()) + self._instance_dict[instance_id]["reward"].append( + result_text.strip()) # Convert metadata to metrics metrics = { @@ -100,7 +111,8 @@ async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) return result_text, 0.0, metrics except Exception as e: - error_result = json.dumps({"result": f"Tool execution failed: {e}"}) + error_result = json.dumps( + {"result": f"Tool execution failed: {e}"}) logger.error(f"[MCPBaseTool] Execution failed: {e}") return error_result, 0.0, {"error": str(e)} @@ -112,5 +124,7 @@ async def release(self, instance_id: str, **kwargs) -> None: del self._instance_dict[instance_id] def _parse_tool_result(self, content: list) -> tuple[str, dict]: - tools_content = [part.text for part in filter(lambda x: x.type == "text", content)] + tools_content = [ + part.text for part in filter(lambda x: x.type == "text", content) + ] return " ".join(tools_content), {} diff --git a/Agent0/executor_train/verl/verl/tools/mcp_search_tool.py b/Agent0/executor_train/verl/verl/tools/mcp_search_tool.py index ac82371..b6fe7d5 100644 --- a/Agent0/executor_train/verl/verl/tools/mcp_search_tool.py +++ b/Agent0/executor_train/verl/verl/tools/mcp_search_tool.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -50,8 +50,10 @@ def _parse_tool_result(self, content: list) -> tuple[str, dict]: title_matches = re.findall(r'"title"\s*:', text) title_count = len(title_matches) - results_match = re.search(r'"results"\s*:\s*(\[.*?\])', text, re.DOTALL) - results_content = results_match.group(1) if results_match else "" + results_match = re.search( + r'"results"\s*:\s*(\[.*?\])', text, re.DOTALL) + results_content = results_match.group( + 1) if results_match else "" res += results_content res_cnt += title_count diff --git a/Agent0/executor_train/verl/verl/tools/sandbox_fusion_tools.py b/Agent0/executor_train/verl/verl/tools/sandbox_fusion_tools.py index c3a2748..4b5206a 100644 --- a/Agent0/executor_train/verl/verl/tools/sandbox_fusion_tools.py +++ b/Agent0/executor_train/verl/verl/tools/sandbox_fusion_tools.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -63,12 +63,15 @@ def get_current_count(self): class ExecutionWorker: def __init__(self, enable_global_rate_limit=True, rate_limit=10): - self.rate_limit_worker = self._init_rate_limit(rate_limit) if enable_global_rate_limit else None + self.rate_limit_worker = (self._init_rate_limit( + rate_limit) if enable_global_rate_limit else None) def _init_rate_limit(self, rate_limit): # TODO validation for rate_limit # A Singleton Rate Limitor - return TokenBucketWorker.options(name="rate-limiter", get_if_exists=True).remote(rate_limit) + return TokenBucketWorker.options( + name="rate-limiter", get_if_exists=True + ).remote(rate_limit) def ping(self): return True @@ -85,14 +88,17 @@ def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T: def init_execution_pool( - num_workers: int, enable_global_rate_limit=True, rate_limit=10, mode: PoolMode = PoolMode.ThreadMode + num_workers: int, + enable_global_rate_limit=True, + rate_limit=10, + mode: PoolMode = PoolMode.ThreadMode, ): if mode == PoolMode.ThreadMode: return ( - ray.remote(ExecutionWorker) - .options(max_concurrency=num_workers) - .remote(enable_global_rate_limit=enable_global_rate_limit, rate_limit=rate_limit) - ) + ray.remote(ExecutionWorker) .options( + max_concurrency=num_workers) .remote( + enable_global_rate_limit=enable_global_rate_limit, + rate_limit=rate_limit)) else: raise NotImplementedError("Process mode is not implemented yet") # return ray.util.multiprocessing.Pool(processes=num_workers) @@ -135,7 +141,8 @@ def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): self.rate_limit = config.get("rate_limit", 10) self.default_timeout = config.get("default_timeout", 30) self.default_language = config.get("default_language", "python") - self.enable_global_rate_limit = config.get("enable_global_rate_limit", True) + self.enable_global_rate_limit = config.get( + "enable_global_rate_limit", True) self.execution_pool = init_execution_pool( num_workers=self.num_workers, enable_global_rate_limit=self.enable_global_rate_limit, @@ -152,7 +159,12 @@ def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: return self.tool_schema - async def create(self, instance_id: Optional[str] = None, ground_truth: Optional[str] = None, **kwargs) -> str: + async def create( + self, + instance_id: Optional[str] = None, + ground_truth: Optional[str] = None, + **kwargs, + ) -> str: if instance_id is None: instance_id = str(uuid4()) self._instance_dict[instance_id] = { @@ -163,25 +175,37 @@ async def create(self, instance_id: Optional[str] = None, ground_truth: Optional return instance_id @rollout_trace_op - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: code = parameters.get("code", "") timeout = parameters.get("timeout", self.default_timeout) language = parameters.get("language", self.default_language) if not isinstance(code, str): code = str(code) - result = await self.execution_pool.execute.remote(self.execute_code, instance_id, code, timeout, language) + result = await self.execution_pool.execute.remote( + self.execute_code, instance_id, code, timeout, language + ) # sandbox has no score or metrics, use Nones return result, None, None def execute_code(self, instance_id, code, timeout=30, language="python"): result_status, metadata = _process_single_case( - 0, None, None, self.sandbox_fusion_url, code, timeout, self.memory_limit_mb, language + 0, + None, + None, + self.sandbox_fusion_url, + code, + timeout, + self.memory_limit_mb, + language, ) # we should always expect this since we don't have correct answer if metadata["run_status"] == "Finished": actual_output = metadata["stdout"] + metadata["stderr"] - logger.debug(f"actual_output from sandbox fusion: {actual_output},{instance_id}") + logger.debug( + f"actual_output from sandbox fusion: {actual_output},{instance_id}") return actual_output else: return "no stdout here" diff --git a/Agent0/executor_train/verl/verl/tools/schemas.py b/Agent0/executor_train/verl/verl/tools/schemas.py index c0c65a3..f8f4f30 100644 --- a/Agent0/executor_train/verl/verl/tools/schemas.py +++ b/Agent0/executor_train/verl/verl/tools/schemas.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -73,12 +73,18 @@ def from_openai_function_parsed_schema( except json.JSONDecodeError: arguments = {} has_decode_error = True - # If the arguments is not a dict, it means the arguments is not a valid JSON string + # If the arguments is not a dict, it means the arguments is not a valid + # JSON string if not isinstance(arguments, dict): arguments = {} has_decode_error = True - return OpenAIFunctionCallSchema(name=parsed_schema.name, arguments=arguments), has_decode_error + return ( + OpenAIFunctionCallSchema( + name=parsed_schema.name, + arguments=arguments), + has_decode_error, + ) class OpenAIFunctionToolCall(BaseModel): diff --git a/Agent0/executor_train/verl/verl/tools/search_tool.py b/Agent0/executor_train/verl/verl/tools/search_tool.py index 3cc6cda..951fa0b 100644 --- a/Agent0/executor_train/verl/verl/tools/search_tool.py +++ b/Agent0/executor_train/verl/verl/tools/search_tool.py @@ -1,278 +1,308 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import json -import logging -import os -import threading -from contextlib import ExitStack -from enum import Enum -from typing import Any, Callable, Optional, TypeVar -from uuid import uuid4 - -import ray -import ray.actor - -from verl.tools.utils.search_r1_like_utils import perform_single_search_batch -from verl.utils.rollout_trace import rollout_trace_op - -from .base_tool import BaseTool -from .schemas import OpenAIFunctionToolSchema - -logger = logging.getLogger(__name__) -logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) - -T = TypeVar("T") - - -# Adapted from verl/tools/sandbox_fusion_tools.py -class PoolMode(Enum): - """Execution pool mode enumeration.""" - - ThreadMode = 1 - ProcessMode = 2 - - -@ray.remote(concurrency_groups={"acquire": 1, "release": 10}) -class TokenBucketWorker: - """Ray actor for rate limiting using token bucket algorithm.""" - - def __init__(self, rate_limit: int): - self.rate_limit = rate_limit - self.current_count = 0 # For observability - self._semaphore = threading.Semaphore(rate_limit) - - @ray.method(concurrency_group="acquire") - def acquire(self): - """Acquire a token from the bucket.""" - self._semaphore.acquire() - self.current_count += 1 - - @ray.method(concurrency_group="release") - def release(self): - """Release a token back to the bucket.""" - self._semaphore.release() - self.current_count -= 1 - - def get_current_count(self): - """Get current number of acquired tokens.""" - return self.current_count - - -class SearchExecutionWorker: - """Worker for executing search operations with optional rate limiting.""" - - def __init__(self, enable_global_rate_limit=True, rate_limit=10): - self.rate_limit_worker = self._init_rate_limit(rate_limit) if enable_global_rate_limit else None - - def _init_rate_limit(self, rate_limit): - """Initialize singleton rate limiter.""" - return TokenBucketWorker.options(name="rate-limiter", get_if_exists=True).remote(rate_limit) - - def ping(self): - """Health check method.""" - return True - - def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T: - """Execute function with optional rate limiting.""" - if self.rate_limit_worker: - with ExitStack() as stack: - stack.callback(self.rate_limit_worker.release.remote) - ray.get(self.rate_limit_worker.acquire.remote()) - try: - return fn(*fn_args, **fn_kwargs) - except Exception as e: - # TODO we should make this available to the tool caller - logger.warning(f"Error when executing search: {e}") - else: - return fn(*fn_args, **fn_kwargs) - - -def init_search_execution_pool( - num_workers: int, enable_global_rate_limit=True, rate_limit=10, mode: PoolMode = PoolMode.ThreadMode -): - """Initialize search execution pool.""" - if mode == PoolMode.ThreadMode: - return ( - ray.remote(SearchExecutionWorker) - .options(max_concurrency=num_workers) - .remote(enable_global_rate_limit=enable_global_rate_limit, rate_limit=rate_limit) - ) - else: - raise NotImplementedError("Process mode is not implemented yet") - - -class SearchTool(BaseTool): - """Search tool for retrieving information using external retrieval services. - - This tool provides search functionality with rate limiting and concurrent execution - support through Ray. It integrates with external retrieval services to perform - semantic search operations. - - Methods: - get_openai_tool_schema: Return the tool schema in OpenAI format - create: Create a tool instance for a trajectory - execute: Execute the search tool - calc_reward: Calculate the reward with respect to tool state - release: Release the tool instance - """ - - def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): - """Initialize SearchTool with configuration and schema. - - Args: - config: Configuration dictionary containing tool settings - tool_schema: OpenAI function tool schema definition - - Example tool_schema: - { - "type": "function", - "function": { - "name": "search", - "description": "Searches for relevant information based on queries.", - "parameters": { - "type": "object", - "properties": { - "query_list": { - "type": "array", - "items": {"type": "string"}, - "description": "List of search queries" - } - }, - "required": ["query_list"] - } - } - } - """ - super().__init__(config, tool_schema) - self._instance_dict = {} - - # Worker and rate limiting configuration - self.num_workers = config.get("num_workers", 120) - self.rate_limit = config.get("rate_limit", 120) - self.timeout = config.get("timeout", 30) - - self.enable_global_rate_limit = config.get("enable_global_rate_limit", True) - self.execution_pool = init_search_execution_pool( - num_workers=self.num_workers, - enable_global_rate_limit=self.enable_global_rate_limit, - rate_limit=self.rate_limit, - mode=PoolMode.ThreadMode, - ) - - # Retrieval service configuration - self.retrieval_service_url = config.get("retrieval_service_url") - assert self.retrieval_service_url, "Configuration must include 'retrieval_service_url'" - self.topk = config.get("topk", 3) - if self.retrieval_service_url == "": - raise ValueError("retrieval_service_url is not set") - - logger.info(f"Initialized SearchTool with config: {config}") - - def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: - """Return the OpenAI tool schema.""" - return self.tool_schema - - async def create(self, instance_id: Optional[str] = None, **kwargs) -> str: - """Create a tool instance. - - Args: - instance_id: The instance id of the tool. - - Returns: - The instance id of the tool. - """ - if instance_id is None: - instance_id = str(uuid4()) - self._instance_dict[instance_id] = { - "response": "", - "reward": [], - } - return instance_id - - def execute_search(self, instance_id: str, query_list: list, retrieval_service_url: str, topk: int, timeout: int): - """Execute search operation using retrieval service. - - Args: - instance_id: Tool instance ID - query_list: List of search queries - retrieval_service_url: URL of the retrieval service - topk: Number of top results to return - timeout: Request timeout in seconds - - Returns: - Tuple of (result_text, metadata) - """ - result_text, metadata = perform_single_search_batch( - retrieval_service_url=retrieval_service_url, - query_list=query_list, - topk=topk, - concurrent_semaphore=None, # Ray handles concurrency control - timeout=timeout, - ) - logger.debug(f"Search result for instance {instance_id}: {result_text}") - return result_text, metadata - - @rollout_trace_op - async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[str, float, dict]: - """Execute the search tool. - - Args: - instance_id: The instance ID of the tool - parameters: Tool parameters containing query_list and optional timeout - - Returns: tool_response, tool_reward_score, tool_metrics - tool_response: The response str of the tool. - tool_reward_score: The step reward score of the tool. - tool_metrics: The metrics of the tool. - """ - timeout = self.timeout - query_list_from_params = parameters.get("query_list") - - if not query_list_from_params or not isinstance(query_list_from_params, list): - error_msg = "Error: 'query_list' is missing, empty, or not a list in parameters." - logger.error(f"[SearchTool] {error_msg} Received parameters: {parameters}") - return json.dumps({"result": error_msg}), 0.0, {} - - # Execute search using Ray execution pool - try: - result_text, metadata = await self.execution_pool.execute.remote( - self.execute_search, instance_id, query_list_from_params, self.retrieval_service_url, self.topk, timeout - ) - - # Store results in instance dictionary - self._instance_dict[instance_id]["reward"].append(result_text.strip()) - - # Convert metadata to metrics - metrics = { - "query_count": metadata.get("query_count", 0), - "status": metadata.get("status", "unknown"), - "total_results": metadata.get("total_results", 0), - "api_request_error": metadata.get("api_request_error"), - } - - return result_text, 0.0, metrics - - except Exception as e: - error_result = json.dumps({"result": f"Search execution failed: {e}"}) - logger.error(f"[SearchTool] Execution failed: {e}") - return error_result, 0.0, {"error": str(e)} - - async def calc_reward(self, instance_id: str, **kwargs) -> str: - return self._instance_dict[instance_id]["reward"] - - async def release(self, instance_id: str, **kwargs) -> None: - if instance_id in self._instance_dict: - del self._instance_dict[instance_id] +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import logging +import os +import threading +from contextlib import ExitStack +from enum import Enum +from typing import Any, Callable, Optional, TypeVar +from uuid import uuid4 + +import ray +import ray.actor + +from verl.tools.utils.search_r1_like_utils import perform_single_search_batch +from verl.utils.rollout_trace import rollout_trace_op + +from .base_tool import BaseTool +from .schemas import OpenAIFunctionToolSchema + +logger = logging.getLogger(__name__) +logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) + +T = TypeVar("T") + + +# Adapted from verl/tools/sandbox_fusion_tools.py +class PoolMode(Enum): + """Execution pool mode enumeration.""" + + ThreadMode = 1 + ProcessMode = 2 + + +@ray.remote(concurrency_groups={"acquire": 1, "release": 10}) +class TokenBucketWorker: + """Ray actor for rate limiting using token bucket algorithm.""" + + def __init__(self, rate_limit: int): + self.rate_limit = rate_limit + self.current_count = 0 # For observability + self._semaphore = threading.Semaphore(rate_limit) + + @ray.method(concurrency_group="acquire") + def acquire(self): + """Acquire a token from the bucket.""" + self._semaphore.acquire() + self.current_count += 1 + + @ray.method(concurrency_group="release") + def release(self): + """Release a token back to the bucket.""" + self._semaphore.release() + self.current_count -= 1 + + def get_current_count(self): + """Get current number of acquired tokens.""" + return self.current_count + + +class SearchExecutionWorker: + """Worker for executing search operations with optional rate limiting.""" + + def __init__(self, enable_global_rate_limit=True, rate_limit=10): + self.rate_limit_worker = (self._init_rate_limit( + rate_limit) if enable_global_rate_limit else None) + + def _init_rate_limit(self, rate_limit): + """Initialize singleton rate limiter.""" + return TokenBucketWorker.options( + name="rate-limiter", get_if_exists=True + ).remote(rate_limit) + + def ping(self): + """Health check method.""" + return True + + def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T: + """Execute function with optional rate limiting.""" + if self.rate_limit_worker: + with ExitStack() as stack: + stack.callback(self.rate_limit_worker.release.remote) + ray.get(self.rate_limit_worker.acquire.remote()) + try: + return fn(*fn_args, **fn_kwargs) + except Exception as e: + # TODO we should make this available to the tool caller + logger.warning(f"Error when executing search: {e}") + else: + return fn(*fn_args, **fn_kwargs) + + +def init_search_execution_pool( + num_workers: int, + enable_global_rate_limit=True, + rate_limit=10, + mode: PoolMode = PoolMode.ThreadMode, +): + """Initialize search execution pool.""" + if mode == PoolMode.ThreadMode: + return ( + ray.remote(SearchExecutionWorker) .options( + max_concurrency=num_workers) .remote( + enable_global_rate_limit=enable_global_rate_limit, + rate_limit=rate_limit)) + else: + raise NotImplementedError("Process mode is not implemented yet") + + +class SearchTool(BaseTool): + """Search tool for retrieving information using external retrieval services. + + This tool provides search functionality with rate limiting and concurrent execution + support through Ray. It integrates with external retrieval services to perform + semantic search operations. + + Methods: + get_openai_tool_schema: Return the tool schema in OpenAI format + create: Create a tool instance for a trajectory + execute: Execute the search tool + calc_reward: Calculate the reward with respect to tool state + release: Release the tool instance + """ + + def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema): + """Initialize SearchTool with configuration and schema. + + Args: + config: Configuration dictionary containing tool settings + tool_schema: OpenAI function tool schema definition + + Example tool_schema: + { + "type": "function", + "function": { + "name": "search", + "description": "Searches for relevant information based on queries.", + "parameters": { + "type": "object", + "properties": { + "query_list": { + "type": "array", + "items": {"type": "string"}, + "description": "List of search queries" + } + }, + "required": ["query_list"] + } + } + } + """ + super().__init__(config, tool_schema) + self._instance_dict = {} + + # Worker and rate limiting configuration + self.num_workers = config.get("num_workers", 120) + self.rate_limit = config.get("rate_limit", 120) + self.timeout = config.get("timeout", 30) + + self.enable_global_rate_limit = config.get( + "enable_global_rate_limit", True) + self.execution_pool = init_search_execution_pool( + num_workers=self.num_workers, + enable_global_rate_limit=self.enable_global_rate_limit, + rate_limit=self.rate_limit, + mode=PoolMode.ThreadMode, + ) + + # Retrieval service configuration + self.retrieval_service_url = config.get("retrieval_service_url") + assert ( + self.retrieval_service_url + ), "Configuration must include 'retrieval_service_url'" + self.topk = config.get("topk", 3) + if self.retrieval_service_url == "": + raise ValueError("retrieval_service_url is not set") + + logger.info(f"Initialized SearchTool with config: {config}") + + def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema: + """Return the OpenAI tool schema.""" + return self.tool_schema + + async def create(self, instance_id: Optional[str] = None, **kwargs) -> str: + """Create a tool instance. + + Args: + instance_id: The instance id of the tool. + + Returns: + The instance id of the tool. + """ + if instance_id is None: + instance_id = str(uuid4()) + self._instance_dict[instance_id] = { + "response": "", + "reward": [], + } + return instance_id + + def execute_search( + self, + instance_id: str, + query_list: list, + retrieval_service_url: str, + topk: int, + timeout: int, + ): + """Execute search operation using retrieval service. + + Args: + instance_id: Tool instance ID + query_list: List of search queries + retrieval_service_url: URL of the retrieval service + topk: Number of top results to return + timeout: Request timeout in seconds + + Returns: + Tuple of (result_text, metadata) + """ + result_text, metadata = perform_single_search_batch( + retrieval_service_url=retrieval_service_url, + query_list=query_list, + topk=topk, + concurrent_semaphore=None, # Ray handles concurrency control + timeout=timeout, + ) + logger.debug( + f"Search result for instance {instance_id}: {result_text}") + return result_text, metadata + + @rollout_trace_op + async def execute( + self, instance_id: str, parameters: dict[str, Any], **kwargs + ) -> tuple[str, float, dict]: + """Execute the search tool. + + Args: + instance_id: The instance ID of the tool + parameters: Tool parameters containing query_list and optional timeout + + Returns: tool_response, tool_reward_score, tool_metrics + tool_response: The response str of the tool. + tool_reward_score: The step reward score of the tool. + tool_metrics: The metrics of the tool. + """ + timeout = self.timeout + query_list_from_params = parameters.get("query_list") + + if not query_list_from_params or not isinstance( + query_list_from_params, list): + error_msg = ( + "Error: 'query_list' is missing, empty, or not a list in parameters." + ) + logger.error( + f"[SearchTool] {error_msg} Received parameters: {parameters}") + return json.dumps({"result": error_msg}), 0.0, {} + + # Execute search using Ray execution pool + try: + result_text, metadata = await self.execution_pool.execute.remote( + self.execute_search, + instance_id, + query_list_from_params, + self.retrieval_service_url, + self.topk, + timeout, + ) + + # Store results in instance dictionary + self._instance_dict[instance_id]["reward"].append( + result_text.strip()) + + # Convert metadata to metrics + metrics = { + "query_count": metadata.get("query_count", 0), + "status": metadata.get("status", "unknown"), + "total_results": metadata.get("total_results", 0), + "api_request_error": metadata.get("api_request_error"), + } + + return result_text, 0.0, metrics + + except Exception as e: + error_result = json.dumps( + {"result": f"Search execution failed: {e}"}) + logger.error(f"[SearchTool] Execution failed: {e}") + return error_result, 0.0, {"error": str(e)} + + async def calc_reward(self, instance_id: str, **kwargs) -> str: + return self._instance_dict[instance_id]["reward"] + + async def release(self, instance_id: str, **kwargs) -> None: + if instance_id in self._instance_dict: + del self._instance_dict[instance_id] diff --git a/Agent0/executor_train/verl/verl/tools/utils/__init__.py b/Agent0/executor_train/verl/verl/tools/utils/__init__.py index c4b932b..72375fe 100644 --- a/Agent0/executor_train/verl/verl/tools/utils/__init__.py +++ b/Agent0/executor_train/verl/verl/tools/utils/__init__.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/tools/utils/mcp_clients/McpClientManager.py b/Agent0/executor_train/verl/verl/tools/utils/mcp_clients/McpClientManager.py index ee5fe31..c3ea4ea 100644 --- a/Agent0/executor_train/verl/verl/tools/utils/mcp_clients/McpClientManager.py +++ b/Agent0/executor_train/verl/verl/tools/utils/mcp_clients/McpClientManager.py @@ -1,97 +1,101 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import asyncio -import json -import logging -from typing import Any - -from fastmcp import Client -from fastmcp.client.transports import SSETransport - -from verl.tools.utils.mcp_clients.utils import TokenBucket, mcp2openai - -logger = logging.getLogger(__name__) - - -class MCPClientManager: - rootServerName = "mcpServers" - initialized = False - clients = [] - tool_client_mapping = {} - rate_limiter = None - - async def initialize(self, config_path, rate_limit: float = 10.0): - if self.initialized: - return - """Initialize the MCP Client Manager and start all clients""" - result = self._load_config(config_path) - servers = result[self.rootServerName] - exclude_sse_servers = {self.rootServerName: {}} - for server_name in servers.keys(): - server = servers[server_name] - if "auth_token" in server: - transport = SSETransport(url=server["url"], headers={"Authorization": f"Bearer {server['auth_token']}"}) - client = Client(transport) - self.clients.append(client) - else: - exclude_sse_servers[self.rootServerName][server_name] = server - - if exclude_sse_servers[self.rootServerName]: - self.clients.append(Client(exclude_sse_servers)) - - # Initialize rate limiter - self.rate_limiter = TokenBucket(rate_limit) - self.initialized = True - - async def call_tool(self, tool_name, parameters, timeout): - # Apply rate limiting - while not self.rate_limiter.acquire(): - await asyncio.sleep(0.1) - - client = self.get_client_with_tool_name(tool_name) - async with client: - return await client.call_tool_mcp(tool_name, parameters) - - async def fetch_tool_schemas(self, tool_selected_list: list[str]) -> list[dict]: - tool_schemas = [] - for client in self.clients: - async with client: - tools = await client.list_tools_mcp() - for tool in tools.tools: - if not tool_selected_list: - self.tool_client_mapping[tool.name] = client - tool_schemas.append(mcp2openai(tool)) - elif tool.name in tool_selected_list: - self.tool_client_mapping[tool.name] = client - tool_schemas.append(mcp2openai(tool)) - - return tool_schemas - - def get_client_with_tool_name(self, tool_name: str): - return self.tool_client_mapping[tool_name] - - def _load_config(self, file: str) -> dict[str, Any]: - try: - with open(file) as f: - return json.load(f) - except FileNotFoundError: - logger.warning(f'the "{file}" file was not found') - except Exception: - logger.error(f'there was an error reading the "{file}" file') - - return {} - - -ClientManager = MCPClientManager() +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import asyncio +import json +import logging +from typing import Any + +from fastmcp import Client +from fastmcp.client.transports import SSETransport + +from verl.tools.utils.mcp_clients.utils import TokenBucket, mcp2openai + +logger = logging.getLogger(__name__) + + +class MCPClientManager: + rootServerName = "mcpServers" + initialized = False + clients = [] + tool_client_mapping = {} + rate_limiter = None + + async def initialize(self, config_path, rate_limit: float = 10.0): + if self.initialized: + return + """Initialize the MCP Client Manager and start all clients""" + result = self._load_config(config_path) + servers = result[self.rootServerName] + exclude_sse_servers = {self.rootServerName: {}} + for server_name in servers.keys(): + server = servers[server_name] + if "auth_token" in server: + transport = SSETransport( + url=server["url"], headers={ + "Authorization": f"Bearer { + server['auth_token']}"}, ) + client = Client(transport) + self.clients.append(client) + else: + exclude_sse_servers[self.rootServerName][server_name] = server + + if exclude_sse_servers[self.rootServerName]: + self.clients.append(Client(exclude_sse_servers)) + + # Initialize rate limiter + self.rate_limiter = TokenBucket(rate_limit) + self.initialized = True + + async def call_tool(self, tool_name, parameters, timeout): + # Apply rate limiting + while not self.rate_limiter.acquire(): + await asyncio.sleep(0.1) + + client = self.get_client_with_tool_name(tool_name) + async with client: + return await client.call_tool_mcp(tool_name, parameters) + + async def fetch_tool_schemas( + self, tool_selected_list: list[str]) -> list[dict]: + tool_schemas = [] + for client in self.clients: + async with client: + tools = await client.list_tools_mcp() + for tool in tools.tools: + if not tool_selected_list: + self.tool_client_mapping[tool.name] = client + tool_schemas.append(mcp2openai(tool)) + elif tool.name in tool_selected_list: + self.tool_client_mapping[tool.name] = client + tool_schemas.append(mcp2openai(tool)) + + return tool_schemas + + def get_client_with_tool_name(self, tool_name: str): + return self.tool_client_mapping[tool_name] + + def _load_config(self, file: str) -> dict[str, Any]: + try: + with open(file) as f: + return json.load(f) + except FileNotFoundError: + logger.warning(f'the "{file}" file was not found') + except Exception: + logger.error(f'there was an error reading the "{file}" file') + + return {} + + +ClientManager = MCPClientManager() diff --git a/Agent0/executor_train/verl/verl/tools/utils/mcp_clients/utils.py b/Agent0/executor_train/verl/verl/tools/utils/mcp_clients/utils.py index 22a5f63..932d991 100644 --- a/Agent0/executor_train/verl/verl/tools/utils/mcp_clients/utils.py +++ b/Agent0/executor_train/verl/verl/tools/utils/mcp_clients/utils.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/tools/utils/search_r1_like_utils.py b/Agent0/executor_train/verl/verl/tools/utils/search_r1_like_utils.py index 23669e4..58e1b24 100644 --- a/Agent0/executor_train/verl/verl/tools/utils/search_r1_like_utils.py +++ b/Agent0/executor_train/verl/verl/tools/utils/search_r1_like_utils.py @@ -1,243 +1,268 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. - -import json -import logging -import threading -import time -import traceback -import uuid -from typing import Any, Optional - -import requests - -DEFAULT_TIMEOUT = 30 # Default search request timeout -MAX_RETRIES = 10 -INITIAL_RETRY_DELAY = 1 -API_TIMEOUT = 10 - -logger = logging.getLogger(__name__) - - -def call_search_api( - retrieval_service_url: str, - query_list: list[str], - topk: int = 3, - return_scores: bool = True, - timeout: int = DEFAULT_TIMEOUT, -) -> tuple[Optional[dict[str, Any]], Optional[str]]: - """ - Calls the remote search API to perform retrieval with retry logic for various errors, - using increasing delay between retries. Logs internal calls with a unique ID. - - Args: - retrieval_service_url: The URL of the retrieval service API. - query_list: List of search queries. - topk: Number of top results to return. - return_scores: Whether to return scores. - timeout: Request timeout in seconds. - - Returns: - A tuple (response_json, error_message). - If successful, response_json is the API's returned JSON object, error_message is None. - If failed after retries, response_json is None, error_message contains the error information. - """ - request_id = str(uuid.uuid4()) - log_prefix = f"[Search Request ID: {request_id}] " - - payload = {"queries": query_list, "topk": topk, "return_scores": return_scores} - - headers = {"Content-Type": "application/json", "Accept": "application/json"} - - last_error = None - - for attempt in range(MAX_RETRIES): - try: - logger.info( - f"{log_prefix}Attempt {attempt + 1}/{MAX_RETRIES}: Calling search API at {retrieval_service_url}" - ) - response = requests.post( - retrieval_service_url, - headers=headers, - json=payload, - timeout=timeout, - ) - - # Check for Gateway Timeout (504) and other server errors for retrying - if response.status_code in [500, 502, 503, 504]: - last_error = ( - f"{log_prefix}API Request Error: Server Error ({response.status_code}) on attempt " - f"{attempt + 1}/{MAX_RETRIES}" - ) - logger.warning(last_error) - if attempt < MAX_RETRIES - 1: - delay = INITIAL_RETRY_DELAY * (attempt + 1) - logger.info(f"{log_prefix}Retrying after {delay} seconds...") - time.sleep(delay) - continue - - # Check for other HTTP errors (e.g., 4xx) - response.raise_for_status() - - # If successful (status code 2xx) - logger.info(f"{log_prefix}Search API call successful on attempt {attempt + 1}") - return response.json(), None - - except requests.exceptions.ConnectionError as e: - last_error = f"{log_prefix}Connection Error: {e}" - logger.warning(last_error) - if attempt < MAX_RETRIES - 1: - delay = INITIAL_RETRY_DELAY * (attempt + 1) - logger.info(f"{log_prefix}Retrying after {delay} seconds...") - time.sleep(delay) - continue - except requests.exceptions.Timeout as e: - last_error = f"{log_prefix}Timeout Error: {e}" - logger.warning(last_error) - if attempt < MAX_RETRIES - 1: - delay = INITIAL_RETRY_DELAY * (attempt + 1) - logger.info(f"{log_prefix}Retrying after {delay} seconds...") - time.sleep(delay) - continue - except requests.exceptions.RequestException as e: - last_error = f"{log_prefix}API Request Error: {e}" - break # Exit retry loop on other request errors - except json.JSONDecodeError as e: - raw_response_text = response.text if "response" in locals() else "N/A" - last_error = f"{log_prefix}API Response JSON Decode Error: {e}, Response: {raw_response_text[:200]}" - break # Exit retry loop on JSON decode errors - except Exception as e: - last_error = f"{log_prefix}Unexpected Error: {e}" - break # Exit retry loop on other unexpected errors - - # If loop finishes without returning success, return the last recorded error - logger.error(f"{log_prefix}Search API call failed. Last error: {last_error}") - return None, last_error.replace(log_prefix, "API Call Failed: ") if last_error else "API Call Failed after retries" - - -def _passages2string(retrieval_result): - """Convert retrieval results to formatted string.""" - 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})\n{text}\n\n" - return format_reference.strip() - - -def perform_single_search_batch( - retrieval_service_url: str, - query_list: list[str], - topk: int = 3, - concurrent_semaphore: Optional[threading.Semaphore] = None, - timeout: int = DEFAULT_TIMEOUT, -) -> tuple[str, dict[str, Any]]: - """ - Performs a single batch search for multiple queries (original search tool behavior). - - Args: - retrieval_service_url: The URL of the retrieval service API. - query_list: List of search queries. - topk: Number of top results to return. - concurrent_semaphore: Optional semaphore for concurrency control. - timeout: Request timeout in seconds. - - Returns: - A tuple (result_text, metadata). - result_text: The search result JSON string. - metadata: Metadata dictionary for the batch search. - """ - logger.info(f"Starting batch search for {len(query_list)} queries.") - - api_response = None - error_msg = None - - try: - if concurrent_semaphore: - with concurrent_semaphore: - api_response, error_msg = call_search_api( - retrieval_service_url=retrieval_service_url, - query_list=query_list, - topk=topk, - return_scores=True, - timeout=timeout, - ) - else: - api_response, error_msg = call_search_api( - retrieval_service_url=retrieval_service_url, - query_list=query_list, - topk=topk, - return_scores=True, - timeout=timeout, - ) - except Exception as e: - error_msg = f"API Request Exception during batch search: {e}" - logger.error(f"Batch search: {error_msg}") - traceback.print_exc() - - metadata = { - "query_count": len(query_list), - "queries": query_list, - "api_request_error": error_msg, - "api_response": None, - "status": "unknown", - "total_results": 0, - "formatted_result": None, - } - - result_text = json.dumps({"result": "Search request failed or timed out after retries."}) - - if error_msg: - metadata["status"] = "api_error" - result_text = json.dumps({"result": f"Search error: {error_msg}"}) - logger.error(f"Batch search: API error occurred: {error_msg}") - elif api_response: - logger.debug(f"Batch search: API Response: {api_response}") - metadata["api_response"] = api_response - - try: - raw_results = api_response.get("result", []) - if raw_results: - pretty_results = [] - total_results = 0 - - for retrieval in raw_results: - formatted = _passages2string(retrieval) - pretty_results.append(formatted) - total_results += len(retrieval) if isinstance(retrieval, list) else 1 - - final_result = "\n---\n".join(pretty_results) - result_text = json.dumps({"result": final_result}) - metadata["status"] = "success" - metadata["total_results"] = total_results - metadata["formatted_result"] = final_result - logger.info(f"Batch search: Successful, got {total_results} total results") - else: - result_text = json.dumps({"result": "No search results found."}) - metadata["status"] = "no_results" - metadata["total_results"] = 0 - logger.info("Batch search: No results found") - except Exception as e: - error_msg = f"Error processing search results: {e}" - result_text = json.dumps({"result": error_msg}) - metadata["status"] = "processing_error" - logger.error(f"Batch search: {error_msg}") - else: - metadata["status"] = "unknown_api_state" - result_text = json.dumps({"result": "Unknown API state (no response and no error message)."}) - logger.error("Batch search: Unknown API state.") - - return result_text, metadata +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import json +import logging +import threading +import time +import traceback +import uuid +from typing import Any, Optional + +import requests + +DEFAULT_TIMEOUT = 30 # Default search request timeout +MAX_RETRIES = 10 +INITIAL_RETRY_DELAY = 1 +API_TIMEOUT = 10 + +logger = logging.getLogger(__name__) + + +def call_search_api( + retrieval_service_url: str, + query_list: list[str], + topk: int = 3, + return_scores: bool = True, + timeout: int = DEFAULT_TIMEOUT, +) -> tuple[Optional[dict[str, Any]], Optional[str]]: + """ + Calls the remote search API to perform retrieval with retry logic for various errors, + using increasing delay between retries. Logs internal calls with a unique ID. + + Args: + retrieval_service_url: The URL of the retrieval service API. + query_list: List of search queries. + topk: Number of top results to return. + return_scores: Whether to return scores. + timeout: Request timeout in seconds. + + Returns: + A tuple (response_json, error_message). + If successful, response_json is the API's returned JSON object, error_message is None. + If failed after retries, response_json is None, error_message contains the error information. + """ + request_id = str(uuid.uuid4()) + log_prefix = f"[Search Request ID: {request_id}] " + + payload = { + "queries": query_list, + "topk": topk, + "return_scores": return_scores} + + headers = { + "Content-Type": "application/json", + "Accept": "application/json"} + + last_error = None + + for attempt in range(MAX_RETRIES): + try: + logger.info( + f"{log_prefix}Attempt { + attempt + 1}/{MAX_RETRIES}: Calling search API at {retrieval_service_url}") + response = requests.post( + retrieval_service_url, + headers=headers, + json=payload, + timeout=timeout, + ) + + # Check for Gateway Timeout (504) and other server errors for + # retrying + if response.status_code in [500, 502, 503, 504]: + last_error = ( + f"{log_prefix}API Request Error: Server Error ({ + response.status_code}) on attempt " f"{ + attempt + 1}/{MAX_RETRIES}") + logger.warning(last_error) + if attempt < MAX_RETRIES - 1: + delay = INITIAL_RETRY_DELAY * (attempt + 1) + logger.info( + f"{log_prefix}Retrying after {delay} seconds...") + time.sleep(delay) + continue + + # Check for other HTTP errors (e.g., 4xx) + response.raise_for_status() + + # If successful (status code 2xx) + logger.info( + f"{log_prefix}Search API call successful on attempt { + attempt + 1}") + return response.json(), None + + except requests.exceptions.ConnectionError as e: + last_error = f"{log_prefix}Connection Error: {e}" + logger.warning(last_error) + if attempt < MAX_RETRIES - 1: + delay = INITIAL_RETRY_DELAY * (attempt + 1) + logger.info(f"{log_prefix}Retrying after {delay} seconds...") + time.sleep(delay) + continue + except requests.exceptions.Timeout as e: + last_error = f"{log_prefix}Timeout Error: {e}" + logger.warning(last_error) + if attempt < MAX_RETRIES - 1: + delay = INITIAL_RETRY_DELAY * (attempt + 1) + logger.info(f"{log_prefix}Retrying after {delay} seconds...") + time.sleep(delay) + continue + except requests.exceptions.RequestException as e: + last_error = f"{log_prefix}API Request Error: {e}" + break # Exit retry loop on other request errors + except json.JSONDecodeError as e: + raw_response_text = response.text if "response" in locals() else "N/A" + last_error = f"{log_prefix}API Response JSON Decode Error: {e}, Response: { + raw_response_text[ + :200]}" + break # Exit retry loop on JSON decode errors + except Exception as e: + last_error = f"{log_prefix}Unexpected Error: {e}" + break # Exit retry loop on other unexpected errors + + # If loop finishes without returning success, return the last recorded + # error + logger.error( + f"{log_prefix}Search API call failed. Last error: {last_error}") + return None, ( + last_error.replace(log_prefix, "API Call Failed: ") + if last_error + else "API Call Failed after retries" + ) + + +def _passages2string(retrieval_result): + """Convert retrieval results to formatted string.""" + 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})\n{text}\n\n" + return format_reference.strip() + + +def perform_single_search_batch( + retrieval_service_url: str, + query_list: list[str], + topk: int = 3, + concurrent_semaphore: Optional[threading.Semaphore] = None, + timeout: int = DEFAULT_TIMEOUT, +) -> tuple[str, dict[str, Any]]: + """ + Performs a single batch search for multiple queries (original search tool behavior). + + Args: + retrieval_service_url: The URL of the retrieval service API. + query_list: List of search queries. + topk: Number of top results to return. + concurrent_semaphore: Optional semaphore for concurrency control. + timeout: Request timeout in seconds. + + Returns: + A tuple (result_text, metadata). + result_text: The search result JSON string. + metadata: Metadata dictionary for the batch search. + """ + logger.info(f"Starting batch search for {len(query_list)} queries.") + + api_response = None + error_msg = None + + try: + if concurrent_semaphore: + with concurrent_semaphore: + api_response, error_msg = call_search_api( + retrieval_service_url=retrieval_service_url, + query_list=query_list, + topk=topk, + return_scores=True, + timeout=timeout, + ) + else: + api_response, error_msg = call_search_api( + retrieval_service_url=retrieval_service_url, + query_list=query_list, + topk=topk, + return_scores=True, + timeout=timeout, + ) + except Exception as e: + error_msg = f"API Request Exception during batch search: {e}" + logger.error(f"Batch search: {error_msg}") + traceback.print_exc() + + metadata = { + "query_count": len(query_list), + "queries": query_list, + "api_request_error": error_msg, + "api_response": None, + "status": "unknown", + "total_results": 0, + "formatted_result": None, + } + + result_text = json.dumps( + {"result": "Search request failed or timed out after retries."} + ) + + if error_msg: + metadata["status"] = "api_error" + result_text = json.dumps({"result": f"Search error: {error_msg}"}) + logger.error(f"Batch search: API error occurred: {error_msg}") + elif api_response: + logger.debug(f"Batch search: API Response: {api_response}") + metadata["api_response"] = api_response + + try: + raw_results = api_response.get("result", []) + if raw_results: + pretty_results = [] + total_results = 0 + + for retrieval in raw_results: + formatted = _passages2string(retrieval) + pretty_results.append(formatted) + total_results += ( + len(retrieval) if isinstance(retrieval, list) else 1 + ) + + final_result = "\n---\n".join(pretty_results) + result_text = json.dumps({"result": final_result}) + metadata["status"] = "success" + metadata["total_results"] = total_results + metadata["formatted_result"] = final_result + logger.info( + f"Batch search: Successful, got {total_results} total results") + else: + result_text = json.dumps( + {"result": "No search results found."}) + metadata["status"] = "no_results" + metadata["total_results"] = 0 + logger.info("Batch search: No results found") + except Exception as e: + error_msg = f"Error processing search results: {e}" + result_text = json.dumps({"result": error_msg}) + metadata["status"] = "processing_error" + logger.error(f"Batch search: {error_msg}") + else: + metadata["status"] = "unknown_api_state" + result_text = json.dumps( + {"result": "Unknown API state (no response and no error message)."} + ) + logger.error("Batch search: Unknown API state.") + + return result_text, metadata diff --git a/Agent0/executor_train/verl/verl/tools/utils/tool_registry.py b/Agent0/executor_train/verl/verl/tools/utils/tool_registry.py index 5c14d10..ba5cff6 100644 --- a/Agent0/executor_train/verl/verl/tools/utils/tool_registry.py +++ b/Agent0/executor_train/verl/verl/tools/utils/tool_registry.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -37,8 +37,14 @@ async def initialize_mcp_tool(tool_cls, tool_config) -> list: tool_list = [] mcp_servers_config_path = tool_config.mcp.mcp_servers_config_path - tool_selected_list = tool_config.mcp.tool_selected_list if "tool_selected_list" in tool_config.mcp else None - await ClientManager.initialize(mcp_servers_config_path, tool_config.config.rate_limit) + tool_selected_list = ( + tool_config.mcp.tool_selected_list + if "tool_selected_list" in tool_config.mcp + else None + ) + await ClientManager.initialize( + mcp_servers_config_path, tool_config.config.rate_limit + ) # Wait for MCP client to be ready max_retries = 10 retry_interval = 2 # seconds @@ -47,10 +53,13 @@ async def initialize_mcp_tool(tool_cls, tool_config) -> list: if tool_schemas: break if i < max_retries - 1: - logger.debug(f"Waiting for MCP client to be ready, attempt {i + 1}/{max_retries}") + logger.debug( + f"Waiting for MCP client to be ready, attempt { + i + 1}/{max_retries}") await asyncio.sleep(retry_interval) else: - raise RuntimeError("Failed to initialize MCP tools after maximum retries") + raise RuntimeError( + "Failed to initialize MCP tools after maximum retries") # mcp registry assert len(tool_schemas), "mcp tool is empty" for tool_schema_dict in tool_schemas: @@ -91,16 +100,24 @@ def initialize_tools_from_config(tools_config_file): if tool_config.get("tool_schema", None) is None: tool_schema = None else: - tool_schema_dict = OmegaConf.to_container(tool_config.tool_schema, resolve=True) - tool_schema = OpenAIFunctionToolSchema.model_validate(tool_schema_dict) + tool_schema_dict = OmegaConf.to_container( + tool_config.tool_schema, resolve=True + ) + tool_schema = OpenAIFunctionToolSchema.model_validate( + tool_schema_dict + ) tool = tool_cls( - config=OmegaConf.to_container(tool_config.config, resolve=True), + config=OmegaConf.to_container( + tool_config.config, + resolve=True), tool_schema=tool_schema, ) tool_list.append(tool) case ToolType.MCP: loop = asyncio.get_event_loop() - mcp_tools = loop.run_until_complete(initialize_mcp_tool(tool_cls, tool_config)) + mcp_tools = loop.run_until_complete( + initialize_mcp_tool(tool_cls, tool_config) + ) tool_list.extend(mcp_tools) case _: raise NotImplementedError diff --git a/Agent0/executor_train/verl/verl/trainer/__init__.py b/Agent0/executor_train/verl/verl/trainer/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/trainer/__init__.py +++ b/Agent0/executor_train/verl/verl/trainer/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/trainer/config/__init__.py b/Agent0/executor_train/verl/verl/trainer/config/__init__.py index f4cc9b8..0590dc8 100644 --- a/Agent0/executor_train/verl/verl/trainer/config/__init__.py +++ b/Agent0/executor_train/verl/verl/trainer/config/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/trainer/config/algorithm.py b/Agent0/executor_train/verl/verl/trainer/config/algorithm.py index e9600a9..c83f9c1 100644 --- a/Agent0/executor_train/verl/verl/trainer/config/algorithm.py +++ b/Agent0/executor_train/verl/verl/trainer/config/algorithm.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/trainer/constants_ppo.py b/Agent0/executor_train/verl/verl/trainer/constants_ppo.py index 84350bb..21a070d 100644 --- a/Agent0/executor_train/verl/verl/trainer/constants_ppo.py +++ b/Agent0/executor_train/verl/verl/trainer/constants_ppo.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/trainer/fsdp_sft_trainer.py b/Agent0/executor_train/verl/verl/trainer/fsdp_sft_trainer.py index 531ebab..d246b33 100644 --- a/Agent0/executor_train/verl/verl/trainer/fsdp_sft_trainer.py +++ b/Agent0/executor_train/verl/verl/trainer/fsdp_sft_trainer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -18,34 +18,20 @@ - Add validation """ -import os - -os.environ["NCCL_DEBUG"] = "WARN" -os.environ["TOKENIZERS_PARALLELISM"] = "true" - -import logging -import re -from contextlib import nullcontext - -import hydra -import torch -import torch.distributed -from peft import LoraConfig, TaskType, get_peft_model -from tensordict import TensorDict -from torch import nn, optim -from torch.distributed.device_mesh import DeviceMesh, init_device_mesh -from torch.distributed.fsdp import CPUOffload, MixedPrecision, ShardingStrategy -from torch.distributed.fsdp import FullyShardedDataParallel as FSDP -from torch.utils.data import DataLoader, Dataset, DistributedSampler -from tqdm import tqdm -from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel - -import verl.utils.hdfs_io as hdfs_io -from verl.utils.dataset import SFTDataset -from verl.utils.dataset.multiturn_sft_dataset import MultiTurnSFTDataset -from verl.utils.device import get_device_id, get_device_name, is_cuda_available, is_npu_available -from verl.utils.distributed import destroy_global_process_group, initialize_global_process_group -from verl.utils.fs import copy_to_local +from verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager +from verl.utils.ulysses import ( + gather_outpus_and_unpad, + get_ulysses_sequence_parallel_world_size, + ulysses_pad_and_slice_inputs, +) +from verl.utils.tracking import Tracking +from verl.utils.torch_functional import ( + get_cosine_schedule_with_warmup, + get_wsd_schedule_with_warmup, +) +from verl.utils.torch_dtypes import PrecisionType +from verl.utils.py_functional import convert_to_regular_types +from verl.utils.profiler import log_gpu_memory_usage from verl.utils.fsdp_utils import ( CPUOffloadPolicy, MixedPrecisionPolicy, @@ -56,22 +42,55 @@ get_init_weight_context_manager, init_fn, ) -from verl.utils.profiler import log_gpu_memory_usage -from verl.utils.py_functional import convert_to_regular_types -from verl.utils.torch_dtypes import PrecisionType -from verl.utils.torch_functional import get_cosine_schedule_with_warmup, get_wsd_schedule_with_warmup -from verl.utils.tracking import Tracking -from verl.utils.ulysses import ( - gather_outpus_and_unpad, - get_ulysses_sequence_parallel_world_size, - ulysses_pad_and_slice_inputs, +from verl.utils.fs import copy_to_local +from verl.utils.distributed import ( + destroy_global_process_group, + initialize_global_process_group, ) -from verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager +from verl.utils.device import ( + get_device_id, + get_device_name, + is_cuda_available, + is_npu_available, +) +from verl.utils.dataset.multiturn_sft_dataset import MultiTurnSFTDataset +from verl.utils.dataset import SFTDataset +import verl.utils.hdfs_io as hdfs_io +from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel +from tqdm import tqdm +from torch.utils.data import DataLoader, Dataset, DistributedSampler +from torch.distributed.fsdp import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp import CPUOffload, MixedPrecision, ShardingStrategy +from torch.distributed.device_mesh import DeviceMesh, init_device_mesh +from torch import nn, optim +from tensordict import TensorDict +from peft import LoraConfig, TaskType, get_peft_model +import torch.distributed +import torch +import hydra +from contextlib import nullcontext +import re +import logging +import os + +os.environ["NCCL_DEBUG"] = "WARN" +os.environ["TOKENIZERS_PARALLELISM"] = "true" + if is_cuda_available: - from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input + from flash_attn.bert_padding import ( + index_first_axis, + pad_input, + rearrange, + unpad_input, + ) elif is_npu_available: - from transformers.integrations.npu_flash_attention import index_first_axis, pad_input, rearrange, unpad_input + from transformers.integrations.npu_flash_attention import ( + index_first_axis, + pad_input, + rearrange, + unpad_input, + ) logger = logging.getLogger(__file__) logger.setLevel(os.getenv("VERL_SFT_LOGGING_LEVEL", "WARN")) @@ -97,19 +116,26 @@ def __init__( self.config = config self.device_mesh = device_mesh self.ulysses_device_mesh = ulysses_device_mesh - self.sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) + self.sharding_manager = FSDPUlyssesShardingManager( + self.ulysses_device_mesh) self.tokenizer = tokenizer if self.config.data.chat_template is not None: - raise ValueError("Apply Chat template from config is not supported yet.") + raise ValueError( + "Apply Chat template from config is not supported yet.") # normalize dp size self._normalize_config_bsz() # Set sequence parallel size - self.config.ulysses_sequence_parallel_size = getattr(self.config, "ulysses_sequence_parallel_size", 1) - self.use_remove_padding = getattr(self.config, "use_remove_padding", False) + self.config.ulysses_sequence_parallel_size = getattr( + self.config, "ulysses_sequence_parallel_size", 1 + ) + self.use_remove_padding = getattr( + self.config, "use_remove_padding", False) if self.device_mesh.get_rank() == 0: - print(f"Using sequence parallel size: {self.config.ulysses_sequence_parallel_size}") + print( + f"Using sequence parallel size: { + self.config.ulysses_sequence_parallel_size}") print(f"Using remove padding: {self.use_remove_padding}") self._build_dataloader(train_dataset, val_dataset) @@ -122,17 +148,26 @@ def __init__( self.device_name = get_device_name() def _normalize_config_bsz(self): - dp_size = self.device_mesh.size(0) if not self.ulysses_device_mesh else self.ulysses_device_mesh.size(0) + dp_size = ( + self.device_mesh.size(0) + if not self.ulysses_device_mesh + else self.ulysses_device_mesh.size(0) + ) 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, ( - f"Global batch size {self.config.data.train_batch_size} is not divisible by dp size {dp_size}" - ) + assert ( + self.config.data.train_batch_size % + dp_size == 0), f"Global batch size { + self.config.data.train_batch_size} is not divisible by dp size {dp_size}" self.config.data.train_batch_size //= dp_size - assert self.config.data.train_batch_size % self.config.data.micro_batch_size_per_gpu == 0 + assert ( + self.config.data.train_batch_size + % self.config.data.micro_batch_size_per_gpu + == 0 + ) def _build_dataloader(self, train_dataset, val_dataset): # build dataset @@ -147,16 +182,24 @@ def _build_dataloader(self, train_dataset, val_dataset): rank = self.ulysses_device_mesh.get_local_rank("dp") world_size = self.ulysses_device_mesh.size(0) if self.ulysses_device_mesh.get_rank() == 0: - print(f"Using SP rank {rank} and size {world_size} for data distribution") - print("Each SP rank gets different data, but the same data WITHIN the same rank") + print( + f"Using SP rank {rank} and size {world_size} for data distribution") + print( + "Each SP rank gets different data, but the same data WITHIN the same rank" + ) else: rank = self.device_mesh.get_rank() world_size = self.device_mesh.size() if self.device_mesh.get_rank() == 0: - print(f"Using FSDP rank {rank} and size {world_size} for data distribution") + print( + f"Using FSDP rank {rank} and size {world_size} for data distribution") self.train_sampler = DistributedSampler( - self.train_dataset, shuffle=True, num_replicas=world_size, rank=rank, drop_last=True + self.train_dataset, + shuffle=True, + num_replicas=world_size, + rank=rank, + drop_last=True, ) self.train_dataloader = DataLoader( dataset=self.train_dataset, @@ -168,7 +211,11 @@ def _build_dataloader(self, train_dataset, val_dataset): ) self.val_sampler = DistributedSampler( - self.val_dataset, shuffle=False, num_replicas=world_size, rank=rank, drop_last=True + self.val_dataset, + shuffle=False, + num_replicas=world_size, + rank=rank, + drop_last=True, ) self.val_dataloader = DataLoader( dataset=self.val_dataset, @@ -183,7 +230,9 @@ 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_to_local(src=self.config.model.partial_pretrain, verbose=True) + local_model_path = copy_to_local( + 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 @@ -197,19 +246,21 @@ def _build_model_optimizer(self): torch_dtype = self.config.model.fsdp_config.get("model_dtype", "fp32") torch_dtype = PrecisionType.to_dtype(torch_dtype) # load config first - config = AutoConfig.from_pretrained(local_model_path, trust_remote_code=trust_remote_code) + config = AutoConfig.from_pretrained( + local_model_path, trust_remote_code=trust_remote_code + ) self.model_config = config if hasattr(self.model_config, "max_position_embeddings"): self.model_config.max_position_embeddings = max( - self.model_config.max_position_embeddings, self.config.data.max_length - ) + self.model_config.max_position_embeddings, self.config.data.max_length) if self.config.ulysses_sequence_parallel_size > 1: - assert self.use_remove_padding, "Sequence parallel is only supported when remove_padding is enabled" + assert ( + self.use_remove_padding + ), "Sequence parallel is only supported when remove_padding is enabled" # This may be very large init_context = get_init_weight_context_manager( - use_meta_tensor=not config.tie_word_embeddings, mesh=self.device_mesh - ) + use_meta_tensor=not config.tie_word_embeddings, mesh=self.device_mesh) with init_context(): self.model: PreTrainedModel = AutoModelForCausalLM.from_pretrained( @@ -220,36 +271,52 @@ def _build_model_optimizer(self): trust_remote_code=trust_remote_code, ) - if self.use_remove_padding or self.config.ulysses_sequence_parallel_size > 1: + if ( + self.use_remove_padding + or self.config.ulysses_sequence_parallel_size > 1 + ): from verl.models.transformers.monkey_patch import apply_monkey_patch - apply_monkey_patch(model=self.model, ulysses_sp_size=self.config.ulysses_sequence_parallel_size) + apply_monkey_patch( + model=self.model, + ulysses_sp_size=self.config.ulysses_sequence_parallel_size, + ) # Apply Liger kernel if use_liger is enabled if self.config.model.get("use_liger", False): - from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance + from liger_kernel.transformers.monkey_patch import ( + _apply_liger_kernel_to_instance, + ) _apply_liger_kernel_to_instance(model=self.model) 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 + # 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), + "target_modules": convert_to_regular_types( + self.config.model.target_modules + ), "bias": "none", } - self.model = get_peft_model(self.model, LoraConfig(**lora_config)) + 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}) + 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 + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + buffer_dtype=torch.float32, ) auto_wrap_policy = get_fsdp_wrap_policy( @@ -263,7 +330,9 @@ def _build_model_optimizer(self): 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) + cpu_offload = CPUOffload( + offload_params=self.config.model.fsdp_config.offload_params + ) fsdp_strategy = self.config.model.strategy if fsdp_strategy == "fsdp": @@ -281,9 +350,13 @@ def _build_model_optimizer(self): forward_prefetch=False, ) elif fsdp_strategy == "fsdp2": - assert CPUOffloadPolicy is not None, "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" + assert ( + CPUOffloadPolicy is not None + ), "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" mp_policy = MixedPrecisionPolicy( - param_dtype=torch.bfloat16, reduce_dtype=torch.float32, cast_forward_inputs=True + param_dtype=torch.bfloat16, + reduce_dtype=torch.float32, + cast_forward_inputs=True, ) fsdp_kwargs = { @@ -294,7 +367,9 @@ def _build_model_optimizer(self): } full_state = self.model.state_dict() apply_fsdp2(self.model, fsdp_kwargs, self.config.model.fsdp_config) - fsdp2_load_full_state_dict(self.model, full_state, self.device_mesh, cpu_offload) + fsdp2_load_full_state_dict( + self.model, full_state, self.device_mesh, cpu_offload + ) self.fsdp_model = self.model else: raise NotImplementedError(f"not implement {fsdp_strategy}") @@ -319,45 +394,64 @@ def _build_model_optimizer(self): f"{self.config.trainer.total_epochs}, total number of steps {self.total_steps}" ) - num_warmup_steps = int(self.total_steps * self.config.optim.warmup_steps_ratio) + num_warmup_steps = int( + self.total_steps * + self.config.optim.warmup_steps_ratio) - if not hasattr(self.config.optim, "lr_scheduler") or self.config.optim.lr_scheduler == "cosine": + if ( + not hasattr(self.config.optim, "lr_scheduler") + or self.config.optim.lr_scheduler == "cosine" + ): self.lr_scheduler = get_cosine_schedule_with_warmup( - optimizer=self.optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=self.total_steps + optimizer=self.optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=self.total_steps, ) elif self.config.optim.lr_scheduler == "wsd": self.lr_scheduler = get_wsd_schedule_with_warmup( - optimizer=self.optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=self.total_steps + optimizer=self.optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=self.total_steps, ) else: - raise ValueError(f"Unknown lr scheduler: {self.config.optim.lr_scheduler}") + raise ValueError( + f"Unknown lr scheduler: { + self.config.optim.lr_scheduler}") def _compute_loss_and_backward(self, batch, do_backward=True): """Compute loss with optional sequence parallelism and remove padding features""" - use_sp = self.use_remove_padding and self.config.ulysses_sequence_parallel_size > 1 + use_sp = ( + self.use_remove_padding and self.config.ulysses_sequence_parallel_size > 1) # Move inputs to GPU and prepare loss mask input_ids = batch["input_ids"].to(self.device_name) attention_mask = batch["attention_mask"].to(self.device_name) position_ids = batch["position_ids"].to(self.device_name) - loss_mask = batch.pop("loss_mask")[:, :-1].reshape(-1).to(self.device_name) + loss_mask = batch.pop("loss_mask")[ + :, :-1].reshape(-1).to(self.device_name) loss_fct = nn.CrossEntropyLoss(reduction="none") # Context manager for sequence parallel if needed context = self.sharding_manager if use_sp else nullcontext() - with context, torch.autocast(device_type=self.device_name, dtype=torch.bfloat16): + with context, torch.autocast( + device_type=self.device_name, dtype=torch.bfloat16 + ): if not use_sp: # Standard forward pass without sequence parallel labels = input_ids[:, 1:].contiguous() output = self.fsdp_model( - input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + use_cache=False, ) logits = output.logits shift_logits = logits[..., :-1, :].contiguous() shift_labels = labels.contiguous() # Flatten the tokens - shift_logits = shift_logits.view(-1, self.model.config.vocab_size) + 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) @@ -375,23 +469,35 @@ def _compute_loss_and_backward(self, batch, do_backward=True): 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) + input_ids_rmpad = input_ids_rmpad.transpose( + 0, 1) # (1, total_nnz) # Unpad position_ids to align rotary position_ids_rmpad = index_first_axis( - rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices + rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), + indices, ).transpose(0, 1) # Pad and slice inputs for sequence parallelism - input_ids_rmpad_sliced, position_ids_rmpad_padded, pad_size = ulysses_pad_and_slice_inputs( - input_ids_rmpad, position_ids_rmpad, sp_size=get_ulysses_sequence_parallel_world_size() + input_ids_rmpad_sliced, position_ids_rmpad_padded, pad_size = ( + ulysses_pad_and_slice_inputs( + input_ids_rmpad, + position_ids_rmpad, + sp_size=get_ulysses_sequence_parallel_world_size(), + ) ) # For computing loss - input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=1) # (1, total_nnz) + input_ids_rmpad_rolled = torch.roll( + input_ids_rmpad, shifts=-1, dims=1 + ) # (1, total_nnz) input_ids_rmpad_rolled, _, _ = ulysses_pad_and_slice_inputs( - input_ids_rmpad_rolled, None, get_ulysses_sequence_parallel_world_size() + input_ids_rmpad_rolled, + None, + get_ulysses_sequence_parallel_world_size(), ) - input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0) # ((total_nnz / sp) + pad) + input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze( + 0 + ) # ((total_nnz / sp) + pad) # Forward pass output = self.fsdp_model( @@ -403,16 +509,23 @@ def _compute_loss_and_backward(self, batch, do_backward=True): # Compute loss locally then aggregate logits_rmpad = output.logits.squeeze(0) - input_ids_rmpad_rolled = input_ids_rmpad_rolled.to(logits_rmpad.device) + input_ids_rmpad_rolled = input_ids_rmpad_rolled.to( + logits_rmpad.device) loss = loss_fct(logits_rmpad, input_ids_rmpad_rolled) # Gather and unpad for sequence parallelism - loss = gather_outpus_and_unpad(loss, gather_dim=0, unpad_dim=0, padding_size=pad_size) + loss = gather_outpus_and_unpad( + loss, gather_dim=0, unpad_dim=0, padding_size=pad_size + ) # This is the loss collected from all ulysses ranks full_loss = pad_input( - hidden_states=loss.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen + hidden_states=loss.unsqueeze(-1), + indices=indices, + batch=batch_size, + seqlen=seqlen, ) - full_loss = full_loss.squeeze(-1)[:, :-1] # Remove last token's loss + # Remove last token's loss + full_loss = full_loss.squeeze(-1)[:, :-1] full_loss = full_loss.reshape(-1) loss_mask = loss_mask.to(full_loss.device) loss = full_loss * loss_mask @@ -421,7 +534,11 @@ def _compute_loss_and_backward(self, batch, do_backward=True): if self.config.data.balance_dp_token: torch.distributed.all_reduce(valid_token_this_rank) - dp_size = self.ulysses_device_mesh.size("dp") if use_sp else torch.distributed.get_world_size() + dp_size = ( + self.ulysses_device_mesh.size("dp") + if use_sp + else torch.distributed.get_world_size() + ) else: dp_size = 1 @@ -444,15 +561,22 @@ def training_step(self, batch: TensorDict): n_micro_batches = len(micro_batches) step_loss = 0 for micro_batch in micro_batches: - loss = self._compute_loss_and_backward(batch=micro_batch) / n_micro_batches + loss = self._compute_loss_and_backward( + batch=micro_batch) / n_micro_batches step_loss += loss.item() if self.config.model.strategy == "fsdp": - grad_norm = self.fsdp_model.clip_grad_norm_(max_norm=self.config.optim.clip_grad) + grad_norm = self.fsdp_model.clip_grad_norm_( + max_norm=self.config.optim.clip_grad + ) elif self.config.model.strategy == "fsdp2": - grad_norm = fsdp2_clip_grad_norm_(self.fsdp_model.parameters(), max_norm=self.config.optim.clip_grad) + grad_norm = fsdp2_clip_grad_norm_( + self.fsdp_model.parameters(), + max_norm=self.config.optim.clip_grad) else: - raise NotImplementedError(f"not implement {self.config.model.strategy}") + raise NotImplementedError( + f"not implement { + self.config.model.strategy}") log_gpu_memory_usage("Before optimizer step", logger=logger) @@ -474,18 +598,22 @@ def training_step(self, batch: TensorDict): step_loss = torch.tensor(step_loss).to(self.device_name) if is_cuda_available: - torch.distributed.all_reduce(step_loss, op=torch.distributed.ReduceOp.AVG) + torch.distributed.all_reduce( + step_loss, op=torch.distributed.ReduceOp.AVG) elif is_npu_available: torch.distributed.all_reduce(step_loss) step_loss /= self.device_mesh.size(0) - return {"train/loss": step_loss.detach().item(), "train/lr(1e-3)": lr * 1e3} + return { + "train/loss": step_loss.detach().item(), + "train/lr(1e-3)": lr * 1e3} def validation_step(self, batch: TensorDict): self.fsdp_model.eval() with torch.no_grad(): loss = self._compute_loss_and_backward(batch, do_backward=False) if is_cuda_available: - torch.distributed.all_reduce(loss, op=torch.distributed.ReduceOp.AVG) + torch.distributed.all_reduce( + loss, op=torch.distributed.ReduceOp.AVG) elif is_npu_available: torch.distributed.all_reduce(loss) loss /= self.device_mesh.size(0) @@ -493,7 +621,9 @@ def validation_step(self, batch: TensorDict): def save_checkpoint(self, step): # save checkpoint - path = os.path.join(self.config.trainer.default_local_dir, f"global_step_{step}") + path = os.path.join( + self.config.trainer.default_local_dir, f"global_step_{step}" + ) fsdp_strategy = self.config.model.strategy if fsdp_strategy == "fsdp": @@ -501,7 +631,9 @@ def save_checkpoint(self, step): 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): + with FSDP.state_dict_type( + self.fsdp_model, StateDictType.FULL_STATE_DICT, cfg + ): state_dict = self.fsdp_model.state_dict() # save huggingface model @@ -511,7 +643,10 @@ def save_checkpoint(self, step): self.tokenizer.save_pretrained(path) elif fsdp_strategy == "fsdp2": # FSDP2 checkpoint saving - from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + ) # Get full state dict with FSDP2 options = StateDictOptions(full_state_dict=True, cpu_offload=True) @@ -528,8 +663,13 @@ def save_checkpoint(self, step): # Copy to HDFS if configured if self.device_mesh.get_rank() == 0 and 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) + 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() @@ -548,7 +688,9 @@ def fit(self): last_valid_metric = None # 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 + 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 @@ -568,7 +710,9 @@ def fit(self): disable=rank != 0, ): global_step += 1 - data = TensorDict(data, batch_size=self.config.data.train_batch_size).to(self.device_name) + data = TensorDict( + data, batch_size=self.config.data.train_batch_size + ).to(self.device_name) metric = self.training_step(data) if rank == 0: tracking.log(data=metric, step=global_step) @@ -578,13 +722,15 @@ def fit(self): is_save_step = global_step % self.config.trainer.save_freq == 0 # early exit or validation step - if is_last_step or (self.config.trainer.test_freq > 0 and is_valid_step): + if is_last_step or ( + self.config.trainer.test_freq > 0 and is_valid_step + ): # Perform validation val_losses = [] for val_data in self.val_dataloader: - val_data = TensorDict(val_data, batch_size=self.config.data.micro_batch_size_per_gpu).to( - self.device_name - ) + val_data = TensorDict( + val_data, batch_size=self.config.data.micro_batch_size_per_gpu, ).to( + self.device_name) val_loss = self.validation_step(val_data) val_losses.append(val_loss) if rank == 0: @@ -594,7 +740,8 @@ def fit(self): last_valid_metric = metric torch.distributed.barrier() - if is_last_step or (self.config.trainer.save_freq > 0 and is_save_step): + if is_last_step or ( + self.config.trainer.save_freq > 0 and is_save_step): self.save_checkpoint(step=global_step) if is_last_step: @@ -607,7 +754,10 @@ def run_sft(config): device_name = get_device_name() local_rank, rank, world_size = initialize_global_process_group() - device_mesh = init_device_mesh(device_type=device_name, mesh_shape=(world_size,), mesh_dim_names=("fsdp",)) + device_mesh = init_device_mesh( + device_type=device_name, mesh_shape=( + world_size,), mesh_dim_names=( + "fsdp",)) dp_size = world_size // config.ulysses_sequence_parallel_size ulysses_device_mesh = init_device_mesh( device_type=device_name, @@ -617,10 +767,15 @@ def run_sft(config): # build tokenizer and datasets first from verl.utils import hf_tokenizer - local_model_path = copy_to_local(src=config.model.partial_pretrain, verbose=True) - tokenizer = hf_tokenizer(local_model_path, trust_remote_code=config.model.trust_remote_code) - train_dataset = create_sft_dataset(config.data.train_files, config.data, tokenizer) - val_dataset = create_sft_dataset(config.data.val_files, config.data, tokenizer) + local_model_path = copy_to_local( + src=config.model.partial_pretrain, verbose=True) + tokenizer = hf_tokenizer( + local_model_path, trust_remote_code=config.model.trust_remote_code + ) + train_dataset = create_sft_dataset( + config.data.train_files, config.data, tokenizer) + val_dataset = create_sft_dataset( + config.data.val_files, config.data, tokenizer) trainer = FSDPSFTTrainer( config=config, @@ -648,7 +803,9 @@ def create_sft_dataset(data_paths, data_config, tokenizer): if data_config.custom_cls.get("path", None): from verl.utils.import_utils import load_extern_type - dataset_cls = load_extern_type(data_config.custom_cls.path, data_config.custom_cls.name) + dataset_cls = load_extern_type( + data_config.custom_cls.path, data_config.custom_cls.name + ) # Then check if multi-turn dataset should be used elif data_config.get("multiturn", {}).get("enable", False): dataset_cls = MultiTurnSFTDataset @@ -657,7 +814,9 @@ def create_sft_dataset(data_paths, data_config, tokenizer): dataset_cls = SFTDataset # Create datasets based on the selected class - dataset = dataset_cls(parquet_files=data_paths, tokenizer=tokenizer, config=data_config) + dataset = dataset_cls( + parquet_files=data_paths, tokenizer=tokenizer, config=data_config + ) return dataset diff --git a/Agent0/executor_train/verl/verl/trainer/main_eval.py b/Agent0/executor_train/verl/verl/trainer/main_eval.py index 0a5c581..4cc7b5d 100644 --- a/Agent0/executor_train/verl/verl/trainer/main_eval.py +++ b/Agent0/executor_train/verl/verl/trainer/main_eval.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -38,7 +38,9 @@ def process_item(reward_fn, data_source, response_lst, reward_data): @hydra.main(config_path="config", config_name="evaluation", version_base=None) def main(config): - local_path = copy_to_local(config.data.path, use_shm=config.data.get("use_shm", False)) + local_path = copy_to_local( + config.data.path, use_shm=config.data.get("use_shm", False) + ) dataset = pd.read_parquet(local_path) responses = dataset[config.data.response_key] data_sources = dataset[config.data.data_source_key] @@ -56,7 +58,10 @@ def main(config): # Create remote tasks remote_tasks = [ - process_item.remote(compute_score, data_sources[i], responses[i], reward_model_data[i]) for i in range(total) + process_item.remote( + compute_score, data_sources[i], responses[i], reward_model_data[i] + ) + for i in range(total) ] # Process results as they come in diff --git a/Agent0/executor_train/verl/verl/trainer/main_generation.py b/Agent0/executor_train/verl/verl/trainer/main_generation.py index b8174ad..3eeb757 100644 --- a/Agent0/executor_train/verl/verl/trainer/main_generation.py +++ b/Agent0/executor_train/verl/verl/trainer/main_generation.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,6 +15,21 @@ Generate responses given a dataset of prompts """ +from verl.workers.fsdp_workers import ActorRolloutRefWorker +from verl.utils.model import compute_position_id_with_mask +from verl.utils.hdfs_io import makedirs +from verl.utils.fs import copy_to_local +from verl.utils import hf_tokenizer +from verl.single_controller.ray import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) +from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto +from verl import DataProto +from omegaconf import OmegaConf +import pandas as pd +from pprint import pprint import os import hydra @@ -25,20 +40,6 @@ os.environ["TOKENIZERS_PARALLELISM"] = "true" # os.environ['TORCH_COMPILE_DISABLE'] = '1' -from pprint import pprint - -import pandas as pd -from omegaconf import OmegaConf - -from verl import DataProto -from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto -from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup -from verl.utils import hf_tokenizer -from verl.utils.fs import copy_to_local -from verl.utils.hdfs_io import makedirs -from verl.utils.model import compute_position_id_with_mask -from verl.workers.fsdp_workers import ActorRolloutRefWorker - @hydra.main(config_path="config", config_name="generation", version_base=None) def main(config): @@ -49,7 +50,10 @@ def run_generation(config) -> None: if not ray.is_initialized(): # this is for local ray cluster ray.init( - runtime_env={"env_vars": {"TOKENIZERS_PARALLELISM": "true", "NCCL_DEBUG": "WARN"}}, + runtime_env={ + "env_vars": { + "TOKENIZERS_PARALLELISM": "true", + "NCCL_DEBUG": "WARN"}}, num_cpus=config.ray_init.num_cpus, ) @@ -58,7 +62,9 @@ def run_generation(config) -> None: @ray.remote(num_cpus=1) def main_task(config): - pprint(OmegaConf.to_container(config, resolve=True)) # resolve=True will eval symbol values + pprint( + OmegaConf.to_container(config, resolve=True) + ) # resolve=True will eval symbol values OmegaConf.resolve(config) local_path = copy_to_local(config.model.path) @@ -69,7 +75,8 @@ def main_task(config): assert config.data.n_samples == 1, "When temperature=0, n_samples must be 1." assert config.data.n_samples >= 1, "n_samples should always >= 1" - # read dataset. Note that the dataset should directly contain chat template format (e.g., a list of dictionary) + # 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() @@ -79,8 +86,13 @@ def main_task(config): 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) + 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, @@ -95,7 +107,9 @@ def main_task(config): 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] + 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, @@ -109,7 +123,11 @@ def main_task(config): 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} + batch_dict = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + } data = DataProto.from_dict(batch_dict) data_padded, pad_size = pad_dataproto_to_divisor(data, wg.world_size) @@ -124,9 +142,15 @@ def main_task(config): for i in range(len(output)): data_item = output[i] prompt_length = data_item.batch["prompts"].shape[-1] - valid_response_length = data_item.batch["attention_mask"][prompt_length:].sum() - valid_response_ids = data_item.batch["responses"][:valid_response_length] - response_str = tokenizer.decode(valid_response_ids, skip_special_tokens=True) + valid_response_length = data_item.batch["attention_mask"][ + prompt_length: + ].sum() + valid_response_ids = data_item.batch["responses"][ + :valid_response_length + ] + response_str = tokenizer.decode( + valid_response_ids, skip_special_tokens=True + ) output_texts.append(response_str) output_lst[n_sample].extend(output_texts) diff --git a/Agent0/executor_train/verl/verl/trainer/main_ppo.py b/Agent0/executor_train/verl/verl/trainer/main_ppo.py index 2a0b21d..3201f10 100644 --- a/Agent0/executor_train/verl/verl/trainer/main_ppo.py +++ b/Agent0/executor_train/verl/verl/trainer/main_ppo.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -61,14 +61,19 @@ def run_ppo(config) -> None: ) # Create a remote instance of the TaskRunner class, and - # Execute the `run` method of the TaskRunner instance remotely and wait for it to complete + # Execute the `run` method of the TaskRunner instance remotely and wait + # for it to complete if ( is_cuda_available and OmegaConf.select(config.trainer, "profile_steps") is not None and len(OmegaConf.select(config.trainer, "profile_steps")) > 0 ): - nsight_options = OmegaConf.to_container(config.trainer.controller_nsight_options) - runner = TaskRunner.options(runtime_env={"nsight": nsight_options}).remote() + nsight_options = OmegaConf.to_container( + config.trainer.controller_nsight_options + ) + runner = TaskRunner.options( + runtime_env={ + "nsight": nsight_options}).remote() else: runner = TaskRunner.remote() ray.get(runner.run.remote(config)) @@ -98,14 +103,18 @@ def run(self, config): config: Training configuration object containing all parameters needed for setting up and running the PPO training process. """ - # Print the initial configuration. `resolve=True` will evaluate symbolic values. + # Print the initial configuration. `resolve=True` will evaluate + # symbolic values. from pprint import pprint from omegaconf import OmegaConf from verl.utils.fs import copy_to_local - print(f"TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}") + print( + f"TaskRunner hostname: { + socket.gethostname()}, PID: { + os.getpid()}") pprint(OmegaConf.to_container(config, resolve=True)) @@ -114,16 +123,20 @@ def run(self, config): # Download the checkpoint from HDFS to the local machine. # `use_shm` determines whether to use shared memory, which could lead to faster model loading if turned on local_path = copy_to_local( - config.actor_rollout_ref.model.path, use_shm=config.actor_rollout_ref.model.get("use_shm", False) + config.actor_rollout_ref.model.path, + use_shm=config.actor_rollout_ref.model.get("use_shm", False), ) # Instantiate the tokenizer and processor. from verl.utils import hf_processor, hf_tokenizer trust_remote_code = config.data.get("trust_remote_code", False) - tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code) + tokenizer = hf_tokenizer( + local_path, trust_remote_code=trust_remote_code) # Used for multimodal LLM, could be None - processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True) + processor = hf_processor( + local_path, trust_remote_code=trust_remote_code, use_fast=True + ) # Version validation for vllm. if config.actor_rollout_ref.rollout.name in ["vllm"]: @@ -131,13 +144,19 @@ def run(self, config): if config.actor_rollout_ref.model.get("lora_rank", 0) > 0: if not is_version_ge(pkg="vllm", minver="0.7.3"): - raise NotImplementedError("PPO LoRA is not supported before vllm 0.7.3") + raise NotImplementedError( + "PPO LoRA is not supported before vllm 0.7.3" + ) # Define worker classes based on the actor strategy. if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"}: assert config.critic.strategy in {"fsdp", "fsdp2"} from verl.single_controller.ray import RayWorkerGroup - from verl.workers.fsdp_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker, CriticWorker + from verl.workers.fsdp_workers import ( + ActorRolloutRefWorker, + AsyncActorRolloutRefWorker, + CriticWorker, + ) actor_rollout_cls = ( AsyncActorRolloutRefWorker @@ -149,7 +168,11 @@ def run(self, config): elif config.actor_rollout_ref.actor.strategy == "megatron": assert config.actor_rollout_ref.actor.strategy == config.critic.strategy from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup - from verl.workers.megatron_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker, CriticWorker + from verl.workers.megatron_workers import ( + ActorRolloutRefWorker, + AsyncActorRolloutRefWorker, + CriticWorker, + ) actor_rollout_cls = ( AsyncActorRolloutRefWorker @@ -173,7 +196,9 @@ def run(self, config): # Map roles to the resource pool. global_pool_id = "global_pool" resource_pool_spec = { - global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes, + global_pool_id: [ + config.trainer.n_gpus_per_node] * + config.trainer.nnodes, } mapping = { Role.ActorRollout: global_pool_id, @@ -193,28 +218,51 @@ def run(self, config): from verl.workers.megatron_workers import RewardModelWorker else: raise NotImplementedError - role_worker_mapping[Role.RewardModel] = ray.remote(RewardModelWorker) + role_worker_mapping[Role.RewardModel] = ray.remote( + RewardModelWorker) mapping[Role.RewardModel] = global_pool_id # Add a reference policy worker if KL loss or KL reward is used. - if config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss: - role_worker_mapping[Role.RefPolicy] = ray.remote(ActorRolloutRefWorker) + if ( + config.algorithm.use_kl_in_reward + or config.actor_rollout_ref.actor.use_kl_loss + ): + role_worker_mapping[Role.RefPolicy] = ray.remote( + ActorRolloutRefWorker) mapping[Role.RefPolicy] = global_pool_id # Load the reward manager for training and validation. reward_fn = load_reward_manager( - config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {}) + config, + tokenizer, + num_examine=0, + **config.reward_model.get("reward_kwargs", {}), ) val_reward_fn = load_reward_manager( - config, tokenizer, num_examine=1, **config.reward_model.get("reward_kwargs", {}) + config, + tokenizer, + num_examine=1, + **config.reward_model.get("reward_kwargs", {}), + ) + resource_pool_manager = ResourcePoolManager( + resource_pool_spec=resource_pool_spec, mapping=mapping ) - resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping) from verl.utils.dataset.rl_dataset import collate_fn # Create training and validation datasets. - train_dataset = create_rl_dataset(config.data.train_files, config.data, tokenizer, processor, is_train=True) - val_dataset = create_rl_dataset(config.data.val_files, config.data, tokenizer, processor, is_train=False) + train_dataset = create_rl_dataset( + config.data.train_files, + config.data, + tokenizer, + processor, + is_train=True) + val_dataset = create_rl_dataset( + config.data.val_files, + config.data, + tokenizer, + processor, + is_train=False) train_sampler = create_rl_sampler(config.data, train_dataset) # Initialize the PPO trainer. @@ -239,7 +287,12 @@ def run(self, config): trainer.fit() -def create_rl_dataset(data_paths, data_config, tokenizer, processor, is_train=True): +def create_rl_dataset( + data_paths, + data_config, + tokenizer, + processor, + is_train=True): """Create a dataset. Arguments: @@ -257,17 +310,28 @@ def create_rl_dataset(data_paths, data_config, tokenizer, processor, is_train=Tr # Check if a custom dataset class is specified in the data configuration # and if the path to the custom class is provided - if "custom_cls" in data_config and data_config.custom_cls.get("path", None) is not None: + if ( + "custom_cls" in data_config + and data_config.custom_cls.get("path", None) is not None + ): # Dynamically load the custom dataset class - dataset_cls = load_extern_type(data_config.custom_cls.path, data_config.custom_cls.name) - # Verify that the custom dataset class inherits from torch.utils.data.Dataset + dataset_cls = load_extern_type( + data_config.custom_cls.path, data_config.custom_cls.name + ) + # Verify that the custom dataset class inherits from + # torch.utils.data.Dataset if not issubclass(dataset_cls, Dataset): raise TypeError( - f"The custom dataset class '{data_config.custom_cls.name}' from " - f"'{data_config.custom_cls.path}' must inherit from torch.utils.data.Dataset" - ) - elif "datagen" in data_config and data_config.datagen.get("path", None) is not None and is_train: - # If a data generation strategy is specified, use the DynamicGenDataset class + f"The custom dataset class '{ + data_config.custom_cls.name}' from " f"'{ + data_config.custom_cls.path}' must inherit from torch.utils.data.Dataset") + elif ( + "datagen" in data_config + and data_config.datagen.get("path", None) is not None + and is_train + ): + # If a data generation strategy is specified, use the DynamicGenDataset + # class from verl.utils.dataset.dynamicgen_dataset import DynamicGenDataset dataset_cls = DynamicGenDataset @@ -302,7 +366,10 @@ def create_rl_sampler(data_config, dataset): import torch from torch.utils.data import RandomSampler, SequentialSampler - if data_config.sampler is not None and data_config.sampler.get("class_path", None) is not None: + if ( + data_config.sampler is not None + and data_config.sampler.get("class_path", None) is not None + ): curriculum_class = load_extern_type( data_config.sampler.class_path, data_config.sampler.class_name, @@ -319,13 +386,17 @@ def create_rl_sampler(data_config, dataset): ) # Use a sampler to facilitate checkpoint resumption. - # If shuffling is enabled in the data configuration, create a random sampler. + # If shuffling is enabled in the data configuration, create a random + # sampler. elif data_config.shuffle: train_dataloader_generator = torch.Generator() train_dataloader_generator.manual_seed(data_config.get("seed", 1)) - sampler = RandomSampler(data_source=dataset, generator=train_dataloader_generator) + sampler = RandomSampler( + data_source=dataset, generator=train_dataloader_generator + ) else: - # If shuffling is disabled, use a sequential sampler to iterate through the dataset in order. + # If shuffling is disabled, use a sequential sampler to iterate through + # the dataset in order. sampler = SequentialSampler(data_source=dataset) return sampler diff --git a/Agent0/executor_train/verl/verl/trainer/ppo/__init__.py b/Agent0/executor_train/verl/verl/trainer/ppo/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/trainer/ppo/__init__.py +++ b/Agent0/executor_train/verl/verl/trainer/ppo/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/trainer/ppo/core_algos.py b/Agent0/executor_train/verl/verl/trainer/ppo/core_algos.py index 5f02675..2f14d73 100644 --- a/Agent0/executor_train/verl/verl/trainer/ppo/core_algos.py +++ b/Agent0/executor_train/verl/verl/trainer/ppo/core_algos.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Copyright 2022 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -63,8 +63,9 @@ def get_policy_loss_fn(name): loss_name = name if loss_name not in POLICY_LOSS_REGISTRY: raise ValueError( - f"Unsupported loss mode: {loss_name}. Supported modes are: {list(POLICY_LOSS_REGISTRY.keys())}" - ) + f"Unsupported loss mode: {loss_name}. Supported modes are: { + list( + POLICY_LOSS_REGISTRY.keys())}") return POLICY_LOSS_REGISTRY[loss_name] @@ -81,11 +82,12 @@ def register_adv_est(name_or_enum): """ def decorator(fn): - name = name_or_enum.value if isinstance(name_or_enum, Enum) else name_or_enum + name = name_or_enum.value if isinstance( + name_or_enum, Enum) else name_or_enum if name in ADV_ESTIMATOR_REGISTRY and ADV_ESTIMATOR_REGISTRY[name] != fn: raise ValueError( - f"Adv estimator {name} has already been registered: {ADV_ESTIMATOR_REGISTRY[name]} vs {fn}" - ) + f"Adv estimator {name} has already been registered: { + ADV_ESTIMATOR_REGISTRY[name]} vs {fn}") ADV_ESTIMATOR_REGISTRY[name] = fn return fn @@ -102,7 +104,8 @@ def get_adv_estimator_fn(name_or_enum): Returns: `(callable)`: The advantage estimator function. """ - name = name_or_enum.value if isinstance(name_or_enum, Enum) else name_or_enum + name = name_or_enum.value if isinstance( + name_or_enum, Enum) else name_or_enum if name not in ADV_ESTIMATOR_REGISTRY: raise ValueError(f"Unknown advantage estimator simply: {name}") return ADV_ESTIMATOR_REGISTRY[name] @@ -184,13 +187,20 @@ def get_kl_controller(kl_ctrl): if kl_ctrl.type == "fixed": return FixedKLController(kl_coef=kl_ctrl.kl_coef) elif kl_ctrl.type == "adaptive": - assert kl_ctrl.horizon > 0, f"horizon must be larger than 0. Got {kl_ctrl.horizon}" - return AdaptiveKLController(init_kl_coef=kl_ctrl.kl_coef, target_kl=kl_ctrl.target_kl, horizon=kl_ctrl.horizon) + assert ( + kl_ctrl.horizon > 0 + ), f"horizon must be larger than 0. Got {kl_ctrl.horizon}" + return AdaptiveKLController( + init_kl_coef=kl_ctrl.kl_coef, + target_kl=kl_ctrl.target_kl, + horizon=kl_ctrl.horizon, + ) else: raise NotImplementedError -@register_adv_est(AdvantageEstimator.GAE) # or simply: @register_adv_est("gae") +# or simply: @register_adv_est("gae") +@register_adv_est(AdvantageEstimator.GAE) def compute_gae_advantage_return( token_level_rewards: torch.Tensor, values: torch.Tensor, @@ -226,12 +236,19 @@ def compute_gae_advantage_return( gen_len = token_level_rewards.shape[-1] for t in reversed(range(gen_len)): - delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t] + delta = token_level_rewards[:, t] + \ + gamma * nextvalues - values[:, t] lastgaelam_ = delta + gamma * lam * lastgaelam # skip values and TD-error on observation tokens - nextvalues = values[:, t] * response_mask[:, t] + (1 - response_mask[:, t]) * nextvalues - lastgaelam = lastgaelam_ * response_mask[:, t] + (1 - response_mask[:, t]) * lastgaelam + nextvalues = ( + values[:, t] * response_mask[:, t] + + (1 - response_mask[:, t]) * nextvalues + ) + lastgaelam = ( + lastgaelam_ * response_mask[:, t] + + (1 - response_mask[:, t]) * lastgaelam + ) advantages_reversed.append(lastgaelam) advantages = torch.stack(advantages_reversed[::-1], dim=1) @@ -241,8 +258,10 @@ def compute_gae_advantage_return( return advantages, returns -# NOTE(sgm): this implementation only consider outcome supervision, where the reward is a scalar. -@register_adv_est(AdvantageEstimator.GRPO) # or simply: @register_adv_est("grpo") +# NOTE(sgm): this implementation only consider outcome supervision, where +# the reward is a scalar. +# or simply: @register_adv_est("grpo") +@register_adv_est(AdvantageEstimator.GRPO) def compute_grpo_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, @@ -300,7 +319,9 @@ def compute_grpo_outcome_advantage( raise ValueError(f"no score in prompt index: {idx}") for i in range(bsz): if norm_adv_by_std_in_grpo: - scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon) + scores[i] = (scores[i] - id2mean[index[i]]) / ( + id2std[index[i]] + epsilon + ) else: scores[i] = scores[i] - id2mean[index[i]] scores = scores.unsqueeze(-1) * response_mask @@ -308,7 +329,9 @@ def compute_grpo_outcome_advantage( return scores, scores -@register_adv_est(AdvantageEstimator.GRPO_PASSK) # or simply: @register_adv_est("grpo_passk") +@register_adv_est( + AdvantageEstimator.GRPO_PASSK +) # or simply: @register_adv_est("grpo_passk") def compute_grpo_passk_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, @@ -418,13 +441,15 @@ def compute_reinforce_plus_plus_baseline_outcome_advantage( for i in range(bsz): scores[i] = scores[i] - id2mean[index[i]] - scores = scores.unsqueeze(-1).tile([1, response_length]) * response_mask + scores = scores.unsqueeze(-1).tile([1, + response_length]) * response_mask scores = verl_F.masked_whiten(scores, response_mask) * response_mask return scores, scores -@register_adv_est(AdvantageEstimator.RLOO) # or simply: @register_adv_est("rloo") +# or simply: @register_adv_est("rloo") +@register_adv_est(AdvantageEstimator.RLOO) def compute_rloo_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, @@ -468,15 +493,16 @@ def compute_rloo_outcome_advantage( for i in range(bsz): response_num = len(id2score[index[i]]) if response_num > 1: - scores[i] = scores[i] * response_num / (response_num - 1) - id2mean[index[i]] * response_num / ( - response_num - 1 - ) + scores[i] = scores[i] * response_num / (response_num - 1) - id2mean[ + index[i] + ] * response_num / (response_num - 1) scores = scores.unsqueeze(-1) * response_mask return scores, scores -@register_adv_est(AdvantageEstimator.OPO) # or simply: @register_adv_est("opo") +# or simply: @register_adv_est("opo") +@register_adv_est(AdvantageEstimator.OPO) def compute_opo_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, @@ -520,7 +546,8 @@ def compute_opo_outcome_advantage( elif len(id2score[idx]) > 1: score_tensor = torch.tensor(id2score[idx]) len_tensor = torch.tensor(id2len[idx]) - id2bsl[idx] = (len_tensor * score_tensor).sum() / len_tensor.sum() + id2bsl[idx] = ( + len_tensor * score_tensor).sum() / len_tensor.sum() else: raise ValueError(f"no score in prompt index: {idx}") for i in range(bsz): @@ -530,9 +557,14 @@ def compute_opo_outcome_advantage( return scores, scores -@register_adv_est(AdvantageEstimator.REINFORCE_PLUS_PLUS) # or simply: @register_adv_est("reinforce_plus_plus") +@register_adv_est( + AdvantageEstimator.REINFORCE_PLUS_PLUS +) # or simply: @register_adv_est("reinforce_plus_plus") def compute_reinforce_plus_plus_outcome_advantage( - token_level_rewards: torch.Tensor, response_mask: torch.Tensor, config: Optional[AlgoConfig] = None, **kwargs + token_level_rewards: torch.Tensor, + response_mask: torch.Tensor, + config: Optional[AlgoConfig] = None, + **kwargs, ) -> tuple[torch.Tensor, torch.Tensor]: """ Compute advantage for REINFORCE++. @@ -569,7 +601,8 @@ def compute_reinforce_plus_plus_outcome_advantage( return advantages, returns -@register_adv_est(AdvantageEstimator.REMAX) # or simply: @register_adv_est("remax") +# or simply: @register_adv_est("remax") +@register_adv_est(AdvantageEstimator.REMAX) def compute_remax_outcome_advantage( token_level_rewards: torch.Tensor, reward_baselines: torch.Tensor, @@ -599,13 +632,19 @@ def compute_remax_outcome_advantage( """ with torch.no_grad(): - returns = (token_level_rewards * response_mask).flip(dims=[-1]).cumsum(dim=-1).flip(dims=[-1]) + returns = ( + (token_level_rewards * response_mask) + .flip(dims=[-1]) + .cumsum(dim=-1) + .flip(dims=[-1]) + ) advantages = returns - reward_baselines.unsqueeze(-1) * response_mask return advantages, returns -@register_adv_est(AdvantageEstimator.GPG) # or simply: @register_adv_est("gpg") +# or simply: @register_adv_est("gpg") +@register_adv_est(AdvantageEstimator.GPG) def compute_gpg_outcome_advantage( token_level_rewards: torch.Tensor, response_mask: torch.Tensor, @@ -683,7 +722,10 @@ def compute_rewards(token_level_scores, old_log_prob, ref_log_prob, kl_ratio): return token_level_scores - kl * kl_ratio -def agg_loss(loss_mat: torch.Tensor, loss_mask: torch.Tensor, loss_agg_mode: str): +def agg_loss( + loss_mat: torch.Tensor, + loss_mask: torch.Tensor, + loss_agg_mode: str): """ Aggregate the loss matrix into a scalar. @@ -704,7 +746,9 @@ def agg_loss(loss_mat: torch.Tensor, loss_mask: torch.Tensor, loss_agg_mode: str seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) # token-sum loss = torch.mean(seq_losses) # seq-mean elif loss_agg_mode == "seq-mean-token-mean": - seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) / torch.sum(loss_mask, dim=-1) # token-mean + seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) / torch.sum( + loss_mask, dim=-1 + ) # token-mean loss = torch.mean(seq_losses) # seq-mean elif loss_agg_mode == "seq-mean-token-sum-norm": seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) @@ -780,22 +824,33 @@ def compute_policy_loss( clip_pg_losses1 = torch.maximum( pg_losses1, pg_losses2 ) # max(-ratio * A, -clip(ratio, 1-cliprange, 1+cliprange) * A) - pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses1).float(), response_mask) + pg_clipfrac = verl_F.masked_mean( + torch.gt(pg_losses2, pg_losses1).float(), response_mask + ) pg_losses3 = -advantages * clip_ratio_c clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1) - pg_clipfrac_lower = verl_F.masked_mean( - torch.gt(clip_pg_losses1, pg_losses3) * (advantages < 0).float(), response_mask - ) + pg_clipfrac_lower = verl_F.masked_mean(torch.gt( + clip_pg_losses1, pg_losses3) * (advantages < 0).float(), response_mask) pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1) - pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode) + pg_loss = agg_loss( + loss_mat=pg_losses, + loss_mask=response_mask, + loss_agg_mode=loss_agg_mode) return pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower @register_policy_loss("gpg") -def compute_policy_loss_gpg(old_log_prob, log_prob, advantages, response_mask, loss_agg_mode="token-mean", config=None): +def compute_policy_loss_gpg( + old_log_prob, + log_prob, + advantages, + response_mask, + loss_agg_mode="token-mean", + config=None, +): """Adapted from https://github.com/AMAP-ML/GPG/blob/main/VisualThinker-R1-Zero/src/open-r1-multimodal/src/open_r1/trainer/grpo_trainer.py#L495 Args: @@ -811,7 +866,10 @@ def compute_policy_loss_gpg(old_log_prob, log_prob, advantages, response_mask, l """ pg_losses = -log_prob * advantages - pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode) + pg_loss = agg_loss( + loss_mat=pg_losses, + loss_mask=response_mask, + loss_agg_mode=loss_agg_mode) return pg_loss, torch.tensor(0.0), torch.tensor(0.0), torch.tensor(0.0) @@ -855,12 +913,26 @@ def compute_policy_loss_clip_cov( clip_cov_ub (float, optional): Upper bound for clipping covariance. Defaults to 5.0. """ - clip_cov_ratio = config.policy_loss.clip_cov_ratio if config.policy_loss.clip_cov_ratio is not None else 0.0002 + clip_cov_ratio = ( + config.policy_loss.clip_cov_ratio + if config.policy_loss.clip_cov_ratio is not None + else 0.0002 + ) cliprange = config.clip_ratio - cliprange_low = config.clip_ratio_low if config.clip_ratio_low is not None else cliprange - cliprange_high = config.clip_ratio_high if config.clip_ratio_high is not None else cliprange - clip_cov_ub = config.policy_loss.clip_cov_ub if config.policy_loss.clip_cov_ub is not None else 5.0 - clip_cov_lb = config.policy_loss.clip_cov_lb if config.policy_loss.clip_cov_lb is not None else 1.0 + cliprange_low = ( + config.clip_ratio_low if config.clip_ratio_low is not None else cliprange) + cliprange_high = ( + config.clip_ratio_high if config.clip_ratio_high is not None else cliprange) + clip_cov_ub = ( + config.policy_loss.clip_cov_ub + if config.policy_loss.clip_cov_ub is not None + else 5.0 + ) + clip_cov_lb = ( + config.policy_loss.clip_cov_lb + if config.policy_loss.clip_cov_lb is not None + else 1.0 + ) assert clip_cov_ratio > 0, "clip_ratio should be larger than 0." @@ -876,7 +948,8 @@ def compute_policy_loss_clip_cov( cliprange_high = cliprange corr = torch.ones_like(advantages) - pg_losses2 = -advantages * torch.clamp(ratio, 1 - cliprange_low, 1 + cliprange_high) + pg_losses2 = -advantages * \ + torch.clamp(ratio, 1 - cliprange_low, 1 + cliprange_high) clip_by_origin = (pg_losses2 > pg_losses1) & (response_mask > 0) cov_all = (advantages - verl_F.masked_mean(advantages, response_mask)) * ( @@ -886,21 +959,28 @@ def compute_policy_loss_clip_cov( cov_all[clip_by_origin] = -torch.inf clip_num = max(int(clip_cov_ratio * response_mask.sum().item()), 1) - top_k_idx = (cov_all < clip_cov_ub) & (cov_all > clip_cov_lb) & (response_mask > 0) + top_k_idx = ( + cov_all < clip_cov_ub) & ( + cov_all > clip_cov_lb) & ( + response_mask > 0) top_k_idx = torch.nonzero(top_k_idx) if len(top_k_idx) > 0: perm = torch.randperm(len(top_k_idx)) top_k_idx = top_k_idx[perm[: min(clip_num, len(top_k_idx))]] else: - top_k_idx = torch.empty((0, 2), device=cov_all.device, dtype=torch.long) + top_k_idx = torch.empty( + (0, 2), device=cov_all.device, dtype=torch.long) corr[top_k_idx[:, 0], top_k_idx[:, 1]] = 0 pg_clipfrac = verl_F.masked_mean((corr == 0).float(), response_mask) pg_losses = torch.maximum(pg_losses1, pg_losses2) * corr - pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode) + pg_loss = agg_loss( + loss_mat=pg_losses, + loss_mask=response_mask, + loss_agg_mode=loss_agg_mode) return pg_loss, pg_clipfrac, ppo_kl, torch.tensor(0.0) @@ -936,8 +1016,16 @@ def compute_policy_loss_kl_cov( ppo_kl_coef (float, optional): Coefficient for the KL penalty term in the loss. Defaults to 1. """ - kl_cov_ratio = config.policy_loss.kl_cov_ratio if config.policy_loss.kl_cov_ratio is not None else 0.0002 - ppo_kl_coef = config.policy_loss.ppo_kl_coef if config.policy_loss.ppo_kl_coef is not None else 1.0 + kl_cov_ratio = ( + config.policy_loss.kl_cov_ratio + if config.policy_loss.kl_cov_ratio is not None + else 0.0002 + ) + ppo_kl_coef = ( + config.policy_loss.ppo_kl_coef + if config.policy_loss.ppo_kl_coef is not None + else 1.0 + ) assert kl_cov_ratio > 0, "kl_cov_ratio should be larger than 0." @@ -957,22 +1045,37 @@ def compute_policy_loss_kl_cov( k = min(kl_cov_ratio, len(all_valid_adv)) if k != 0: - cov_lst_all = (all_valid_adv - all_valid_adv.mean()) * (all_valid_logp - all_valid_logp.mean()) + cov_lst_all = (all_valid_adv - all_valid_adv.mean()) * ( + all_valid_logp - all_valid_logp.mean() + ) k_percent_nums = max(1, int(len(cov_lst_all) * kl_cov_ratio)) - large_cov_idxs = torch.topk(cov_lst_all, k_percent_nums, largest=True).indices + large_cov_idxs = torch.topk( + cov_lst_all, + k_percent_nums, + largest=True).indices if len(large_cov_idxs) != 0: large_cov_idxs = all_valid_idx[large_cov_idxs] - pg_losses[large_cov_idxs // advantages.shape[1], large_cov_idxs % advantages.shape[1]] = pg_losses_kl[ - large_cov_idxs // advantages.shape[1], large_cov_idxs % advantages.shape[1] + pg_losses[ + large_cov_idxs // advantages.shape[1], + large_cov_idxs % advantages.shape[1], + ] = pg_losses_kl[ + large_cov_idxs // advantages.shape[1], + large_cov_idxs % advantages.shape[1], ] - pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode) + pg_loss = agg_loss( + loss_mat=pg_losses, + loss_mask=response_mask, + loss_agg_mode=loss_agg_mode) return pg_loss, torch.tensor(0.0), ppo_kl_abs, torch.tensor(0.0) -def compute_entropy_loss(logits, response_mask, loss_agg_mode: str = "token-mean"): +def compute_entropy_loss( + logits, + response_mask, + loss_agg_mode: str = "token-mean"): """Compute categorical entropy loss (For backward compatibility) Args: @@ -985,7 +1088,10 @@ def compute_entropy_loss(logits, response_mask, loss_agg_mode: str = "token-mean """ # compute entropy token_entropy = verl_F.entropy_from_logits(logits) # (bs, response_len) - entropy_loss = agg_loss(loss_mat=token_entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode) + entropy_loss = agg_loss( + loss_mat=token_entropy, + loss_mask=response_mask, + loss_agg_mode=loss_agg_mode) return entropy_loss @@ -1022,16 +1128,24 @@ def compute_value_loss( vf_clipfrac (float): Fraction of elements where the clipped loss was used. """ - vpredclipped = verl_F.clip_by_value(vpreds, values - cliprange_value, values + cliprange_value) + vpredclipped = verl_F.clip_by_value( + vpreds, values - cliprange_value, values + cliprange_value + ) vf_losses1 = (vpreds - returns) ** 2 vf_losses2 = (vpredclipped - returns) ** 2 clipped_vf_losses = torch.max(vf_losses1, vf_losses2) - vf_loss = 0.5 * agg_loss(loss_mat=clipped_vf_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode) - vf_clipfrac = verl_F.masked_mean(torch.gt(vf_losses2, vf_losses1).float(), response_mask) + vf_loss = 0.5 * agg_loss(loss_mat=clipped_vf_losses, + loss_mask=response_mask, + loss_agg_mode=loss_agg_mode) + vf_clipfrac = verl_F.masked_mean( + torch.gt(vf_losses2, vf_losses1).float(), response_mask + ) return vf_loss, vf_clipfrac -def kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty) -> torch.FloatTensor: +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 See more description in http://joschu.net/blog/kl-approx.html @@ -1063,7 +1177,8 @@ def kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_pe 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 + # so, here logprob and ref_logprob should contain the logits for every + # token in vocabulary raise NotImplementedError raise NotImplementedError @@ -1086,7 +1201,9 @@ def compute_pf_ppo_reweight_data( """ @torch.no_grad() - def compute_weights(scores: torch.Tensor, reweight_method: str, weight_pow: float) -> torch.Tensor: + def compute_weights( + scores: torch.Tensor, reweight_method: str, weight_pow: float + ) -> torch.Tensor: """Compute importance weights for resampling based on scores. Args: @@ -1105,7 +1222,9 @@ def compute_weights(scores: torch.Tensor, reweight_method: str, weight_pow: floa elif reweight_method == "max_min": max_score = torch.max(scores) min_score = torch.min(scores) - weights = torch.where((scores == max_score) | (scores == min_score), 1.0, 0.0) + weights = torch.where( + (scores == max_score) | (scores == min_score), 1.0, 0.0 + ) elif reweight_method == "max_random": max_score = torch.max(scores) weights = torch.where(scores == max_score, 0.4, 0.1) @@ -1120,7 +1239,9 @@ def compute_weights(scores: torch.Tensor, reweight_method: str, weight_pow: floa batch_size = scores.shape[0] sample_indices = torch.multinomial(weights, batch_size, replacement=True) - resampled_batch = {key: tensor[sample_indices] for key, tensor in data.batch.items()} + resampled_batch = { + key: tensor[sample_indices] for key, tensor in data.batch.items() + } sample_indices_np = sample_indices.numpy() resampled_non_tensor_batch = {} @@ -1128,7 +1249,8 @@ def compute_weights(scores: torch.Tensor, reweight_method: str, weight_pow: floa if isinstance(array, np.ndarray): resampled_non_tensor_batch[key] = array[sample_indices_np] else: - resampled_non_tensor_batch[key] = [array[i] for i in sample_indices_np] + resampled_non_tensor_batch[key] = [array[i] + for i in sample_indices_np] resampled_meta_info = {} for key, value in data.meta_info.items(): diff --git a/Agent0/executor_train/verl/verl/trainer/ppo/metric_utils.py b/Agent0/executor_train/verl/verl/trainer/ppo/metric_utils.py index 3b6b47b..91a5692 100644 --- a/Agent0/executor_train/verl/verl/trainer/ppo/metric_utils.py +++ b/Agent0/executor_train/verl/verl/trainer/ppo/metric_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -77,7 +77,8 @@ def _compute_response_info(batch: DataProto) -> dict[str, Any]: ) -def compute_data_metrics(batch: DataProto, use_critic: bool = True) -> dict[str, Any]: +def compute_data_metrics( + batch: DataProto, use_critic: bool = True) -> dict[str, Any]: """ Computes various metrics from a batch of data for PPO training. @@ -109,7 +110,8 @@ def compute_data_metrics(batch: DataProto, use_critic: bool = True) -> dict[str, max_response_length = batch.batch["responses"].shape[-1] - prompt_mask = batch.batch["attention_mask"][:, :-max_response_length].bool() + prompt_mask = batch.batch["attention_mask"][:, + :-max_response_length].bool() response_mask = batch.batch["response_mask"].bool() max_prompt_length = prompt_mask.size(-1) @@ -151,7 +153,9 @@ def compute_data_metrics(batch: DataProto, use_critic: bool = True) -> dict[str, "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(), + "critic/vf_explained_var": (1.0 - return_diff_var / (return_var + 1e-5)) + .detach() + .item(), } if use_critic else {} @@ -160,14 +164,20 @@ def compute_data_metrics(batch: DataProto, use_critic: bool = True) -> dict[str, "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()) + "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(), + "prompt_length/clip_ratio": torch.mean( + torch.eq(prompt_length, max_prompt_length).float() + ) + .detach() + .item(), } # multi-turn conversation @@ -180,7 +190,9 @@ def compute_data_metrics(batch: DataProto, use_critic: bool = True) -> dict[str, return metrics -def compute_timing_metrics(batch: DataProto, timing_raw: dict[str, float]) -> dict[str, Any]: +def compute_timing_metrics( + batch: DataProto, timing_raw: dict[str, float] +) -> dict[str, Any]: """ Computes timing metrics for different processing stages in PPO training. @@ -208,21 +220,28 @@ def compute_timing_metrics(batch: DataProto, timing_raw: dict[str, float]) -> di 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"]}, - } + num_tokens_of_section = {"gen": num_response_tokens, + **{name: num_overall_tokens for name in ["ref", + "values", + "adv", + "update_critic", + "update_actor"]}, + } 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] + 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()) }, } -def compute_throughout_metrics(batch: DataProto, timing_raw: dict[str, float], n_gpus: int) -> dict[str, Any]: +def compute_throughout_metrics( + batch: DataProto, timing_raw: dict[str, float], n_gpus: int +) -> dict[str, Any]: """ Computes throughput metrics for PPO training. @@ -292,14 +311,16 @@ def bootstrap_metric( bootstrap_metric_lsts = [[] for _ in range(len(reduce_fns))] for _ in range(n_bootstrap): - bootstrap_idxs = np.random.choice(len(data), size=subset_size, replace=True) + bootstrap_idxs = np.random.choice( + len(data), size=subset_size, replace=True) bootstrap_data = [data[i] for i in bootstrap_idxs] for i, reduce_fn in enumerate(reduce_fns): bootstrap_metric_lsts[i].append(reduce_fn(bootstrap_data)) return [(np.mean(lst), np.std(lst)) for lst in bootstrap_metric_lsts] -def calc_maj_val(data: list[dict[str, Any]], vote_key: str, val_key: str) -> float: +def calc_maj_val(data: list[dict[str, Any]], + vote_key: str, val_key: str) -> float: """ Calculate a value based on majority voting. @@ -336,7 +357,10 @@ def calc_maj_val(data: list[dict[str, Any]], vote_key: str, val_key: str) -> flo def process_validation_metrics( - data_sources: list[str], sample_inputs: list[str], infos_dict: dict[str, list[Any]], seed: int = 42 + data_sources: list[str], + sample_inputs: list[str], + infos_dict: dict[str, list[Any]], + seed: int = 42, ) -> dict[str, dict[str, dict[str, float]]]: """ Process validation metrics into a structured format with statistical analysis. @@ -380,7 +404,9 @@ def process_validation_metrics( >>> # result will contain statistics for each data source and variable """ # Group metrics by data source, prompt and variable - data_src2prompt2var2vals = defaultdict(lambda: defaultdict(lambda: defaultdict(list))) + data_src2prompt2var2vals = defaultdict( + lambda: defaultdict(lambda: defaultdict(list)) + ) for sample_idx, data_source in enumerate(data_sources): prompt = sample_inputs[sample_idx] var2vals = data_src2prompt2var2vals[data_source][prompt] @@ -388,7 +414,9 @@ def process_validation_metrics( var2vals[var_name].append(var_vals[sample_idx]) # Calculate metrics for each group - data_src2prompt2var2metric = defaultdict(lambda: defaultdict(lambda: defaultdict(dict))) + data_src2prompt2var2metric = defaultdict( + lambda: defaultdict(lambda: defaultdict(dict)) + ) for data_source, prompt2var2vals in data_src2prompt2var2vals.items(): for prompt, var2vals in prompt2var2vals.items(): for var_name, var_vals in var2vals.items(): @@ -411,36 +439,62 @@ def process_validation_metrics( for n in ns: [(bon_mean, bon_std), (won_mean, won_std)] = bootstrap_metric( - data=var_vals, subset_size=n, reduce_fns=[np.max, np.min], seed=seed + data=var_vals, + subset_size=n, + reduce_fns=[np.max, np.min], + seed=seed, + ) + metric[f"best@{n}/mean"], metric[f"best@{n}/std"] = ( + bon_mean, + bon_std, + ) + metric[f"worst@{n}/mean"], metric[f"worst@{n}/std"] = ( + won_mean, + won_std, ) - metric[f"best@{n}/mean"], metric[f"best@{n}/std"] = bon_mean, bon_std - metric[f"worst@{n}/mean"], metric[f"worst@{n}/std"] = won_mean, won_std if var2vals.get("pred", None) is not None: vote_data = [ - {"val": val, "pred": pred} for val, pred in zip(var_vals, var2vals["pred"], strict=True) + {"val": val, "pred": pred} + for val, pred in zip( + var_vals, var2vals["pred"], strict=True + ) ] [(maj_n_mean, maj_n_std)] = bootstrap_metric( data=vote_data, subset_size=n, - reduce_fns=[partial(calc_maj_val, vote_key="pred", val_key="val")], + reduce_fns=[ + partial( + calc_maj_val, vote_key="pred", val_key="val" + ) + ], seed=seed, ) - metric[f"maj@{n}/mean"], metric[f"maj@{n}/std"] = maj_n_mean, maj_n_std + metric[f"maj@{n}/mean"], metric[f"maj@{n}/std"] = ( + maj_n_mean, + maj_n_std, + ) data_src2prompt2var2metric[data_source][prompt][var_name] = metric # Aggregate metrics across prompts - data_src2var2metric2prompt_vals = defaultdict(lambda: defaultdict(lambda: defaultdict(list))) + data_src2var2metric2prompt_vals = defaultdict( + lambda: defaultdict(lambda: defaultdict(list)) + ) for data_source, prompt2var2metric in data_src2prompt2var2metric.items(): for prompt, var2metric in prompt2var2metric.items(): for var_name, metric in var2metric.items(): for metric_name, metric_val in metric.items(): - data_src2var2metric2prompt_vals[data_source][var_name][metric_name].append(metric_val) + data_src2var2metric2prompt_vals[data_source][var_name][ + metric_name + ].append(metric_val) - data_src2var2metric2val = defaultdict(lambda: defaultdict(lambda: defaultdict(float))) + data_src2var2metric2val = defaultdict( + lambda: defaultdict(lambda: defaultdict(float)) + ) for data_source, var2metric2prompt_vals in data_src2var2metric2prompt_vals.items(): for var_name, metric2prompt_vals in var2metric2prompt_vals.items(): for metric_name, prompt_vals in metric2prompt_vals.items(): - data_src2var2metric2val[data_source][var_name][metric_name] = np.mean(prompt_vals) + data_src2var2metric2val[data_source][var_name][metric_name] = np.mean( + prompt_vals) return data_src2var2metric2val diff --git a/Agent0/executor_train/verl/verl/trainer/ppo/ray_trainer.py b/Agent0/executor_train/verl/verl/trainer/ppo/ray_trainer.py index 5ba32ac..8f38ad5 100644 --- a/Agent0/executor_train/verl/verl/trainer/ppo/ray_trainer.py +++ b/Agent0/executor_train/verl/verl/trainer/ppo/ray_trainer.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -40,7 +40,11 @@ from verl.experimental.dataset.sampler import AbstractCurriculumSampler from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto from verl.single_controller.base import Worker -from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup +from verl.single_controller.ray import ( + RayClassWithInitArgs, + RayResourcePool, + RayWorkerGroup, +) from verl.single_controller.ray.base import create_colocated_worker_cls from verl.trainer.config import AlgoConfig from verl.trainer.ppo import core_algos @@ -52,12 +56,18 @@ process_validation_metrics, ) from verl.trainer.ppo.reward import compute_reward, compute_reward_async -from verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path, should_save_ckpt_esi +from verl.utils.checkpoint.checkpoint_manager import ( + find_latest_ckpt_path, + should_save_ckpt_esi, +) from verl.utils.debug import marked_timer from verl.utils.metric import ( reduce_metrics, ) -from verl.utils.seqlen_balancing import get_seqlen_balanced_partitions, log_seqlen_unbalance +from verl.utils.seqlen_balancing import ( + get_seqlen_balanced_partitions, + log_seqlen_unbalance, +) from verl.utils.torch_functional import masked_mean from verl.utils.tracking import ValidationGenerationsLogger @@ -86,7 +96,8 @@ class ResourcePoolManager: resource_pool_spec: dict[str, list[int]] mapping: dict[Role, str] - resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict) + resource_pool_dict: dict[str, RayResourcePool] = field( + default_factory=dict) def create_resource_pool(self): """Create Ray resource pools for distributed training. @@ -102,7 +113,10 @@ def create_resource_pool(self): # 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 + 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 @@ -114,27 +128,41 @@ def get_resource_pool(self, role: Role) -> RayResourcePool: def get_n_gpus(self) -> int: """Get the number of gpus in this cluster.""" - return sum([n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes]) + return sum( + [ + n_gpus + for process_on_nodes in self.resource_pool_spec.values() + for n_gpus in process_on_nodes + ] + ) def _check_resource_available(self): """Check if the resource pool can be satisfied in this ray cluster.""" node_available_resources = ray.state.available_resources_per_node() node_available_gpus = { - node: node_info.get("GPU", 0) if "GPU" in node_info else node_info.get("NPU", 0) + node: ( + node_info.get("GPU", 0) + if "GPU" in node_info + else node_info.get("NPU", 0) + ) for node, node_info in node_available_resources.items() } # check total required gpus can be satisfied total_available_gpus = sum(node_available_gpus.values()) total_required_gpus = sum( - [n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes] + [ + n_gpus + for process_on_nodes in self.resource_pool_spec.values() + for n_gpus in process_on_nodes + ] ) if total_available_gpus < total_required_gpus: raise ValueError( - f"Total available GPUs {total_available_gpus} is less than total desired GPUs {total_required_gpus}" - ) + f"Total available GPUs {total_available_gpus} is less than total desired GPUs {total_required_gpus}") - # check each resource pool can be satisfied, O(#resource_pools * #nodes) + # check each resource pool can be satisfied, O(#resource_pools * + # #nodes) for resource_pool_name, process_on_nodes in self.resource_pool_spec.items(): num_gpus, num_nodes = process_on_nodes[0], len(process_on_nodes) for node, available_gpus in node_available_gpus.items(): @@ -150,7 +178,9 @@ def _check_resource_available(self): ) -def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty="kl"): +def apply_kl_penalty( + data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty="kl" +): """Apply KL penalty to the token-level rewards. This function computes the KL divergence between the reference policy and current policy, @@ -172,23 +202,33 @@ def apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, batch_size = data.batch.batch_size[0] # compute kl between ref_policy and current policy - # When apply_kl_penalty, algorithm.use_kl_in_reward=True, so the reference model has been enabled. + # When apply_kl_penalty, algorithm.use_kl_in_reward=True, so the reference + # model has been enabled. kld = core_algos.kl_penalty( - data.batch["old_log_probs"], data.batch["ref_log_prob"], kl_penalty=kl_penalty - ) # (batch_size, response_length) + 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 token_level_rewards = token_level_scores - beta * kld - current_kl = masked_mean(kld, mask=response_mask, axis=-1) # average over sequence + 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 + # 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 = {"actor/reward_kl_penalty": current_kl, "actor/reward_kl_penalty_coeff": beta} + metrics = { + "actor/reward_kl_penalty": current_kl, + "actor/reward_kl_penalty_coeff": beta, + } return data, metrics @@ -243,7 +283,8 @@ def compute_advantage( data.batch["response_mask"] = compute_response_mask(data) # prepare response group if adv_estimator == AdvantageEstimator.GAE: - # Compute advantages and returns using Generalized Advantage Estimation (GAE) + # Compute advantages and returns using Generalized Advantage Estimation + # (GAE) advantages, returns = core_algos.compute_gae_advantage_return( token_level_rewards=data.batch["token_level_rewards"], values=data.batch["values"], @@ -262,7 +303,8 @@ def compute_advantage( elif adv_estimator == AdvantageEstimator.GRPO: # Initialize the mask for GRPO calculation grpo_calculation_mask = data.batch["response_mask"] - # Call compute_grpo_outcome_advantage with parameters matching its definition + # Call compute_grpo_outcome_advantage with parameters matching its + # definition advantages, returns = core_algos.compute_grpo_outcome_advantage( token_level_rewards=data.batch["token_level_rewards"], response_mask=grpo_calculation_mask, @@ -352,7 +394,9 @@ def __init__( 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()=}" + 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 @@ -362,13 +406,17 @@ def __init__( self.device_name = device_name self.validation_generations_logger = ValidationGenerationsLogger() - # if ref_in_actor is True, the reference policy will be actor without lora applied - self.ref_in_actor = config.actor_rollout_ref.model.get("lora_rank", 0) > 0 + # if ref_in_actor is True, the reference policy will be actor without + # lora applied + self.ref_in_actor = config.actor_rollout_ref.model.get( + "lora_rank", 0) > 0 # define in-reward KL control # kl loss control currently not suppoorted if self.config.algorithm.use_kl_in_reward: - self.kl_ctrl_in_reward = core_algos.get_kl_controller(self.config.algorithm.kl_ctrl) + self.kl_ctrl_in_reward = core_algos.get_kl_controller( + self.config.algorithm.kl_ctrl + ) if self.config.algorithm.adv_estimator == AdvantageEstimator.GAE: self.use_critic = True @@ -387,7 +435,11 @@ def __init__( raise NotImplementedError self._validate_config() - self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler) + self._create_dataloader( + train_dataset, + val_dataset, + collate_fn, + train_sampler) def _validate_config(self): config = self.config @@ -395,31 +447,42 @@ def _validate_config(self): n_gpus = config.trainer.n_gpus_per_node * config.trainer.nnodes if config.actor_rollout_ref.actor.strategy == "megatron": model_parallel_size = ( - config.actor_rollout_ref.actor.megatron.tensor_model_parallel_size - * config.actor_rollout_ref.actor.megatron.pipeline_model_parallel_size - ) + config.actor_rollout_ref.actor.megatron.tensor_model_parallel_size * + config.actor_rollout_ref.actor.megatron.pipeline_model_parallel_size) assert ( - n_gpus % (model_parallel_size * config.actor_rollout_ref.actor.megatron.context_parallel_size) == 0 + n_gpus + % ( + model_parallel_size + * config.actor_rollout_ref.actor.megatron.context_parallel_size + ) + == 0 ), ( f"n_gpus ({n_gpus}) must be divisible by model_parallel_size ({model_parallel_size}) times " f"context_parallel_size ({config.actor_rollout_ref.actor.megatron.context_parallel_size})" ) megatron_dp = n_gpus // ( - model_parallel_size * config.actor_rollout_ref.actor.megatron.context_parallel_size + model_parallel_size + * config.actor_rollout_ref.actor.megatron.context_parallel_size + ) + minimal_bsz = ( + megatron_dp + * config.actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu ) - minimal_bsz = megatron_dp * config.actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu else: minimal_bsz = n_gpus # 1. Check total batch size for data correctness - real_train_batch_size = config.data.train_batch_size * config.actor_rollout_ref.rollout.n + real_train_batch_size = ( + config.data.train_batch_size * config.actor_rollout_ref.rollout.n + ) assert real_train_batch_size % minimal_bsz == 0, ( f"real_train_batch_size ({real_train_batch_size}) must be divisible by minimal possible batch size " f"({minimal_bsz})" ) # A helper function to check "micro_batch_size" vs "micro_batch_size_per_gpu" - # We throw an error if the user sets both. The new convention is "..._micro_batch_size_per_gpu". + # We throw an error if the user sets both. The new convention is + # "..._micro_batch_size_per_gpu". def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): """Validate mutually exclusive micro batch size configuration options. @@ -466,14 +529,16 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): ) if self.use_reference_policy: - # reference: log_prob_micro_batch_size vs. log_prob_micro_batch_size_per_gpu + # reference: log_prob_micro_batch_size vs. + # log_prob_micro_batch_size_per_gpu check_mutually_exclusive( config.actor_rollout_ref.ref.log_prob_micro_batch_size, config.actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu, "actor_rollout_ref.ref", ) - # The rollout section also has log_prob_micro_batch_size vs. log_prob_micro_batch_size_per_gpu + # The rollout section also has log_prob_micro_batch_size vs. + # log_prob_micro_batch_size_per_gpu check_mutually_exclusive( config.actor_rollout_ref.rollout.log_prob_micro_batch_size, config.actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu, @@ -483,13 +548,17 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): if self.use_critic and not config.critic.use_dynamic_bsz: # Check for critic micro-batch size conflicts check_mutually_exclusive( - config.critic.ppo_micro_batch_size, config.critic.ppo_micro_batch_size_per_gpu, "critic" + config.critic.ppo_micro_batch_size, + config.critic.ppo_micro_batch_size_per_gpu, + "critic", ) # Check for reward model micro-batch size conflicts if config.reward_model.enable and not config.reward_model.use_dynamic_bsz: check_mutually_exclusive( - config.reward_model.micro_batch_size, config.reward_model.micro_batch_size_per_gpu, "reward_model" + config.reward_model.micro_batch_size, + config.reward_model.micro_batch_size_per_gpu, + "reward_model", ) # Actor @@ -498,15 +567,22 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): # ppo_mini_batch_size is divisible by ppo_micro_batch_size # ppo_micro_batch_size * sequence_parallel_size >= n_gpus if not config.actor_rollout_ref.actor.use_dynamic_bsz: - assert config.data.train_batch_size >= config.actor_rollout_ref.actor.ppo_mini_batch_size - sp_size = config.actor_rollout_ref.actor.get("ulysses_sequence_parallel_size", 1) + assert ( + config.data.train_batch_size + >= config.actor_rollout_ref.actor.ppo_mini_batch_size + ) + sp_size = config.actor_rollout_ref.actor.get( + "ulysses_sequence_parallel_size", 1 + ) if config.actor_rollout_ref.actor.ppo_micro_batch_size is not None: assert ( config.actor_rollout_ref.actor.ppo_mini_batch_size % config.actor_rollout_ref.actor.ppo_micro_batch_size == 0 ) - assert config.actor_rollout_ref.actor.ppo_micro_batch_size * sp_size >= n_gpus + assert ( + config.actor_rollout_ref.actor.ppo_micro_batch_size * + sp_size >= n_gpus) assert config.actor_rollout_ref.actor.loss_agg_mode in [ "token-mean", @@ -515,7 +591,10 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): "seq-mean-token-sum-norm", ], f"Invalid loss_agg_mode: {config.actor_rollout_ref.actor.loss_agg_mode}" - if self.config.algorithm.use_kl_in_reward and config.actor_rollout_ref.actor.use_kl_loss: + if ( + self.config.algorithm.use_kl_in_reward + and config.actor_rollout_ref.actor.use_kl_loss + ): print("NOTICE: You have both enabled in-reward kl and kl loss.") # critic @@ -523,42 +602,51 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): assert config.data.train_batch_size >= config.critic.ppo_mini_batch_size sp_size = config.critic.get("ulysses_sequence_parallel_size", 1) if config.critic.ppo_micro_batch_size is not None: - assert config.critic.ppo_mini_batch_size % config.critic.ppo_micro_batch_size == 0 + assert ( + config.critic.ppo_mini_batch_size + % config.critic.ppo_micro_batch_size + == 0 + ) assert config.critic.ppo_micro_batch_size * sp_size >= n_gpus - # Check if use_remove_padding is enabled when using sequence parallelism for fsdp - if config.actor_rollout_ref.actor.strategy in {"fsdp", "fsdp2"} and ( - config.actor_rollout_ref.actor.get("ulysses_sequence_parallel_size", 1) > 1 - or config.actor_rollout_ref.ref.get("ulysses_sequence_parallel_size", 1) > 1 - ): - assert config.actor_rollout_ref.model.use_remove_padding, ( - "When using sequence parallelism for actor/ref policy, you must enable `use_remove_padding`." - ) + # Check if use_remove_padding is enabled when using sequence + # parallelism for fsdp + if config.actor_rollout_ref.actor.strategy in { + "fsdp", + "fsdp2"} and ( + config.actor_rollout_ref.actor.get( + "ulysses_sequence_parallel_size", + 1) > 1 or config.actor_rollout_ref.ref.get( + "ulysses_sequence_parallel_size", + 1) > 1): + assert ( + config.actor_rollout_ref.model.use_remove_padding + ), "When using sequence parallelism for actor/ref policy, you must enable `use_remove_padding`." if self.use_critic and config.critic.strategy in {"fsdp", "fsdp2"}: if config.critic.get("ulysses_sequence_parallel_size", 1) > 1: - assert config.critic.model.use_remove_padding, ( - "When using sequence parallelism for critic, you must enable `use_remove_padding`." - ) + assert ( + config.critic.model.use_remove_padding + ), "When using sequence parallelism for critic, you must enable `use_remove_padding`." if config.data.get("val_batch_size", None) is not None: print( - "WARNING: val_batch_size is deprecated." - + " Validation datasets are sent to inference engines as a whole batch," - + " which will schedule the memory themselves." - ) + "WARNING: val_batch_size is deprecated." + + " Validation datasets are sent to inference engines as a whole batch," + + " which will schedule the memory themselves.") # check eval config if config.actor_rollout_ref.rollout.val_kwargs.do_sample: - assert config.actor_rollout_ref.rollout.temperature > 0, ( - "validation gen temperature should be greater than 0 when enabling do_sample" - ) + assert ( + config.actor_rollout_ref.rollout.temperature > 0 + ), "validation gen temperature should be greater than 0 when enabling do_sample" # check multi_turn with tool config if config.actor_rollout_ref.rollout.multi_turn.enable: assert ( config.actor_rollout_ref.rollout.multi_turn.tool_config_path is not None - or config.actor_rollout_ref.rollout.multi_turn.interaction_config_path is not None + or config.actor_rollout_ref.rollout.multi_turn.interaction_config_path + is not None ), ( "tool_config_path or interaction_config_path must be set when enabling multi_turn with tool, " "due to no role-playing support" @@ -566,7 +654,12 @@ def check_mutually_exclusive(mbs, mbs_per_gpu, name: str): print("[validate_config] All configuration checks passed successfully!") - def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampler: Optional[Sampler]): + def _create_dataloader( + self, + train_dataset, + val_dataset, + collate_fn, + train_sampler: Optional[Sampler]): """ Creates the train and validation dataloaders. """ @@ -575,16 +668,23 @@ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampl if train_dataset is None: train_dataset = create_rl_dataset( - self.config.data.train_files, self.config.data, self.tokenizer, self.processor + self.config.data.train_files, + self.config.data, + self.tokenizer, + self.processor, ) if val_dataset is None: val_dataset = create_rl_dataset( - self.config.data.val_files, self.config.data, self.tokenizer, self.processor + self.config.data.val_files, + self.config.data, + self.tokenizer, + self.processor, ) self.train_dataset, self.val_dataset = train_dataset, val_dataset if train_sampler is None: - train_sampler = create_rl_sampler(self.config.data, self.train_dataset) + train_sampler = create_rl_sampler( + self.config.data, self.train_dataset) if collate_fn is None: from verl.utils.dataset.rl_dataset import collate_fn as default_collate_fn @@ -594,7 +694,9 @@ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampl self.train_dataloader = StatefulDataLoader( dataset=self.train_dataset, - batch_size=self.config.data.get("gen_batch_size", self.config.data.train_batch_size), + batch_size=self.config.data.get( + "gen_batch_size", self.config.data.train_batch_size + ), num_workers=num_workers, drop_last=True, collate_fn=collate_fn, @@ -622,7 +724,9 @@ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampl f"{len(self.val_dataloader)}" ) - total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs + 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 @@ -633,14 +737,21 @@ def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampl try: OmegaConf.set_struct(self.config, True) with open_dict(self.config): - if OmegaConf.select(self.config, "actor_rollout_ref.actor.optim"): - self.config.actor_rollout_ref.actor.optim.total_training_steps = total_training_steps + if OmegaConf.select( + self.config, + "actor_rollout_ref.actor.optim"): + self.config.actor_rollout_ref.actor.optim.total_training_steps = ( + total_training_steps) if OmegaConf.select(self.config, "critic.optim"): self.config.critic.optim.total_training_steps = total_training_steps except Exception as e: - print(f"Warning: Could not set total_training_steps in config. Structure missing? Error: {e}") + print( + f"Warning: Could not set total_training_steps in config. Structure missing? Error: {e}" + ) - def _dump_generations(self, inputs, outputs, scores, reward_extra_infos_dict, dump_path): + def _dump_generations( + self, inputs, outputs, scores, reward_extra_infos_dict, dump_path + ): """Dump rollout/validation samples as JSONL.""" os.makedirs(dump_path, exist_ok=True) filename = os.path.join(dump_path, f"{self.global_steps}.jsonl") @@ -689,7 +800,9 @@ def _maybe_log_val_generations(self, inputs, outputs, scores): samples = samples[:generations_to_log] # Log to each configured logger - self.validation_generations_logger.log(self.config.trainer.logger, samples, self.global_steps) + self.validation_generations_logger.log( + self.config.trainer.logger, samples, self.global_steps + ) def _validate(self): data_source_lst = [] @@ -706,17 +819,24 @@ def _validate(self): # repeat test batch test_batch = test_batch.repeat( - repeat_times=self.config.actor_rollout_ref.rollout.val_kwargs.n, interleave=True + repeat_times=self.config.actor_rollout_ref.rollout.val_kwargs.n, + interleave=True, ) # 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": + if ( + self.config.reward_model.enable + and test_batch[0].non_tensor_batch["reward_model"]["style"] == "model" + ): return {} # Store original inputs input_ids = test_batch.batch["input_ids"] # TODO: Can we keep special tokens except for padding tokens? - input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids] + input_texts = [ + self.tokenizer.decode(ids, skip_special_tokens=True) + for ids in input_ids + ] sample_inputs.extend(input_texts) batch_keys_to_pop = ["input_ids", "attention_mask", "position_ids"] @@ -751,20 +871,29 @@ def _validate(self): if not self.async_rollout_mode else self.config.actor_rollout_ref.rollout.agent.num_workers ) - test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, size_divisor) + test_gen_batch_padded, pad_size = pad_dataproto_to_divisor( + test_gen_batch, size_divisor + ) if not self.async_rollout_mode: - test_output_gen_batch_padded = self.actor_rollout_wg.generate_sequences(test_gen_batch_padded) + test_output_gen_batch_padded = self.actor_rollout_wg.generate_sequences( + test_gen_batch_padded) else: - test_output_gen_batch_padded = self.async_rollout_manager.generate_sequences(test_gen_batch_padded) + test_output_gen_batch_padded = ( + self.async_rollout_manager.generate_sequences(test_gen_batch_padded)) # unpad - test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size) + test_output_gen_batch = unpad_dataproto( + test_output_gen_batch_padded, pad_size=pad_size + ) print("validation generation end") # Store generated outputs output_ids = test_output_gen_batch.batch["responses"] - output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids] + output_texts = [ + self.tokenizer.decode(ids, skip_special_tokens=True) + for ids in output_ids + ] sample_outputs.extend(output_texts) test_batch = test_batch.union(test_output_gen_batch) @@ -777,19 +906,30 @@ def _validate(self): sample_scores.extend(scores) reward_extra_infos_dict["reward"].extend(scores) - print(f"len reward_extra_infos_dict['reward']: {len(reward_extra_infos_dict['reward'])}") + print( + f"len reward_extra_infos_dict['reward']: {len(reward_extra_infos_dict['reward'])}" + ) if "reward_extra_info" in result: for key, lst in result["reward_extra_info"].items(): reward_extra_infos_dict[key].extend(lst) - print(f"len reward_extra_infos_dict['{key}']: {len(reward_extra_infos_dict[key])}") + print( + f"len reward_extra_infos_dict['{key}']: {len(reward_extra_infos_dict[key])}" + ) # collect num_turns of each prompt if "__num_turns__" in test_batch.non_tensor_batch: - sample_turns.append(test_batch.non_tensor_batch["__num_turns__"]) + sample_turns.append( + test_batch.non_tensor_batch["__num_turns__"]) - data_source_lst.append(test_batch.non_tensor_batch.get("data_source", ["unknown"] * reward_tensor.shape[0])) + data_source_lst.append( + test_batch.non_tensor_batch.get( + "data_source", ["unknown"] * reward_tensor.shape[0] + ) + ) - self._maybe_log_val_generations(inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores) + self._maybe_log_val_generations( + inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores + ) # dump generations val_data_dir = self.config.trainer.get("validation_data_dir", None) @@ -803,20 +943,32 @@ def _validate(self): ) for key_info, lst in reward_extra_infos_dict.items(): - assert len(lst) == 0 or len(lst) == len(sample_scores), f"{key_info}: {len(lst)=}, {len(sample_scores)=}" + assert len(lst) == 0 or len(lst) == len( + sample_scores + ), f"{key_info}: {len(lst)=}, {len(sample_scores)=}" data_sources = np.concatenate(data_source_lst, axis=0) - data_src2var2metric2val = process_validation_metrics(data_sources, sample_inputs, reward_extra_infos_dict) + data_src2var2metric2val = process_validation_metrics( + data_sources, sample_inputs, reward_extra_infos_dict + ) metric_dict = {} for data_source, var2metric2val in data_src2var2metric2val.items(): core_var = "acc" if "acc" in var2metric2val else "reward" for var_name, metric2val in var2metric2val.items(): - n_max = max([int(name.split("@")[-1].split("/")[0]) for name in metric2val.keys()]) + n_max = max( + [ + int(name.split("@")[-1].split("/")[0]) + for name in metric2val.keys() + ] + ) for metric_name, metric_val in metric2val.items(): if ( (var_name == core_var) - and any(metric_name.startswith(pfx) for pfx in ["mean", "maj", "best"]) + and any( + metric_name.startswith(pfx) + for pfx in ["mean", "maj", "best"] + ) and (f"@{n_max}" in metric_name) ): metric_sec = "val-core" @@ -842,30 +994,39 @@ def init_workers(self): """ self.resource_pool_manager.create_resource_pool() - self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()} + 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) + 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", profile_option=self.config.trainer.npu_profile.options, ) - self.resource_pool_to_cls[resource_pool]["actor_rollout"] = actor_rollout_cls + self.resource_pool_to_cls[resource_pool][ + "actor_rollout" + ] = actor_rollout_cls else: raise NotImplementedError # create critic if self.use_critic: - 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) + 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 # create reference policy if needed if self.use_reference_policy: - resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy) + 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, @@ -877,30 +1038,44 @@ def init_workers(self): # 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) + 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. + # See + # https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb + # for more information. all_wg = {} wg_kwargs = {} # Setting up kwargs for RayWorkerGroup - if OmegaConf.select(self.config.trainer, "ray_wait_register_center_timeout") is not None: - wg_kwargs["ray_wait_register_center_timeout"] = self.config.trainer.ray_wait_register_center_timeout + if (OmegaConf.select(self.config.trainer, + "ray_wait_register_center_timeout") is not None): + wg_kwargs["ray_wait_register_center_timeout"] = ( + self.config.trainer.ray_wait_register_center_timeout + ) if OmegaConf.select(self.config.trainer, "profile_steps") is not None: - wg_kwargs["profile_steps"] = OmegaConf.select(self.config.trainer, "profile_steps") - assert OmegaConf.select(self.config.trainer, "worker_nsight_options") is not None, ( - "worker_nsight_options must be set when profile_steps is set" + wg_kwargs["profile_steps"] = OmegaConf.select( + self.config.trainer, "profile_steps" ) + assert ( + OmegaConf.select(self.config.trainer, "worker_nsight_options") + is not None + ), "worker_nsight_options must be set when profile_steps is set" wg_kwargs["worker_nsight_options"] = OmegaConf.to_container( OmegaConf.select(self.config.trainer, "worker_nsight_options") ) for resource_pool, class_dict in self.resource_pool_to_cls.items(): - worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict) + 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, @@ -922,7 +1097,8 @@ def init_workers(self): 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 + # 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() @@ -942,8 +1118,9 @@ def _save_checkpoint(self): # path: given_path + `/global_step_{global_steps}` + `/actor` local_global_step_folder = os.path.join( - self.config.trainer.default_local_dir, f"global_step_{self.global_steps}" - ) + self.config.trainer.default_local_dir, + f"global_step_{ + self.global_steps}") print(f"local_global_step_folder: {local_global_step_folder}") actor_local_path = os.path.join(local_global_step_folder, "actor") @@ -951,47 +1128,68 @@ def _save_checkpoint(self): actor_remote_path = ( None if self.config.trainer.default_hdfs_dir is None - else os.path.join(self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", "actor") + else os.path.join( + self.config.trainer.default_hdfs_dir, + f"global_step_{self.global_steps}", + "actor", + ) ) - remove_previous_ckpt_in_save = self.config.trainer.get("remove_previous_ckpt_in_save", False) + remove_previous_ckpt_in_save = self.config.trainer.get( + "remove_previous_ckpt_in_save", False + ) if remove_previous_ckpt_in_save: print( - "Warning: remove_previous_ckpt_in_save is deprecated," - + " set max_actor_ckpt_to_keep=1 and max_critic_ckpt_to_keep=1 instead" - ) + "Warning: remove_previous_ckpt_in_save is deprecated," + + " set max_actor_ckpt_to_keep=1 and max_critic_ckpt_to_keep=1 instead") max_actor_ckpt_to_keep = ( - self.config.trainer.get("max_actor_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1 + self.config.trainer.get("max_actor_ckpt_to_keep", None) + if not remove_previous_ckpt_in_save + else 1 ) max_critic_ckpt_to_keep = ( - self.config.trainer.get("max_critic_ckpt_to_keep", None) if not remove_previous_ckpt_in_save else 1 + self.config.trainer.get("max_critic_ckpt_to_keep", None) + if not remove_previous_ckpt_in_save + else 1 ) self.actor_rollout_wg.save_checkpoint( - actor_local_path, actor_remote_path, self.global_steps, max_ckpt_to_keep=max_actor_ckpt_to_keep + actor_local_path, + actor_remote_path, + self.global_steps, + max_ckpt_to_keep=max_actor_ckpt_to_keep, ) if self.use_critic: - critic_local_path = os.path.join(local_global_step_folder, "critic") + critic_local_path = os.path.join( + local_global_step_folder, "critic") critic_remote_path = ( None if self.config.trainer.default_hdfs_dir is None - else os.path.join(self.config.trainer.default_hdfs_dir, f"global_step_{self.global_steps}", "critic") + else os.path.join( + self.config.trainer.default_hdfs_dir, + f"global_step_{self.global_steps}", + "critic", + ) ) self.critic_wg.save_checkpoint( - critic_local_path, critic_remote_path, self.global_steps, max_ckpt_to_keep=max_critic_ckpt_to_keep + critic_local_path, + critic_remote_path, + self.global_steps, + max_ckpt_to_keep=max_critic_ckpt_to_keep, ) # save dataloader local_mkdir_safe(local_global_step_folder) - dataloader_local_path = os.path.join(local_global_step_folder, "data.pt") + dataloader_local_path = os.path.join( + local_global_step_folder, "data.pt") dataloader_state_dict = self.train_dataloader.state_dict() torch.save(dataloader_state_dict, dataloader_local_path) # latest checkpointed iteration tracker (for atomic usage) local_latest_checkpointed_iteration = os.path.join( - self.config.trainer.default_local_dir, "latest_checkpointed_iteration.txt" - ) + self.config.trainer.default_local_dir, + "latest_checkpointed_iteration.txt") with open(local_latest_checkpointed_iteration, "w") as f: f.write(str(self.global_steps)) @@ -1003,11 +1201,16 @@ def _load_checkpoint(self): if self.config.trainer.default_hdfs_dir is not None: raise NotImplementedError("load from hdfs is not implemented yet") else: - checkpoint_folder = self.config.trainer.default_local_dir # TODO: check path + checkpoint_folder = ( + self.config.trainer.default_local_dir + ) # TODO: check path if not os.path.isabs(checkpoint_folder): working_dir = os.getcwd() - checkpoint_folder = os.path.join(working_dir, checkpoint_folder) - global_step_folder = find_latest_ckpt_path(checkpoint_folder) # None if no latest + checkpoint_folder = os.path.join( + working_dir, checkpoint_folder) + global_step_folder = find_latest_ckpt_path( + checkpoint_folder + ) # None if no latest # find global_step_folder if self.config.trainer.resume_mode == "auto": @@ -1016,14 +1219,17 @@ def _load_checkpoint(self): return 0 else: if self.config.trainer.resume_mode == "resume_path": - assert isinstance(self.config.trainer.resume_from_path, str), "resume ckpt must be str type" - assert "global_step_" in self.config.trainer.resume_from_path, ( - "resume ckpt must specify the global_steps" - ) + assert isinstance( + self.config.trainer.resume_from_path, str + ), "resume ckpt must be str type" + assert ( + "global_step_" in self.config.trainer.resume_from_path + ), "resume ckpt must specify the global_steps" global_step_folder = self.config.trainer.resume_from_path if not os.path.isabs(global_step_folder): working_dir = os.getcwd() - global_step_folder = os.path.join(working_dir, global_step_folder) + global_step_folder = os.path.join( + working_dir, global_step_folder) print(f"Load from checkpoint folder: {global_step_folder}") # set global step self.global_steps = int(global_step_folder.split("global_step_")[-1]) @@ -1035,37 +1241,51 @@ def _load_checkpoint(self): critic_path = os.path.join(global_step_folder, "critic") # load actor self.actor_rollout_wg.load_checkpoint( - actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load + actor_path, + del_local_after_load=self.config.trainer.del_local_ckpt_after_load, ) # load critic if self.use_critic: self.critic_wg.load_checkpoint( - critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load + critic_path, + del_local_after_load=self.config.trainer.del_local_ckpt_after_load, ) # load dataloader, # TODO: from remote not implemented yet dataloader_local_path = os.path.join(global_step_folder, "data.pt") if os.path.exists(dataloader_local_path): - dataloader_state_dict = torch.load(dataloader_local_path, weights_only=False) + dataloader_state_dict = torch.load( + dataloader_local_path, weights_only=False + ) self.train_dataloader.load_state_dict(dataloader_state_dict) else: - print(f"Warning: No dataloader state found at {dataloader_local_path}, will start from scratch") + print( + f"Warning: No dataloader state found at {dataloader_local_path}, will start from scratch" + ) - def _balance_batch(self, batch: DataProto, metrics, logging_prefix="global_seqlen"): + 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 = batch.batch["attention_mask"].view(batch_size, -1).sum(-1).tolist() # (train_batch_size,) + global_seqlen_lst = ( + batch.batch["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]) + # 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 + seqlen_list=global_seqlen_lst, + partitions=global_partition_lst, + prefix=logging_prefix, ) metrics.update(global_balance_stats) @@ -1094,7 +1314,9 @@ def fit(self): # 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): + if self.val_reward_fn is not None and self.config.trainer.get( + "val_before_train", True + ): val_metrics = self._validate() assert val_metrics, f"{val_metrics=}" pprint(f"Initial validation metrics: {val_metrics}") @@ -1103,7 +1325,11 @@ def fit(self): return # add tqdm - progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc="Training Progress") + progress_bar = tqdm( + total=self.total_training_steps, + initial=self.global_steps, + desc="Training Progress", + ) # we start from step 1 self.global_steps += 1 @@ -1122,7 +1348,9 @@ def fit(self): ) with marked_timer("start_profile", timing_raw): if do_profile: - self.actor_rollout_wg.start_profile(role="e2e", profile_step=self.global_steps) + self.actor_rollout_wg.start_profile( + role="e2e", profile_step=self.global_steps + ) if self.use_reference_policy: self.ref_policy_wg.start_profile() if self.use_critic: @@ -1133,7 +1361,8 @@ def fit(self): batch: DataProto = DataProto.from_single_dict(batch_dict) # pop those keys for generation - batch_keys_to_pop = ["input_ids", "attention_mask", "position_ids"] + batch_keys_to_pop = [ + "input_ids", "attention_mask", "position_ids"] non_tensor_batch_keys_to_pop = ["raw_prompt_ids"] if "multi_modal_data" in batch.non_tensor_batch: non_tensor_batch_keys_to_pop.append("multi_modal_data") @@ -1155,7 +1384,10 @@ def fit(self): # pass global_steps to trace gen_batch.meta_info["global_steps"] = self.global_steps - gen_batch = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True) + gen_batch = gen_batch.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) is_last_step = self.global_steps >= self.total_training_steps @@ -1163,9 +1395,11 @@ def fit(self): # generate a batch with marked_timer("gen", timing_raw, color="red"): if not self.async_rollout_mode: - gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch) + gen_batch_output = self.actor_rollout_wg.generate_sequences( + gen_batch) else: - gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch) + gen_batch_output = ( + self.async_rollout_manager.generate_sequences(gen_batch)) timing_raw.update(gen_batch_output.meta_info["timing"]) gen_batch_output.meta_info.pop("timing", None) @@ -1173,27 +1407,39 @@ def fit(self): with marked_timer("gen_max", timing_raw, color="purple"): gen_baseline_batch = deepcopy(gen_batch) gen_baseline_batch.meta_info["do_sample"] = False - gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch) + gen_baseline_output = ( + self.actor_rollout_wg.generate_sequences( + gen_baseline_batch + ) + ) batch = batch.union(gen_baseline_output) reward_baseline_tensor = self.reward_fn(batch) - reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1) + reward_baseline_tensor = reward_baseline_tensor.sum( + dim=-1) - batch.pop(batch_keys=list(gen_baseline_output.batch.keys())) + batch.pop( + batch_keys=list( + gen_baseline_output.batch.keys())) batch.batch["reward_baselines"] = reward_baseline_tensor del gen_baseline_batch, gen_baseline_output batch.non_tensor_batch["uid"] = np.array( - [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object + [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.repeat( + repeat_times=self.config.actor_rollout_ref.rollout.n, + interleave=True, + ) batch = batch.union(gen_batch_output) if "response_mask" not in batch.batch.keys(): - batch.batch["response_mask"] = compute_response_mask(batch) + batch.batch["response_mask"] = compute_response_mask( + batch) # Balance the number of valid tokens across DP ranks. # NOTE: This usually changes the order of data in the `batch`, # which won't affect the advantage calculation (since it's based on uid), @@ -1203,7 +1449,9 @@ def fit(self): 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_token_num"] = torch.sum( + batch.batch["attention_mask"], dim=-1 + ).tolist() with marked_timer("reward", timing_raw, color="yellow"): # compute reward model score @@ -1212,18 +1460,30 @@ def fit(self): batch = batch.union(reward_tensor) if self.config.reward_model.launch_reward_fn_async: - future_reward = compute_reward_async.remote(batch, self.config, self.tokenizer) + future_reward = compute_reward_async.remote( + batch, self.config, self.tokenizer + ) else: - reward_tensor, reward_extra_infos_dict = compute_reward(batch, self.reward_fn) + reward_tensor, reward_extra_infos_dict = compute_reward( + batch, self.reward_fn) # recompute old_log_probs with marked_timer("old_log_prob", timing_raw, color="blue"): - old_log_prob = self.actor_rollout_wg.compute_log_prob(batch) + old_log_prob = self.actor_rollout_wg.compute_log_prob( + batch) entropys = old_log_prob.batch["entropys"] response_masks = batch.batch["response_mask"] - loss_agg_mode = self.config.actor_rollout_ref.actor.loss_agg_mode - entropy_agg = agg_loss(loss_mat=entropys, loss_mask=response_masks, loss_agg_mode=loss_agg_mode) - old_log_prob_metrics = {"actor/entropy": entropy_agg.detach().item()} + loss_agg_mode = ( + self.config.actor_rollout_ref.actor.loss_agg_mode + ) + entropy_agg = agg_loss( + loss_mat=entropys, + loss_mask=response_masks, + loss_agg_mode=loss_agg_mode, + ) + old_log_prob_metrics = { + "actor/entropy": entropy_agg.detach().item() + } metrics.update(old_log_prob_metrics) old_log_prob.batch.pop("entropys") batch = batch.union(old_log_prob) @@ -1235,30 +1495,38 @@ def fit(self): attention_mask = batch.batch["attention_mask"] responses = batch.batch["responses"] response_length = responses.size(1) - response_mask = attention_mask[:, -response_length:] + response_mask = attention_mask[:, - + response_length:] rollout_probs = torch.exp(rollout_old_log_probs) actor_probs = torch.exp(actor_old_log_probs) - rollout_probs_diff = torch.abs(rollout_probs - actor_probs) - rollout_probs_diff = torch.masked_select(rollout_probs_diff, response_mask.bool()) - rollout_probs_diff_max = torch.max(rollout_probs_diff) - rollout_probs_diff_mean = torch.mean(rollout_probs_diff) - rollout_probs_diff_std = torch.std(rollout_probs_diff) + rollout_probs_diff = torch.abs( + rollout_probs - actor_probs) + rollout_probs_diff = torch.masked_select( + rollout_probs_diff, response_mask.bool() + ) + rollout_probs_diff_max = torch.max( + rollout_probs_diff) + rollout_probs_diff_mean = torch.mean( + rollout_probs_diff) + rollout_probs_diff_std = torch.std( + rollout_probs_diff) metrics.update( { "training/rollout_probs_diff_max": rollout_probs_diff_max.detach().item(), "training/rollout_probs_diff_mean": rollout_probs_diff_mean.detach().item(), "training/rollout_probs_diff_std": rollout_probs_diff_std.detach().item(), - } - ) + }) if self.use_reference_policy: # compute reference log_prob with marked_timer("ref", timing_raw, color="olive"): if not self.ref_in_actor: - ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch) + ref_log_prob = self.ref_policy_wg.compute_ref_log_prob( + batch) else: - ref_log_prob = self.actor_rollout_wg.compute_ref_log_prob(batch) + ref_log_prob = ( + self.actor_rollout_wg.compute_ref_log_prob(batch)) batch = batch.union(ref_log_prob) # compute values @@ -1271,20 +1539,31 @@ def fit(self): # we combine with rule-based rm reward_extra_infos_dict: dict[str, list] if self.config.reward_model.launch_reward_fn_async: - reward_tensor, reward_extra_infos_dict = ray.get(future_reward) + reward_tensor, reward_extra_infos_dict = ray.get( + future_reward + ) batch.batch["token_level_scores"] = reward_tensor if reward_extra_infos_dict: - batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()}) + batch.non_tensor_batch.update( + { + k: np.array(v) + for k, v in reward_extra_infos_dict.items() + } + ) # compute rewards. apply_kl_penalty if available if self.config.algorithm.use_kl_in_reward: batch, kl_metrics = apply_kl_penalty( - batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty + batch, + kl_ctrl=self.kl_ctrl_in_reward, + kl_penalty=self.config.algorithm.kl_penalty, ) metrics.update(kl_metrics) else: - batch.batch["token_level_rewards"] = batch.batch["token_level_scores"] + batch.batch["token_level_rewards"] = batch.batch[ + "token_level_scores" + ] # compute advantages, executed on the driver process @@ -1306,26 +1585,39 @@ def fit(self): if self.use_critic: with marked_timer("update_critic", timing_raw, color="pink"): critic_output = self.critic_wg.update_critic(batch) - critic_output_metrics = reduce_metrics(critic_output.meta_info["metrics"]) + critic_output_metrics = reduce_metrics( + critic_output.meta_info["metrics"] + ) metrics.update(critic_output_metrics) # implement critic warmup if self.config.trainer.critic_warmup <= self.global_steps: # update actor with marked_timer("update_actor", timing_raw, color="red"): - batch.meta_info["multi_turn"] = self.config.actor_rollout_ref.rollout.multi_turn.enable - actor_output = self.actor_rollout_wg.update_actor(batch) - actor_output_metrics = reduce_metrics(actor_output.meta_info["metrics"]) + batch.meta_info["multi_turn"] = ( + self.config.actor_rollout_ref.rollout.multi_turn.enable) + actor_output = self.actor_rollout_wg.update_actor( + batch) + actor_output_metrics = reduce_metrics( + actor_output.meta_info["metrics"] + ) metrics.update(actor_output_metrics) # Log rollout generations if enabled - rollout_data_dir = self.config.trainer.get("rollout_data_dir", None) + rollout_data_dir = self.config.trainer.get( + "rollout_data_dir", None) if rollout_data_dir: - with marked_timer("dump_rollout_generations", timing_raw, color="green"): + with marked_timer( + "dump_rollout_generations", timing_raw, color="green" + ): print(batch.batch.keys()) - inputs = self.tokenizer.batch_decode(batch.batch["prompts"], skip_special_tokens=True) - outputs = self.tokenizer.batch_decode(batch.batch["responses"], skip_special_tokens=True) - scores = batch.batch["token_level_scores"].sum(-1).cpu().tolist() + inputs = self.tokenizer.batch_decode( + batch.batch["prompts"], skip_special_tokens=True) + outputs = self.tokenizer.batch_decode( + batch.batch["responses"], skip_special_tokens=True) + scores = ( + batch.batch["token_level_scores"].sum(-1).cpu().tolist() + ) self._dump_generations( inputs=inputs, outputs=outputs, @@ -1338,7 +1630,10 @@ def fit(self): if ( self.val_reward_fn is not None and self.config.trainer.test_freq > 0 - and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0) + and ( + is_last_step + or self.global_steps % self.config.trainer.test_freq == 0 + ) ): with marked_timer("testing", timing_raw, color="green"): val_metrics: dict = self._validate() @@ -1346,7 +1641,8 @@ def fit(self): last_val_metrics = val_metrics metrics.update(val_metrics) - # Check if the ESI (Elastic Server Instance)/training plan is close to expiration. + # Check if the ESI (Elastic Server Instance)/training plan + # is close to expiration. esi_close_to_expiration = should_save_ckpt_esi( max_steps_duration=self.max_steps_duration, redundant_time=self.config.trainer.esi_redundant_time, @@ -1357,14 +1653,15 @@ def fit(self): # 1. The save frequency is set to a positive value. # 2. It's the last training step. # 3. The current step number is a multiple of the save frequency. - # 4. The ESI(Elastic Server Instance)/training plan is close to expiration. + # 4. The ESI(Elastic Server Instance)/training plan is + # close to expiration. if self.config.trainer.save_freq > 0 and ( - is_last_step - or self.global_steps % self.config.trainer.save_freq == 0 - or esi_close_to_expiration - ): + is_last_step or self.global_steps % + self.config.trainer.save_freq == 0 or esi_close_to_expiration): if esi_close_to_expiration: - print("Force saving checkpoint: ESI instance expiration approaching.") + print( + "Force saving checkpoint: ESI instance expiration approaching." + ) with marked_timer("save_checkpoint", timing_raw, color="green"): self._save_checkpoint() @@ -1379,7 +1676,8 @@ def fit(self): self.rm_wg.stop_profile() steps_duration = timing_raw["step"] - self.max_steps_duration = max(self.max_steps_duration, steps_duration) + self.max_steps_duration = max( + self.max_steps_duration, steps_duration) # training metrics metrics.update( @@ -1389,14 +1687,26 @@ def fit(self): } ) # collect metrics - metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic)) - metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw)) + metrics.update( + compute_data_metrics( + batch=batch, + use_critic=self.use_critic)) + metrics.update( + compute_timing_metrics(batch=batch, timing_raw=timing_raw) + ) # TODO: implement actual tflpo and theoretical tflpo n_gpus = self.resource_pool_manager.get_n_gpus() - metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus)) + metrics.update( + compute_throughout_metrics( + batch=batch, timing_raw=timing_raw, n_gpus=n_gpus + ) + ) - # this is experimental and may be changed/removed in the future in favor of a general-purpose one - if isinstance(self.train_dataloader.sampler, AbstractCurriculumSampler): + # this is experimental and may be changed/removed in the future + # in favor of a general-purpose one + if isinstance( + self.train_dataloader.sampler, + AbstractCurriculumSampler): self.train_dataloader.sampler.update(batch=batch) # TODO: make a canonical logger that supports various backend diff --git a/Agent0/executor_train/verl/verl/trainer/ppo/reward.py b/Agent0/executor_train/verl/verl/trainer/ppo/reward.py index 143b631..ff32775 100644 --- a/Agent0/executor_train/verl/verl/trainer/ppo/reward.py +++ b/Agent0/executor_train/verl/verl/trainer/ppo/reward.py @@ -1,4 +1,4 @@ -# Copyright 2025 Individual Contributor: Thibaut Barroyer +# Copyright 2025-2026 Individual Contributor: Thibaut Barroyer # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -59,7 +59,8 @@ def get_custom_reward_fn(config): return None if not os.path.exists(file_path): - raise FileNotFoundError(f"Reward function file '{file_path}' not found.") + raise FileNotFoundError( + f"Reward function file '{file_path}' not found.") spec = importlib.util.spec_from_file_location("custom_module", file_path) module = importlib.util.module_from_spec(spec) @@ -67,13 +68,17 @@ def get_custom_reward_fn(config): sys.modules["custom_module"] = module spec.loader.exec_module(module) except Exception as e: - raise RuntimeError(f"Error loading module from '{file_path}': {e}") from e + raise RuntimeError( + f"Error loading module from '{file_path}': {e}") from e function_name = reward_fn_config.get("name") if not hasattr(module, function_name): - raise AttributeError(f"Reward function '{function_name}' not found in '{file_path}'.") + raise AttributeError( + f"Reward function '{function_name}' not found in '{file_path}'." + ) - print(f"using customized reward function '{function_name}' from '{file_path}'") + print( + f"using customized reward function '{function_name}' from '{file_path}'") raw_fn = getattr(module, function_name) reward_kwargs = dict(reward_fn_config.get("reward_kwargs", {})) @@ -118,7 +123,9 @@ def load_reward_manager(config, tokenizer, num_examine, **reward_kwargs): if sandbox_url: sandbox_manager = multiprocessing.Manager() # Create a semaphore to control concurrent access to the sandbox - _concurrent_semaphore = sandbox_manager.Semaphore(sandbox_config.get("max_concurrent", 64)) + _concurrent_semaphore = sandbox_manager.Semaphore( + sandbox_config.get("max_concurrent", 64) + ) final_compute_score = partial( default_compute_score, sandbox_fusion_url=sandbox_url, @@ -165,5 +172,7 @@ def compute_reward_async(data: DataProto, config, tokenizer): Load the reward manager and compute the reward for a batch of data. This is meant to be run in a separate Ray worker. """ - reward_fn = load_reward_manager(config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {})) + reward_fn = load_reward_manager( + config, tokenizer, num_examine=0, **config.reward_model.get("reward_kwargs", {}) + ) return compute_reward(data, reward_fn) diff --git a/Agent0/executor_train/verl/verl/utils/__init__.py b/Agent0/executor_train/verl/verl/utils/__init__.py index 0345849..c11673d 100644 --- a/Agent0/executor_train/verl/verl/utils/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -16,4 +16,8 @@ from .config import omega_conf_to_dataclass from .tokenizer import hf_processor, hf_tokenizer -__all__ = tokenizer.__all__ + config.__all__ + ["hf_processor", "hf_tokenizer", "omega_conf_to_dataclass"] +__all__ = ( + tokenizer.__all__ + + config.__all__ + + ["hf_processor", "hf_tokenizer", "omega_conf_to_dataclass"] +) diff --git a/Agent0/executor_train/verl/verl/utils/activation_offload.py b/Agent0/executor_train/verl/verl/utils/activation_offload.py index 73e2e83..73743ee 100644 --- a/Agent0/executor_train/verl/verl/utils/activation_offload.py +++ b/Agent0/executor_train/verl/verl/utils/activation_offload.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -33,7 +33,8 @@ def _get_unique_tensor_key(tensor): - key = (tensor.untyped_storage().data_ptr() + tensor.storage_offset(), tensor.dtype) + key = (tensor.untyped_storage().data_ptr() + + tensor.storage_offset(), tensor.dtype) return key @@ -72,18 +73,24 @@ def __init__( def __enter__(self): self.inside_context = True - torch._C._autograd._push_saved_tensors_default_hooks(self.on_save_for_backward, self.on_get_saved_tensor) + torch._C._autograd._push_saved_tensors_default_hooks( + self.on_save_for_backward, self.on_get_saved_tensor + ) def __exit__(self, *args: Any): self.inside_context = False torch._C._autograd._pop_saved_tensors_default_hooks() def on_save_for_backward(self, tensor: torch.Tensor) -> Any: - retrieve_identifier = self.offload_handler.tensor_push(tensor, **self.handler_extra_kwargs) + retrieve_identifier = self.offload_handler.tensor_push( + tensor, **self.handler_extra_kwargs + ) return retrieve_identifier def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor: - tensor = self.offload_handler.tensor_pop(saved_state, **self.handler_extra_kwargs) + tensor = self.offload_handler.tensor_pop( + saved_state, **self.handler_extra_kwargs + ) return tensor @@ -97,15 +104,13 @@ def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: """Tensor push.""" raise NotImplementedError( "`tensor_push is not implented in OffloadHandler class. Inherit this class and implement your " - "custom tensor_push." - ) + "custom tensor_push.") def tensor_pop(self, tensor_tag: Any, **kwargs): """Tensor pop.""" raise NotImplementedError( "`tensor_pop is not implented in OffloadHandler class. Inherit this class and implement your " - "custom tensor_pop." - ) + "custom tensor_pop.") class GroupCommitFunction(torch.autograd.Function): @@ -140,7 +145,11 @@ class SynchronizedGroupOffloadHandler(OffloadHandler): as the computation kernels, thus the copying will block computation. """ - def __init__(self, num_offload_group, tensor_need_offloading_checker=(lambda _: True)) -> None: + def __init__( + self, + num_offload_group, + tensor_need_offloading_checker=( + lambda _: True)) -> None: super().__init__() self.num_offload_group = num_offload_group @@ -198,7 +207,10 @@ def tensor_push(self, tensor: torch.Tensor, **kwargs): tensor_tag = (self.current_group, self.tensor_count_current_group) self.tensor_count_current_group += 1 assert tensor_tag not in self.tensor_tag_to_state - if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker(tensor): + if ( + self.current_group < self.num_offload_group + and self.tensor_need_offloading_checker(tensor) + ): state = SynchronizedGroupOffloadHandler.offload(tensor) self.tensor_tag_to_state[tensor_tag] = state else: @@ -227,7 +239,8 @@ class AsyncDoubleBufferGroupOffloadHandler(SynchronizedGroupOffloadHandler): def __init__( self, - num_offload_group, # must be <= actual number of groups (number of commits) + num_offload_group, + # must be <= actual number of groups (number of commits) num_model_group, tensor_need_offloading_checker=(lambda t: True), ) -> None: @@ -249,7 +262,9 @@ def __init__( # for optimal CPU/GPU interconnect usage constant = 0 for i in range(self.num_offload_group): - self.layer_window_map[i] = ((self.num_layers // self.num_offload_group) * (i + 1)) - 1 + self.layer_window_map[i] = ( + (self.num_layers // self.num_offload_group) * (i + 1) + ) - 1 if i < (self.num_layers % self.num_offload_group): self.layer_window_map[i] += i + 1 constant = i + 1 @@ -263,10 +278,12 @@ def __init__( def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any: torch_stray_tensor = isinstance( tensor, - torch._subclasses.fake_tensor.FakeTensor | torch._subclasses.functional_tensor.FunctionalTensor, + torch._subclasses.fake_tensor.FakeTensor + | torch._subclasses.functional_tensor.FunctionalTensor, ) need_offload = not torch_stray_tensor - need_offload = need_offload and self.tensor_need_offloading_checker(tensor) + need_offload = need_offload and self.tensor_need_offloading_checker( + tensor) if need_offload: # obtain a unique tensor tag @@ -361,10 +378,12 @@ def bulk_reload_group(self, group_to_reload): offload_mapping = self.group_offload_mapping.pop(group_to_reload) assert offload_mapping is not None for key, state in offload_mapping.items(): - offload_mapping[key] = SynchronizedGroupOffloadHandler.reload(state) + offload_mapping[key] = SynchronizedGroupOffloadHandler.reload( + state) for tensor_label, state in self.tensor_tag_to_state.items(): group_id, _ = tensor_label - if group_id == group_to_reload and not isinstance(state, torch.Tensor): + if group_id == group_to_reload and not isinstance( + state, torch.Tensor): assert isinstance(state, tuple), f"{group_id} {state}" key, shape = state recovered_tensor = offload_mapping[key].view(shape) @@ -378,7 +397,8 @@ def on_group_commit_backward(self): assert self.current_group >= 0 # Layer window data structure helps us to reload at right times - if self.layer_window_map[self.offloaded_group_count - 1] == self.current_group: + if self.layer_window_map[self.offloaded_group_count - + 1] == self.current_group: # Stream synchronization both ways self.h2d_stream.wait_stream(get_torch_device().current_stream()) get_torch_device().current_stream().wait_stream(self.h2d_stream) @@ -396,7 +416,9 @@ def on_group_commit_backward(self): def get_activation_offload_context( - num_layers: int = 1, model_layers: int = 1, tensor_need_offloading_checker=(lambda t: True) + num_layers: int = 1, + model_layers: int = 1, + tensor_need_offloading_checker=(lambda t: True), ): cpu_offload_handler = AsyncDoubleBufferGroupOffloadHandler( num_offload_group=num_layers, @@ -444,11 +466,14 @@ def _pack_kwargs(self, *args, **kwargs): return tuple(flat_args), tuple(kwarg_keys) def _unpack_kwargs(self, flat_args, kwarg_keys): - assert len(kwarg_keys) <= len(flat_args), f"too many keys {len(kwarg_keys)} vs. {len(flat_args)}" + assert len(kwarg_keys) <= len( + flat_args + ), f"too many keys {len(kwarg_keys)} vs. {len(flat_args)}" if len(kwarg_keys) == 0: return flat_args, {} args = flat_args[: -len(kwarg_keys)] - kwargs = dict(zip(kwarg_keys, flat_args[-len(kwarg_keys) :], strict=True)) + kwargs = dict( + zip(kwarg_keys, flat_args[-len(kwarg_keys):], strict=True)) return args, kwargs def _ckpt_forward(self, forward_method, *args, **kwargs): @@ -457,7 +482,8 @@ def _ckpt_forward(self, forward_method, *args, **kwargs): def my_function(*inputs): # unpack back into args and kwargs nonlocal forward_method, kwarg_keys - unpacked_args, unpacked_kwargs = self._unpack_kwargs(inputs, kwarg_keys) + unpacked_args, unpacked_kwargs = self._unpack_kwargs( + inputs, kwarg_keys) # run original module return forward_method(*unpacked_args, **unpacked_kwargs) @@ -518,7 +544,9 @@ def enable_activation_offloading(model, strategy, enable_ckpt=False): """ - assert strategy == "fsdp" or strategy == "fsdp2", "activation offloading only supports fsdp strategy" + assert ( + strategy == "fsdp" or strategy == "fsdp2" + ), "activation offloading only supports fsdp strategy" layers = [] def get_layers(module): @@ -530,17 +558,22 @@ def get_layers(module): if isinstance(child, FSDP): wrapped_module = child._fsdp_wrapped_module # In some cases, torch.nn.Embedding is wrapped with FSDP alone. However, the activation - # size of torch.nn.Embedding is small, so it's not necessary to offload it. + # size of torch.nn.Embedding is small, so it's not necessary to + # offload it. if not isinstance(wrapped_module, torch.nn.Embedding): layers.append(child) get_layers(model) if len(layers) < 3: - logger.warning(f"Find only {len(layers)} fsdp layers, not neccessary to enable async activation offloading") + logger.warning( + f"Find only { + len(layers)} fsdp layers, not neccessary to enable async activation offloading") return tensor_filter = FSDPParameterFilter() - context, sync_func = get_activation_offload_context(len(layers) - 1, len(layers), tensor_filter) + context, sync_func = get_activation_offload_context( + len(layers) - 1, len(layers), tensor_filter + ) if enable_ckpt: # The implementation of activation checkpointing in transformers library is incompatible with # activation offloading, diff --git a/Agent0/executor_train/verl/verl/utils/checkpoint/__init__.py b/Agent0/executor_train/verl/verl/utils/checkpoint/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/utils/checkpoint/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/checkpoint/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/checkpoint/checkpoint_manager.py b/Agent0/executor_train/verl/verl/utils/checkpoint/checkpoint_manager.py index ff861ab..e116670 100644 --- a/Agent0/executor_train/verl/verl/utils/checkpoint/checkpoint_manager.py +++ b/Agent0/executor_train/verl/verl/utils/checkpoint/checkpoint_manager.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -49,8 +49,14 @@ def __init__( checkpoint_config: DictConfig = None, ): self.checkpoint_config = checkpoint_config - checkpoint_load_contents = checkpoint_config.get("load_contents", None) if checkpoint_config else None - checkpoint_save_contents = checkpoint_config.get("save_contents", None) if checkpoint_config else None + checkpoint_load_contents = ( + checkpoint_config.get( + "load_contents", + None) if checkpoint_config else None) + checkpoint_save_contents = ( + checkpoint_config.get( + "save_contents", + None) if checkpoint_config else None) if checkpoint_load_contents is None: checkpoint_load_contents = ["model", "optimizer", "extra"] if checkpoint_save_contents is None: @@ -118,25 +124,38 @@ def should_load_extra(self) -> bool: """ return "extra" in self.checkpoint_load_contents - def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_after_load: bool = False): + def load_checkpoint( + self, + local_path: str, + hdfs_path: str = None, + del_local_after_load: bool = False): raise NotImplementedError def save_checkpoint( - self, local_path: str, hdfs_path: str = None, global_step: int = 0, max_ckpt_to_keep: int = None + self, + local_path: str, + hdfs_path: str = None, + global_step: int = 0, + max_ckpt_to_keep: int = None, ): raise NotImplementedError @staticmethod def checkpath(local_path: str, hdfs_path: str): - assert local_path is not None or hdfs_path is not None, "local_path and hdfs_path cannot be both None" - return local_path is not None, local_path if local_path is not None else hdfs_path + assert ( + local_path is not None or hdfs_path is not None + ), "local_path and hdfs_path cannot be both None" + return local_path is not None, ( + local_path if local_path is not None else hdfs_path + ) def remove_previous_save_local_path(self, path): if isinstance(path, str): path = [path] for p in path: abs_path = os.path.abspath(p) - print(f"Checkpoint manager remove previous save local path: {abs_path}") + print( + f"Checkpoint manager remove previous save local path: {abs_path}") if not os.path.exists(abs_path): continue shutil.rmtree(abs_path, ignore_errors=True) @@ -203,7 +222,10 @@ def get_checkpoint_tracker_filename(root_path: str): return os.path.join(root_path, "latest_checkpointed_iteration.txt") -def should_save_ckpt_esi(max_steps_duration: float, save_ckpt_duration: float = 60, redundant_time: float = 0) -> bool: +def should_save_ckpt_esi( + max_steps_duration: float, + save_ckpt_duration: float = 60, + redundant_time: float = 0) -> bool: """ Determine if checkpoint should be saved based on capacity esi expiration. @@ -212,8 +234,11 @@ def should_save_ckpt_esi(max_steps_duration: float, save_ckpt_duration: float = save_ckpt_duration: Estimated time (seconds) required to save checkpoint (default: 60) redundant_time: Additional buffer time (seconds) for unexpected delays (default: 0) """ - exp_ts_mlp = os.getenv("MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP") # vemlp - exp_ts_aws = os.getenv("SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP") # aws + exp_ts_mlp = os.getenv( + "MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP") # vemlp + exp_ts_aws = os.getenv( + "SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP" + ) # aws if exp_ts_mlp: try: import time @@ -221,17 +246,16 @@ def should_save_ckpt_esi(max_steps_duration: float, save_ckpt_duration: float = remaining = float(exp_ts_mlp) - time.time() except ValueError: return False - return ( - remaining > 0 - and max_steps_duration > 0 - and remaining <= save_ckpt_duration + max_steps_duration + redundant_time - ) + return (remaining > 0 and max_steps_duration > 0 and remaining <= + save_ckpt_duration + max_steps_duration + redundant_time) elif exp_ts_aws: from datetime import datetime, timedelta expiration_time = datetime.fromtimestamp(int(exp_ts_aws)) time_difference = expiration_time - datetime.now() - threshold_minutes = (save_ckpt_duration + max_steps_duration + redundant_time) / 60 + threshold_minutes = ( + save_ckpt_duration + max_steps_duration + redundant_time + ) / 60 return time_difference < timedelta(minutes=threshold_minutes) else: return False diff --git a/Agent0/executor_train/verl/verl/utils/checkpoint/fsdp_checkpoint_manager.py b/Agent0/executor_train/verl/verl/utils/checkpoint/fsdp_checkpoint_manager.py index e042ae8..d4b406c 100644 --- a/Agent0/executor_train/verl/verl/utils/checkpoint/fsdp_checkpoint_manager.py +++ b/Agent0/executor_train/verl/verl/utils/checkpoint/fsdp_checkpoint_manager.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,12 +24,20 @@ from accelerate import init_empty_weights from omegaconf import DictConfig from torch.distributed.fsdp import FullyShardedDataParallel as FSDP -from torch.distributed.fsdp import ShardedOptimStateDictConfig, ShardedStateDictConfig, StateDictType +from torch.distributed.fsdp import ( + ShardedOptimStateDictConfig, + ShardedStateDictConfig, + StateDictType, +) from transformers import GenerationConfig, PreTrainedTokenizer, ProcessorMixin from verl.utils.device import is_cuda_available from verl.utils.fs import copy_to_local, is_non_local, local_mkdir_safe -from verl.utils.fsdp_utils import fsdp_version, get_fsdp_full_state_dict, get_fsdp_state_ctx +from verl.utils.fsdp_utils import ( + fsdp_version, + get_fsdp_full_state_dict, + get_fsdp_state_ctx, +) from verl.utils.logger import log_with_rank from .checkpoint_manager import BaseCheckpointManager @@ -80,7 +88,9 @@ def __init__( if processing_class is None: assert "tokenizer" in kwargs, "tokenizer or processor must be provided" warnings.warn( - "`tokenizer` is deprecated. use `processing_class` instead.", DeprecationWarning, stacklevel=2 + "`tokenizer` is deprecated. use `processing_class` instead.", + DeprecationWarning, + stacklevel=2, ) processing_class = kwargs.pop("tokenizer") @@ -92,7 +102,11 @@ def __init__( checkpoint_config=checkpoint_config, ) - def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_after_load=False): + def load_checkpoint( + self, + local_path: str, + hdfs_path: str = None, + del_local_after_load=False): """ Load an FSDP checkpoint for this rank. @@ -110,11 +124,13 @@ def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_afte # check if the checkpoint_load_contents is valid if self.should_load_model: - assert self.model is not None, "model must be provided when checkpoint_contents.load includes ['model']" + assert ( + self.model is not None + ), "model must be provided when checkpoint_contents.load includes ['model']" if self.should_load_optimizer: - assert self.optimizer is not None, ( - "optimizer must be provided when checkpoint_contents.load includes ['optimizer']" - ) + assert ( + self.optimizer is not None + ), "optimizer must be provided when checkpoint_contents.load includes ['optimizer']" # every rank download its own checkpoint state_dict_cfg = ( @@ -123,47 +139,83 @@ def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_afte else None ) optim_cfg = ( - ShardedOptimStateDictConfig(offload_to_cpu=True if is_cuda_available else False) + ShardedOptimStateDictConfig( + offload_to_cpu=True if is_cuda_available else False + ) if self.should_load_optimizer else None ) - with get_fsdp_state_ctx(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg): + with get_fsdp_state_ctx( + self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg + ): if self.should_load_model: - remote_model_path = os.path.join(local_path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt") + remote_model_path = os.path.join( + local_path, + f"model_world_size_{self.world_size}_rank_{self.rank}.pt", + ) local_model_path = copy_to_local(remote_model_path) - model_state_dict = torch.load(local_model_path, weights_only=False) + model_state_dict = torch.load( + local_model_path, weights_only=False) self.model.load_state_dict(model_state_dict) - log_with_rank(f"Loaded model from {remote_model_path}", rank=self.rank, logger=logger) + log_with_rank( + f"Loaded model from {remote_model_path}", + rank=self.rank, + logger=logger, + ) if self.should_load_optimizer: - remote_optim_path = os.path.join(local_path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt") + remote_optim_path = os.path.join( + local_path, + f"optim_world_size_{self.world_size}_rank_{self.rank}.pt", + ) local_optim_path = copy_to_local(remote_optim_path) - optimizer_state_dict = torch.load(local_optim_path, weights_only=False) + optimizer_state_dict = torch.load( + local_optim_path, weights_only=False) self.optimizer.load_state_dict(optimizer_state_dict) - log_with_rank(f"Loaded optimizer from {remote_optim_path}", rank=self.rank, logger=logger) + log_with_rank( + f"Loaded optimizer from {remote_optim_path}", + rank=self.rank, + logger=logger, + ) if self.should_load_extra: remote_extra_state_path = os.path.join( - local_path, f"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt" - ) + local_path, f"extra_state_world_size_{ + self.world_size}_rank_{ + self.rank}.pt", ) local_extra_state_path = copy_to_local(remote_extra_state_path) - extra_state_dict = torch.load(local_extra_state_path, weights_only=False) + extra_state_dict = torch.load( + local_extra_state_path, weights_only=False) # recover random state if "rng" in extra_state_dict: # 'rng' may not exist for backward compatibility self.load_rng_state(extra_state_dict["rng"]) - log_with_rank(f"Loaded rng from {remote_extra_state_path}", rank=self.rank, logger=logger) + log_with_rank( + f"Loaded rng from {remote_extra_state_path}", + rank=self.rank, + logger=logger, + ) lr_scheduler_state_dict = extra_state_dict["lr_scheduler"] if lr_scheduler_state_dict is not None and self.lr_scheduler is not None: self.lr_scheduler.load_state_dict(lr_scheduler_state_dict) - log_with_rank(f"Loaded lr_scheduler from {remote_extra_state_path}", rank=self.rank, logger=logger) + log_with_rank( + f"Loaded lr_scheduler from {remote_extra_state_path}", + rank=self.rank, + logger=logger, + ) if self.rank == 0 and del_local_after_load: try: - os.remove(local_model_path) if is_non_local(local_model_path) else None - os.remove(local_optim_path) if is_non_local(local_optim_path) else None - os.remove(local_extra_state_path) if is_non_local(local_extra_state_path) else None + os.remove(local_model_path) if is_non_local( + local_model_path) else None + os.remove(local_optim_path) if is_non_local( + local_optim_path) else None + ( + os.remove(local_extra_state_path) + if is_non_local(local_extra_state_path) + else None + ) except Exception as e: log_with_rank( f"remove local resume ckpt file after loading failed, exception {e} will be ignored", @@ -174,7 +226,13 @@ def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_afte # wait for everyone to load checkpoints torch.distributed.barrier() - def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: int = 0, max_ckpt_to_keep=None): + def save_checkpoint( + self, + local_path: str, + hdfs_path: str = None, + global_step: int = 0, + max_ckpt_to_keep=None, + ): """ Save an FSDP checkpoint for this rank. @@ -207,7 +265,8 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i and len(self.previous_saved_paths) >= max_ckpt_to_keep ): keep_start = len(self.previous_saved_paths) - max_ckpt_to_keep + 1 - self.remove_previous_save_local_path(self.previous_saved_paths[:keep_start]) + self.remove_previous_save_local_path( + self.previous_saved_paths[:keep_start]) self.previous_saved_paths = self.previous_saved_paths[keep_start:] local_path = local_mkdir_safe(local_path) @@ -215,40 +274,73 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i # check if the checkpoint_save_contents is valid if self.should_save_model: - assert self.model is not None, "model must be provided when checkpoint_contents.save includes ['model']" + assert ( + self.model is not None + ), "model must be provided when checkpoint_contents.save includes ['model']" if self.should_save_optimizer: - assert self.optimizer is not None, ( - "optimizer must be provided when checkpoint_contents.save includes ['optimizer']" - ) + assert ( + self.optimizer is not None + ), "optimizer must be provided when checkpoint_contents.save includes ['optimizer']" # every rank will save its own model and optim shard - state_dict_cfg = ShardedStateDictConfig(offload_to_cpu=True if is_cuda_available else False) - optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=True if is_cuda_available else False) + state_dict_cfg = ShardedStateDictConfig( + offload_to_cpu=True if is_cuda_available else False + ) + optim_cfg = ShardedOptimStateDictConfig( + offload_to_cpu=True if is_cuda_available else False + ) with warnings.catch_warnings(): warnings.simplefilter("ignore") - with get_fsdp_state_ctx(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg): - model_path = os.path.join(local_path, f"model_world_size_{self.world_size}_rank_{self.rank}.pt") - optim_path = os.path.join(local_path, f"optim_world_size_{self.world_size}_rank_{self.rank}.pt") - extra_path = os.path.join(local_path, f"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt") + with get_fsdp_state_ctx( + self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg + ): + model_path = os.path.join( + local_path, + f"model_world_size_{self.world_size}_rank_{self.rank}.pt", + ) + optim_path = os.path.join( + local_path, + f"optim_world_size_{self.world_size}_rank_{self.rank}.pt", + ) + extra_path = os.path.join( + local_path, f"extra_state_world_size_{ + self.world_size}_rank_{ + self.rank}.pt", ) if self.should_save_model: model_state_dict = self.model.state_dict() torch.save(model_state_dict, model_path) - log_with_rank(f"Saved model to {os.path.abspath(model_path)}", rank=self.rank, logger=logger) + log_with_rank( + f"Saved model to {os.path.abspath(model_path)}", + rank=self.rank, + logger=logger, + ) if self.should_save_optimizer: optimizer_state_dict = self.optimizer.state_dict() torch.save(optimizer_state_dict, optim_path) - log_with_rank(f"Saved optim to {os.path.abspath(optim_path)}", rank=self.rank, logger=logger) + log_with_rank( + f"Saved optim to {os.path.abspath(optim_path)}", + rank=self.rank, + logger=logger, + ) if self.should_save_extra: - lr_scheduler_state_dict = self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None + lr_scheduler_state_dict = ( + self.lr_scheduler.state_dict() + if self.lr_scheduler is not None + else None + ) extra_state_dict = { "lr_scheduler": lr_scheduler_state_dict, "rng": self.get_rng_state(), } torch.save(extra_state_dict, extra_path) - log_with_rank(f"Saved extra_state to {os.path.abspath(extra_path)}", rank=self.rank, logger=logger) + log_with_rank( + f"Saved extra_state to {os.path.abspath(extra_path)}", + rank=self.rank, + logger=logger, + ) if self.rank == 0: # Save HF tokenizer/processor and model config on rank 0 to huggingface/ directory, no matter whether @@ -262,10 +354,16 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i hf_config_tokenizer_path = os.path.join(local_path, "huggingface") local_mkdir_safe(hf_config_tokenizer_path) model_config = unwrap_model.config - if unwrap_model.can_generate() and hasattr(model_config, "name_or_path") and model_config.name_or_path: + if ( + unwrap_model.can_generate() + and hasattr(model_config, "name_or_path") + and model_config.name_or_path + ): # Some model's name_or_path is empty if not initialized from pretrained, # in this cases, we don't save generation config. - generation_config = GenerationConfig.from_pretrained(model_config.name_or_path) + generation_config = GenerationConfig.from_pretrained( + model_config.name_or_path + ) generation_config.save_pretrained(hf_config_tokenizer_path) else: generation_config = None @@ -273,7 +371,8 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i model_config.save_pretrained(hf_config_tokenizer_path) self.processing_class.save_pretrained(hf_config_tokenizer_path) log_with_rank( - f"Saved model config and tokenizer class to {os.path.abspath(hf_config_tokenizer_path)}", + f"Saved model config and tokenizer class to { + os.path.abspath(hf_config_tokenizer_path)}", rank=self.rank, logger=logger, log_only_rank_0=True, @@ -293,8 +392,11 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i if self.should_save_hf_model: # Only rank 0 will save hf model and, - # offload to cpu to save LLMs which may be too large to fit in one GPU - state_dict = get_fsdp_full_state_dict(self.model, offload_to_cpu=True, rank0_only=True) + # offload to cpu to save LLMs which may be too large to fit in one + # GPU + state_dict = get_fsdp_full_state_dict( + self.model, offload_to_cpu=True, rank0_only=True + ) if self.rank == 0: hf_local_path = os.path.join(local_path, "huggingface") @@ -313,10 +415,14 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i auto_model_cls = AutoModelForVision2Seq else: - raise NotImplementedError(f"Unknown architecture {model_config['architectures']}") + raise NotImplementedError( + f"Unknown architecture {model_config['architectures']}" + ) with init_empty_weights(): - save_model = auto_model_cls.from_config(model_config, torch_dtype=torch.bfloat16) + save_model = auto_model_cls.from_config( + model_config, torch_dtype=torch.bfloat16 + ) save_model.to_empty(device="cpu") if save_model.can_generate(): @@ -328,7 +434,8 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i f"in, using a generation config created from the model config when saving hf_model." ) - save_model.save_pretrained(hf_local_path, state_dict=state_dict) + save_model.save_pretrained( + hf_local_path, state_dict=state_dict) log_with_rank( f"Saved hf_model to {os.path.abspath(hf_local_path)}", rank=self.rank, diff --git a/Agent0/executor_train/verl/verl/utils/checkpoint/megatron_checkpoint_manager.py b/Agent0/executor_train/verl/verl/utils/checkpoint/megatron_checkpoint_manager.py index f0071b8..87617e2 100644 --- a/Agent0/executor_train/verl/verl/utils/checkpoint/megatron_checkpoint_manager.py +++ b/Agent0/executor_train/verl/verl/utils/checkpoint/megatron_checkpoint_manager.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,7 +31,10 @@ from verl.utils.device import get_device_name, get_torch_device from verl.utils.fs import is_non_local, local_mkdir_safe from verl.utils.logger import log_with_rank -from verl.utils.megatron.dist_checkpointing import load_dist_checkpointing, save_dist_checkpointing +from verl.utils.megatron.dist_checkpointing import ( + load_dist_checkpointing, + save_dist_checkpointing, +) from verl.utils.megatron_utils import ( get_dist_checkpoint_path, get_hf_model_checkpoint_path, @@ -143,12 +146,15 @@ def __init__( self.use_checkpoint_opt_param_scheduler = use_checkpoint_opt_param_scheduler self.bridge = bridge self.rank = torch.distributed.get_rank() - self.use_dist_checkpointing = use_dist_checkpointing or not self.bridge or self.is_value_model + self.use_dist_checkpointing = ( + use_dist_checkpointing or not self.bridge or self.is_value_model + ) self.use_hf_checkpoint = not self.use_dist_checkpointing self.weight_saver = get_weight_saver(self.arch) - def get_rng_state(self, use_dist_ckpt: bool = True, data_parallel_random_init: bool = False): + def get_rng_state(self, use_dist_ckpt: bool = True, + data_parallel_random_init: bool = False): """collect rng state across data parallel ranks""" rng_state = { "random_rng_state": random.getstate(), @@ -158,12 +164,22 @@ def get_rng_state(self, use_dist_ckpt: bool = True, data_parallel_random_init: b } if get_device_name() != "cpu": - rng_state[f"{get_device_name()}_rng_state"] = get_torch_device().get_rng_state() + rng_state[f"{get_device_name()}_rng_state"] = ( + get_torch_device().get_rng_state() + ) rng_state_list = None - if torch.distributed.is_initialized() and mpu.get_data_parallel_world_size() > 1 and data_parallel_random_init: - rng_state_list = [None for i in range(mpu.get_data_parallel_world_size())] - torch.distributed.all_gather_object(rng_state_list, rng_state, group=mpu.get_data_parallel_group()) + if ( + torch.distributed.is_initialized() + and mpu.get_data_parallel_world_size() > 1 + and data_parallel_random_init + ): + rng_state_list = [ + None for i in range( + mpu.get_data_parallel_world_size())] + torch.distributed.all_gather_object( + rng_state_list, rng_state, group=mpu.get_data_parallel_group() + ) else: rng_state_list = [rng_state] @@ -213,11 +229,17 @@ def get_checkpoint_name( # optimizer, then the optimizer's path must additionally include the # data parallel rank. - # due to the fact that models are identical across cp ranks, cp rank is not used in the checkpoint path + # due to the fact that models are identical across cp ranks, cp rank is + # not used in the checkpoint path if not pipeline_parallel: - common_path = os.path.join(checkpoints_path, f"mp_rank_{tensor_rank:02d}") + common_path = os.path.join( + checkpoints_path, f"mp_rank_{ + tensor_rank:02d}") else: - common_path = os.path.join(checkpoints_path, f"mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}") + common_path = os.path.join( + checkpoints_path, f"mp_rank_{ + tensor_rank:02d}_{ + pipeline_rank:03d}") if expert_parallel: common_path = common_path + f"_{expert_rank:03d}" @@ -234,11 +256,13 @@ def generate_state_dict(self): # All ranks Save Model to reduce memory pressure if self.should_save_model or self.should_load_model: - # Get sharded state dict, notice that state_dict will collect among dp groups, causing memory pressure + # Get sharded state dict, notice that state_dict will collect among + # dp groups, causing memory pressure for vpp_rank, model in enumerate(self.model): if len(self.model) > 1: mpu.set_virtual_pipeline_model_parallel_rank(vpp_rank) - key = f"model{vpp_rank}" if len(self.model) > 1 else "model" + key = f"model{vpp_rank}" if len( + self.model) > 1 else "model" else: key = "model" if hasattr(model, "module"): @@ -248,7 +272,8 @@ def generate_state_dict(self): # Optimizer State Dict if self.should_save_optimizer or self.should_load_optimizer: torch.distributed.barrier() - optimizer_sharded_states = self.optimizer.sharded_state_dict(state_dict) + optimizer_sharded_states = self.optimizer.sharded_state_dict( + state_dict) state_dict["optimizer"] = optimizer_sharded_states if self.lr_scheduler is not None: @@ -263,7 +288,9 @@ def generate_state_dict(self): return state_dict - def load_rng_states(self, rng_states, data_parallel_random_init=False, use_dist_ckpt=True): + def load_rng_states( + self, rng_states, data_parallel_random_init=False, use_dist_ckpt=True + ): # access rng_state for data parallel rank if data_parallel_random_init: rng_states = rng_states[mpu.get_data_parallel_rank()] @@ -274,29 +301,48 @@ def load_rng_states(self, rng_states, data_parallel_random_init=False, use_dist_ torch.set_rng_state(rng_states["torch_rng_state"]) if get_device_name() != "cpu": - get_torch_device().set_rng_state(rng_states[f"{get_device_name()}_rng_state"]) + get_torch_device().set_rng_state( + rng_states[f"{get_device_name()}_rng_state"] + ) # Check for empty states array if not rng_states["rng_tracker_states"]: raise KeyError - tensor_parallel.get_cuda_rng_tracker().set_states(rng_states["rng_tracker_states"]) + tensor_parallel.get_cuda_rng_tracker().set_states( + rng_states["rng_tracker_states"] + ) - def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_after_load=False): + def load_checkpoint( + self, + local_path: str, + hdfs_path: str = None, + del_local_after_load=False): if local_path is not None: - assert os.path.exists(local_path), f"Checkpoint path {local_path} does not exist." + assert os.path.exists( + local_path + ), f"Checkpoint path {local_path} does not exist." dist_checkpoint_path = get_dist_checkpoint_path(local_path) # Get State Dict for loading sharded_state_dict = self.generate_state_dict() - log_with_rank(f"Generated state dict for saving: {sharded_state_dict.keys()}", rank=self.rank, logger=logger) + log_with_rank( + f"Generated state dict for saving: {sharded_state_dict.keys()}", + rank=self.rank, + logger=logger, + ) for vpp_rank, model in enumerate(self.model): if len(self.model) > 1: model_i_keys = sharded_state_dict[f"model{vpp_rank}"].keys() - log_with_rank(f"Generated state dict for saving: {model_i_keys}", rank=self.rank, logger=logger) + log_with_rank( + f"Generated state dict for saving: {model_i_keys}", + rank=self.rank, + logger=logger, + ) else: log_with_rank( - f"Generated state dict for saving: {sharded_state_dict['model'].keys()}", + f"Generated state dict for saving: { + sharded_state_dict['model'].keys()}", rank=self.rank, logger=logger, ) @@ -309,46 +355,67 @@ def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_afte if self.should_load_model and self.use_dist_checkpointing: assert "model" in state_dict or any( - f"model{vpp_rank}" in state_dict for vpp_rank in range(len(self.model)) - ), f"Model state dict not found in {state_dict.keys()}. Please check the checkpoint file {local_path}." + f"model{vpp_rank}" in state_dict for vpp_rank in range( + len( + self.model))), f"Model state dict not found in { + state_dict.keys()}. Please check the checkpoint file {local_path}." for vpp_rank, model in enumerate(self.model): if len(self.model) == 1: model_state_dict = state_dict["model"] else: - assert f"model{vpp_rank}" in state_dict, f"model{vpp_rank} not found in state_dict" + assert ( + f"model{vpp_rank}" in state_dict + ), f"model{vpp_rank} not found in state_dict" model_state_dict = state_dict[f"model{vpp_rank}"] mpu.set_virtual_pipeline_model_parallel_rank(vpp_rank) self.model[vpp_rank].load_state_dict(model_state_dict) - log_with_rank(f"Loaded sharded model checkpoint from {local_path}", rank=self.rank, logger=logger) + log_with_rank( + f"Loaded sharded model checkpoint from {local_path}", + rank=self.rank, + logger=logger, + ) elif self.should_load_model and self.use_hf_checkpoint: hf_model_path = get_hf_model_checkpoint_path(local_path) self.bridge.load_weights(self.model, hf_model_path) - log_with_rank(f"Loaded HF model checkpoint from {hf_model_path} with bridge", rank=self.rank, logger=logger) + log_with_rank( + f"Loaded HF model checkpoint from {hf_model_path} with bridge", + rank=self.rank, + logger=logger, + ) if self.should_load_optimizer: - assert "optimizer" in state_dict, ( - f"Optimizer state dict not found in {state_dict.keys()}. Please check the checkpoint file {local_path}." - ) + assert ( + "optimizer" in state_dict), f"Optimizer state dict not found in { + state_dict.keys()}. Please check the checkpoint file {local_path}." optimizer_state_dict = state_dict["optimizer"] self.optimizer.load_state_dict(optimizer_state_dict) - log_with_rank(f"Loaded optimizer checkpoint from {local_path}", rank=self.rank, logger=logger) + log_with_rank( + f"Loaded optimizer checkpoint from {local_path}", + rank=self.rank, + logger=logger, + ) if self.use_checkpoint_opt_param_scheduler: - assert "lr_scheduler" in state_dict, ( - f"LR scheduler state dict not found in {state_dict.keys()}. Please check the checkpoint file " - f"{local_path}." - ) + assert "lr_scheduler" in state_dict, (f"LR scheduler state dict not found in { + state_dict.keys()}. Please check the checkpoint file " f"{local_path}.") lr_scheduler_state_dict = state_dict["lr_scheduler"] if self.lr_scheduler is not None: self.lr_scheduler.load_state_dict(lr_scheduler_state_dict) - log_with_rank(f"Loaded LR scheduler checkpoint from {local_path}", rank=self.rank, logger=logger) + log_with_rank( + f"Loaded LR scheduler checkpoint from {local_path}", + rank=self.rank, + logger=logger, + ) if self.should_load_extra: - assert "rng_state" in state_dict, ( - f"RNG state dict not found in {state_dict.keys()}. Please check the checkpoint file {local_path}." - ) + assert ( + "rng_state" in state_dict), f"RNG state dict not found in { + state_dict.keys()}. Please check the checkpoint file {local_path}." rng_state = state_dict["rng_state"] self.load_rng_states(rng_state) - log_with_rank(f"Loaded RNG states from {local_path}", rank=self.rank, logger=logger) + log_with_rank( + f"Loaded RNG states from {local_path}", + rank=self.rank, + logger=logger) if del_local_after_load: try: @@ -360,7 +427,13 @@ def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_afte logger=logger, ) - def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: int = 0, max_ckpt_to_keep=None): + def save_checkpoint( + self, + local_path: str, + hdfs_path: str = None, + global_step: int = 0, + max_ckpt_to_keep=None, + ): # record the previous global step self.previous_global_step = global_step @@ -372,7 +445,8 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i and len(self.previous_saved_paths) >= max_ckpt_to_keep ): keep_start = len(self.previous_saved_paths) - max_ckpt_to_keep + 1 - self.remove_previous_save_local_path(self.previous_saved_paths[:keep_start]) + self.remove_previous_save_local_path( + self.previous_saved_paths[:keep_start]) self.previous_saved_paths = self.previous_saved_paths[keep_start:] local_path = local_mkdir_safe(local_path) @@ -381,14 +455,25 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i if self.use_dist_checkpointing: # Generate state dict for saving state_dict = self.generate_state_dict() - log_with_rank(f"Generated state dict for saving: {state_dict.keys()}", rank=self.rank, logger=logger) + log_with_rank( + f"Generated state dict for saving: {state_dict.keys()}", + rank=self.rank, + logger=logger, + ) for vpp_rank, model in enumerate(self.model): if len(self.model) > 1: model_i_keys = state_dict[f"model{vpp_rank}"].keys() - log_with_rank(f"Generated state dict for saving: {model_i_keys}", rank=self.rank, logger=logger) + log_with_rank( + f"Generated state dict for saving: {model_i_keys}", + rank=self.rank, + logger=logger, + ) else: log_with_rank( - f"Generated state dict for saving: {state_dict['model'].keys()}", rank=self.rank, logger=logger + f"Generated state dict for saving: { + state_dict['model'].keys()}", + rank=self.rank, + logger=logger, ) # Start Async save if enabled async_save_request = save_dist_checkpointing( @@ -399,30 +484,50 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i # Synchronize all async save requests if not self.checkpoint_config.async_save: - assert async_save_request is None, "Async save request should be None when not using async save." + assert ( + async_save_request is None + ), "Async save request should be None when not using async save." torch.distributed.barrier() else: - assert self.use_hf_checkpoint, "use_hf_checkpoint should be True when not using dist checkpointing" - log_with_rank(f"Saving HF model checkpoint to {local_path} with bridge", rank=self.rank, logger=logger) + assert ( + self.use_hf_checkpoint + ), "use_hf_checkpoint should be True when not using dist checkpointing" + log_with_rank( + f"Saving HF model checkpoint to {local_path} with bridge", + rank=self.rank, + logger=logger, + ) hf_ckpt_path = get_hf_model_checkpoint_path(local_path) self.bridge.save_weights(self.model, hf_ckpt_path) - log_with_rank(f"Saved bridge checkpoint to {hf_ckpt_path}", rank=self.rank, logger=logger) + log_with_rank( + f"Saved bridge checkpoint to {hf_ckpt_path}", + rank=self.rank, + logger=logger, + ) if self.should_save_model: # Only rank 0 saves the hf config and tokenizer to huggingface path # No matter whether we save hf model or not if self.rank == 0: # Save tokenizer - hf_config_tokenizer_path = get_hf_model_checkpoint_path(local_path) + hf_config_tokenizer_path = get_hf_model_checkpoint_path( + local_path) self.processing_class.save_pretrained(hf_config_tokenizer_path) # Save huggingface config self.hf_config.save_pretrained(hf_config_tokenizer_path) - if hasattr(self.hf_config, "name_or_path") and self.hf_config.name_or_path: + if ( + hasattr(self.hf_config, "name_or_path") + and self.hf_config.name_or_path + ): try: - generation_config = GenerationConfig.from_pretrained(self.hf_config.name_or_path) - generation_config.save_pretrained(hf_config_tokenizer_path) + generation_config = GenerationConfig.from_pretrained( + self.hf_config.name_or_path + ) + generation_config.save_pretrained( + hf_config_tokenizer_path) except Exception: - # if the generation config isn't available, we don't save it + # if the generation config isn't available, we don't + # save it pass log_with_rank( f"Saved Huggingface config and tokenizer to {hf_config_tokenizer_path}", @@ -441,14 +546,16 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i pop_keys = [] for key, value in transformer_config_dict.items(): if type(value) in to_convert_types: - transformer_config_dict[key] = to_convert_types[type(value)](value) + transformer_config_dict[key] = to_convert_types[type(value)]( + value) if type(value) in ignore_types: pop_keys.append(key) if callable(value): pop_keys.append(key) for key in pop_keys: transformer_config_dict.pop(key) - transformer_config_path = get_transformer_config_checkpoint_path(local_path) + transformer_config_path = get_transformer_config_checkpoint_path( + local_path) with open(transformer_config_path, "w") as f: json.dump(transformer_config_dict, f, indent=2) @@ -481,8 +588,11 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i else: from transformers import AutoModelForCausalLM - model = AutoModelForCausalLM.from_pretrained(self.config.model.path, torch_dtype="auto") - model.save_pretrained(hf_model_ckpt_path, state_dict=state_dict) + model = AutoModelForCausalLM.from_pretrained( + self.config.model.path, torch_dtype="auto" + ) + model.save_pretrained( + hf_model_ckpt_path, state_dict=state_dict) log_with_rank( f"Saved Huggingface config and tokenizer to {hf_model_ckpt_path}", rank=self.rank, @@ -492,32 +602,55 @@ def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: i if hdfs_path is not None: log_with_rank( - f"Uploading checkpoint to {hdfs_path}", rank=self.rank, logger=logger, log_only_rank_0=True + f"Uploading checkpoint to {hdfs_path}", + rank=self.rank, + logger=logger, + log_only_rank_0=True, ) from verl.utils import hdfs_io hdfs_io.makedirs(hdfs_path, exist_ok=True) - hdfs_io.copy(src=hf_model_ckpt_path, dst=hdfs_path, dirs_exist_ok=True) + hdfs_io.copy( + src=hf_model_ckpt_path, + dst=hdfs_path, + dirs_exist_ok=True) log_with_rank( - f"HDFS checkpoint uploaded to {hdfs_path}", rank=self.rank, logger=logger, log_only_rank_0=True + f"HDFS checkpoint uploaded to {hdfs_path}", + rank=self.rank, + logger=logger, + log_only_rank_0=True, ) def finalize_save_fn(): # Rank 0 uploads checkpoint to HDFS if hdfs_path is provided log_with_rank( - f"Dist checkpointing save completed for {dist_checkpoint_path}", rank=self.rank, logger=logger + f"Dist checkpointing save completed for {dist_checkpoint_path}", + rank=self.rank, + logger=logger, ) if self.rank == 0: if hdfs_path is not None: - log_with_rank(f"Uploading checkpoint to {hdfs_path}", rank=self.rank, logger=logger) + log_with_rank( + f"Uploading checkpoint to {hdfs_path}", + rank=self.rank, + logger=logger, + ) from verl.utils import hdfs_io hdfs_io.makedirs(hdfs_path, exist_ok=True) - hdfs_io.copy(src=dist_checkpoint_path, dst=hdfs_path, dirs_exist_ok=True) - hdfs_io.copy(src=hf_config_tokenizer_path, dst=hdfs_path, dirs_exist_ok=True) + hdfs_io.copy( + src=dist_checkpoint_path, + dst=hdfs_path, + dirs_exist_ok=True) + hdfs_io.copy( + src=hf_config_tokenizer_path, + dst=hdfs_path, + dirs_exist_ok=True) if self.checkpoint_config.async_save: - assert async_save_request is not None, "Async save request should not be None when using async save." + assert ( + async_save_request is not None + ), "Async save request should not be None when using async save." async_save_request.add_finalize_fn(finalize_save_fn) else: finalize_save_fn() diff --git a/Agent0/executor_train/verl/verl/utils/config.py b/Agent0/executor_train/verl/verl/utils/config.py index f1c301f..8e110a8 100644 --- a/Agent0/executor_train/verl/verl/utils/config.py +++ b/Agent0/executor_train/verl/verl/utils/config.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -20,7 +20,9 @@ __all__ = ["omega_conf_to_dataclass"] -def omega_conf_to_dataclass(config: DictConfig | dict, dataclass_type: Optional[type[Any]] = None) -> Any: +def omega_conf_to_dataclass( + config: DictConfig | dict, dataclass_type: Optional[type[Any]] = None +) -> Any: """ Convert an OmegaConf DictConfig to a dataclass. diff --git a/Agent0/executor_train/verl/verl/utils/dataset/__init__.py b/Agent0/executor_train/verl/verl/utils/dataset/__init__.py index 6032d68..19ce563 100644 --- a/Agent0/executor_train/verl/verl/utils/dataset/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/dataset/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/dataset/multiturn_sft_dataset.py b/Agent0/executor_train/verl/verl/utils/dataset/multiturn_sft_dataset.py index e3eed0f..198402d 100644 --- a/Agent0/executor_train/verl/verl/utils/dataset/multiturn_sft_dataset.py +++ b/Agent0/executor_train/verl/verl/utils/dataset/multiturn_sft_dataset.py @@ -1,5 +1,5 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,14 +32,18 @@ def convert_nested_value_to_list_recursive(data_item): if isinstance(data_item, dict): - return {k: convert_nested_value_to_list_recursive(v) for k, v in data_item.items()} + return {k: convert_nested_value_to_list_recursive( + v) for k, v in data_item.items()} elif isinstance(data_item, list): - return [convert_nested_value_to_list_recursive(elem) for elem in data_item] + return [convert_nested_value_to_list_recursive( + elem) for elem in data_item] elif isinstance(data_item, np.ndarray): - # Convert to list, then recursively process the elements of the new list + # Convert to list, then recursively process the elements of the new + # list return convert_nested_value_to_list_recursive(data_item.tolist()) else: - # Base case: item is already a primitive type (int, str, float, bool, etc.) + # Base case: item is already a primitive type (int, str, float, bool, + # etc.) return data_item @@ -57,7 +61,9 @@ def __init__(self, parquet_files: str | list[str], tokenizer, config=None): multiturn_config = config.get("multiturn", {}) self.messages_key = multiturn_config.get("messages_key", "messages") self.tools_key = multiturn_config.get("tools_key", "tools") - self.enable_thinking_key = multiturn_config.get("enable_thinking_key", "enable_thinking") + self.enable_thinking_key = multiturn_config.get( + "enable_thinking_key", "enable_thinking" + ) assert self.truncation in ["error", "left", "right"] if not isinstance(parquet_files, list): @@ -73,14 +79,19 @@ def __init__(self, parquet_files: str | list[str], tokenizer, config=None): def _download(self): for i, parquet_file in enumerate(self.parquet_files): - self.parquet_files[i] = copy_local_path_from_hdfs(parquet_file, verbose=True) + self.parquet_files[i] = copy_local_path_from_hdfs( + parquet_file, verbose=True + ) def _read_files_and_process(self): def series_to_item(ls): import numpy import pandas - while isinstance(ls, pandas.core.series.Series | numpy.ndarray) and len(ls) == 1: + while ( + isinstance(ls, pandas.core.series.Series | numpy.ndarray) + and len(ls) == 1 + ): ls = ls[0] return ls @@ -91,16 +102,22 @@ def series_to_item(ls): self.dataframe = pd.concat(dataframes) # Extract messages list from dataframe - self.messages = self.dataframe[self.messages_key].apply(series_to_item).tolist() + self.messages = self.dataframe[self.messages_key].apply( + series_to_item).tolist() # Extract tools list from dataframe if self.tools_key in self.dataframe.columns: - self.tools = self.dataframe[self.tools_key].apply(convert_nested_value_to_list_recursive).tolist() + self.tools = ( + self.dataframe[self.tools_key] + .apply(convert_nested_value_to_list_recursive) + .tolist() + ) else: self.tools = None # Extract enable_thinking list from dataframe if self.enable_thinking_key in self.dataframe.columns: - self.enable_thinking = self.dataframe[self.enable_thinking_key].tolist() + self.enable_thinking = self.dataframe[self.enable_thinking_key].tolist( + ) else: self.enable_thinking = None @@ -138,12 +155,14 @@ def _process_message_tokens( tools=tools, ) if is_assistant: - prev_applied_text_w_generation_prompt = self.tokenizer.apply_chat_template( - messages[:start_idx], - tokenize=False, - add_generation_prompt=True, - enable_thinking=enable_thinking, - tools=tools, + prev_applied_text_w_generation_prompt = ( + self.tokenizer.apply_chat_template( + messages[:start_idx], + tokenize=False, + add_generation_prompt=True, + enable_thinking=enable_thinking, + tools=tools, + ) ) else: @@ -158,13 +177,15 @@ def _process_message_tokens( ) # Get tokens for the current message only if is_assistant: - generation_prompt_text = prev_applied_text_w_generation_prompt[len(prev_applied_text) :] + generation_prompt_text = prev_applied_text_w_generation_prompt[ + len(prev_applied_text): + ] generation_prompt_tokens = self.tokenizer.encode( generation_prompt_text, add_special_tokens=False, ) _message_tokens = self.tokenizer.encode( - cur_applied_text[len(prev_applied_text_w_generation_prompt) :], + cur_applied_text[len(prev_applied_text_w_generation_prompt):], add_special_tokens=False, ) message_tokens = generation_prompt_tokens + _message_tokens @@ -173,7 +194,7 @@ def _process_message_tokens( ) else: message_tokens = self.tokenizer.encode( - cur_applied_text[len(prev_applied_text) :], + cur_applied_text[len(prev_applied_text):], add_special_tokens=False, ) loss_mask = [0] * len(message_tokens) @@ -207,7 +228,8 @@ def _validate_and_convert_tokens( a == b for a, b in zip(concat_tokens, full_tokens_list, strict=True) ): logging.warning( - f"Token mismatch detected! Full tokenization length: {len(full_tokens_list)}, Concatenated tokens " + f"Token mismatch detected! Full tokenization length: { + len(full_tokens_list)}, Concatenated tokens " f"length: {len(concat_tokens)}. Using concatenated version." # f"full tokens text: {self.tokenizer.decode(full_tokens_list)}" # f"concat tokens text: {self.tokenizer.decode(concat_tokens)}" @@ -228,7 +250,8 @@ def __getitem__(self, item): tokenizer = self.tokenizer messages = self.messages[item] tools = self.tools[item] if self.tools is not None else None - enable_thinking = self.enable_thinking[item] if self.enable_thinking is not None else None + enable_thinking = ( + self.enable_thinking[item] if self.enable_thinking is not None else None) if self.tools is not None: tools = json.loads(self.tools[item]) @@ -263,7 +286,12 @@ def __getitem__(self, item): if cur_messages["role"] == "assistant": # Process assistant message tokens, loss_mask, attention_mask = self._process_message_tokens( - messages, i, i + 1, is_assistant=True, enable_thinking=enable_thinking, tools=tools + messages, + i, + i + 1, + is_assistant=True, + enable_thinking=enable_thinking, + tools=tools, ) concat_tokens.extend(tokens) concat_loss_mask.extend(loss_mask) @@ -285,7 +313,8 @@ def __getitem__(self, item): elif cur_messages["role"] in ["user", "system"]: # Process user or system message if cur_messages["role"] == "system" and i != 0: - raise ValueError("System message should be the first message") + raise ValueError( + "System message should be the first message") tokens, loss_mask, attention_mask = self._process_message_tokens( messages, i, i + 1, enable_thinking=enable_thinking, tools=tools ) @@ -305,27 +334,42 @@ def __getitem__(self, item): sequence_length = input_ids.shape[0] if sequence_length < self.max_length: # Pad sequences - pad_token_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0 - padded_input_ids = torch.full((self.max_length - sequence_length,), pad_token_id, dtype=input_ids.dtype) - padded_attention_mask = torch.zeros((self.max_length - sequence_length,), dtype=attention_mask.dtype) - padded_loss_mask = torch.zeros((self.max_length - sequence_length,), dtype=loss_mask.dtype) + pad_token_id = ( + self.tokenizer.pad_token_id + if self.tokenizer.pad_token_id is not None + else 0 + ) + padded_input_ids = torch.full( + (self.max_length - sequence_length,), + pad_token_id, + dtype=input_ids.dtype, + ) + padded_attention_mask = torch.zeros( + (self.max_length - sequence_length,), dtype=attention_mask.dtype) + padded_loss_mask = torch.zeros( + (self.max_length - sequence_length,), dtype=loss_mask.dtype + ) input_ids = torch.cat((input_ids, padded_input_ids)) attention_mask = torch.cat((attention_mask, padded_attention_mask)) loss_mask = torch.cat((loss_mask, padded_loss_mask)) elif sequence_length > self.max_length: if self.truncation == "left": - input_ids = input_ids[-self.max_length :] - attention_mask = attention_mask[-self.max_length :] - loss_mask = loss_mask[-self.max_length :] + input_ids = input_ids[-self.max_length:] + attention_mask = attention_mask[-self.max_length:] + loss_mask = loss_mask[-self.max_length:] elif self.truncation == "right": input_ids = input_ids[: self.max_length] attention_mask = attention_mask[: self.max_length] loss_mask = loss_mask[: self.max_length] elif self.truncation == "error": - raise ValueError(f"{sequence_length=} is larger than {self.max_length=}") + raise ValueError( + f"{sequence_length=} is larger than {self.max_length=}" + ) else: - raise ValueError(f"Unknown truncation method {self.truncation}") + raise ValueError( + f"Unknown truncation method { + self.truncation}") # Create position IDs position_ids = torch.arange(len(input_ids), dtype=torch.long) diff --git a/Agent0/executor_train/verl/verl/utils/dataset/rl_dataset.py b/Agent0/executor_train/verl/verl/utils/dataset/rl_dataset.py index e053a67..012238c 100644 --- a/Agent0/executor_train/verl/verl/utils/dataset/rl_dataset.py +++ b/Agent0/executor_train/verl/verl/utils/dataset/rl_dataset.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -43,7 +43,7 @@ def collate_fn(data_list: list[dict]) -> dict: Returns: Dict where tensor entries are stacked into a torch.Tensor of shape - (batch_size, \*dims) and non-tensor entries are converted to + (batch_size, \\*dims) and non-tensor entries are converted to np.ndarray of dtype object with shape (batch_size,). """ tensors = defaultdict(list) @@ -98,7 +98,9 @@ def __init__( self.processor = processor self.config = config - self.cache_dir = os.path.expanduser(config.get("cache_dir", "~/.cache/verl/rlhf")) + self.cache_dir = os.path.expanduser( + config.get("cache_dir", "~/.cache/verl/rlhf") + ) self.prompt_key = config.get("prompt_key", "prompt") self.image_key = config.get("image_key", "images") self.video_key = config.get("video_key", "videos") @@ -106,16 +108,20 @@ def __init__( self.return_raw_chat = config.get("return_raw_chat", False) self.return_full_prompt = config.get("return_full_prompt", False) self.truncation = config.get("truncation", "error") - self.filter_overlong_prompts = config.get("filter_overlong_prompts", True) + self.filter_overlong_prompts = config.get( + "filter_overlong_prompts", True) - self.num_workers = config.get("filter_overlong_prompts_workers", max(1, os.cpu_count() // 4)) + self.num_workers = config.get( + "filter_overlong_prompts_workers", max(1, os.cpu_count() // 4) + ) self.num_workers = min(self.num_workers, os.cpu_count()) self.use_shm = config.get("use_shm", False) self.chat_template_func = config.get("chat_template_func", None) self.need_tools_kwargs = config.get("need_tools_kwargs", False) self.filter_prompts = config.get("filter_prompts", True) self.serialize_dataset = False - self.return_multi_modal_inputs = config.get("return_multi_modal_inputs", True) + self.return_multi_modal_inputs = config.get( + "return_multi_modal_inputs", True) self._download() self._read_files_and_tokenize() @@ -123,23 +129,30 @@ def __init__( def _download(self, use_origin_parquet=False): from verl.utils.fs import copy_to_local - data_files = self.data_files if not use_origin_parquet else self.original_data_files + data_files = ( + self.data_files if not use_origin_parquet else self.original_data_files) for i, parquet_file in enumerate(data_files): - self.data_files[i] = copy_to_local(src=parquet_file, cache_dir=self.cache_dir, use_shm=self.use_shm) + self.data_files[i] = copy_to_local( + src=parquet_file, + cache_dir=self.cache_dir, + use_shm=self.use_shm) def _read_files_and_tokenize(self): dataframes = [] for parquet_file in self.data_files: # read parquet files and cache - dataframe = datasets.load_dataset("parquet", data_files=parquet_file)["train"] + dataframe = datasets.load_dataset( + "parquet", data_files=parquet_file)["train"] dataframes.append(dataframe) - self.dataframe: datasets.Dataset = datasets.concatenate_datasets(dataframes) + self.dataframe: datasets.Dataset = datasets.concatenate_datasets( + dataframes) print(f"dataset len: {len(self.dataframe)}") self.dataframe = self.maybe_filter_out_long_prompts(self.dataframe) - def maybe_filter_out_long_prompts(self, dataframe: datasets.Dataset = None): + def maybe_filter_out_long_prompts( + self, dataframe: datasets.Dataset = None): # filter out too long prompts if self.filter_overlong_prompts: tokenizer = self.tokenizer @@ -157,23 +170,36 @@ def doc2len(doc) -> int: messages, add_generation_prompt=True, tokenize=False ) images = ( - [process_image(image) for image in messages.pop(image_key)] if image_key in messages else None + [process_image(image) for image in messages.pop(image_key)] + if image_key in messages + else None ) videos = ( - [process_video(video) for video in messages.pop(video_key)] if video_key in messages else None + [process_video(video) for video in messages.pop(video_key)] + if video_key in messages + else None ) - return len(processor(text=[raw_prompt], images=images, videos=videos)["input_ids"][0]) + return len( + processor( + text=[raw_prompt], + images=images, + videos=videos)["input_ids"][0]) else: def doc2len(doc) -> int: - return len(tokenizer.apply_chat_template(doc[prompt_key], add_generation_prompt=True)) + return len( + tokenizer.apply_chat_template( + doc[prompt_key], add_generation_prompt=True + ) + ) dataframe = dataframe.filter( lambda doc: doc2len(doc) <= self.max_prompt_length, num_proc=self.num_workers, - desc=f"Filtering prompts longer than {self.max_prompt_length} tokens", + desc=f"Filtering prompts longer than { + self.max_prompt_length} tokens", ) print(f"filter dataset len: {len(dataframe)}") @@ -183,10 +209,14 @@ def resume_dataset_state(self): self.serialize_dataset = not hasattr(self, "original_data_files") # resume dataframe if not it's serialized in data.pt if not self.serialize_dataset: - self._download(use_origin_parquet=True) # download and resume from original parquet files + self._download( + use_origin_parquet=True + ) # download and resume from original parquet files self._read_files_and_tokenize() else: - print(r"old dataloader ckpt file is used, please train from scratch for better ckpt performance") + print( + r"old dataloader ckpt file is used, please train from scratch for better ckpt performance" + ) def __len__(self): return len(self.dataframe) @@ -223,26 +253,44 @@ def __getitem__(self, item): if self.processor is not None: from verl.utils.dataset.vision_utils import process_image, process_video - raw_prompt = self.processor.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) + raw_prompt = self.processor.apply_chat_template( + messages, add_generation_prompt=True, tokenize=False + ) multi_modal_data = {} images = None - if self.image_key in row_dict and row_dict.get(self.image_key, None) is not None: - images = [process_image(image) for image in row_dict.pop(self.image_key)] + if ( + self.image_key in row_dict + and row_dict.get(self.image_key, None) is not None + ): + images = [ + process_image(image) for image in row_dict.pop( + self.image_key)] # due to the image key is "image" instead of "images" in vllm, we need to use "image" here - # link: https://github.com/vllm-project/vllm/blob/3c545c0c3b98ee642373a308197d750d0e449403/vllm/multimodal/parse.py#L205 + # link: + # https://github.com/vllm-project/vllm/blob/3c545c0c3b98ee642373a308197d750d0e449403/vllm/multimodal/parse.py#L205 multi_modal_data["image"] = images videos = None - if self.video_key in row_dict and row_dict.get(self.video_key, None) is not None: - videos = [process_video(video) for video in row_dict.pop(self.video_key)] + if ( + self.video_key in row_dict + and row_dict.get(self.video_key, None) is not None + ): + videos = [ + process_video(video) for video in row_dict.pop( + self.video_key)] # due to the video key is "video" instead of "videos" in vllm, we need to use "video" here - # link: https://github.com/vllm-project/vllm/blob/3c545c0c3b98ee642373a308197d750d0e449403/vllm/multimodal/parse.py#L205 + # link: + # https://github.com/vllm-project/vllm/blob/3c545c0c3b98ee642373a308197d750d0e449403/vllm/multimodal/parse.py#L205 multi_modal_data["video"] = [video.numpy() for video in videos] - model_inputs = self.processor(text=[raw_prompt], images=images, videos=videos, return_tensors="pt") + model_inputs = self.processor( + text=[raw_prompt], + images=images, + videos=videos, + return_tensors="pt") input_ids = model_inputs.pop("input_ids") attention_mask = model_inputs.pop("attention_mask") @@ -250,11 +298,13 @@ def __getitem__(self, item): if "second_per_grid_ts" in model_inputs: model_inputs.pop("second_per_grid_ts") - # There's a trap here, multi_modal_inputs has to be a dict, not BatchFeature + # There's a trap here, multi_modal_inputs has to be a dict, not + # BatchFeature row_dict["multi_modal_data"] = multi_modal_data # We will do batch.union() in the trainer, - # so we cannot have "multi_modal_inputs" in row_dict if rollout generates new multi_modal_inputs + # so we cannot have "multi_modal_inputs" in row_dict if rollout + # generates new multi_modal_inputs if self.return_multi_modal_inputs: row_dict["multi_modal_inputs"] = dict(model_inputs) @@ -262,8 +312,12 @@ def __getitem__(self, item): row_dict["multi_modal_inputs"].pop("second_per_grid_ts", None) else: - raw_prompt = self.tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=False) - model_inputs = self.tokenizer(raw_prompt, return_tensors="pt", add_special_tokens=False) + raw_prompt = self.tokenizer.apply_chat_template( + messages, add_generation_prompt=True, tokenize=False + ) + model_inputs = self.tokenizer( + raw_prompt, return_tensors="pt", add_special_tokens=False + ) input_ids = model_inputs.pop("input_ids") attention_mask = model_inputs.pop("attention_mask") @@ -276,7 +330,11 @@ def __getitem__(self, item): truncation=self.truncation, ) - if self.processor is not None and "Qwen2VLImageProcessor" in self.processor.image_processor.__class__.__name__: + if ( + self.processor is not None + and "Qwen2VLImageProcessor" + in self.processor.image_processor.__class__.__name__ + ): from verl.models.transformers.qwen2_vl import get_rope_index position_ids = [ @@ -297,18 +355,24 @@ def __getitem__(self, item): row_dict["attention_mask"] = attention_mask[0] row_dict["position_ids"] = position_ids[0] - raw_prompt_ids = self.tokenizer.encode(raw_prompt, add_special_tokens=False) + raw_prompt_ids = self.tokenizer.encode( + raw_prompt, add_special_tokens=False) if len(raw_prompt_ids) > self.max_prompt_length: if self.truncation == "left": - raw_prompt_ids = raw_prompt_ids[-self.max_prompt_length :] + raw_prompt_ids = raw_prompt_ids[-self.max_prompt_length:] elif self.truncation == "right": raw_prompt_ids = raw_prompt_ids[: self.max_prompt_length] elif self.truncation == "middle": left_half = self.max_prompt_length // 2 right_half = self.max_prompt_length - left_half - raw_prompt_ids = raw_prompt_ids[:left_half] + raw_prompt_ids[-right_half:] + raw_prompt_ids = ( + raw_prompt_ids[:left_half] + raw_prompt_ids[-right_half:] + ) elif self.truncation == "error": - raise RuntimeError(f"Prompt length {len(raw_prompt_ids)} is longer than {self.max_prompt_length}.") + raise RuntimeError( + f"Prompt length { + len(raw_prompt_ids)} is longer than { + self.max_prompt_length}.") row_dict["raw_prompt_ids"] = raw_prompt_ids # encode prompts without chat template @@ -322,10 +386,18 @@ def __getitem__(self, item): # add index for each prompt index = row_dict.get("extra_info", {}).get("index", 0) tools_kwargs = row_dict.get("extra_info", {}).get("tools_kwargs", {}) - interaction_kwargs = row_dict.get("extra_info", {}).get("interaction_kwargs", {}) - need_tools_kwargs = row_dict.get("extra_info", {}).get("need_tools_kwargs", self.need_tools_kwargs) + interaction_kwargs = row_dict.get("extra_info", {}).get( + "interaction_kwargs", {} + ) + need_tools_kwargs = row_dict.get("extra_info", {}).get( + "need_tools_kwargs", self.need_tools_kwargs + ) if need_tools_kwargs and not tools_kwargs: - logger.warning("tools_kwargs is empty for index {}, data source: {}", index, row_dict["data_source"]) + logger.warning( + "tools_kwargs is empty for index {}, data source: {}", + index, + row_dict["data_source"], + ) row_dict["index"] = index row_dict["tools_kwargs"] = tools_kwargs row_dict["interaction_kwargs"] = interaction_kwargs diff --git a/Agent0/executor_train/verl/verl/utils/dataset/rm_dataset.py b/Agent0/executor_train/verl/verl/utils/dataset/rm_dataset.py index 7af7923..ba7519a 100644 --- a/Agent0/executor_train/verl/verl/utils/dataset/rm_dataset.py +++ b/Agent0/executor_train/verl/verl/utils/dataset/rm_dataset.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -74,7 +74,8 @@ def _download_files(): 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)) + 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 @@ -100,10 +101,26 @@ def _pad_to_length(self, input_ids, attention_mask): 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 + (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 + ( + 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] @@ -116,27 +133,41 @@ def __getitem__(self, 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] + 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) + 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 + (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_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) + 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((chosen_attention_mask, rejected_attention_mask), dim=0) + attention_mask = torch.stack( + (chosen_attention_mask, rejected_attention_mask), dim=0 + ) return { "input_ids": input_ids, diff --git a/Agent0/executor_train/verl/verl/utils/dataset/sft_dataset.py b/Agent0/executor_train/verl/verl/utils/dataset/sft_dataset.py index 2aa7b20..1ef485e 100644 --- a/Agent0/executor_train/verl/verl/utils/dataset/sft_dataset.py +++ b/Agent0/executor_train/verl/verl/utils/dataset/sft_dataset.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -58,8 +58,14 @@ def __init__(self, parquet_files: str | ListConfig, tokenizer, config): tokenizer = hf_tokenizer(tokenizer) self.tokenizer: PreTrainedTokenizer = tokenizer - self.prompt_key = prompt_key if isinstance(prompt_key, tuple | list) else [prompt_key] - self.response_key = response_key if isinstance(response_key, tuple | list) else [response_key] + self.prompt_key = ( + prompt_key if isinstance( + prompt_key, + tuple | list) else [prompt_key]) + self.response_key = ( + response_key if isinstance( + response_key, + tuple | list) else [response_key]) self.prompt_dict_keys = prompt_dict_keys if prompt_dict_keys else [] self.response_dict_keys = response_dict_keys if response_dict_keys else [] @@ -70,14 +76,19 @@ def __init__(self, parquet_files: str | ListConfig, tokenizer, config): def _download(self): for i, parquet_file in enumerate(self.parquet_files): - self.parquet_files[i] = copy_to_local(parquet_file, verbose=True, use_shm=self.use_shm) + self.parquet_files[i] = copy_to_local( + parquet_file, verbose=True, use_shm=self.use_shm + ) def _read_files_and_tokenize(self): def series_to_item(ls): import numpy import pandas - while isinstance(ls, pandas.core.series.Series | numpy.ndarray) and len(ls) == 1: + while ( + isinstance(ls, pandas.core.series.Series | numpy.ndarray) + and len(ls) == 1 + ): ls = ls[0] return ls @@ -93,7 +104,9 @@ def series_to_item(ls): # type(x[0]): numpy.ndarray # type(x[0][0]): dict try: - self.prompts = self.prompts.apply(lambda x: series_to_item(x)[key], axis=1) # noqa: B023 + self.prompts = self.prompts.apply( + lambda x: series_to_item(x)[key], axis=1 + ) # noqa: B023 except Exception: print(f"self.prompts={self.prompts}") raise @@ -103,7 +116,9 @@ def series_to_item(ls): self.responses = self.dataframe[self.response_key] for key in self.response_dict_keys: try: - self.responses = self.responses.apply(lambda x: series_to_item(x)[key], axis=1) # noqa: B023 + self.responses = self.responses.apply( + lambda x: series_to_item(x)[key], axis=1 + ) # noqa: B023 except Exception: print(f"self.responses={self.responses}") raise @@ -124,15 +139,21 @@ def __getitem__(self, item): prompt_chat = [{"role": "user", "content": prompt}] # string - prompt_chat_str = tokenizer.apply_chat_template(prompt_chat, add_generation_prompt=True, tokenize=False) + prompt_chat_str = tokenizer.apply_chat_template( + prompt_chat, add_generation_prompt=True, tokenize=False + ) response_chat_str = response + tokenizer.eos_token # tokenize - prompt_ids_output = tokenizer(prompt_chat_str, return_tensors="pt", add_special_tokens=False) + prompt_ids_output = tokenizer( + prompt_chat_str, return_tensors="pt", add_special_tokens=False + ) prompt_ids = prompt_ids_output["input_ids"][0] prompt_attention_mask = prompt_ids_output["attention_mask"][0] - response_ids_output = tokenizer(response_chat_str, return_tensors="pt", add_special_tokens=False) + response_ids_output = tokenizer( + response_chat_str, return_tensors="pt", add_special_tokens=False + ) response_ids = response_ids_output["input_ids"][0] response_attention_mask = response_ids_output["attention_mask"][0] @@ -140,31 +161,46 @@ def __getitem__(self, item): response_length = response_ids.shape[0] input_ids = torch.cat((prompt_ids, response_ids), dim=-1) - attention_mask = torch.cat((prompt_attention_mask, response_attention_mask), dim=-1) + attention_mask = torch.cat( + (prompt_attention_mask, response_attention_mask), dim=-1 + ) # padding to max length sequence_length = input_ids.shape[0] if sequence_length < self.max_length: padded_input_ids = ( - torch.ones(size=(self.max_length - sequence_length,), dtype=input_ids.dtype) - * self.tokenizer.pad_token_id - ) - padded_attention_mask = torch.zeros(size=(self.max_length - sequence_length,), dtype=attention_mask.dtype) + torch.ones( + size=( + self.max_length - + sequence_length, + ), + dtype=input_ids.dtype) * + self.tokenizer.pad_token_id) + padded_attention_mask = torch.zeros( + size=( + self.max_length - + sequence_length, + ), + dtype=attention_mask.dtype) input_ids = torch.cat((input_ids, padded_input_ids)) attention_mask = torch.cat((attention_mask, padded_attention_mask)) elif sequence_length > self.max_length: if self.truncation == "left": # actually, left truncation may not be reasonable - input_ids = input_ids[-self.max_length :] - attention_mask = attention_mask[-self.max_length :] + input_ids = input_ids[-self.max_length:] + attention_mask = attention_mask[-self.max_length:] elif self.truncation == "right": input_ids = input_ids[: self.max_length] attention_mask = attention_mask[: self.max_length] elif self.truncation == "error": - raise NotImplementedError(f"{sequence_length=} is larger than {self.max_length=}") + raise NotImplementedError( + f"{sequence_length=} is larger than {self.max_length=}" + ) else: - raise NotImplementedError(f"Unknown truncation method {self.truncation}") + raise NotImplementedError( + f"Unknown truncation method {self.truncation}" + ) position_ids = compute_position_id_with_mask(attention_mask) @@ -173,7 +209,8 @@ def __getitem__(self, item): # mask out prompt for SFT. loss_mask[: min(prompt_length, loss_mask.size(0)) - 1] = 0 # mask out the last token in response - loss_mask[min(prompt_length + response_length, loss_mask.size(0)) - 1] = 0 + loss_mask[min(prompt_length + response_length, + loss_mask.size(0)) - 1] = 0 return { "input_ids": input_ids, diff --git a/Agent0/executor_train/verl/verl/utils/dataset/vision_utils.py b/Agent0/executor_train/verl/verl/utils/dataset/vision_utils.py index 75cce7f..6c63cce 100644 --- a/Agent0/executor_train/verl/verl/utils/dataset/vision_utils.py +++ b/Agent0/executor_train/verl/verl/utils/dataset/vision_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -92,14 +92,18 @@ def process_video( return fetch_video(video) -def process_multi_modal_inputs_for_minicpmo(input_ids, attention_mask, position_ids, cu_seqlens, multi_modal_inputs): +def process_multi_modal_inputs_for_minicpmo( + input_ids, attention_mask, position_ids, cu_seqlens, multi_modal_inputs +): # Adjust image bounds based on left padding and cumulative sequence lengths # This is necessary for MiniCPM-o's vision-language alignment left_padding_length = torch.argmax(attention_mask, dim=1) image_bounds = [] for i in range(len(multi_modal_inputs["image_bound"])): image_bound = ( - multi_modal_inputs["image_bound"][i].to(left_padding_length.device) - left_padding_length[i] + cu_seqlens[i] + multi_modal_inputs["image_bound"][i].to(left_padding_length.device) + - left_padding_length[i] + + cu_seqlens[i] ) image_bounds.append(image_bound) @@ -110,7 +114,8 @@ def process_multi_modal_inputs_for_minicpmo(input_ids, attention_mask, position_ multi_modal_inputs["pixel_values"] = [pixel_values] multi_modal_inputs["image_bound"] = [torch.vstack(image_bounds)] - multi_modal_inputs["tgt_sizes"] = [torch.vstack(multi_modal_inputs["tgt_sizes"])] + multi_modal_inputs["tgt_sizes"] = [ + torch.vstack(multi_modal_inputs["tgt_sizes"])] multi_modal_inputs["input_ids"] = input_ids multi_modal_inputs["attention_mask"] = attention_mask multi_modal_inputs["position_ids"] = position_ids diff --git a/Agent0/executor_train/verl/verl/utils/debug/__init__.py b/Agent0/executor_train/verl/verl/utils/debug/__init__.py index eb67df1..4716788 100644 --- a/Agent0/executor_train/verl/verl/utils/debug/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/debug/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/debug/performance.py b/Agent0/executor_train/verl/verl/utils/debug/performance.py index 9186e12..a3dac23 100644 --- a/Agent0/executor_train/verl/verl/utils/debug/performance.py +++ b/Agent0/executor_train/verl/verl/utils/debug/performance.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -13,5 +13,6 @@ # limitations under the License. # APIs kept for backward compatibility purpose -# This file is deprecated, for new features please develop in profiler/performance.py +# This file is deprecated, for new features please develop in +# profiler/performance.py from verl.utils.profiler.performance import simple_timer, reduce_timing # noqa diff --git a/Agent0/executor_train/verl/verl/utils/debug/trajectory_tracker.py b/Agent0/executor_train/verl/verl/utils/debug/trajectory_tracker.py index 73afb85..a65481f 100644 --- a/Agent0/executor_train/verl/verl/utils/debug/trajectory_tracker.py +++ b/Agent0/executor_train/verl/verl/utils/debug/trajectory_tracker.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -58,7 +58,12 @@ def __init__(self, hdfs_dir, verbose) -> None: 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)) + self.handle.append( + save_to_hdfs.remote( + data, + name, + self.hdfs_dir, + self.verbose)) def wait_for_hdfs(self): while len(self.handle) != 0: @@ -80,9 +85,9 @@ 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 - ) + tracker = TrajectoryTracker.options( + name="global_tracker", get_if_exists=True, lifetime="detached" + ).remote(hdfs_dir, verbose) return tracker diff --git a/Agent0/executor_train/verl/verl/utils/device.py b/Agent0/executor_train/verl/verl/utils/device.py index ed85b0d..1f14cba 100644 --- a/Agent0/executor_train/verl/verl/utils/device.py +++ b/Agent0/executor_train/verl/verl/utils/device.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # This code is inspired by the torchtune. # https://github.com/pytorch/torchtune/blob/main/torchtune/utils/_device.py @@ -6,7 +6,8 @@ # Copyright (c) Meta Platforms, Inc. and affiliates. # All rights reserved. # -# This source code is licensed under the BSD-style license in https://github.com/pytorch/torchtune/blob/main/LICENSE +# This source code is licensed under the BSD-style license in +# https://github.com/pytorch/torchtune/blob/main/LICENSE import logging @@ -61,7 +62,8 @@ def get_torch_device() -> any: try: return getattr(torch, device_name) except AttributeError: - logger.warning(f"Device namespace '{device_name}' not found in torch, try to load torch.cuda.") + logger.warning( + f"Device namespace '{device_name}' not found in torch, try to load torch.cuda.") return torch.cuda @@ -83,4 +85,6 @@ def get_nccl_backend() -> str: elif is_npu_available: return "hccl" else: - raise RuntimeError(f"No available nccl backend found on device type {get_device_name()}.") + raise RuntimeError( + f"No available nccl backend found on device type { + get_device_name()}.") diff --git a/Agent0/executor_train/verl/verl/utils/distributed.py b/Agent0/executor_train/verl/verl/utils/distributed.py index 610b5d4..46e563b 100644 --- a/Agent0/executor_train/verl/verl/utils/distributed.py +++ b/Agent0/executor_train/verl/verl/utils/distributed.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/experimental/__init__.py b/Agent0/executor_train/verl/verl/utils/experimental/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/utils/experimental/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/experimental/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/experimental/torch_functional.py b/Agent0/executor_train/verl/verl/utils/experimental/torch_functional.py index 0b4ce5c..8026d0a 100644 --- a/Agent0/executor_train/verl/verl/utils/experimental/torch_functional.py +++ b/Agent0/executor_train/verl/verl/utils/experimental/torch_functional.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,7 +32,8 @@ def _fused_linear_for_ppo_fwd( log_probs = logits.log_softmax(dim=-1) token_log_probs = log_probs.gather(-1, input_ids.unsqueeze(-1)).squeeze(-1) - entropy = torch.logsumexp(logits, dim=-1) - torch.sum(probs * logits, dim=-1) + entropy = torch.logsumexp(logits, dim=-1) - \ + torch.sum(probs * logits, dim=-1) return token_log_probs.to(orig_dtype), entropy.to(orig_dtype) @@ -55,14 +56,19 @@ def _fused_linear_for_ppo_bwd( # Gradient from log_probs if dlog_probs is not None: - one_hot_input = torch.zeros_like(logits).scatter_(-1, input_ids.unsqueeze(-1), 1) - dlogits += dlog_probs.to(torch.float32).unsqueeze(-1) * (one_hot_input - probs) + one_hot_input = torch.zeros_like(logits).scatter_( + -1, input_ids.unsqueeze(-1), 1 + ) + dlogits += dlog_probs.to(torch.float32).unsqueeze(-1) * \ + (one_hot_input - probs) # Gradient from entropy if dentropy is not None: log_probs = logits.log_softmax(dim=-1) - entropy = torch.logsumexp(logits, dim=-1) - torch.sum(probs * logits, dim=-1) - dlogits += probs * (log_probs + entropy.unsqueeze(-1)) * (-dentropy.unsqueeze(-1)) + entropy = torch.logsumexp(logits, dim=-1) - \ + torch.sum(probs * logits, dim=-1) + dlogits += (probs * (log_probs + entropy.unsqueeze(-1)) + * (-dentropy.unsqueeze(-1))) dlogits = dlogits.to(orig_dtype) / temperature @@ -86,11 +92,17 @@ def forward( # Cast to a 2D tensor of the shape [T, D] for ease of working orig_ndim = hidden_states.ndim - assert orig_ndim in (2, 3), f"Invalid hidden_states shape, received {hidden_states.shape}" + assert orig_ndim in ( + 2, + 3, + ), f"Invalid hidden_states shape, received {hidden_states.shape}" orig_batch_size = -1 if orig_ndim == 3: - assert input_ids.ndim == 2, f"input_ids shape doesn't match, {hidden_states.shape} {input_ids.shape}" + assert ( + input_ids.ndim == 2), f"input_ids shape doesn't match, { + hidden_states.shape} { + input_ids.shape}" orig_batch_size = hidden_states.shape[0] hidden_states = hidden_states.flatten(0, 1) input_ids = input_ids.flatten(0, 1) @@ -98,9 +110,13 @@ def forward( T = hidden_states.shape[0] # Allocate memory for outputs - output_requires_grad = hidden_states.requires_grad or vocab_weights.requires_grad - log_probs = hidden_states.new_zeros(T, requires_grad=output_requires_grad) - entropy = hidden_states.new_zeros(T, requires_grad=output_requires_grad) + output_requires_grad = ( + hidden_states.requires_grad or vocab_weights.requires_grad + ) + log_probs = hidden_states.new_zeros( + T, requires_grad=output_requires_grad) + entropy = hidden_states.new_zeros( + T, requires_grad=output_requires_grad) # Perform forward one chunk at a time for chunk_start in range(0, T, chunk_size): @@ -129,7 +145,11 @@ def forward( return log_probs, entropy @staticmethod - def backward(ctx, dlog_probs: Optional[torch.FloatTensor], dentropy: Optional[torch.FloatTensor]): + def backward( + ctx, + dlog_probs: Optional[torch.FloatTensor], + dentropy: Optional[torch.FloatTensor], + ): assert dlog_probs is not None or dentropy is not None hidden_states, vocab_weights, input_ids = ctx.saved_tensors @@ -182,7 +202,8 @@ def backward(ctx, dlog_probs: Optional[torch.FloatTensor], dentropy: Optional[to # Cast the output back to the original input dimension if orig_ndim == 3 and hidden_states.requires_grad: hidden_size = hidden_states.shape[-1] - dhidden_states = dhidden_states.view(orig_batch_size, -1, hidden_size) + dhidden_states = dhidden_states.view( + orig_batch_size, -1, hidden_size) return ( dhidden_states, # hidden_states diff --git a/Agent0/executor_train/verl/verl/utils/flops_counter.py b/Agent0/executor_train/verl/verl/utils/flops_counter.py index 1bed929..74734cc 100644 --- a/Agent0/executor_train/verl/verl/utils/flops_counter.py +++ b/Agent0/executor_train/verl/verl/utils/flops_counter.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -76,9 +76,8 @@ class FlopsCounter: def __init__(self, config: PretrainedConfig): if config.model_type not in VALID_CONFIG_TYPE: print( - f"Only support config type of {VALID_CONFIG_TYPE}, but got {config.model_type}. MFU will always be " - f"zero." - ) + f"Only support config type of {VALID_CONFIG_TYPE}, but got { + config.model_type}. MFU will always be " f"zero.") self.estimate_func = { "qwen2": self._estimate_qwen2_flops, @@ -105,7 +104,11 @@ def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): num_attention_heads = self.config.num_attention_heads intermediate_size = self.config.intermediate_size - head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) + head_dim = getattr( + self.config, + "head_dim", + self.config.hidden_size // self.config.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 @@ -113,10 +116,13 @@ def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): # 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) + 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 + 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 @@ -124,14 +130,23 @@ def _estimate_qwen2_flops(self, tokens_sum, batch_seqlens, delta_time): 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 + 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_deepseek_v3_flops(self, tokens_sum, batch_seqlens, delta_time): + def _estimate_deepseek_v3_flops( + self, + tokens_sum, + batch_seqlens, + delta_time): hidden_size = self.config.hidden_size vocab_size = self.config.vocab_size moe_intermediate_size = self.config.moe_intermediate_size @@ -145,8 +160,10 @@ def _estimate_deepseek_v3_flops(self, tokens_sum, batch_seqlens, delta_time): # non-attn per layer parm moe_gata_N = hidden_size * moe_num_expert - # moe has fc1_1, fc1_2 and fc2 using SwiGLU in ExpertMlp layer & shared experts - moe_expertmlp_N = hidden_size * moe_intermediate_size * (moe_topk + share_expert_num) * 3 + # moe has fc1_1, fc1_2 and fc2 using SwiGLU in ExpertMlp layer & shared + # experts + moe_expertmlp_N = (hidden_size * moe_intermediate_size * + (moe_topk + share_expert_num) * 3) # MLA attn attn_linear_N = 0 q_head_dim = self.config.qk_nope_head_dim + self.config.qk_rope_head_dim @@ -156,18 +173,22 @@ def _estimate_deepseek_v3_flops(self, tokens_sum, batch_seqlens, delta_time): attn_linear_N += hidden_size * self.config.q_lora_rank attn_linear_N += num_query_heads * q_head_dim * self.config.q_lora_rank - attn_linear_N += hidden_size * (self.config.kv_lora_rank + self.config.qk_rope_head_dim) - attn_linear_N += ( - num_query_heads - * (q_head_dim - self.config.qk_rope_head_dim + self.config.v_head_dim) - * self.config.kv_lora_rank + attn_linear_N += hidden_size * ( + self.config.kv_lora_rank + self.config.qk_rope_head_dim ) + attn_linear_N += (num_query_heads * + (q_head_dim - + self.config.qk_rope_head_dim + + self.config.v_head_dim) * + self.config.kv_lora_rank) attn_linear_N += num_query_heads * self.config.v_head_dim * hidden_size emd_and_lm_head_N = vocab_size * hidden_size * 2 # non-attn all_layer parm moe_N = ( - (moe_gata_N + moe_expertmlp_N + attn_linear_N) * (num_hidden_layers - first_k_dense_replace) - + (hidden_size * self.config.intermediate_size * 3 + attn_linear_N) * first_k_dense_replace + (moe_gata_N + moe_expertmlp_N + attn_linear_N) + * (num_hidden_layers - first_k_dense_replace) + + (hidden_size * self.config.intermediate_size * 3 + attn_linear_N) + * first_k_dense_replace + emd_and_lm_head_N ) # non-attn all_layer & all_token fwd & bwd flops @@ -195,18 +216,28 @@ def _estimate_qwen2_moe_flops(self, tokens_sum, batch_seqlens, delta_time): moe_topk = self.config.num_experts_per_tok num_experts = self.config.num_experts - head_dim = getattr(self.config, "head_dim", self.config.hidden_size // self.config.num_attention_heads) + head_dim = getattr( + self.config, + "head_dim", + self.config.hidden_size // self.config.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 # gate + moe export - moe_mlp_N = hidden_size * moe_topk * moe_intermediate_size * 3 + hidden_size * num_experts - attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim) + moe_mlp_N = ( + hidden_size * moe_topk * moe_intermediate_size * 3 + + hidden_size * num_experts + ) + 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 = (moe_mlp_N + attn_linear_N) * num_hidden_layers + emd_and_lm_head_N + dense_N = (moe_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 @@ -214,7 +245,12 @@ def _estimate_qwen2_moe_flops(self, tokens_sum, batch_seqlens, delta_time): 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 + 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 @@ -235,7 +271,9 @@ def estimate_flops(self, batch_seqlens, delta_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) + 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 diff --git a/Agent0/executor_train/verl/verl/utils/fs.py b/Agent0/executor_train/verl/verl/utils/fs.py index 7cc1130..83063f3 100644 --- a/Agent0/executor_train/verl/verl/utils/fs.py +++ b/Agent0/executor_train/verl/verl/utils/fs.py @@ -1,5 +1,5 @@ #!/usr/bin/env python -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -144,7 +144,9 @@ def copy_to_shm(src: str): """ shm_model_root = "/dev/shm/verl-cache/" src_abs = os.path.abspath(os.path.normpath(src)) - dest = os.path.join(shm_model_root, hashlib.md5(src_abs.encode("utf-8")).hexdigest()) + dest = os.path.join( + shm_model_root, hashlib.md5(src_abs.encode("utf-8")).hexdigest() + ) os.makedirs(dest, exist_ok=True) dest = os.path.join(dest, os.path.basename(src_abs)) if os.path.exists(dest) and verify_copy(src, dest): @@ -166,11 +168,15 @@ def _record_directory_structure(folder_path): with open(record_file, "w") as f: for root, dirs, files in os.walk(folder_path): for dir_name in dirs: - relative_dir = os.path.relpath(os.path.join(root, dir_name), folder_path) + relative_dir = os.path.relpath( + os.path.join(root, dir_name), folder_path + ) f.write(f"dir:{relative_dir}\n") for file_name in files: if file_name != ".directory_record.txt": - relative_file = os.path.relpath(os.path.join(root, file_name), folder_path) + relative_file = os.path.relpath( + os.path.join(root, file_name), folder_path + ) f.write(f"file:{relative_file}\n") return record_file @@ -181,11 +187,14 @@ def _check_directory_structure(folder_path, record_file): existing_entries = set() for root, dirs, files in os.walk(folder_path): for dir_name in dirs: - relative_dir = os.path.relpath(os.path.join(root, dir_name), folder_path) + relative_dir = os.path.relpath( + os.path.join(root, dir_name), folder_path) existing_entries.add(f"dir:{relative_dir}") for file_name in files: if file_name != ".directory_record.txt": - relative_file = os.path.relpath(os.path.join(root, file_name), folder_path) + relative_file = os.path.relpath( + os.path.join(root, file_name), folder_path + ) existing_entries.add(f"file:{relative_file}") with open(record_file) as f: recorded_entries = set(f.read().splitlines()) @@ -193,7 +202,12 @@ def _check_directory_structure(folder_path, record_file): def copy_to_local( - src: str, cache_dir=None, filelock=".file.lock", verbose=False, always_recopy=False, use_shm: bool = False + src: str, + cache_dir=None, + filelock=".file.lock", + verbose=False, + always_recopy=False, + use_shm: bool = False, ) -> str: """Copy files/directories from HDFS to local cache with validation. @@ -209,7 +223,9 @@ def copy_to_local( str: Local filesystem path to copied resource """ # Save to a local path for persistence. - local_path = copy_local_path_from_hdfs(src, cache_dir, filelock, verbose, always_recopy) + local_path = copy_local_path_from_hdfs( + src, cache_dir, filelock, verbose, always_recopy + ) # Load into shm to improve efficiency. if use_shm: return copy_to_shm(local_path) @@ -217,12 +233,17 @@ def copy_to_local( def copy_local_path_from_hdfs( - src: str, cache_dir=None, filelock=".file.lock", verbose=False, always_recopy=False -) -> str: + src: str, + cache_dir=None, + filelock=".file.lock", + verbose=False, + always_recopy=False) -> str: """Deprecated. Please use copy_to_local instead.""" from filelock import FileLock - assert src[-1] != "/", f"Make sure the last char in src is not / because it will cause error. Got {src}" + 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 @@ -248,11 +269,13 @@ def copy_local_path_from_hdfs( if os.path.isdir(local_path): _record_directory_structure(local_path) elif os.path.isdir(local_path): - # always_recopy=False, local path exists, and it is a folder: check whether there is anything missed + # always_recopy=False, local path exists, and it is a folder: + # check whether there is anything missed record_file = os.path.join(local_path, ".directory_record.txt") if not _check_directory_structure(local_path, record_file): if verbose: - print(f"Recopy from {src} to {local_path} due to missing files or directories.") + print( + f"Recopy from {src} to {local_path} due to missing files or directories.") shutil.rmtree(local_path, ignore_errors=True) copy(src, local_path) _record_directory_structure(local_path) diff --git a/Agent0/executor_train/verl/verl/utils/fsdp_utils.py b/Agent0/executor_train/verl/verl/utils/fsdp_utils.py index 7465b40..cd56f20 100644 --- a/Agent0/executor_train/verl/verl/utils/fsdp_utils.py +++ b/Agent0/executor_train/verl/verl/utils/fsdp_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,17 +27,35 @@ from torch.distributed import DeviceMesh from torch.distributed.fsdp import FullyShardedDataParallel as FSDP from torch.distributed.fsdp._runtime_utils import _lazy_init -from torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy +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 from verl.utils.device import get_device_id, get_device_name, get_torch_device if version.parse(torch.__version__) >= version.parse("2.6"): - from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, fully_shard + from torch.distributed.fsdp import ( + CPUOffloadPolicy, + FSDPModule, + MixedPrecisionPolicy, + fully_shard, + ) elif version.parse(torch.__version__) >= version.parse("2.4"): - from torch.distributed._composable.fsdp import CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, fully_shard + from torch.distributed._composable.fsdp import ( + CPUOffloadPolicy, + FSDPModule, + MixedPrecisionPolicy, + fully_shard, + ) else: - fully_shard, MixedPrecisionPolicy, FSDPModule, CPUOffloadPolicy = None, None, None, None + fully_shard, MixedPrecisionPolicy, FSDPModule, CPUOffloadPolicy = ( + None, + None, + None, + None, + ) def init_fn(x: torch.nn.Module): @@ -47,22 +65,33 @@ def init_fn(x: torch.nn.Module): return x -def get_init_weight_context_manager(use_meta_tensor=True, mesh: DeviceMesh = None): +def get_init_weight_context_manager( + use_meta_tensor=True, + mesh: DeviceMesh = None): from accelerate import init_empty_weights - cpu_init_weights = lambda: torch.device("cpu") + def cpu_init_weights(): return torch.device("cpu") if use_meta_tensor: if mesh is None: - init_context = init_empty_weights if torch.distributed.get_rank() != 0 else cpu_init_weights + init_context = ( + init_empty_weights + if torch.distributed.get_rank() != 0 + else cpu_init_weights + ) else: - init_context = init_empty_weights if mesh.get_coordinate()[-1] != 0 else cpu_init_weights + init_context = ( + init_empty_weights + if mesh.get_coordinate()[-1] != 0 + else cpu_init_weights + ) else: init_context = cpu_init_weights return init_context # Copyright 2020-present the HuggingFace Inc. team. -# Adapted from https://github.com/huggingface/transformers/src/transformers/trainer.py +# Adapted from +# https://github.com/huggingface/transformers/src/transformers/trainer.py def get_fsdp_wrap_policy(module, config=None, is_lora=False): """Get FSDP wrap policy for the module. @@ -85,7 +114,8 @@ def _get_attr(attr_name, default_value=None): if _get_attr("disable", False): return None - default_transformer_cls_names_to_wrap = getattr(module, "_no_split_modules", None) + default_transformer_cls_names_to_wrap = getattr( + module, "_no_split_modules", None) fsdp_transformer_layer_cls_to_wrap = _get_attr( "transformer_layer_cls_to_wrap", default_transformer_cls_names_to_wrap ) @@ -106,18 +136,24 @@ def lambda_policy_fn(module): and module.weight.requires_grad ) - lambda_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=lambda_policy_fn) + 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) + 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.") + raise Exception( + "Could not find the transformer layer class to wrap in the model." + ) else: transformer_cls_to_wrap.add(transformer_cls) @@ -183,7 +219,9 @@ def load_fsdp_model_to_gpu(model: FSDP): if handle._offload_params: continue flat_param = handle.flat_param - handle.flat_param_to(torch.device(f"{get_device_name()}:{device_id}"), non_blocking=True) + handle.flat_param_to( + torch.device(f"{get_device_name()}:{device_id}"), non_blocking=True + ) # the following still keeps id(._local_shard) != id(.data) flat_param._local_shard = flat_param.data @@ -240,7 +278,9 @@ def register_empty_parameter(module, name, param): 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) + module._parameters[name] = param_cls( + module._parameters[name].to(device), **kwargs + ) registered.add(module._parameters[name]) try: @@ -282,14 +322,16 @@ def parallel_load_safetensors(filepath): 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) + 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)] + ckpt_chunks = [ckpt_chunks[rank * size: rank * size + size] + for rank in range(world_size)] shard_states = {} device = get_device_id() @@ -307,7 +349,9 @@ def parallel_load_safetensors(filepath): return shard_states -def parallel_init_module_fn(module: torch.nn.Module, shard_states: dict[str, torch.nn.Parameter]): +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. @@ -322,7 +366,8 @@ def parallel_init_module_fn(module: torch.nn.Module, shard_states: dict[str, tor state2fqn = {} for name, state in itertools.chain( - module.named_parameters(remove_duplicate=False), module.named_buffers(remove_duplicate=False) + module.named_parameters(remove_duplicate=False), + module.named_buffers(remove_duplicate=False), ): state2fqn.setdefault(state, []).append(name) # remove standalone parameters and buffers @@ -334,7 +379,10 @@ def create_and_sync_state(param_name, state, is_param): assert param_name in shard_states, f"{param_name} not loaded" device = get_device_id() if is_param: - param = torch.nn.Parameter(torch.empty_like(state.data, device=device), requires_grad=state.requires_grad) + 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] @@ -350,14 +398,16 @@ def create_and_sync_state(param_name, state, is_param): 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 = 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 + # 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( @@ -368,7 +418,9 @@ def init_fn(sub_mod: torch.nn.Module, recurse: bool = True): # 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) + materialized_states[state] = create_and_sync_state( + fqn, state, is_param + ) else: if fqn in shard_states: shard_states.pop(fqn) @@ -407,7 +459,10 @@ def get_fsdp_state_ctx(model, state_type, state_cfg, optim_cfg): return nullcontext() -def get_fsdp_full_state_dict(model: torch.nn.Module, offload_to_cpu: bool = True, rank0_only: bool = True): +def get_fsdp_full_state_dict( + model: torch.nn.Module, + offload_to_cpu: bool = True, + rank0_only: bool = True): """ Get the full state dict from an FSDP model. @@ -425,17 +480,27 @@ def get_fsdp_full_state_dict(model: torch.nn.Module, offload_to_cpu: bool = True if fsdp_version(model) == 1: from torch.distributed.fsdp import FullStateDictConfig, StateDictType - state_dict_config = FullStateDictConfig(offload_to_cpu=offload_to_cpu, rank0_only=rank0_only) + state_dict_config = FullStateDictConfig( + offload_to_cpu=offload_to_cpu, rank0_only=rank0_only + ) with get_fsdp_state_ctx( - model, state_type=StateDictType.FULL_STATE_DICT, state_cfg=state_dict_config, optim_cfg=None + model, + state_type=StateDictType.FULL_STATE_DICT, + state_cfg=state_dict_config, + optim_cfg=None, ): state_dict = model.state_dict() return state_dict elif fsdp_version(model) == 2: - from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, + ) state_dict_config = StateDictOptions( - full_state_dict=True, cpu_offload=offload_to_cpu, broadcast_from_rank0=not rank0_only + full_state_dict=True, + cpu_offload=offload_to_cpu, + broadcast_from_rank0=not rank0_only, ) state_dict = get_model_state_dict(model, options=state_dict_config) return state_dict @@ -443,7 +508,11 @@ def get_fsdp_full_state_dict(model: torch.nn.Module, offload_to_cpu: bool = True raise NotImplementedError(f"Unknown FSDP version {fsdp_version}") -def fsdp2_load_full_state_dict(model: torch.nn.Module, full_state: dict, device_mesh=None, cpu_offload=None): +def fsdp2_load_full_state_dict( + model: torch.nn.Module, + full_state: dict, + device_mesh=None, + cpu_offload=None): """ Loads the full state dict (could be only on rank 0) into the sharded model. This is done by broadcasting the parameters from rank 0 to all other ranks. This function modifies the model in-place. @@ -452,7 +521,10 @@ def fsdp2_load_full_state_dict(model: torch.nn.Module, full_state: dict, device_ model (`torch.nn.Module`): The model to load the state dict into full_state (`dict`): The full state dict to load, can only be on rank 0 """ - from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + set_model_state_dict, + ) # To broadcast, it needs to be instantiated in the GPU. if dist.get_rank() == 0: @@ -461,7 +533,10 @@ def fsdp2_load_full_state_dict(model: torch.nn.Module, full_state: dict, device_ model = model.to_empty(device=get_device_id()) cpu_offload = cpu_offload is not None - options = StateDictOptions(full_state_dict=True, cpu_offload=cpu_offload, broadcast_from_rank0=True) + options = StateDictOptions( + full_state_dict=True, + cpu_offload=cpu_offload, + broadcast_from_rank0=True) set_model_state_dict(model, full_state, options=options) # rotary_emb is not in state_dict, so we need to broadcast it manually @@ -476,17 +551,24 @@ def fsdp2_load_full_state_dict(model: torch.nn.Module, full_state: dict, device_ def apply_fsdp2(model, fsdp_kwargs, config): """model: AutoModelForCausalLM""" - assert CPUOffloadPolicy is not None, "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" + assert ( + CPUOffloadPolicy is not None + ), "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" - default_transformer_cls_names_to_wrap = getattr(model, "_no_split_modules", None) + default_transformer_cls_names_to_wrap = getattr( + model, "_no_split_modules", None) fsdp_transformer_layer_cls_to_wrap = config.get("wrap_policy", {}).get( "transformer_layer_cls_to_wrap", default_transformer_cls_names_to_wrap ) if isinstance(fsdp_transformer_layer_cls_to_wrap, str): - fsdp_transformer_layer_cls_to_wrap = [fsdp_transformer_layer_cls_to_wrap] + fsdp_transformer_layer_cls_to_wrap = [ + fsdp_transformer_layer_cls_to_wrap] - assert len(fsdp_transformer_layer_cls_to_wrap) > 0 and fsdp_transformer_layer_cls_to_wrap[0] is not None + assert ( + len(fsdp_transformer_layer_cls_to_wrap) > 0 + and fsdp_transformer_layer_cls_to_wrap[0] is not None + ) modules = [] for name, module in model.named_modules(): @@ -497,10 +579,14 @@ def apply_fsdp2(model, fsdp_kwargs, config): for idx, module in enumerate(modules): fully_shard(module, **fsdp_kwargs) - fully_shard(model, **fsdp_kwargs) # fsdp2 will not reshard_after_forward for root module + fully_shard( + model, **fsdp_kwargs + ) # fsdp2 will not reshard_after_forward for root module -def fsdp2_clip_grad_norm_(parameters, max_norm, norm_type=2.0, error_if_nonfinite=False, foreach=None): +def fsdp2_clip_grad_norm_( + parameters, max_norm, norm_type=2.0, error_if_nonfinite=False, foreach=None +): """torch.nn.utils.clip_grad_norm_ cann't run on cpu parameter DTensor""" from torch.nn.utils.clip_grad import _clip_grads_with_norm_, _get_total_norm @@ -521,7 +607,7 @@ def layered_summon_lora_params(fsdp_module) -> OrderedDict: def __prefix_submodules(module, prefix): for name, submodule in module.named_modules(): - if name.startswith(prefix) and "." not in name[len(prefix) :]: + if name.startswith(prefix) and "." not in name[len(prefix):]: yield name, submodule lora_params = OrderedDict() @@ -538,16 +624,22 @@ def __prefix_submodules(module, prefix): peft_model = getattr(fsdp_module, "_fsdp_wrapped_module", fsdp_module) for prefix in prefix_list: for name, submodule in __prefix_submodules(fsdp_module, prefix): - prefix = name.replace("_fsdp_wrapped_module.base_model.model.", "base_model.model.") + prefix = name.replace( + "_fsdp_wrapped_module.base_model.model.", "base_model.model." + ) if name.endswith(".model") or name.endswith(".layers"): continue if fsdp_version(submodule) > 0: with FSDP.summon_full_params(submodule, writeback=False): - sub_lora_params = get_peft_model_state_dict(peft_model, state_dict=submodule.state_dict()) + sub_lora_params = get_peft_model_state_dict( + peft_model, state_dict=submodule.state_dict() + ) sub_lora_params = { - f"{prefix}.{name}": param.full_tensor().detach().cpu() - if hasattr(param, "full_tensor") - else param.detach().cpu() + f"{prefix}.{name}": ( + param.full_tensor().detach().cpu() + if hasattr(param, "full_tensor") + else param.detach().cpu() + ) for name, param in sub_lora_params.items() } lora_params.update(sub_lora_params) diff --git a/Agent0/executor_train/verl/verl/utils/hdfs_io.py b/Agent0/executor_train/verl/verl/utils/hdfs_io.py index 31edda1..d4c6115 100644 --- a/Agent0/executor_train/verl/verl/utils/hdfs_io.py +++ b/Agent0/executor_train/verl/verl/utils/hdfs_io.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -113,9 +113,13 @@ def copy(src: str, dst: str, **kwargs) -> bool: 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) + 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) + returncode = _run_cmd( + _hdfs_cmd(f"-put -f {from_path} {to_path}"), timeout=timeout + ) else: if from_path.startswith("hdfs"): returncode = _run_cmd( diff --git a/Agent0/executor_train/verl/verl/utils/import_utils.py b/Agent0/executor_train/verl/verl/utils/import_utils.py index fc75541..93fac92 100644 --- a/Agent0/executor_train/verl/verl/utils/import_utils.py +++ b/Agent0/executor_train/verl/verl/utils/import_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -80,7 +80,9 @@ def import_external_libs(external_libs=None): importlib.import_module(external_lib) -def load_extern_type(file_path: Optional[str], type_name: Optional[str]) -> type: +def load_extern_type( + file_path: Optional[str], + type_name: Optional[str]) -> type: """Load a external data type based on the file path and type name""" if not file_path: return None @@ -99,17 +101,21 @@ def load_extern_type(file_path: Optional[str], type_name: Optional[str]) -> type file_path = file_path[7:] if not os.path.exists(file_path): - raise FileNotFoundError(f"Custom type file '{file_path}' not found.") + raise FileNotFoundError( + f"Custom type file '{file_path}' not found.") - spec = importlib.util.spec_from_file_location("custom_module", file_path) + spec = importlib.util.spec_from_file_location( + "custom_module", file_path) module = importlib.util.module_from_spec(spec) try: spec.loader.exec_module(module) except Exception as e: - raise RuntimeError(f"Error loading module from '{file_path}'") from e + raise RuntimeError( + f"Error loading module from '{file_path}'") from e if not hasattr(module, type_name): - raise AttributeError(f"Custom type '{type_name}' not found in '{file_path}'.") + raise AttributeError( + f"Custom type '{type_name}' not found in '{file_path}'.") return getattr(module, type_name) diff --git a/Agent0/executor_train/verl/verl/utils/kernel/__init__.py b/Agent0/executor_train/verl/verl/utils/kernel/__init__.py index e32d583..ac310b9 100644 --- a/Agent0/executor_train/verl/verl/utils/kernel/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/kernel/__init__.py @@ -15,7 +15,7 @@ # limitations under the License. # -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -28,4 +28,3 @@ # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. - diff --git a/Agent0/executor_train/verl/verl/utils/kernel/kernels.py b/Agent0/executor_train/verl/verl/utils/kernel/kernels.py index a125bac..4fa275d 100644 --- a/Agent0/executor_train/verl/verl/utils/kernel/kernels.py +++ b/Agent0/executor_train/verl/verl/utils/kernel/kernels.py @@ -15,7 +15,7 @@ # limitations under the License. # -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -92,10 +92,10 @@ class BackwardEnum: Enum for the backward method. """ - _Total_Fuse_MN = ( - 0 # Fuse d_logits & d_hidden & d_weight, no intermediate storage, requires fp32 for d_hidden & d_weight + _Total_Fuse_MN = 0 # Fuse d_logits & d_hidden & d_weight, no intermediate storage, requires fp32 for d_hidden & d_weight + _Total_Separate = ( + 1 # Store d_logits, no special requirements for d_hidden & d_weight ) - _Total_Separate = 1 # Store d_logits, no special requirements for d_hidden & d_weight _Split_Dlogits_N = 2 # split d_logits along its N dimension, aka. vocab_size _Split_Dlogits_M = 3 # split d_logits along its M dimension, aka. num_tokens @@ -118,7 +118,13 @@ def set_backward_method(backward_method: BackwardEnum): @triton.autotune( - configs=[triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32}, num_stages=3, num_warps=8)], + configs=[ + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32}, + num_stages=3, + num_warps=8, + ) + ], key=["num_tokens", "hidden_size", "vocab_size"], ) @triton.jit @@ -169,7 +175,9 @@ def efficient_entropy_kernel_general_mainloop( # create pointers for the first blocks of hidden offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_k = tl.arange(0, BLOCK_SIZE_K) - hidden_ptrs = hidden_ptr + (offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k) + hidden_ptrs = hidden_ptr + ( + offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k + ) # load labels for this block labels = tl.load(labels_ptr + offs_am, mask=offs_am < num_tokens) @@ -181,9 +189,18 @@ def efficient_entropy_kernel_general_mainloop( _entropy_b = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32) _logprobs = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32) for n in range(0, num_pid_n): - offs_bn = pid_n * vocab_per_split + n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_bn = ( + pid_n * + vocab_per_split + + n * + BLOCK_SIZE_N + + tl.arange( + 0, + BLOCK_SIZE_N)) # weight_ptrs = weight_ptr + (offs_k[:, None] * stride_weight_k + offs_bn[None, :] * stride_weight_n) - weight_ptrs = weight_ptr + (offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k) + weight_ptrs = weight_ptr + ( + offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k + ) # iterate over K dimension logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) @@ -191,7 +208,8 @@ def efficient_entropy_kernel_general_mainloop( # load the next block of hidden and weight _hidden = tl.load( hidden_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_am[:, None] < num_tokens), other=0.0, ) # _weight = tl.load(weight_ptrs, @@ -236,17 +254,34 @@ def efficient_entropy_kernel_general_mainloop( offs_max_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) offs_max_n = pid_n maximum_ptrs = max_ptr + offs_max_n * stride_max_n + offs_max_m * stride_max_m - tl.store(maximum_ptrs, _max, mask=(offs_max_m < num_tokens) & (offs_max_n < num_splits)) + tl.store( + maximum_ptrs, _max, mask=( + offs_max_m < num_tokens) & ( + offs_max_n < num_splits)) # store entropy accu_ptrs = accu_ptr + offs_max_n * stride_accu_n + offs_max_m * stride_accu_m - tl.store(accu_ptrs, _accu, mask=(offs_max_m < num_tokens) & (offs_max_n[None] < num_splits)) - entropy_b_ptrs = entropy_b_ptr + offs_max_n * stride_entropy_b_n + offs_max_m * stride_entropy_b_m - tl.store(entropy_b_ptrs, _entropy_b, mask=(offs_max_m < num_tokens) & (offs_max_n < num_splits)) + tl.store( + accu_ptrs, + _accu, + mask=(offs_max_m < num_tokens) & (offs_max_n[None] < num_splits), + ) + entropy_b_ptrs = ( + entropy_b_ptr + + offs_max_n * stride_entropy_b_n + + offs_max_m * stride_entropy_b_m + ) + tl.store( + entropy_b_ptrs, + _entropy_b, + mask=(offs_max_m < num_tokens) & (offs_max_n < num_splits), + ) # store logprobs vocab_left_idx = pid_n * vocab_per_split + rank * vocab_size - vocab_right_idx = min((pid_n + 1) * vocab_per_split, vocab_size) + rank * vocab_size + vocab_right_idx = min( + (pid_n + 1) * vocab_per_split, + vocab_size) + rank * vocab_size mask = (labels >= vocab_left_idx) & (labels < vocab_right_idx) mask &= offs_am < num_tokens global_logprobs_ptrs = global_logprobs_ptr + offs_am * stride_global_logprobs @@ -254,7 +289,10 @@ def efficient_entropy_kernel_general_mainloop( tl.store(global_logprobs_ptrs, _logprobs, mask=mask) -@triton.autotune(configs=[triton.Config({"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64})], key=["num_tokens", "num_splits"]) +@triton.autotune( + configs=[triton.Config({"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64})], + key=["num_tokens", "num_splits"], +) @triton.jit def efficient_entropy_triton_kernel_epilogue( max_ptr, @@ -294,16 +332,32 @@ def efficient_entropy_triton_kernel_epilogue( global_entropy_b = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32) for pid_n in range(0, tl.cdiv(num_splits, BLOCK_SIZE_N)): offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - max_ptrs = max_ptr + offs_m[:, None] * stride_max_m + offs_n[None, :] * stride_max_n - - _max = tl.load(max_ptrs, mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), other=0.0) - - accu_ptrs = accu_ptr + offs_m[:, None] * stride_accu_m + offs_n[None, :] * stride_accu_n - _accu = tl.load(accu_ptrs, mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), other=0.0) - - entropy_b_ptrs = entropy_b_ptr + offs_m[:, None] * stride_entropy_b_m + offs_n[None, :] * stride_entropy_b_n + max_ptrs = (max_ptr + + offs_m[:, None] * + stride_max_m + + offs_n[None, :] * + stride_max_n) + + _max = tl.load(max_ptrs, mask=(offs_m[:, None] < num_tokens) & ( + offs_n[None, :] < num_splits), other=0.0, ) + + accu_ptrs = (accu_ptr + + offs_m[:, None] * + stride_accu_m + + offs_n[None, :] * + stride_accu_n) + _accu = tl.load(accu_ptrs, mask=(offs_m[:, None] < num_tokens) & ( + offs_n[None, :] < num_splits), other=0.0, ) + + entropy_b_ptrs = ( + entropy_b_ptr + + offs_m[:, None] * stride_entropy_b_m + + offs_n[None, :] * stride_entropy_b_n + ) _entropy_b = tl.load( - entropy_b_ptrs, mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), other=0.0 + entropy_b_ptrs, + mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), + other=0.0, ) # local reduction @@ -314,7 +368,9 @@ def efficient_entropy_triton_kernel_epilogue( _scale = tl.exp(_max - global_max[:, None]) _coeff = tl.exp(_max_old - global_max) global_accu = _coeff * global_accu + tl.sum(_scale * _accu, axis=1) - global_entropy_b = _coeff * global_entropy_b + tl.sum(_scale * _entropy_b, axis=1) + global_entropy_b = _coeff * global_entropy_b + tl.sum( + _scale * _entropy_b, axis=1 + ) # store maximum_ptrs = global_max_ptr + offs_m * stride_global_max @@ -322,12 +378,17 @@ def efficient_entropy_triton_kernel_epilogue( # store entropy_b global_entropy_b = tl.fdiv(global_entropy_b, global_accu) # entropy_b - tl.store(global_entropy_b_ptr + offs_m * stride_global_entropy_b, global_entropy_b, mask=offs_m < num_tokens) + tl.store( + global_entropy_b_ptr + offs_m * stride_global_entropy_b, + global_entropy_b, + mask=offs_m < num_tokens, + ) # store entropy global_accu_ptrs = global_accu_ptr + offs_m * stride_global_accu tl.store(global_accu_ptrs, global_accu, mask=offs_m < num_tokens) - global_entropy = tl.log(global_accu) + global_max - global_entropy_b # entropy_a + global_entropy = tl.log(global_accu) + global_max - \ + global_entropy_b # entropy_a global_entropy_ptrs = global_entropy_ptr + offs_m * stride_global_entropy tl.store(global_entropy_ptrs, global_entropy, mask=offs_m < num_tokens) # update logprobs @@ -337,16 +398,23 @@ def efficient_entropy_triton_kernel_epilogue( global_logprobs = -1 * global_logprobs if reduction == 0: - tl.store(global_logprobs_ptrs, global_logprobs, mask=offs_m < num_tokens) + tl.store( + global_logprobs_ptrs, + global_logprobs, + mask=offs_m < num_tokens) elif reduction == 1: global_logprobs_scalar = tl.sum(global_logprobs, axis=0) tl.atomic_add(global_logprobs_scalar_ptr, global_logprobs_scalar) elif reduction == 2: - global_logprobs_scalar = tl.sum(global_logprobs, axis=0) / num_tokens.to(tl.float32) + global_logprobs_scalar = tl.sum( + global_logprobs, axis=0) / num_tokens.to(tl.float32) tl.atomic_add(global_logprobs_scalar_ptr, global_logprobs_scalar) -@triton.autotune(configs=[triton.Config({"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64})], key=["num_tokens", "num_splits"]) +@triton.autotune( + configs=[triton.Config({"BLOCK_SIZE_M": 16, "BLOCK_SIZE_N": 64})], + key=["num_tokens", "num_splits"], +) @triton.jit def efficient_entropy_triton_kernel_epilogue_tp( num_tokens, @@ -382,21 +450,29 @@ def efficient_entropy_triton_kernel_epilogue_tp( for pid_n in range(0, tl.cdiv(num_splits, BLOCK_SIZE_N)): offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - _reduced_max = tl.load( - reduced_max_ptr + offs_m[:, None] * stride_reduced_max_m + offs_n[None, :] * stride_reduced_max_n, - mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), - other=0.0, - ) + _reduced_max = tl.load(reduced_max_ptr + offs_m[:, + None] * stride_reduced_max_m + offs_n[None, + :] * stride_reduced_max_n, + mask=(offs_m[:, + None] < num_tokens) & (offs_n[None, + :] < num_splits), + other=0.0, + ) _original_max = tl.load( - original_max_ptr + offs_m[:, None] * stride_original_max_m + offs_n[None, :] * stride_original_max_n, - mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), - other=0.0, - ) - _accu = tl.load( - accu_ptr + offs_m[:, None] * stride_accu_m + offs_n[None, :] * stride_accu_n, + original_max_ptr + + offs_m[:, None] * stride_original_max_m + + offs_n[None, :] * stride_original_max_n, mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), other=0.0, ) + _accu = tl.load(accu_ptr + offs_m[:, + None] * stride_accu_m + offs_n[None, + :] * stride_accu_n, + mask=(offs_m[:, + None] < num_tokens) & (offs_n[None, + :] < num_splits), + other=0.0, + ) # local reduce-max _max_old = global_max @@ -409,20 +485,38 @@ def efficient_entropy_triton_kernel_epilogue_tp( global_accu = _coeff * global_accu + tl.sum(_scale * _accu, axis=1) # update entropy_b - _entropy_b = tl.load( - entropy_b_ptr + offs_m[:, None] * stride_entropy_b_m + offs_n[None, :] * stride_entropy_b_n, - mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), - other=0.0, + _entropy_b = tl.load(entropy_b_ptr + offs_m[:, + None] * stride_entropy_b_m + offs_n[None, + :] * stride_entropy_b_n, + mask=(offs_m[:, + None] < num_tokens) & (offs_n[None, + :] < num_splits), + other=0.0, + ) + global_entropy_b = _coeff * global_entropy_b + tl.sum( + _scale * _entropy_b, axis=1 ) - global_entropy_b = _coeff * global_entropy_b + tl.sum(_scale * _entropy_b, axis=1) # store - tl.store(global_max_ptr + offs_m * stride_global_max, global_max, mask=offs_m < num_tokens) - tl.store(global_accu_ptr + offs_m * stride_global_accu, global_accu, mask=offs_m < num_tokens) - tl.store(global_entropy_b_ptr + offs_m * stride_global_entropy_b, global_entropy_b, mask=offs_m < num_tokens) + tl.store( + global_max_ptr + offs_m * stride_global_max, + global_max, + mask=offs_m < num_tokens, + ) + tl.store( + global_accu_ptr + offs_m * stride_global_accu, + global_accu, + mask=offs_m < num_tokens, + ) + tl.store( + global_entropy_b_ptr + offs_m * stride_global_entropy_b, + global_entropy_b, + mask=offs_m < num_tokens, + ) -@triton.autotune(configs=[triton.Config({"BLOCK_SIZE_M": 16})], key=["num_tokens"]) +@triton.autotune(configs=[triton.Config( + {"BLOCK_SIZE_M": 16})], key=["num_tokens"]) @triton.jit def efficient_entropy_triton_epilogue_tp_update( num_tokens, @@ -444,22 +538,47 @@ def efficient_entropy_triton_epilogue_tp_update( offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - maximum = tl.load(maximum_ptr + offs_m * stride_maximum, mask=offs_m < num_tokens) - accumulate = tl.load(accumulate_ptr + offs_m * stride_accumulate, mask=offs_m < num_tokens) + maximum = tl.load( + maximum_ptr + + offs_m * + stride_maximum, + mask=offs_m < num_tokens) + accumulate = tl.load( + accumulate_ptr + offs_m * stride_accumulate, mask=offs_m < num_tokens + ) - entropy_b = tl.load(entropy_b_ptr + offs_m * stride_entropy_b, mask=offs_m < num_tokens) + entropy_b = tl.load( + entropy_b_ptr + offs_m * stride_entropy_b, mask=offs_m < num_tokens + ) entropy_b = tl.fdiv(entropy_b, accumulate) - tl.store(entropy_b_ptr + offs_m * stride_entropy_b, entropy_b, mask=offs_m < num_tokens) + tl.store( + entropy_b_ptr + + offs_m * + stride_entropy_b, + entropy_b, + mask=offs_m < num_tokens) entropy = tl.log(accumulate) + maximum - entropy_b - tl.store(entropy_ptr + offs_m * stride_entropy, entropy, mask=offs_m < num_tokens) - - logprobs = tl.load(logprobs_ptr + offs_m * stride_logprobs, mask=offs_m < num_tokens) + tl.store( + entropy_ptr + + offs_m * + stride_entropy, + entropy, + mask=offs_m < num_tokens) + + logprobs = tl.load( + logprobs_ptr + offs_m * stride_logprobs, mask=offs_m < num_tokens + ) logprobs = maximum + tl.log(accumulate) - logprobs logprobs = -1 * logprobs if reduction == 0: - tl.store(logprobs_ptr + offs_m * stride_logprobs, logprobs, mask=offs_m < num_tokens) + tl.store( + logprobs_ptr + + offs_m * + stride_logprobs, + logprobs, + mask=offs_m < num_tokens) elif reduction == 1: logprobs_scalar = tl.sum(logprobs, axis=0) tl.atomic_add(logprobs_scalar_ptr, logprobs_scalar) @@ -489,10 +608,14 @@ def efficient_entropy_forward( assert hidden.shape[0] == labels.shape[0] and hidden.shape[1] == weight.shape[1] - _rank = 0 if dist_process_group is None else dist.get_rank(dist_process_group) - _world_size = 1 if dist_process_group is None else dist.get_world_size(dist_process_group) + _rank = 0 if dist_process_group is None else dist.get_rank( + dist_process_group) + _world_size = ( + 1 if dist_process_group is None else dist.get_world_size(dist_process_group)) - if dist_process_group is not None and not hasattr(efficient_entropy_forward, "_initialized"): + if dist_process_group is not None and not hasattr( + efficient_entropy_forward, "_initialized" + ): global _dedicated_stream, _dedicated_events _dedicated_stream = get_torch_device().Stream(hidden.device) _dedicated_events = [get_torch_device().Event() for _ in range(2)] @@ -507,36 +630,56 @@ def efficient_entropy_forward( if REDUCTION == EntropyReductionEnum._None: if dist_process_group is None: - logprobs = torch.empty((num_tokens,), device=hidden.device, dtype=torch.float32) + logprobs = torch.empty( + (num_tokens,), device=hidden.device, dtype=torch.float32 + ) else: - logprobs = torch.zeros((num_tokens,), device=hidden.device, dtype=torch.float32) + logprobs = torch.zeros( + (num_tokens,), device=hidden.device, dtype=torch.float32 + ) elif REDUCTION in (EntropyReductionEnum._Sum, EntropyReductionEnum._Mean): logprobs = torch.empty((), device=hidden.device, dtype=torch.float32) else: raise ValueError(f"Invalid reduction: {reduction}") - entropy = torch.empty((num_tokens,), device=hidden.device, dtype=torch.float32) + entropy = torch.empty( + (num_tokens,), device=hidden.device, dtype=torch.float32) assert logprobs.is_contiguous() and entropy.is_contiguous() maximum = torch.empty_like(entropy) - accumulate_and_entropy_b = torch.empty((num_tokens * 2,), device=hidden.device, dtype=torch.float32) - accumulate_and_entropy_b_view = accumulate_and_entropy_b.view(2, num_tokens) + accumulate_and_entropy_b = torch.empty( + (num_tokens * 2,), device=hidden.device, dtype=torch.float32 + ) + accumulate_and_entropy_b_view = accumulate_and_entropy_b.view( + 2, num_tokens) accumulate = accumulate_and_entropy_b_view[0, :] entropy_b = accumulate_and_entropy_b_view[1, :] - assert maximum.is_contiguous() and accumulate.is_contiguous() and entropy_b.is_contiguous() + assert ( + maximum.is_contiguous() + and accumulate.is_contiguous() + and entropy_b.is_contiguous() + ) vocab_per_split = 1024 assert vocab_per_split % 128 == 0 num_splits = (vocab_size + vocab_per_split - 1) // vocab_per_split - _max = torch.empty((num_tokens, num_splits), device=hidden.device, dtype=torch.float32) - _accu = torch.empty((num_tokens, num_splits), device=hidden.device, dtype=torch.float32) - _entropy_b = torch.empty((num_tokens, num_splits), device=hidden.device, dtype=torch.float32) + _max = torch.empty( + (num_tokens, num_splits), device=hidden.device, dtype=torch.float32 + ) + _accu = torch.empty( + (num_tokens, num_splits), device=hidden.device, dtype=torch.float32 + ) + _entropy_b = torch.empty( + (num_tokens, num_splits), device=hidden.device, dtype=torch.float32 + ) if REDUCTION == EntropyReductionEnum._None: _logprobs = logprobs else: - _logprobs = torch.empty((num_tokens,), device=hidden.device, dtype=torch.float32) + _logprobs = torch.empty( + (num_tokens,), device=hidden.device, dtype=torch.float32 + ) assert _accu.is_contiguous() and _entropy_b.is_contiguous() and _max.is_contiguous() assert _accu.is_cuda and _entropy_b.is_cuda and _max.is_cuda @@ -544,7 +687,12 @@ def efficient_entropy_forward( if _config._use_triton: # 1D kernel launch, then split the tile def mainloop_grid(meta): - return (triton.cdiv(num_tokens, meta["BLOCK_SIZE_M"]) * num_splits,) + return ( + triton.cdiv( + num_tokens, + meta["BLOCK_SIZE_M"]) * + num_splits, + ) efficient_entropy_kernel_general_mainloop[mainloop_grid]( _rank, @@ -614,7 +762,10 @@ def epilogue_grid(meta): get_torch_device().current_stream().record_event(_dedicated_events[0]) with get_torch_device().stream(_dedicated_stream): _dedicated_stream.wait_event(_dedicated_events[0]) - dist.all_reduce(_logprobs, op=dist.ReduceOp.SUM, group=dist_process_group) + dist.all_reduce( + _logprobs, + op=dist.ReduceOp.SUM, + group=dist_process_group) _dedicated_stream.record_event(_dedicated_events[1]) efficient_entropy_triton_kernel_epilogue_tp[epilogue_grid]( @@ -641,7 +792,10 @@ def epilogue_grid(meta): ) get_torch_device().current_stream().wait_event(_dedicated_events[1]) - dist.all_reduce(accumulate_and_entropy_b, op=dist.ReduceOp.SUM, group=dist_process_group) + dist.all_reduce( + accumulate_and_entropy_b, + op=dist.ReduceOp.SUM, + group=dist_process_group) # update logprobs & entropy efficient_entropy_triton_epilogue_tp_update[epilogue_grid]( @@ -667,7 +821,12 @@ def epilogue_grid(meta): @triton.autotune( configs=[ triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 16}, + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + }, num_stages=3, num_warps=8, ) @@ -737,40 +896,61 @@ def efficient_entropy_backward_kernel_general_mainloop_MN( maximum_ptrs = maximum_ptr + offs_am * stride_maximum maximum = tl.load(maximum_ptrs, mask=offs_am < num_tokens, other=0.0) accu_ptrs = accu_ptr + offs_am * stride_accu - accu = tl.load(accu_ptrs, mask=offs_am < num_tokens, other=1e-6) # epsilon to avoid division by zero + accu = tl.load( + accu_ptrs, mask=offs_am < num_tokens, other=1e-6 + ) # epsilon to avoid division by zero accu_rcp = tl.fdiv(1.0, accu) d_entropy_ptrs = d_entropy_ptr + offs_am * stride_d_entropy d_entropy = tl.load(d_entropy_ptrs, mask=offs_am < num_tokens, other=0.0) if reduction == 0: # none d_logprobs_ptrs = d_logprobs_ptr + offs_am * stride_d_logprobs - d_logprobs = tl.load(d_logprobs_ptrs, mask=offs_am < num_tokens, other=0.0) + d_logprobs = tl.load( + d_logprobs_ptrs, + mask=offs_am < num_tokens, + other=0.0) elif reduction == 1: # sum d_logprobs = tl.load(d_logprobs_ptr) d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,)) else: # mean - d_logprobs = tl.fdiv(tl.load(d_logprobs_ptr), num_tokens.to(tl.float32)) + d_logprobs = tl.fdiv( + tl.load(d_logprobs_ptr), + num_tokens.to( + tl.float32)) d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,)) d_logprobs = -1 * d_logprobs entropy_b_ptrs = entropy_b_ptr + offs_am * stride_entropy_b entropy_b = tl.load(entropy_b_ptrs, mask=offs_am < num_tokens, other=0.0) - hidden_ptrs = hidden_ptr + (offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k) + hidden_ptrs = hidden_ptr + ( + offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k + ) # weight_ptrs = weight_ptr + (offs_k[:, None] * stride_weight_k + offs_bn[None, :] * stride_weight_n) - weight_ptrs = weight_ptr + (offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k) + weight_ptrs = weight_ptr + ( + offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k + ) labels_ptrs = labels_ptr + offs_am * stride_labels labels = tl.load(labels_ptrs, mask=offs_am < num_tokens, other=0) - d_hidden_ptrs = d_hidden_ptr + offs_am[:, None] * stride_d_hidden_m + offs_k[None, :] * stride_d_hidden_k + d_hidden_ptrs = ( + d_hidden_ptr + + offs_am[:, None] * stride_d_hidden_m + + offs_k[None, :] * stride_d_hidden_k + ) # d_weight_ptrs = d_weight_ptr + offs_k[:, None] * stride_d_weight_k + offs_bn[None, :] * stride_d_weight_n - d_weight_ptrs = d_weight_ptr + offs_bn[:, None] * stride_d_weight_n + offs_k[None, :] * stride_d_weight_k + d_weight_ptrs = ( + d_weight_ptr + + offs_bn[:, None] * stride_d_weight_n + + offs_k[None, :] * stride_d_weight_k + ) logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)): _hidden = tl.load( hidden_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_am[:, None] < num_tokens), other=0.0, ) # _weight = tl.load(weight_ptrs, @@ -778,7 +958,8 @@ def efficient_entropy_backward_kernel_general_mainloop_MN( # other=0.0) _weight = tl.load( weight_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[:, None] < vocab_size), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_bn[:, None] < vocab_size), other=0.0, ) @@ -796,7 +977,11 @@ def efficient_entropy_backward_kernel_general_mainloop_MN( mask = (offs_bn + rank * vocab_size)[None, :] == labels[:, None] d_logits = d_logprobs[:, None] * (exp_logits * accu_rcp[:, None] - mask) - d_logits += d_entropy[:, None] * (-exp_logits * accu_rcp[:, None]) * (logits - entropy_b[:, None]) + d_logits += ( + d_entropy[:, None] + * (-exp_logits * accu_rcp[:, None]) + * (logits - entropy_b[:, None]) + ) # scale d_logits by temperature d_logits *= rcp_temperature @@ -805,18 +990,21 @@ def efficient_entropy_backward_kernel_general_mainloop_MN( for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)): _hidden = tl.load( hidden_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_am[:, None] < num_tokens), other=0.0, ) # _d_weight = tl.dot(tl.trans(_hidden).to(tl.float32), d_logits) # tl.atomic_add(d_weight_ptrs, # _d_weight, - # mask=(offs_k[:, None] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[None, :] < vocab_size)) + # mask=(offs_k[:, None] < hidden_size - k * BLOCK_SIZE_K) & + # (offs_bn[None, :] < vocab_size)) _d_weight = tl.dot(d_logits.trans(), _hidden.to(tl.float32)) tl.atomic_add( d_weight_ptrs, _d_weight, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[:, None] < vocab_size), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_bn[:, None] < vocab_size), ) # _weight = tl.load(weight_ptrs, @@ -825,14 +1013,16 @@ def efficient_entropy_backward_kernel_general_mainloop_MN( # _d_hidden = tl.dot(d_logits, tl.trans(_weight).to(tl.float32)) _weight = tl.load( weight_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[:, None] < vocab_size), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_bn[:, None] < vocab_size), other=0.0, ) _d_hidden = tl.dot(d_logits, _weight.to(tl.float32)) tl.atomic_add( d_hidden_ptrs, _d_hidden, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_am[:, None] < num_tokens), ) hidden_ptrs += BLOCK_SIZE_K * stride_hidden_k @@ -844,7 +1034,12 @@ def efficient_entropy_backward_kernel_general_mainloop_MN( @triton.autotune( configs=[ triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 16}, + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + }, num_stages=3, num_warps=8, ), @@ -897,42 +1092,74 @@ def efficient_entropy_backward_kernel_d_hidden( offs_k = tl.arange(0, BLOCK_SIZE_K) result_offs_k = pid_k * BLOCK_SIZE_K + offs_k - maximum = tl.load(maximum_ptr + offs_m * stride_maximum, mask=offs_m < num_tokens, other=0.0) - accu = tl.load(accu_ptr + offs_m * stride_accu, mask=offs_m < num_tokens, other=1e-6) + maximum = tl.load( + maximum_ptr + + offs_m * + stride_maximum, + mask=offs_m < num_tokens, + other=0.0) + accu = tl.load( + accu_ptr + offs_m * stride_accu, mask=offs_m < num_tokens, other=1e-6 + ) accu_rcp = tl.fdiv(1.0, accu) - d_entropy = tl.load(d_entropy_ptr + offs_m * stride_d_entropy, mask=offs_m < num_tokens, other=0.0) + d_entropy = tl.load( + d_entropy_ptr + + offs_m * + stride_d_entropy, + mask=offs_m < num_tokens, + other=0.0) if reduction == 0: - d_logprobs = tl.load(d_logprobs_ptr + offs_m * stride_d_logprobs, mask=offs_m < num_tokens, other=0.0) + d_logprobs = tl.load( + d_logprobs_ptr + offs_m * stride_d_logprobs, + mask=offs_m < num_tokens, + other=0.0, + ) elif reduction == 1: d_logprobs = tl.load(d_logprobs_ptr) d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,)) else: - d_logprobs = tl.fdiv(tl.load(d_logprobs_ptr), num_tokens.to(tl.float32)) + d_logprobs = tl.fdiv( + tl.load(d_logprobs_ptr), + num_tokens.to( + tl.float32)) d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,)) d_logprobs = -1 * d_logprobs - entropy_b = tl.load(entropy_b_ptr + offs_m * stride_entropy_b, mask=offs_m < num_tokens, other=0.0) - labels = tl.load(labels_ptr + offs_m * stride_labels, mask=offs_m < num_tokens, other=0) + entropy_b = tl.load( + entropy_b_ptr + + offs_m * + stride_entropy_b, + mask=offs_m < num_tokens, + other=0.0) + labels = tl.load( + labels_ptr + offs_m * stride_labels, mask=offs_m < num_tokens, other=0 + ) # iterate over vocab_size d_hidden = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K), dtype=tl.float32) for n in range(0, tl.cdiv(vocab_size, BLOCK_SIZE_N)): offs_n = n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - hidden_ptrs = hidden_ptr + (offs_m[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k) - weight_ptrs = weight_ptr + (offs_n[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k) + hidden_ptrs = hidden_ptr + ( + offs_m[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k + ) + weight_ptrs = weight_ptr + ( + offs_n[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k + ) # iterate over hidden_size to get logits logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)): _hidden = tl.load( hidden_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_m[:, None] < num_tokens), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_m[:, None] < num_tokens), other=0.0, ) _weight = tl.load( weight_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_n[:, None] < vocab_size), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_n[:, None] < vocab_size), other=0.0, ) @@ -947,31 +1174,53 @@ def efficient_entropy_backward_kernel_d_hidden( exp_logits = tl.exp(logits - maximum[:, None]) mask = (offs_n + rank * vocab_size)[None, :] == labels[:, None] - d_logits = d_logprobs[:, None] * (exp_logits * accu_rcp[:, None] - mask) - d_logits += d_entropy[:, None] * (-exp_logits * accu_rcp[:, None]) * (logits - entropy_b[:, None]) + d_logits = d_logprobs[:, None] * \ + (exp_logits * accu_rcp[:, None] - mask) + d_logits += ( + d_entropy[:, None] + * (-exp_logits * accu_rcp[:, None]) + * (logits - entropy_b[:, None]) + ) # scale d_logits d_logits *= rcp_temperature # calculate d_hidden - weight_ptrs = weight_ptr + (offs_n[:, None] * stride_weight_n + result_offs_k[None, :] * stride_weight_k) + weight_ptrs = weight_ptr + ( + offs_n[:, None] * stride_weight_n + result_offs_k[None, :] * stride_weight_k + ) _weight = tl.load( - weight_ptrs, mask=(result_offs_k[None, :] < hidden_size) & (offs_n[:, None] < vocab_size), other=0.0 + weight_ptrs, + mask=(result_offs_k[None, :] < hidden_size) + & (offs_n[:, None] < vocab_size), + other=0.0, ) - d_hidden = tl.dot(d_logits.to(weight_ptr.dtype.element_ty), _weight, d_hidden) + d_hidden = tl.dot( + d_logits.to( + weight_ptr.dtype.element_ty), + _weight, + d_hidden) # write back - tl.store( - d_hidden_ptr + offs_m[:, None] * stride_d_hidden_m + result_offs_k[None, :] * stride_d_hidden_k, - d_hidden, - mask=(offs_m[:, None] < num_tokens) & (result_offs_k[None, :] < hidden_size), - ) + tl.store(d_hidden_ptr + offs_m[:, + None] * stride_d_hidden_m + result_offs_k[None, + :] * stride_d_hidden_k, + d_hidden, + mask=(offs_m[:, + None] < num_tokens) & (result_offs_k[None, + :] < hidden_size), + ) @triton.autotune( configs=[ triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 16}, + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 128, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + }, num_stages=3, num_warps=8, ), @@ -1025,36 +1274,72 @@ def efficient_entropy_backward_kernel_d_weight( for m in range(0, tl.cdiv(num_tokens, BLOCK_SIZE_M)): offs_m = m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - maximum = tl.load(maximum_ptr + offs_m * stride_maximum, mask=offs_m < num_tokens, other=0.0) - accu = tl.load(accu_ptr + offs_m * stride_accu, mask=offs_m < num_tokens, other=1e-6) + maximum = tl.load( + maximum_ptr + + offs_m * + stride_maximum, + mask=offs_m < num_tokens, + other=0.0) + accu = tl.load( + accu_ptr + + offs_m * + stride_accu, + mask=offs_m < num_tokens, + other=1e-6) accu_rcp = tl.fdiv(1.0, accu) - d_entropy = tl.load(d_entropy_ptr + offs_m * stride_d_entropy, mask=offs_m < num_tokens, other=0.0) + d_entropy = tl.load( + d_entropy_ptr + offs_m * stride_d_entropy, + mask=offs_m < num_tokens, + other=0.0, + ) if reduction == 0: - d_logprobs = tl.load(d_logprobs_ptr + offs_m * stride_d_logprobs, mask=offs_m < num_tokens, other=0.0) + d_logprobs = tl.load( + d_logprobs_ptr + offs_m * stride_d_logprobs, + mask=offs_m < num_tokens, + other=0.0, + ) elif reduction == 1: d_logprobs = tl.load(d_logprobs_ptr) d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,)) else: - d_logprobs = tl.fdiv(tl.load(d_logprobs_ptr), num_tokens.to(tl.float32)) + d_logprobs = tl.fdiv( + tl.load(d_logprobs_ptr), + num_tokens.to( + tl.float32)) d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,)) d_logprobs = -1 * d_logprobs - entropy_b = tl.load(entropy_b_ptr + offs_m * stride_entropy_b, mask=offs_m < num_tokens, other=0.0) - labels = tl.load(labels_ptr + offs_m * stride_labels, mask=offs_m < num_tokens, other=0) - - hidden_ptrs = hidden_ptr + (offs_m[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k) - weight_ptrs = weight_ptr + (offs_n[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k) + entropy_b = tl.load( + entropy_b_ptr + offs_m * stride_entropy_b, + mask=offs_m < num_tokens, + other=0.0, + ) + labels = tl.load( + labels_ptr + + offs_m * + stride_labels, + mask=offs_m < num_tokens, + other=0) + + hidden_ptrs = hidden_ptr + ( + offs_m[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k + ) + weight_ptrs = weight_ptr + ( + offs_n[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k + ) logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)): _hidden = tl.load( hidden_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_m[:, None] < num_tokens), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_m[:, None] < num_tokens), other=0.0, ) _weight = tl.load( weight_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_n[:, None] < vocab_size), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_n[:, None] < vocab_size), other=0.0, ) @@ -1068,30 +1353,52 @@ def efficient_entropy_backward_kernel_d_weight( exp_logits = tl.exp(logits - maximum[:, None]) mask = (offs_n + rank * vocab_size)[None, :] == labels[:, None] - d_logits = d_logprobs[:, None] * (exp_logits * accu_rcp[:, None] - mask) - d_logits += d_entropy[:, None] * (-exp_logits * accu_rcp[:, None]) * (logits - entropy_b[:, None]) + d_logits = d_logprobs[:, None] * \ + (exp_logits * accu_rcp[:, None] - mask) + d_logits += ( + d_entropy[:, None] + * (-exp_logits * accu_rcp[:, None]) + * (logits - entropy_b[:, None]) + ) d_logits *= rcp_temperature - hidden_ptrs = hidden_ptr + (offs_m[:, None] * stride_hidden_m + result_offs_k[None, :] * stride_hidden_k) + hidden_ptrs = hidden_ptr + ( + offs_m[:, None] * stride_hidden_m + result_offs_k[None, :] * stride_hidden_k + ) _hidden = tl.load( - hidden_ptrs, mask=(result_offs_k[None, :] < hidden_size) & (offs_m[:, None] < num_tokens), other=0.0 + hidden_ptrs, + mask=(result_offs_k[None, :] < hidden_size) + & (offs_m[:, None] < num_tokens), + other=0.0, ) - d_weight = tl.dot(d_logits.to(d_weight_ptr.dtype.element_ty).trans(), _hidden, d_weight) + d_weight = tl.dot( + d_logits.to( + d_weight_ptr.dtype.element_ty).trans(), + _hidden, + d_weight) # write back - tl.store( - d_weight_ptr + offs_n[:, None] * stride_d_weight_n + result_offs_k[None, :] * stride_d_weight_k, - d_weight, - mask=(offs_n[:, None] < vocab_size) & (result_offs_k[None, :] < hidden_size), - ) + tl.store(d_weight_ptr + offs_n[:, + None] * stride_d_weight_n + result_offs_k[None, + :] * stride_d_weight_k, + d_weight, + mask=(offs_n[:, + None] < vocab_size) & (result_offs_k[None, + :] < hidden_size), + ) # NOTE: split tile from d_logits' perspective @triton.autotune( configs=[ triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 16}, + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + }, num_stages=3, num_warps=8, ), @@ -1158,28 +1465,40 @@ def efficient_entropy_backward_kernel_general_d_logits( maximum_ptrs = maximum_ptr + offs_am * stride_maximum maximum = tl.load(maximum_ptrs, mask=offs_am < num_tokens, other=0.0) accu_ptrs = accu_ptr + offs_am * stride_accu - accu = tl.load(accu_ptrs, mask=offs_am < num_tokens, other=1e-6) # epsilon to avoid division by zero + accu = tl.load( + accu_ptrs, mask=offs_am < num_tokens, other=1e-6 + ) # epsilon to avoid division by zero accu_rcp = tl.fdiv(1.0, accu) d_entropy_ptrs = d_entropy_ptr + offs_am * stride_d_entropy d_entropy = tl.load(d_entropy_ptrs, mask=offs_am < num_tokens, other=0.0) if reduction == 0: # none d_logprobs_ptrs = d_logprobs_ptr + offs_am * stride_d_logprobs - d_logprobs = tl.load(d_logprobs_ptrs, mask=offs_am < num_tokens, other=0.0) + d_logprobs = tl.load( + d_logprobs_ptrs, + mask=offs_am < num_tokens, + other=0.0) elif reduction == 1: # sum d_logprobs = tl.load(d_logprobs_ptr) d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,)) else: # mean - d_logprobs = tl.fdiv(tl.load(d_logprobs_ptr), num_tokens.to(tl.float32)) + d_logprobs = tl.fdiv( + tl.load(d_logprobs_ptr), + num_tokens.to( + tl.float32)) d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,)) d_logprobs = -1 * d_logprobs entropy_b_ptrs = entropy_b_ptr + offs_am * stride_entropy_b entropy_b = tl.load(entropy_b_ptrs, mask=offs_am < num_tokens, other=0.0) - hidden_ptrs = hidden_ptr + (offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k) + hidden_ptrs = hidden_ptr + ( + offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k + ) # weight_ptrs = weight_ptr + (offs_k[:, None] * stride_weight_k + offs_bn[None, :] * stride_weight_n) - weight_ptrs = weight_ptr + (offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k) + weight_ptrs = weight_ptr + ( + offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k + ) labels_ptrs = labels_ptr + offs_am * stride_labels labels = tl.load(labels_ptrs, mask=offs_am < num_tokens, other=0) @@ -1187,7 +1506,8 @@ def efficient_entropy_backward_kernel_general_d_logits( for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)): _hidden = tl.load( hidden_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_am[:, None] < num_tokens), other=0.0, ) # _weight = tl.load(weight_ptrs, @@ -1195,7 +1515,8 @@ def efficient_entropy_backward_kernel_general_d_logits( # other=0.0) _weight = tl.load( weight_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[:, None] < vocab_size), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_bn[:, None] < vocab_size), other=0.0, ) @@ -1213,13 +1534,21 @@ def efficient_entropy_backward_kernel_general_d_logits( mask = (offs_bn + rank * vocab_size)[None, :] == labels[:, None] d_logits = d_logprobs[:, None] * (exp_logits * accu_rcp[:, None] - mask) - d_logits += d_entropy[:, None] * (-exp_logits * accu_rcp[:, None]) * (logits - entropy_b[:, None]) + d_logits += ( + d_entropy[:, None] + * (-exp_logits * accu_rcp[:, None]) + * (logits - entropy_b[:, None]) + ) # scale d_logits by temperature d_logits *= rcp_temperature # store d_logits - d_logits_ptrs = d_logits_ptr + offs_am[:, None] * stride_d_logits_m + offs_bn[None, :] * stride_d_logits_n + d_logits_ptrs = ( + d_logits_ptr + + offs_am[:, None] * stride_d_logits_m + + offs_bn[None, :] * stride_d_logits_n + ) tl.store( d_logits_ptrs, d_logits, # will be implicitly converted to d_logits_ptrs.dtype.element_ty @@ -1230,7 +1559,12 @@ def efficient_entropy_backward_kernel_general_d_logits( @triton.autotune( configs=[ triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 16}, + { + "BLOCK_SIZE_M": 128, + "BLOCK_SIZE_N": 256, + "BLOCK_SIZE_K": 32, + "GROUP_SIZE_M": 16, + }, num_stages=3, num_warps=8, ), @@ -1284,39 +1618,81 @@ def efficient_entropy_backward_kernel_general_d_logits_split_N( pid_n = (pid % num_pid_in_group) // group_size_m offs_am = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_bn = split_idx * vocab_per_split + pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + offs_bn = ( + split_idx * + vocab_per_split + + pid_n * + BLOCK_SIZE_N + + tl.arange( + 0, + BLOCK_SIZE_N)) offs_k = tl.arange(0, BLOCK_SIZE_K) - maximum = tl.load(maximum_ptr + offs_am * stride_maximum, mask=offs_am < num_tokens, other=0.0) - accu = tl.load(accu_ptr + offs_am * stride_accu, mask=offs_am < num_tokens, other=1e-6) + maximum = tl.load( + maximum_ptr + + offs_am * + stride_maximum, + mask=offs_am < num_tokens, + other=0.0) + accu = tl.load( + accu_ptr + offs_am * stride_accu, mask=offs_am < num_tokens, other=1e-6 + ) accu_rcp = tl.fdiv(1.0, accu) - d_entropy = tl.load(d_entropy_ptr + offs_am * stride_d_entropy, mask=offs_am < num_tokens, other=0.0) + d_entropy = tl.load( + d_entropy_ptr + + offs_am * + stride_d_entropy, + mask=offs_am < num_tokens, + other=0.0) if reduction == 0: - d_logprobs = tl.load(d_logprobs_ptr + offs_am * stride_d_logprobs, mask=offs_am < num_tokens, other=0.0) + d_logprobs = tl.load( + d_logprobs_ptr + offs_am * stride_d_logprobs, + mask=offs_am < num_tokens, + other=0.0, + ) elif reduction == 1: d_logprobs = tl.load(d_logprobs_ptr) d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,)) else: - d_logprobs = tl.fdiv(tl.load(d_logprobs_ptr), num_tokens.to(tl.float32)) + d_logprobs = tl.fdiv( + tl.load(d_logprobs_ptr), + num_tokens.to( + tl.float32)) d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,)) d_logprobs = -1 * d_logprobs - entropy_b = tl.load(entropy_b_ptr + offs_am * stride_entropy_b, mask=offs_am < num_tokens, other=0.0) - labels = tl.load(labels_ptr + offs_am * stride_labels, mask=offs_am < num_tokens, other=0) - - hidden_ptrs = hidden_ptr + (offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k) - weight_ptrs = weight_ptr + (offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k) + entropy_b = tl.load( + entropy_b_ptr + + offs_am * + stride_entropy_b, + mask=offs_am < num_tokens, + other=0.0) + labels = tl.load( + labels_ptr + + offs_am * + stride_labels, + mask=offs_am < num_tokens, + other=0) + + hidden_ptrs = hidden_ptr + ( + offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k + ) + weight_ptrs = weight_ptr + ( + offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k + ) vocab_right_bound = min((split_idx + 1) * vocab_per_split, vocab_size) logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)): _hidden = tl.load( hidden_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_am[:, None] < num_tokens), other=0.0, ) _weight = tl.load( weight_ptrs, - mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[:, None] < vocab_right_bound), + mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) + & (offs_bn[:, None] < vocab_right_bound), other=0.0, ) logits = tl.dot(_hidden, _weight.trans(), logits) @@ -1329,16 +1705,25 @@ def efficient_entropy_backward_kernel_general_d_logits_split_N( mask = (offs_bn + rank * vocab_size)[None, :] == labels[:, None] d_logits = d_logprobs[:, None] * (exp_logits * accu_rcp[:, None] - mask) - d_logits += d_entropy[:, None] * (-exp_logits * accu_rcp[:, None]) * (logits - entropy_b[:, None]) + d_logits += ( + d_entropy[:, None] + * (-exp_logits * accu_rcp[:, None]) + * (logits - entropy_b[:, None]) + ) d_logits *= rcp_temperature # filter d_logits with mask result_offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - mask = (offs_am[:, None] < num_tokens) & (result_offs_n[None, :] < vocab_per_split) + mask = (offs_am[:, None] < num_tokens) & ( + result_offs_n[None, :] < vocab_per_split) tl.store( - d_logits_ptr + offs_am[:, None] * stride_d_logits_m + result_offs_n[None, :] * stride_d_logits_n, d_logits, mask + d_logits_ptr + + offs_am[:, None] * stride_d_logits_m + + result_offs_n[None, :] * stride_d_logits_n, + d_logits, + mask, ) @@ -1365,8 +1750,10 @@ def efficient_entropy_backward( assert hidden.is_contiguous() and weight.is_contiguous() and labels.is_contiguous() assert hidden.shape[0] == labels.shape[0] and hidden.shape[1] == weight.shape[1] - _rank = 0 if dist_process_group is None else dist.get_rank(dist_process_group) - _world_size = 1 if dist_process_group is None else dist.get_world_size(dist_process_group) + _rank = 0 if dist_process_group is None else dist.get_rank( + dist_process_group) + _world_size = ( + 1 if dist_process_group is None else dist.get_world_size(dist_process_group)) num_tokens, hidden_size = hidden.shape num_tokens = labels.shape[0] @@ -1387,11 +1774,15 @@ def efficient_entropy_backward( d_hidden, d_weight = None, None if _config._backward == BackwardEnum._Total_Fuse_MN or should_return_fp32_grad: - d_hidden = torch.zeros_like(hidden, dtype=torch.float32, device=hidden.device) - d_weight = torch.zeros_like(weight, dtype=torch.float32, device=weight.device) + d_hidden = torch.zeros_like( + hidden, dtype=torch.float32, device=hidden.device) + d_weight = torch.zeros_like( + weight, dtype=torch.float32, device=weight.device) else: - d_hidden = torch.empty_like(hidden, dtype=hidden.dtype, device=hidden.device) - d_weight = torch.empty_like(weight, dtype=hidden.dtype, device=weight.device) + d_hidden = torch.empty_like( + hidden, dtype=hidden.dtype, device=hidden.device) + d_weight = torch.empty_like( + weight, dtype=hidden.dtype, device=weight.device) assert d_hidden.is_contiguous() and d_weight.is_contiguous() assert maximum.is_contiguous() and acc.is_contiguous() @@ -1409,7 +1800,10 @@ def efficient_entropy_backward( if _config._backward == BackwardEnum._Total_Fuse_MN: # --- Triton doesn't materialize d_logits at all. Split tiles at the perspective of d_logits. def mainloop_grid(meta): - return (triton.cdiv(num_tokens, meta["BLOCK_SIZE_M"]) * triton.cdiv(vocab_size, meta["BLOCK_SIZE_N"]),) + return ( + triton.cdiv(num_tokens, meta["BLOCK_SIZE_M"]) + * triton.cdiv(vocab_size, meta["BLOCK_SIZE_N"]), + ) efficient_entropy_backward_kernel_general_mainloop_MN[mainloop_grid]( num_tokens, @@ -1445,13 +1839,18 @@ def mainloop_grid(meta): ) elif _config._backward == BackwardEnum._Total_Separate: - _d_logits = torch.empty((num_tokens, vocab_size), device=hidden.device, dtype=hidden.dtype).contiguous() + _d_logits = torch.empty( + (num_tokens, vocab_size), device=hidden.device, dtype=hidden.dtype + ).contiguous() assert _d_logits.is_contiguous() if _config._use_triton: def d_logits_grid(meta): - return (triton.cdiv(num_tokens, meta["BLOCK_SIZE_M"]) * triton.cdiv(vocab_size, meta["BLOCK_SIZE_N"]),) + return ( + triton.cdiv(num_tokens, meta["BLOCK_SIZE_M"]) + * triton.cdiv(vocab_size, meta["BLOCK_SIZE_N"]), + ) efficient_entropy_backward_kernel_general_d_logits[d_logits_grid]( num_tokens, @@ -1486,17 +1885,25 @@ def d_logits_grid(meta): torch.matmul(_d_logits, weight, out=d_hidden) torch.matmul(_d_logits.T, hidden, out=d_weight) else: - raise AssertionError("Triton is required for efficient entropy kernel") + raise AssertionError( + "Triton is required for efficient entropy kernel") elif _config._backward == BackwardEnum._Split_Dlogits_N: vocab_per_split = 9504 num_splits = (vocab_size + vocab_per_split - 1) // vocab_per_split - _d_logits = torch.empty((num_tokens, vocab_per_split), device=hidden.device, dtype=hidden.dtype).contiguous() + _d_logits = torch.empty( + (num_tokens, + vocab_per_split), + device=hidden.device, + dtype=hidden.dtype).contiguous() assert _d_logits.is_contiguous() def d_logits_grid(meta): - return (triton.cdiv(num_tokens, meta["BLOCK_SIZE_M"]) * triton.cdiv(vocab_per_split, meta["BLOCK_SIZE_N"]),) + return ( + triton.cdiv(num_tokens, meta["BLOCK_SIZE_M"]) + * triton.cdiv(vocab_per_split, meta["BLOCK_SIZE_N"]), + ) for split_idx in range(num_splits): efficient_entropy_backward_kernel_general_d_logits_split_N[d_logits_grid]( @@ -1532,22 +1939,34 @@ def d_logits_grid(meta): ) if split_idx == (num_splits - 1): - vocab_right_bound = min((split_idx + 1) * vocab_per_split, vocab_size) - split_idx * vocab_per_split + vocab_right_bound = ( + min((split_idx + 1) * vocab_per_split, vocab_size) + - split_idx * vocab_per_split + ) _d_logits = _d_logits[:, :vocab_right_bound].contiguous() if split_idx == 0: - torch.matmul( - _d_logits, weight[split_idx * vocab_per_split : (split_idx + 1) * vocab_per_split, :], out=d_hidden - ) + torch.matmul(_d_logits, weight[split_idx * vocab_per_split: ( + split_idx + 1) * vocab_per_split, :, ], out=d_hidden, ) else: d_hidden += torch.matmul( - _d_logits, weight[split_idx * vocab_per_split : (split_idx + 1) * vocab_per_split, :] + _d_logits, + weight[ + split_idx * vocab_per_split: (split_idx + 1) * vocab_per_split, + :, + ], ) torch.matmul( - _d_logits.T, hidden, out=d_weight[split_idx * vocab_per_split : (split_idx + 1) * vocab_per_split, :] + _d_logits.T, + hidden, + out=d_weight[ + split_idx * vocab_per_split: (split_idx + 1) * vocab_per_split, : + ], ) elif _config._backward == BackwardEnum._Split_Dlogits_M: - raise NotImplementedError("BackwardEnum._Split_Dlogits_M is not implemented yet") + raise NotImplementedError( + "BackwardEnum._Split_Dlogits_M is not implemented yet" + ) return d_hidden, d_weight diff --git a/Agent0/executor_train/verl/verl/utils/kernel/linear_cross_entropy.py b/Agent0/executor_train/verl/verl/utils/kernel/linear_cross_entropy.py index 733a815..a613025 100644 --- a/Agent0/executor_train/verl/verl/utils/kernel/linear_cross_entropy.py +++ b/Agent0/executor_train/verl/verl/utils/kernel/linear_cross_entropy.py @@ -15,7 +15,7 @@ # limitations under the License. # -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -63,22 +63,31 @@ def forward( typing.List[torch.Tensor]: _description_ """ - assert isinstance(temperature, float), f"temperature must be a float, but got {type(temperature)}" - assert isinstance(reduction, str), f"reduction must be a str, but got {type(reduction)}" + assert isinstance( + temperature, float + ), f"temperature must be a float, but got {type(temperature)}" + assert isinstance( + reduction, str + ), f"reduction must be a str, but got {type(reduction)}" with torch.cuda.nvtx.range("LinearCrossEntropy-forward"): - REDUCTION = kernels.get_entropy_reduction_enum_number(reduction.lower()) + REDUCTION = kernels.get_entropy_reduction_enum_number( + reduction.lower()) original_hidden_shape = hidden.shape if len(hidden.shape) != 2: - hidden = hidden.view(-1, hidden.shape[-1]) # (batch_size * num_tokens, hidden_size) + hidden = hidden.view( + -1, hidden.shape[-1] + ) # (batch_size * num_tokens, hidden_size) if len(labels.shape) != 1: labels = labels.view(-1) - logprobs, entropy, _maximum, _accumulate, _entropy_b = kernels.efficient_entropy_forward( - hidden, weight, labels, REDUCTION, temperature, dist_process_group - ) + logprobs, entropy, _maximum, _accumulate, _entropy_b = ( + kernels.efficient_entropy_forward( + hidden, weight, labels, REDUCTION, temperature, dist_process_group)) - ctx.save_for_backward(hidden, weight, labels, _maximum, _accumulate, _entropy_b) + ctx.save_for_backward( + hidden, weight, labels, _maximum, _accumulate, _entropy_b + ) ctx.original_hidden_shape = original_hidden_shape ctx.REDUCTION = REDUCTION ctx.dist_process_group = dist_process_group @@ -87,9 +96,13 @@ def forward( return logprobs, entropy @staticmethod - def backward(ctx, dlogprobs: torch.Tensor, dentropy: torch.Tensor) -> list[torch.Tensor]: + def backward( + ctx, dlogprobs: torch.Tensor, dentropy: torch.Tensor + ) -> list[torch.Tensor]: with torch.cuda.nvtx.range("LinearCrossEntropy-backward"): - (hidden, weight, labels, _maximum, _accumulate, _entropy_b) = ctx.saved_tensors + (hidden, weight, labels, _maximum, _accumulate, _entropy_b) = ( + ctx.saved_tensors + ) REDUCTION = ctx.REDUCTION dist_process_group = ctx.dist_process_group should_return_fp32_grad = ctx.should_return_fp32_grad diff --git a/Agent0/executor_train/verl/verl/utils/logger/__init__.py b/Agent0/executor_train/verl/verl/utils/logger/__init__.py index e318436..1d03993 100644 --- a/Agent0/executor_train/verl/verl/utils/logger/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/logger/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/logger/aggregate_logger.py b/Agent0/executor_train/verl/verl/utils/logger/aggregate_logger.py index d29698a..5baa780 100644 --- a/Agent0/executor_train/verl/verl/utils/logger/aggregate_logger.py +++ b/Agent0/executor_train/verl/verl/utils/logger/aggregate_logger.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -64,7 +64,12 @@ class DecoratorLoggerBase: """ def __init__( - self, role: str, logger: logging.Logger = None, level=logging.DEBUG, rank: int = 0, log_only_rank_0: bool = True + self, + role: str, + logger: logging.Logger = None, + level=logging.DEBUG, + rank: int = 0, + log_only_rank_0: bool = True, ): self.role = role self.logger = logger @@ -95,7 +100,10 @@ def print_rank_0(message): print(message, flush=True) -def print_with_rank(message: str, rank: int = 0, log_only_rank_0: bool = False): +def print_with_rank( + message: str, + rank: int = 0, + log_only_rank_0: bool = False): """_summary_ Print a message with rank information. This function prints the message only if `log_only_rank_0` is False or if the rank is 0. @@ -109,7 +117,9 @@ def print_with_rank(message: str, rank: int = 0, log_only_rank_0: bool = False): print(f"[Rank {rank}] {message}", flush=True) -def print_with_rank_and_timer(message: str, rank: int = 0, log_only_rank_0: bool = False): +def print_with_rank_and_timer( + message: str, rank: int = 0, log_only_rank_0: bool = False +): """_summary_ Print a message with rank information and a timestamp. This function prints the message only if `log_only_rank_0` is False or if the rank is 0. @@ -125,7 +135,13 @@ def print_with_rank_and_timer(message: str, rank: int = 0, log_only_rank_0: bool print(message, flush=True) -def log_with_rank(message: str, rank, logger: logging.Logger, level=logging.INFO, log_only_rank_0: bool = False): +def log_with_rank( + message: str, + rank, + logger: logging.Logger, + level=logging.INFO, + log_only_rank_0: bool = False, +): """_summary_ Log a message with rank information using a logger. This function logs the message only if `log_only_rank_0` is False or if the rank is 0. diff --git a/Agent0/executor_train/verl/verl/utils/logging_utils.py b/Agent0/executor_train/verl/verl/utils/logging_utils.py index 13fa917..9e4634c 100644 --- a/Agent0/executor_train/verl/verl/utils/logging_utils.py +++ b/Agent0/executor_train/verl/verl/utils/logging_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -22,7 +22,9 @@ 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) + logging.basicConfig( + format="%(levelname)s:%(asctime)s:%(message)s", + level=level) def log_to_file(string): diff --git a/Agent0/executor_train/verl/verl/utils/megatron/__init__.py b/Agent0/executor_train/verl/verl/utils/megatron/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/utils/megatron/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/megatron/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/megatron/dist_checkpointing.py b/Agent0/executor_train/verl/verl/utils/megatron/dist_checkpointing.py index d95752a..e22acbd 100644 --- a/Agent0/executor_train/verl/verl/utils/megatron/dist_checkpointing.py +++ b/Agent0/executor_train/verl/verl/utils/megatron/dist_checkpointing.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -51,6 +51,8 @@ def load_dist_checkpointing(sharded_state_dict, ckpt_dir): ) # Load model sharded state dicts - state_dict = dist_checkpointing.load(sharded_state_dict, ckpt_dir, sharded_strategy=load_strategy) + state_dict = dist_checkpointing.load( + sharded_state_dict, ckpt_dir, sharded_strategy=load_strategy + ) return state_dict diff --git a/Agent0/executor_train/verl/verl/utils/megatron/memory.py b/Agent0/executor_train/verl/verl/utils/megatron/memory.py index bc62d42..88d59ae 100644 --- a/Agent0/executor_train/verl/verl/utils/megatron/memory.py +++ b/Agent0/executor_train/verl/verl/utils/megatron/memory.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -22,7 +22,12 @@ 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=get_device_id(), requires_grad=False) + self.data = torch.zeros( + self.numel_padded, + dtype=self.dtype, + device=get_device_id(), + requires_grad=False, + ) def zero(self): """Reset the buffer to zero.""" diff --git a/Agent0/executor_train/verl/verl/utils/megatron/optimizer.py b/Agent0/executor_train/verl/verl/utils/megatron/optimizer.py index 100c161..6075664 100644 --- a/Agent0/executor_train/verl/verl/utils/megatron/optimizer.py +++ b/Agent0/executor_train/verl/verl/utils/megatron/optimizer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -14,7 +14,9 @@ # limitations under the License. from megatron.core.optimizer import OptimizerConfig -from megatron.core.optimizer import get_megatron_optimizer as get_megatron_optimizer_native +from megatron.core.optimizer import ( + get_megatron_optimizer as get_megatron_optimizer_native, +) from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler @@ -47,10 +49,15 @@ def get_megatron_optimizer_param_scheduler( wsd_decay_steps = None if config.get("lr_wsd_decay_steps", None) is not None: wsd_decay_steps = config.lr_wsd_decay_steps - if config.get("lr_warmup_steps_ratio", None) is not None and ( - config.get("lr_warmup_steps", None) is None or config.lr_warmup_steps <= 0 - ): - config.lr_warmup_steps = int(config.lr_warmup_steps_ratio * config.lr_decay_steps) + if config.get( + "lr_warmup_steps_ratio", + None) is not None and ( + config.get( + "lr_warmup_steps", + None) is None or config.lr_warmup_steps <= 0): + config.lr_warmup_steps = int( + config.lr_warmup_steps_ratio * config.lr_decay_steps + ) opt_param_scheduler = OptimizerParamScheduler( optimizer, @@ -65,7 +72,8 @@ def get_megatron_optimizer_param_scheduler( wd_incr_steps=config.total_training_steps, wd_incr_style=config.weight_decay_incr_style, use_checkpoint_opt_param_scheduler=config.use_checkpoint_opt_param_scheduler, - override_opt_param_scheduler=(not config.use_checkpoint_opt_param_scheduler), + override_opt_param_scheduler=( + not config.use_checkpoint_opt_param_scheduler), wsd_decay_steps=wsd_decay_steps, lr_wsd_decay_style=config.lr_wsd_decay_style, ) diff --git a/Agent0/executor_train/verl/verl/utils/megatron/pipeline_parallel.py b/Agent0/executor_train/verl/verl/utils/megatron/pipeline_parallel.py index 50ba697..13d305e 100644 --- a/Agent0/executor_train/verl/verl/utils/megatron/pipeline_parallel.py +++ b/Agent0/executor_train/verl/verl/utils/megatron/pipeline_parallel.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -27,14 +27,16 @@ def compute_transformers_input_shapes(batches, meta_info): 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) + 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(), + input_ids_rmpad.shape[0] + // mpu.get_tensor_model_parallel_world_size(), 1, meta_info["hidden_size"], ] @@ -42,7 +44,8 @@ def compute_transformers_input_shapes(batches, meta_info): ) else: # compute shapes for model_inputs - input_shapes.append(torch.Size([input_ids_rmpad.shape[0], 1, meta_info["hidden_size"]])) + input_shapes.append(torch.Size( + [input_ids_rmpad.shape[0], 1, meta_info["hidden_size"]])) return input_shapes diff --git a/Agent0/executor_train/verl/verl/utils/megatron/sequence_parallel.py b/Agent0/executor_train/verl/verl/utils/megatron/sequence_parallel.py index 52fda9b..fcbc5ad 100644 --- a/Agent0/executor_train/verl/verl/utils/megatron/sequence_parallel.py +++ b/Agent0/executor_train/verl/verl/utils/megatron/sequence_parallel.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -39,7 +39,11 @@ def pad_to_sequence_parallel(unpad_tokens: torch.Tensor): total_nnz = unpad_tokens.shape[0] sp_world_size = mpu.get_tensor_model_parallel_world_size() - pad_size = 0 if total_nnz % sp_world_size == 0 else sp_world_size - total_nnz % sp_world_size + pad_size = ( + 0 + if total_nnz % sp_world_size == 0 + else sp_world_size - total_nnz % sp_world_size + ) if pad_size > 0: if unpad_tokens.ndim == 1: @@ -47,6 +51,8 @@ def pad_to_sequence_parallel(unpad_tokens: torch.Tensor): 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") + raise NotImplementedError( + f"Padding dim {unpad_tokens.ndim()} is not supported" + ) return unpad_tokens diff --git a/Agent0/executor_train/verl/verl/utils/megatron/tensor_parallel.py b/Agent0/executor_train/verl/verl/utils/megatron/tensor_parallel.py index d4a99b9..64c8dbe 100644 --- a/Agent0/executor_train/verl/verl/utils/megatron/tensor_parallel.py +++ b/Agent0/executor_train/verl/verl/utils/megatron/tensor_parallel.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -46,7 +46,8 @@ def get_default_kwargs_for_model_parallel_config(): def get_default_model_parallel_config(): from megatron.core import ModelParallelConfig - return ModelParallelConfig(**get_default_kwargs_for_model_parallel_config()) + return ModelParallelConfig( + **get_default_kwargs_for_model_parallel_config()) def get_common_default_kwargs_for_parallel_linear(): @@ -93,7 +94,9 @@ def get_default_kwargs_for_parallel_embedding(): def is_tensor_parallel_param(param): - return hasattr(param, "tensor_model_parallel") and param.tensor_model_parallel + return hasattr( + param, + "tensor_model_parallel") and param.tensor_model_parallel def get_tensor_parallel_partition_dim(param): @@ -114,21 +117,36 @@ def mul_reduce(a, b): return (a * b).sum(dim=-1, keepdim=True) 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()) + 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()) + 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.div_(normalized_sum_exp_logits) - sum_softmax_times_logits = mul_reduce(softmax_logits, vocab_parallel_logits) - 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) + sum_softmax_times_logits = mul_reduce( + softmax_logits, vocab_parallel_logits) + 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 + vocab_parallel_logits, softmax_logits, sum_softmax_times_logits = ( + ctx.saved_tensors + ) # reuse softmax_logits as grad vocab_parallel_logits.sub_(sum_softmax_times_logits) softmax_logits.mul_(vocab_parallel_logits) @@ -139,7 +157,8 @@ def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor: return softmax_logits -def vocab_parallel_entropy(vocab_parallel_logits: torch.Tensor) -> torch.Tensor: +def vocab_parallel_entropy( + vocab_parallel_logits: torch.Tensor) -> torch.Tensor: """Compute entropy when the logits are sharded in tp ranks Args: @@ -155,10 +174,14 @@ def vocab_parallel_log_probs_from_logits(logits, labels): """TODO(zhangchi.usc1992): We may change the implementation later""" from megatron.core import tensor_parallel - return -tensor_parallel.vocab_parallel_cross_entropy(vocab_parallel_logits=logits, target=labels) + return -tensor_parallel.vocab_parallel_cross_entropy( + vocab_parallel_logits=logits, target=labels + ) -def vocab_parallel_log_probs_from_logits_response_rmpad(input_ids, attention_mask, logits_rmpad, response_length): +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 @@ -173,14 +196,21 @@ def vocab_parallel_log_probs_from_logits_response_rmpad(input_ids, attention_mas 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, 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 + 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] + output = full_output.squeeze(-1)[ + :, -response_length - 1: -1 + ] # [batch_size, response_length] return output diff --git a/Agent0/executor_train/verl/verl/utils/megatron_utils.py b/Agent0/executor_train/verl/verl/utils/megatron_utils.py index 2fc7437..3b25522 100644 --- a/Agent0/executor_train/verl/verl/utils/megatron_utils.py +++ b/Agent0/executor_train/verl/verl/utils/megatron_utils.py @@ -1,7 +1,7 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -57,16 +57,18 @@ def get_model( 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" - ) + 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_provider_func( + pre_process=pre_process, post_process=post_process + ) this_model.model_type = model_type model.append(this_model) mpu.set_virtual_pipeline_model_parallel_rank(0) @@ -77,21 +79,32 @@ def get_model( 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" - ) + 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)) + 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 + 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_provider_func( + pre_process=pre_process, post_process=post_process + ) model.model_type = model_type if not isinstance(model, list): @@ -103,7 +116,8 @@ def get_model( # 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) + tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes( + param) # Print number of parameters. if mpu.get_data_parallel_rank() == 0: @@ -111,13 +125,19 @@ def get_model( " > 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]), + sum( + [ + sum([p.nelement() for p in model_module.parameters()]) + for model_module in model + ] + ), ), flush=True, ) # GPU allocation. - if transformer_config is None or (not transformer_config.use_cpu_initialization): + if transformer_config is None or ( + not transformer_config.use_cpu_initialization): for model_module in model: model_module.to(f"{get_device_name()}:{get_device_id()}") @@ -138,7 +158,8 @@ def get_model( ddp_config=DistributedDataParallelConfig( overlap_grad_reduce=False, use_distributed_optimizer=use_distributed_optimizer, - grad_reduce_in_fp32=True, # [old] accumulate_allreduce_grads_in_fp32=True, + grad_reduce_in_fp32=True, + # [old] accumulate_allreduce_grads_in_fp32=True, ), ) ddp_models.append(ddp_model) @@ -168,11 +189,17 @@ def unwrap_model(model, module_instances=ALL_MODULE_WRAPPER_CLASSNAMES): return unwrapped_model -def convert_config(hf_config: PretrainedConfig, megatron_config) -> TransformerConfig: +def convert_config( + hf_config: PretrainedConfig, + megatron_config) -> TransformerConfig: print(f"megatron config {megatron_config}") dt = PrecisionType.to_dtype(megatron_config.params_dtype) print(f"pipeline_dtype=megatron_config {dt}") - qkv_bias = True if "Qwen2ForCausalLM" in hf_config.architectures else getattr(hf_config, "attention_bias", False) + qkv_bias = ( + True + if "Qwen2ForCausalLM" in hf_config.architectures + else getattr(hf_config, "attention_bias", False) + ) overlap_p2p_comm = ( mpu.get_virtual_pipeline_model_parallel_world_size() is not None and mpu.get_virtual_pipeline_model_parallel_world_size() > 1 @@ -232,7 +259,8 @@ def mcore_model_parallel_config( params_dtype: torch.dtype, ) -> ModelParallelConfig: # WARNING: Code should not reach this point. This function is deprecated and will be removed. - # Please use hf_to_mcore_config_dense() from verl.models.mcore.config_converter instead. + # Please use hf_to_mcore_config_dense() from + # verl.models.mcore.config_converter instead. warnings.warn( "Code should not reach this point. This function is deprecated and will be removed. Please use " "hf_to_mcore_config_dense() from verl.models.mcore.config_converter instead.", @@ -264,19 +292,28 @@ def offload_megatron_model_to_cpu(models): """ for model_chunk in models: if isinstance(model_chunk, DDP): - model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers] + model_chunk_all_buffers = [ + model_chunk.buffers, + model_chunk.expert_parallel_buffers, + ] for buffers in model_chunk_all_buffers: for buffer in buffers: # offload parameters if buffer.param_data.storage().size() > 0: - buffer.param_data.cpu_data = buffer.param_data.data.cpu().pin_memory() + buffer.param_data.cpu_data = ( + buffer.param_data.data.cpu().pin_memory() + ) buffer.param_data_size = buffer.param_data.storage().size() buffer.param_data.storage().resize_(0) - assert buffer.param_data_size == buffer.param_data.cpu_data.storage().size() + assert ( + buffer.param_data_size + == buffer.param_data.cpu_data.storage().size() + ) if buffer.grad_data.storage().size() > 0: - # if the grad_data size is already zero, we assume that it is already offloaded + # if the grad_data size is already zero, we assume that + # it is already offloaded buffer.grad_data_size = buffer.grad_data.storage().size() buffer.grad_data.storage().resize_(0) else: @@ -293,7 +330,10 @@ def offload_megatron_model_to_cpu(models): def load_megatron_model_to_gpu(models, load_grad=True): for model_chunk in models: if isinstance(model_chunk, DDP): - model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers] + model_chunk_all_buffers = [ + model_chunk.buffers, + model_chunk.expert_parallel_buffers, + ] for buffers in model_chunk_all_buffers: for buffer in buffers: # sometimes, we don't want to load grad for pure inference @@ -304,7 +344,9 @@ def load_megatron_model_to_gpu(models, load_grad=True): if buffer.param_data.storage().size() == 0: buffer.param_data.storage().resize_(buffer.param_data_size) # copy data from cpu to cuda - buffer.param_data.copy_(buffer.param_data.cpu_data, non_blocking=True) + buffer.param_data.copy_( + buffer.param_data.cpu_data, non_blocking=True + ) else: # we need this for ref module device_id = get_device_id() @@ -429,9 +471,11 @@ def _iter_opts(opt): opt_state_dict_values = _opt.optimizer.state.values() for v in opt_state_dict_values: if "exp_avg" in v: - v["exp_avg"] = v["exp_avg"].to(get_device_id(), non_blocking=True) + v["exp_avg"] = v["exp_avg"].to( + get_device_id(), non_blocking=True) if "exp_avg_sq" in v: - v["exp_avg_sq"] = v["exp_avg_sq"].to(get_device_id(), non_blocking=True) + v["exp_avg_sq"] = v["exp_avg_sq"].to( + get_device_id(), non_blocking=True) gc.collect() get_torch_device().empty_cache() @@ -472,7 +516,11 @@ def convert_qkv_shard(full_tensor, q_name, k_name, v_name): q_shard_list = [] k_shard_list = [] v_shard_list = [] - hidden_size_per_head = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads) + hidden_size_per_head = getattr( + config, + "head_dim", + config.hidden_size // + config.num_attention_heads) if config.num_key_value_heads >= tp_size: q_size_tp = hidden_size_per_head * config.num_attention_heads // tp_size @@ -480,13 +528,14 @@ def convert_qkv_shard(full_tensor, q_name, k_name, v_name): total_size = q_size_tp + 2 * kv_size_tp for i in range(tp_size): num_query_groups_per_partition = num_query_groups // tp_size - qkv_part = full_tensor[i * total_size : (i + 1) * total_size] + qkv_part = full_tensor[i * total_size: (i + 1) * total_size] q_size_chunk = q_size_tp // num_query_groups_per_partition kv_size_chunk = kv_size_tp // num_query_groups_per_partition - for qkv_part_chunk in qkv_part.chunk(num_query_groups_per_partition): + for qkv_part_chunk in qkv_part.chunk( + num_query_groups_per_partition): q_part = qkv_part_chunk[:q_size_chunk] - k_part = qkv_part_chunk[q_size_chunk : q_size_chunk + kv_size_chunk] - v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk :] + k_part = qkv_part_chunk[q_size_chunk: q_size_chunk + kv_size_chunk] + v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk:] q_shard_list.append(q_part) k_shard_list.append(k_part) v_shard_list.append(v_part) @@ -496,13 +545,14 @@ def convert_qkv_shard(full_tensor, q_name, k_name, v_name): total_size = q_size_tp + 2 * kv_size_tp for i in range(tp_size): num_query_groups_per_partition = num_query_groups // tp_size - qkv_part = full_tensor[i * total_size : (i + 1) * total_size] + qkv_part = full_tensor[i * total_size: (i + 1) * total_size] q_size_chunk = q_size_tp // num_query_groups_per_partition kv_size_chunk = kv_size_tp // num_query_groups_per_partition - for qkv_part_chunk in qkv_part.chunk(num_query_groups_per_partition): + for qkv_part_chunk in qkv_part.chunk( + num_query_groups_per_partition): q_part = qkv_part_chunk[:q_size_chunk] - k_part = qkv_part_chunk[q_size_chunk : q_size_chunk + kv_size_chunk] - v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk :] + k_part = qkv_part_chunk[q_size_chunk: q_size_chunk + kv_size_chunk] + v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk:] q_shard_list.append(q_part) if i * config.num_key_value_heads % tp_size == 0: k_shard_list.append(k_part) @@ -520,7 +570,8 @@ def convert_gate_up_shard(full_tensor, gate_name, up_name): 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_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) @@ -540,7 +591,8 @@ def convert_gate_up_shard(full_tensor, gate_name, up_name): new_params[f"model.layers.{layer_number}.self_attn.o_proj.weight"] = param elif component == "linear_qkv" and not isinstance(param, list): if param_type == "layer_norm_weight": - new_params[f"model.layers.{layer_number}.input_layernorm.weight"] = param + new_params[f"model.layers.{layer_number}.input_layernorm.weight"] = ( + param) else: if convert_qkv_gate_up_by_trunk_concat: convert_qkv_shard( @@ -550,16 +602,23 @@ def convert_gate_up_shard(full_tensor, gate_name, up_name): f"model.layers.{layer_number}.self_attn.v_proj.{param_type}", ) else: - new_params[f"model.layers.{layer_number}.self_attn.qkv_proj.{param_type}"] = param + new_params[ + f"model.layers.{layer_number}.self_attn.qkv_proj.{param_type}" + ] = param elif component == "q_layernorm" or component == "k_layernorm": hf_component = component.replace("layer", "") - new_params[f"model.layers.{layer_number}.self_attn.{hf_component}.weight"] = param + new_params[ + f"model.layers.{layer_number}.self_attn.{hf_component}.weight" + ] = param else: assert isinstance(param, list) and len(param) == 3 assert param_type == "weight" or param_type == "bias" - new_params[f"model.layers.{layer_number}.self_attn.q_proj.{param_type}"] = param[0] - new_params[f"model.layers.{layer_number}.self_attn.k_proj.{param_type}"] = param[1] - new_params[f"model.layers.{layer_number}.self_attn.v_proj.{param_type}"] = param[2] + new_params[f"model.layers.{layer_number}.self_attn.q_proj.{param_type}"] = ( + param[0]) + new_params[f"model.layers.{layer_number}.self_attn.k_proj.{param_type}"] = ( + param[1]) + new_params[f"model.layers.{layer_number}.self_attn.v_proj.{param_type}"] = ( + param[2]) elif "mlp" in name: splitted_name = name.split(".") layer_number = splitted_name[2] @@ -567,7 +626,9 @@ def convert_gate_up_shard(full_tensor, gate_name, up_name): param_type = splitted_name[5] if component == "linear_fc1" and not isinstance(param, list): if param_type == "layer_norm_weight": - new_params[f"model.layers.{layer_number}.post_attention_layernorm.weight"] = param + new_params[ + f"model.layers.{layer_number}.post_attention_layernorm.weight" + ] = param elif param_type == "weight": if convert_qkv_gate_up_by_trunk_concat: convert_gate_up_shard( @@ -576,7 +637,9 @@ def convert_gate_up_shard(full_tensor, gate_name, up_name): f"model.layers.{layer_number}.mlp.up_proj.weight", ) else: - new_params[f"model.layers.{layer_number}.mlp.gate_up_proj.weight"] = param + new_params[ + f"model.layers.{layer_number}.mlp.gate_up_proj.weight" + ] = param elif component == "linear_fc1" and isinstance(param, list): assert len(param) == 2 assert param_type == "weight" or param_type == "bias" @@ -605,7 +668,9 @@ def broadcast_from_megatron_pp(tensor: torch.Tensor): tensor_spec = None tensor_spec_output = [None] * mpu.get_pipeline_model_parallel_world_size() torch.distributed.all_gather_object( - object_list=tensor_spec_output, obj=tensor_spec, group=mpu.get_pipeline_model_parallel_group() + object_list=tensor_spec_output, + obj=tensor_spec, + group=mpu.get_pipeline_model_parallel_group(), ) # find the src rank target_tensor_spec = None @@ -619,20 +684,32 @@ def broadcast_from_megatron_pp(tensor: torch.Tensor): src_rank = rank assert target_tensor_spec is not None if tensor is None: - tensor = torch.empty(size=target_tensor_spec[0], dtype=target_tensor_spec[1], device=get_device_id()) + tensor = torch.empty( + size=target_tensor_spec[0], + dtype=target_tensor_spec[1], + device=get_device_id(), + ) if target_tensor_spec[2] is not None: tensor.tensor_model_parallel = target_tensor_spec[2] if target_tensor_spec[3] is not None: tensor.partition_dim = target_tensor_spec[3] - global_rank = torch.distributed.get_global_rank(group=mpu.get_pipeline_model_parallel_group(), group_rank=src_rank) - torch.distributed.broadcast(tensor=tensor, src=global_rank, group=mpu.get_pipeline_model_parallel_group()) + global_rank = torch.distributed.get_global_rank( + group=mpu.get_pipeline_model_parallel_group(), group_rank=src_rank + ) + torch.distributed.broadcast( + tensor=tensor, + src=global_rank, + group=mpu.get_pipeline_model_parallel_group()) return tensor def broadcast_str_from_megatron_pp(obj: Any): obj_output = [None] * mpu.get_pipeline_model_parallel_world_size() - torch.distributed.all_gather_object(object_list=obj_output, obj=obj, group=mpu.get_pipeline_model_parallel_group()) + torch.distributed.all_gather_object( + object_list=obj_output, + obj=obj, + group=mpu.get_pipeline_model_parallel_group()) src_rank = None target_obj = None @@ -645,12 +722,18 @@ def broadcast_str_from_megatron_pp(obj: Any): assert target_obj is not None, "No valid object found to broadcast." - global_rank = torch.distributed.get_global_rank(group=mpu.get_pipeline_model_parallel_group(), group_rank=src_rank) + global_rank = torch.distributed.get_global_rank( + group=mpu.get_pipeline_model_parallel_group(), group_rank=src_rank + ) - obj_output = [None] * torch.distributed.get_world_size(group=mpu.get_pipeline_model_parallel_group()) + obj_output = [None] * torch.distributed.get_world_size( + group=mpu.get_pipeline_model_parallel_group() + ) obj_output[0] = target_obj torch.distributed.broadcast_object_list( - object_list=obj_output, src=global_rank, group=mpu.get_pipeline_model_parallel_group() + object_list=obj_output, + src=global_rank, + group=mpu.get_pipeline_model_parallel_group(), ) return obj_output[0] @@ -677,7 +760,8 @@ def default_tp_concat_fn( from megatron.core import mpu train_tp_size = mpu.get_tensor_model_parallel_world_size() - if layer_name_mapping.get("qkv_layer_name") in name and "layer_norm" not in name: + if layer_name_mapping.get( + "qkv_layer_name") in name and "layer_norm" not in name: # if the tensor is qkv, for each param on tp, split into q, k, v # concat q, k, v separately. q_lst = [] @@ -690,18 +774,28 @@ def default_tp_concat_fn( num_key_value_heads = hf_config.vision_config.num_heads assert num_attention_heads % num_key_value_heads == 0 num_q_per_kv = num_attention_heads // num_key_value_heads - assert infer_params[0].shape[0] % (num_q_per_kv + 2) == 0, ( - f"param '{name}' shape '{infer_params[0].shape}' dim0 is not divisible by {num_q_per_kv + 2}" - ) + assert ( + infer_params[0].shape[0] % + (num_q_per_kv + 2) == 0), f"param '{name}' shape '{ + infer_params[0].shape}' dim0 is not divisible by { + num_q_per_kv + 2}" 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] + split_size = [ + kv_size_per_tp * + num_q_per_kv, + kv_size_per_tp, + kv_size_per_tp] for infer_param in infer_params: num_query_groups_per_partition = num_key_value_heads // train_tp_size for chunk in infer_param.chunk(num_query_groups_per_partition): split_size = [ - kv_size_per_tp * num_q_per_kv // num_query_groups_per_partition, - kv_size_per_tp // num_query_groups_per_partition, - kv_size_per_tp // num_query_groups_per_partition, + kv_size_per_tp * + num_q_per_kv // + num_query_groups_per_partition, + kv_size_per_tp // + num_query_groups_per_partition, + kv_size_per_tp // + num_query_groups_per_partition, ] q, k, v = chunk.split(split_size) q_lst.append(q) @@ -710,7 +804,11 @@ def default_tp_concat_fn( 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) if not convert_qkv_gate_up_by_simple_split else [q, k, v] + infer_params = ( + torch.cat((q, k, v), dim=0) + if not convert_qkv_gate_up_by_simple_split + else [q, k, v] + ) elif ( layer_name_mapping.get("gate_proj_layer_name") in name @@ -726,14 +824,20 @@ def default_tp_concat_fn( 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) if not convert_qkv_gate_up_by_simple_split else [gate, up] + infer_params = ( + torch.cat((gate, up), dim=0) + if not convert_qkv_gate_up_by_simple_split + else [gate, up] + ) elif "mlp.experts.linear_fc2.weight" in name: # moe infer_params = torch.cat(infer_params, dim=1) else: # concat tensor - infer_params = torch.cat(infer_params, dim=tp_utils.get_tensor_parallel_partition_dim(train_params)) + infer_params = torch.cat( + infer_params, + dim=tp_utils.get_tensor_parallel_partition_dim(train_params)) return infer_params @@ -755,7 +859,8 @@ def per_tensor_generator( etp_group = mpu.get_expert_tensor_parallel_group() vpp_size = len(actor_module) all_gather_group = mpu.get_tensor_model_parallel_group() - all_gather_group_size = torch.distributed.get_world_size(group=all_gather_group) + all_gather_group_size = torch.distributed.get_world_size( + group=all_gather_group) def tensor_generator(): for scan_vpp_idx in range(vpp_size): @@ -767,8 +872,13 @@ def tensor_generator(): # note # there is a bug in megatron GPTModel # decoder.layers[n].mlp.router.expert_bias" in GPTModel is not registered in named_parameter, but in - # state_dict(). for now we patch it by adding those keys to extra_keys. - extra_keys = [x for x in model.state_dict().keys() if "_extra_state" not in x and x not in existing_keys] + # state_dict(). for now we patch it by adding those keys to + # extra_keys. + extra_keys = [ + x + for x in model.state_dict().keys() + if "_extra_state" not in x and x not in existing_keys + ] for name in extra_keys: yield name, model.state_dict()[name].to(get_device_id()) @@ -780,13 +890,19 @@ def tensor_generator(): for idx, (name, _) in enumerate(model.named_parameters()): existing_keys.add(name) meta_info.append((pp_rank, scan_vpp_idx, idx, name)) - extra_keys = [x for x in model.state_dict().keys() if "_extra_state" not in x and x not in existing_keys] + extra_keys = [ + x + for x in model.state_dict().keys() + if "_extra_state" not in x and x not in existing_keys + ] for name in extra_keys: meta_info.append((pp_rank, scan_vpp_idx, idx, name)) obj_spec_output = [None] * mpu.get_pipeline_model_parallel_world_size() torch.distributed.all_gather_object( - object_list=obj_spec_output, obj=meta_info, group=mpu.get_pipeline_model_parallel_group() + object_list=obj_spec_output, + obj=meta_info, + group=mpu.get_pipeline_model_parallel_group(), ) layer_list_meta = [item for sublist in obj_spec_output for item in sublist] @@ -798,7 +914,8 @@ def tensor_generator(): import warnings warnings.warn( - "Current model sharing word and embedding weights, skip output layer conversion", stacklevel=2 + "Current model sharing word and embedding weights, skip output layer conversion", + stacklevel=2, ) continue @@ -807,7 +924,9 @@ def tensor_generator(): cur_name, cur_tensor = next(gen_func) except StopIteration: cur_name, cur_tensor = None, None - cur_name = normalize_model_name(name, cur_pp_rank, scan_vpp_idx, transformer_config) + cur_name = normalize_model_name( + name, cur_pp_rank, scan_vpp_idx, transformer_config + ) else: cur_tensor, cur_name = None, None @@ -817,25 +936,34 @@ def tensor_generator(): # (xya): this is a hack to fix the name of the parameters while cur_name.startswith("module."): - cur_name = cur_name[len("module.") :] + cur_name = cur_name[len("module."):] # EP if ".mlp.experts.linear_fc" in cur_name and ep_size > 1: num_experts = weight_converter.mcore_config.num_moe_experts num_experts_per_rank = num_experts // ep_size - infer_params = [torch.empty_like(broad_pp_tensor) for _ in range(ep_size)] - torch.distributed.all_gather(infer_params, broad_pp_tensor, group=ep_group) + infer_params = [torch.empty_like( + broad_pp_tensor) for _ in range(ep_size)] + torch.distributed.all_gather( + infer_params, broad_pp_tensor, group=ep_group) name_prefix, local_expert_id = cur_name.split(".weight") local_expert_id = int(local_expert_id) - global_expert_ids = [num_experts_per_rank * ep_rank + local_expert_id for ep_rank in range(ep_size)] - global_expert_names = [f"{name_prefix}.weight{expert_id}" for expert_id in global_expert_ids] - - for name, param in zip(global_expert_names, infer_params, strict=True): + global_expert_ids = [ + num_experts_per_rank * ep_rank + local_expert_id + for ep_rank in range(ep_size) + ] + global_expert_names = [ + f"{name_prefix}.weight{expert_id}" for expert_id in global_expert_ids] + + for name, param in zip( + global_expert_names, infer_params, strict=True): if etp_size > 1: # gather etp - etp_params = [torch.empty_like(param) for _ in range(etp_size)] - torch.distributed.all_gather(etp_params, param, group=etp_group) + etp_params = [ + torch.empty_like(param) for _ in range(etp_size)] + torch.distributed.all_gather( + etp_params, param, group=etp_group) params = etp_params else: params = [param] @@ -851,7 +979,8 @@ def tensor_generator(): ) if not isinstance(merge_params, list): merge_params = [merge_params] - converted_names, converted_params = weight_converter.convert_param(name, merge_params) + converted_names, converted_params = weight_converter.convert_param( + name, merge_params) yield from zip(converted_names, converted_params, strict=True) continue @@ -862,8 +991,15 @@ def tensor_generator(): if all_gather_group_size <= 1: infer_params = [broad_pp_tensor] else: - infer_params = [torch.empty_like(broad_pp_tensor) for _ in range(all_gather_group_size)] - torch.distributed.all_gather(infer_params, broad_pp_tensor, group=mpu.get_tensor_model_parallel_group()) + infer_params = [ + torch.empty_like(broad_pp_tensor) + for _ in range(all_gather_group_size) + ] + torch.distributed.all_gather( + infer_params, + broad_pp_tensor, + group=mpu.get_tensor_model_parallel_group(), + ) infer_params = default_tp_concat_fn( layer_name_mapping, cur_name, @@ -878,12 +1014,17 @@ def tensor_generator(): if not isinstance(infer_params, list): infer_params = [infer_params] - converted_names, converted_params = weight_converter.convert_param(cur_name, infer_params) + converted_names, converted_params = weight_converter.convert_param( + cur_name, infer_params + ) yield from zip(converted_names, converted_params, strict=True) -def get_transformer_layer_offset(pipeline_rank, vp_rank, config: TransformerConfig): +def get_transformer_layer_offset( + pipeline_rank, + vp_rank, + config: TransformerConfig): ''' Get the index offset of any pipeline stage, given the level of pipelining. @@ -899,7 +1040,8 @@ def get_transformer_layer_offset(pipeline_rank, vp_rank, config: TransformerConf or config.num_layers_in_last_pipeline_stage is not None ): # Calculate number of pipeline stages to distribute the remaining Transformer - # layers after deducting the Transformer layers in the first or the last stages + # layers after deducting the Transformer layers in the first or the + # last stages middle_pipeline_stages = config.pipeline_model_parallel_size middle_pipeline_stages -= sum( [ @@ -916,14 +1058,20 @@ def get_transformer_layer_offset(pipeline_rank, vp_rank, config: TransformerConf # are not set, we will not enable uneven pipeline. All layers will be treated # as middle layers. num_layers_in_first_pipeline_stage = ( - 0 if config.num_layers_in_first_pipeline_stage is None else config.num_layers_in_first_pipeline_stage + 0 + if config.num_layers_in_first_pipeline_stage is None + else config.num_layers_in_first_pipeline_stage ) num_layers_in_last_pipeline_stage = ( - 0 if config.num_layers_in_last_pipeline_stage is None else config.num_layers_in_last_pipeline_stage + 0 + if config.num_layers_in_last_pipeline_stage is None + else config.num_layers_in_last_pipeline_stage ) middle_num_layers = ( - config.num_layers - num_layers_in_first_pipeline_stage - num_layers_in_last_pipeline_stage + config.num_layers + - num_layers_in_first_pipeline_stage + - num_layers_in_last_pipeline_stage ) if mpu.get_virtual_pipeline_model_parallel_world_size() is not None: @@ -945,39 +1093,48 @@ def get_transformer_layer_offset(pipeline_rank, vp_rank, config: TransformerConf else config.num_layers_in_last_pipeline_stage // vp_size ) - num_layers_per_vritual_model_chunk_in_middle_pipeline_stage = middle_num_layers // vp_size + num_layers_per_vritual_model_chunk_in_middle_pipeline_stage = ( + middle_num_layers // vp_size + ) # First stage + middle stage + last stage total_virtual_chunks = ( - num_layers_per_virtual_model_chunk_in_first_pipeline_stage - + num_layers_per_vritual_model_chunk_in_middle_pipeline_stage - + num_layers_per_virtual_model_chunk_in_last_pipeline_stage - ) + num_layers_per_virtual_model_chunk_in_first_pipeline_stage + + num_layers_per_vritual_model_chunk_in_middle_pipeline_stage + + num_layers_per_virtual_model_chunk_in_last_pipeline_stage) - # Calculate the layer offset with interleaved uneven pipeline parallelism + # Calculate the layer offset with interleaved uneven pipeline + # parallelism if pipeline_rank == 0: offset = vp_rank * total_virtual_chunks else: - offset = ( - vp_rank * total_virtual_chunks - + num_layers_per_virtual_model_chunk_in_first_pipeline_stage - + (pipeline_rank - 1) - * (num_layers_per_vritual_model_chunk_in_middle_pipeline_stage // middle_pipeline_stages) - ) + offset = (vp_rank * + total_virtual_chunks + + num_layers_per_virtual_model_chunk_in_first_pipeline_stage + + (pipeline_rank - + 1) * + (num_layers_per_vritual_model_chunk_in_middle_pipeline_stage // + middle_pipeline_stages)) else: if middle_pipeline_stages > 0: - num_layers_per_pipeline_rank = middle_num_layers // middle_pipeline_stages + num_layers_per_pipeline_rank = ( + middle_num_layers // middle_pipeline_stages + ) else: num_layers_per_pipeline_rank = 0 middle_pipeline_rank = ( - pipeline_rank if config.num_layers_in_first_pipeline_stage is None else pipeline_rank - 1 + pipeline_rank + if config.num_layers_in_first_pipeline_stage is None + else pipeline_rank - 1 ) if pipeline_rank == 0: offset = 0 else: - offset = (middle_pipeline_rank * num_layers_per_pipeline_rank) + num_layers_in_first_pipeline_stage + offset = ( + middle_pipeline_rank * num_layers_per_pipeline_rank + ) + num_layers_in_first_pipeline_stage else: num_layers = config.num_layers @@ -989,23 +1146,35 @@ def get_transformer_layer_offset(pipeline_rank, vp_rank, config: TransformerConf if config.account_for_loss_in_pipeline_split: num_layers += 1 - num_layers_per_pipeline_rank = num_layers // config.pipeline_model_parallel_size + num_layers_per_pipeline_rank = ( + num_layers // config.pipeline_model_parallel_size + ) if mpu.get_virtual_pipeline_model_parallel_world_size() is not None: vp_size = mpu.get_virtual_pipeline_model_parallel_world_size() num_layers_per_virtual_rank = num_layers_per_pipeline_rank // vp_size total_virtual_chunks = num_layers // vp_size - offset = vp_rank * total_virtual_chunks + (pipeline_rank * num_layers_per_virtual_rank) + offset = vp_rank * total_virtual_chunks + ( + pipeline_rank * num_layers_per_virtual_rank + ) - # Reduce the offset of embedding layer from the total layer number - if config.account_for_embedding_in_pipeline_split and not mpu.is_pipeline_first_stage(): + # Reduce the offset of embedding layer from the total layer + # number + if ( + config.account_for_embedding_in_pipeline_split + and not mpu.is_pipeline_first_stage() + ): offset -= 1 else: offset = pipeline_rank * num_layers_per_pipeline_rank - # Reduce the offset of embedding layer from the total layer number - if config.account_for_embedding_in_pipeline_split and not mpu.is_pipeline_first_stage(): + # Reduce the offset of embedding layer from the total layer + # number + if ( + config.account_for_embedding_in_pipeline_split + and not mpu.is_pipeline_first_stage() + ): offset -= 1 else: offset = 0 diff --git a/Agent0/executor_train/verl/verl/utils/memory_buffer.py b/Agent0/executor_train/verl/verl/utils/memory_buffer.py index 9386f0d..f0a1b1c 100644 --- a/Agent0/executor_train/verl/verl/utils/memory_buffer.py +++ b/Agent0/executor_train/verl/verl/utils/memory_buffer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -29,14 +29,25 @@ class MemoryBuffer: memory. It must have a unique type to support this behavior. """ - def __init__(self, numel: int, numel_padded: int, dtype: torch.dtype, source: Optional[torch.Tensor] = None): + def __init__( + self, + numel: int, + numel_padded: int, + dtype: torch.dtype, + source: Optional[torch.Tensor] = None, + ): self.numel = numel self.numel_padded = numel_padded self.dtype = dtype if source is not None: self.data = source else: - self.data = torch.zeros(self.numel_padded, dtype=self.dtype, device=get_device_name(), requires_grad=False) + self.data = torch.zeros( + self.numel_padded, + dtype=self.dtype, + device=get_device_name(), + requires_grad=False, + ) def zero(self): """Reset the buffer to zero.""" @@ -69,7 +80,9 @@ def get_weight_buffer_meta_from_module(module: nn.Module) -> dict[str, dict]: return weight_buffer_meta -def build_memory_buffer(weight_buffer_meta: dict[str, dict]) -> dict[torch.dtype, MemoryBuffer]: +def build_memory_buffer( + weight_buffer_meta: dict[str, dict], +) -> dict[torch.dtype, MemoryBuffer]: """Build the memory buffer given weight_buffer_meta Args: @@ -99,14 +112,18 @@ def build_memory_buffer(weight_buffer_meta: dict[str, dict]) -> dict[torch.dtype def build_memory_reference_from_module( - module: torch.nn.Module, memory_buffers: dict[torch.dtype, MemoryBuffer], maintain_weight=True + module: torch.nn.Module, + memory_buffers: dict[torch.dtype, MemoryBuffer], + maintain_weight=True, ): start_index = {} for dtype in memory_buffers: 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]) + 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, param.dtype) if maintain_weight: @@ -114,7 +131,9 @@ def build_memory_reference_from_module( param.data = buffer -def build_memory_reference(weight_buffer_meta: dict[str, dict], memory_buffers: dict[torch.dtype, MemoryBuffer]): +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. @@ -150,7 +169,8 @@ class MemoryBufferModuleWrapper: def __init__(self, module: nn.Module): super().__init__() self.module = module - self.weight_buffer_meta = get_weight_buffer_meta_from_module(self.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) @@ -182,7 +202,8 @@ def __init__(self, transform_memory_param_fn): 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]]): + 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. @@ -202,8 +223,11 @@ def initialize_weight_buffer(self, weight_buffer_meta_pp: list[dict[str, dict]]) 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) + 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): diff --git a/Agent0/executor_train/verl/verl/utils/metric/__init__.py b/Agent0/executor_train/verl/verl/utils/metric/__init__.py index 1e19d3f..d2284e9 100644 --- a/Agent0/executor_train/verl/verl/utils/metric/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/metric/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/metric/utils.py b/Agent0/executor_train/verl/verl/utils/metric/utils.py index f9e7cd5..0b646b6 100644 --- a/Agent0/executor_train/verl/verl/utils/metric/utils.py +++ b/Agent0/executor_train/verl/verl/utils/metric/utils.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/model.py b/Agent0/executor_train/verl/verl/utils/model.py index 04cc34f..c7df28b 100644 --- a/Agent0/executor_train/verl/verl/utils/model.py +++ b/Agent0/executor_train/verl/verl/utils/model.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -64,13 +64,17 @@ def update_model_config(module_config, override_config_kwargs): setattr(module_config, key, val) -def get_huggingface_actor_config(model_name: str, override_config_kwargs=None, trust_remote_code=False) -> dict: +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)}" + 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 ) - module_config = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code) update_model_config(module_config, override_config_kwargs) return module_config @@ -93,7 +97,9 @@ def get_generation_config( return None -def create_huggingface_actor(model_name: str, override_config_kwargs=None, automodel_kwargs=None) -> nn.Module: +def create_huggingface_actor( + model_name: str, override_config_kwargs=None, automodel_kwargs=None +) -> nn.Module: """ Args: @@ -107,17 +113,23 @@ def create_huggingface_actor(model_name: str, override_config_kwargs=None, autom 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)}" - ) + 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) + 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 ) - module: nn.Module = AutoModelForCausalLM.from_config(module_config, **automodel_kwargs) return module -def create_huggingface_critic(model_name: str, override_config_kwargs=None, automodel_kwargs=None) -> nn.Module: +def create_huggingface_critic( + model_name: str, override_config_kwargs=None, automodel_kwargs=None +) -> nn.Module: """ Args: @@ -128,13 +140,16 @@ def create_huggingface_critic(model_name: str, override_config_kwargs=None, auto """ critic_module: nn.Module = create_huggingface_actor( - model_name, override_config_kwargs=override_config_kwargs, automodel_kwargs=automodel_kwargs + 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) + nn.Linear(critic_module.config.hidden_size, 1, dtype=torch_dtype), + LambdaLayer(fn=squeeze), ) return critic_module @@ -198,15 +213,21 @@ def create_random_mask( 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)) + 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) + 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 @@ -220,16 +241,22 @@ def compute_position_id_with_mask(mask): return torch.clip(torch.cumsum(mask, dim=-1) - 1, min=0, max=None) -def convert_weight_keys(state_dict: dict[str, torch.Tensor], model: PreTrainedModel): - # convert state dict keys: https://github.com/huggingface/transformers/pull/38385 +def convert_weight_keys( + state_dict: dict[str, torch.Tensor], model: PreTrainedModel): + # convert state dict keys: + # https://github.com/huggingface/transformers/pull/38385 if not hasattr(model, "_checkpoint_conversion_mapping"): return state_dict - reverse_key_mapping = {v: k for k, v in model._checkpoint_conversion_mapping.items()} + reverse_key_mapping = { + v: k for k, v in model._checkpoint_conversion_mapping.items() + } original_weights = {} for key, value in state_dict.items(): for pattern, replacement in reverse_key_mapping.items(): - replacement = replacement.lstrip("^") # strip off un-needed chars and patterns + replacement = replacement.lstrip( + "^" + ) # strip off un-needed chars and patterns replacement = re.sub(r"\(.*\)", "", replacement) key, n_replace = re.subn(pattern, replacement, key) # Early exit of the loop @@ -259,7 +286,9 @@ def check_exclude_modules(config, key: str) -> bool: return True elif key in config.exclude_modules: return True - elif any(key.endswith(f".{exclude_key}") for exclude_key in config.exclude_modules): + elif any( + key.endswith(f".{exclude_key}") for exclude_key in config.exclude_modules + ): return True return False @@ -282,7 +311,9 @@ def check_target_modules(config, key: str) -> bool: # this module is specified directly in target_modules target_module_found = True else: - target_module_found = any(key.endswith(f".{target_key}") for target_key in config.target_modules) + target_module_found = any( + key.endswith( + f".{target_key}") for target_key in config.target_modules) layer_indexes = getattr(config, "layers_to_transform", None) layers_pattern = getattr(config, "layers_pattern", None) @@ -297,7 +328,11 @@ def check_target_modules(config, key: str) -> bool: if layers_pattern is None or len(layers_pattern) == 0: layer_index = re.match(r".*\.[^.]*\.(\d+)\.", key) else: - layers_pattern = [layers_pattern] if isinstance(layers_pattern, str) else layers_pattern + layers_pattern = ( + [layers_pattern] + if isinstance(layers_pattern, str) + else layers_pattern + ) for pattern in layers_pattern: layer_index = re.match(rf".*\.{pattern}\.(\d+)\.", key) if layer_index is not None: @@ -315,13 +350,16 @@ def check_target_modules(config, key: str) -> bool: return target_module_found -def normalize_model_name(name, pp_rank, vpp_rank, transformer_config, layer_name="layers"): +def normalize_model_name( + name, pp_rank, vpp_rank, transformer_config, layer_name="layers" +): """ Transform the model name in each model_chunk in each pp stage into the name in inference engine """ from verl.utils.megatron_utils import get_transformer_layer_offset - layer_offset = get_transformer_layer_offset(pp_rank, vpp_rank, transformer_config) + layer_offset = get_transformer_layer_offset( + pp_rank, vpp_rank, transformer_config) if layer_name in name: # belong to an intermediate layer split_name = name.split(".") @@ -331,10 +369,13 @@ def normalize_model_name(name, pp_rank, vpp_rank, transformer_config, 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}" + 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) + 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 @@ -355,13 +396,24 @@ def normalize_pp_vpp_params(params, num_hidden_layers, layer_name="layers"): 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, layer_name=layer_name + name, + pp_rank, + vpp_rank, + pp_size, + vpp_size, + num_hidden_layers, + layer_name=layer_name, ) yield normalized_name, param def get_parallel_model_from_config( - config, megatron_config, pre_process=None, post_process=None, share_embeddings_and_output_weights=False, value=False + config, + megatron_config, + pre_process=None, + post_process=None, + share_embeddings_and_output_weights=False, + value=False, ): from megatron.core import ModelParallelConfig @@ -378,7 +430,9 @@ def get_parallel_model_from_config( return model -def _get_parallel_model_architecture_from_config(config: PretrainedConfig, value=False) -> type[nn.Module]: +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) @@ -398,7 +452,9 @@ def _load_hf_model(config, model_config, is_value_model, local_cache_path): from verl.models.mcore.saver import _megatron_calc_global_rank - assert hasattr(model_config, "architectures"), "architectures cannot be empty when load weight!" + 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) @@ -407,16 +463,27 @@ def _load_hf_model(config, model_config, is_value_model, local_cache_path): print(f"start download from {config.model.path}") local_model_path = copy_to_local( - src=config.model.path, cache_dir=local_cache_path, use_shm=config.model.get("use_shm", False) + src=config.model.path, + cache_dir=local_cache_path, + use_shm=config.model.get("use_shm", False), ) print("finish download") else: local_model_path = config.model.path print(f"load from local dir {local_model_path}") - src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=0, cp_rank=mpu.get_context_parallel_rank()) - cpu_init_weights = lambda: torch.device("cpu") - init_context = init_empty_weights if torch.distributed.get_rank() != src_rank else cpu_init_weights + src_rank = _megatron_calc_global_rank( + tp_rank=0, + dp_rank=0, + pp_rank=0, + cp_rank=mpu.get_context_parallel_rank()) + + def cpu_init_weights(): return torch.device("cpu") + init_context = ( + init_empty_weights + if torch.distributed.get_rank() != src_rank + else cpu_init_weights + ) with init_context(), warnings.catch_warnings(): warnings.simplefilter("ignore") # TODO: to find a better way to load mistral7b-rm lm_head @@ -429,7 +496,9 @@ def _load_hf_model(config, model_config, is_value_model, local_cache_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"][ + state_dict["model.embed_tokens.weight"] = state_dict[ + "model.embed_tokens.weight" + ][ :32000 ] # workaround, 32001 -> 32000 is_value_model = True @@ -451,7 +520,9 @@ def get_hf_model_path(config, local_cache_path="~/.cache/verl/rlhf"): from verl.utils.fs import copy_to_local local_model_path = copy_to_local( - src=config.model.path, cache_dir=local_cache_path, use_shm=config.model.get("use_shm", False) + src=config.model.path, + cache_dir=local_cache_path, + use_shm=config.model.get("use_shm", False), ) else: local_model_path = config.model.path @@ -459,7 +530,12 @@ def get_hf_model_path(config, local_cache_path="~/.cache/verl/rlhf"): def load_megatron_model_weights( - config, model_config, parallel_model, params_dtype, is_value_model=False, local_cache_path="~/.cache/verl/rlhf" + config, + model_config, + parallel_model, + params_dtype, + is_value_model=False, + local_cache_path="~/.cache/verl/rlhf", ): """Load weights for verl customized model.""" architectures, model, state_dict, is_value_model = _load_hf_model( @@ -470,7 +546,9 @@ def load_megatron_model_weights( print(f"before weight loader: architectures = {architectures}...") for arch in architectures: - print(f"call weight loader arch = {arch}, model config = {model.config}") + print( + f"call weight loader arch = {arch}, model config = { + model.config}") weight_loader = get_weight_loader(arch) weight_loader( state_dict=state_dict, @@ -484,10 +562,17 @@ def load_megatron_model_weights( def load_megatron_gptmodel_weights( - config, model_config, parallel_model, params_dtype, is_value_model=False, local_cache_path="~/.cache/verl/rlhf" + config, + model_config, + parallel_model, + params_dtype, + is_value_model=False, + local_cache_path="~/.cache/verl/rlhf", ): """Load weights for mcore GPT model.""" - _, model, state_dict, is_value_model = _load_hf_model(config, model_config, is_value_model, local_cache_path) + _, model, state_dict, is_value_model = _load_hf_model( + config, model_config, is_value_model, local_cache_path + ) from verl.models.mcore.loader import load_state_dict_to_megatron_gptmodel @@ -502,7 +587,9 @@ def load_megatron_gptmodel_weights( # pad input_ids_rmpad, cu_seqlens and max_seqlen_in_batch to be divisible by tp -def pad_packed_inputs(unpad_tokens: torch.Tensor, cu_seqlens, max_seqlen_in_batch, size): +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 @@ -527,7 +614,9 @@ def pad_packed_inputs(unpad_tokens: torch.Tensor, cu_seqlens, max_seqlen_in_batc 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") + 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) @@ -535,7 +624,10 @@ def pad_packed_inputs(unpad_tokens: torch.Tensor, cu_seqlens, max_seqlen_in_batc return unpad_tokens, cu_seqlens, max_seqlen_in_batch -def load_mcore_dist_weights(parallel_model, dist_weight_path, is_value_model=False): +def load_mcore_dist_weights( + parallel_model, + dist_weight_path, + is_value_model=False): from megatron.core import dist_checkpointing from megatron.core.dist_checkpointing.serialization import StrictHandling @@ -555,18 +647,29 @@ def load_mcore_dist_weights(parallel_model, dist_weight_path, is_value_model=Fal def get_parallel_gptmodel_from_config( - tfconfig, hf_config, pre_process=None, post_process=None, share_embeddings_and_output_weights=False, value=False + tfconfig, + hf_config, + pre_process=None, + post_process=None, + share_embeddings_and_output_weights=False, + value=False, ): from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec from megatron.core.models.gpt.gpt_model import GPTModel use_te = True assert tfconfig.normalization == "RMSNorm", "only RMSNorm is supported for now" - transformer_layer_spec = get_gpt_decoder_block_spec(tfconfig, use_transformer_engine=use_te) + transformer_layer_spec = get_gpt_decoder_block_spec( + tfconfig, use_transformer_engine=use_te + ) rope_scaling_args = {} if hf_config.rope_scaling is not None: - assert hf_config.rope_scaling["type"] == "linear", "only linear scaling is supported for now" - rope_scaling_args["seq_len_interpolation_factor"] = hf_config.rope_scaling["factor"] + assert ( + hf_config.rope_scaling["type"] == "linear" + ), "only linear scaling is supported for now" + rope_scaling_args["seq_len_interpolation_factor"] = hf_config.rope_scaling[ + "factor" + ] parallel_model = GPTModel( config=tfconfig, transformer_layer_spec=transformer_layer_spec, @@ -600,28 +703,43 @@ def tie_weights(self: "AutoModelForCausalLMWithValueHead") -> None: if isinstance(self.pretrained_model, PreTrainedModel): self.pretrained_model.tie_weights() - def get_input_embeddings(self: "AutoModelForCausalLMWithValueHead") -> torch.nn.Module: + def get_input_embeddings( + self: "AutoModelForCausalLMWithValueHead", + ) -> torch.nn.Module: if isinstance(self.pretrained_model, PreTrainedModel): return self.pretrained_model.get_input_embeddings() - def get_output_embeddings(self: "AutoModelForCausalLMWithValueHead") -> torch.nn.Module: + def get_output_embeddings( + self: "AutoModelForCausalLMWithValueHead", + ) -> torch.nn.Module: if isinstance(self.pretrained_model, PreTrainedModel): return self.pretrained_model.get_output_embeddings() def can_generate(self): return False - ignore_modules = [name for name, _ in model.named_parameters() if "pretrained_model" in name] + ignore_modules = [ + name for name, + _ in model.named_parameters() if "pretrained_model" in name] model._keys_to_ignore_on_save = ignore_modules model.tie_weights = MethodType(tie_weights, model) model.get_input_embeddings = MethodType(get_input_embeddings, model) model.get_output_embeddings = MethodType(get_output_embeddings, model) model.can_generate = MethodType(can_generate, model) - model._no_split_modules = getattr(model.pretrained_model, "_no_split_modules", []) - - -def load_valuehead_model(local_path, torch_dtype, model_config, trust_remote_code): - from transformers import AutoModelForCausalLM, AutoModelForTokenClassification, AutoModelForVision2Seq + model._no_split_modules = getattr( + model.pretrained_model, "_no_split_modules", []) + + +def load_valuehead_model( + local_path, + torch_dtype, + model_config, + trust_remote_code): + from transformers import ( + AutoModelForCausalLM, + AutoModelForTokenClassification, + AutoModelForVision2Seq, + ) try: model = AutoModelForTokenClassification.from_pretrained( diff --git a/Agent0/executor_train/verl/verl/utils/net_utils.py b/Agent0/executor_train/verl/verl/utils/net_utils.py index 138821c..494c145 100644 --- a/Agent0/executor_train/verl/verl/utils/net_utils.py +++ b/Agent0/executor_train/verl/verl/utils/net_utils.py @@ -1,4 +1,4 @@ -# Copyright 2023-2024 SGLang Team +# Copyright 2023-2026 SGLang Team # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at @@ -11,7 +11,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/profiler/__init__.py b/Agent0/executor_train/verl/verl/utils/profiler/__init__.py index 2242c24..c57439f 100644 --- a/Agent0/executor_train/verl/verl/utils/profiler/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/profiler/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,10 +19,20 @@ if is_nvtx_available(): from .nvtx_profile import NsightSystemsProfiler as DistProfiler - from .nvtx_profile import mark_annotate, mark_end_range, mark_start_range, marked_timer + from .nvtx_profile import ( + mark_annotate, + mark_end_range, + mark_start_range, + marked_timer, + ) elif is_npu_available: from .mstx_profile import NPUProfiler as DistProfiler - from .mstx_profile import mark_annotate, mark_end_range, mark_start_range, marked_timer + from .mstx_profile import ( + mark_annotate, + mark_end_range, + mark_start_range, + marked_timer, + ) else: from .performance import marked_timer from .profile import DistProfiler, mark_annotate, mark_end_range, mark_start_range diff --git a/Agent0/executor_train/verl/verl/utils/profiler/config.py b/Agent0/executor_train/verl/verl/utils/profiler/config.py index 8acf075..908d9fd 100644 --- a/Agent0/executor_train/verl/verl/utils/profiler/config.py +++ b/Agent0/executor_train/verl/verl/utils/profiler/config.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -52,6 +52,6 @@ def intersect(self, other: "ProfilerConfig") -> "ProfilerConfig": def __post_init__(self) -> None: """config validation logics go here""" - assert isinstance(self.ranks, set | list | tuple), ( - f"Profiler ranks must be of type list, got {type(self.ranks)}" - ) + assert isinstance( + self.ranks, set | list | tuple + ), f"Profiler ranks must be of type list, got {type(self.ranks)}" diff --git a/Agent0/executor_train/verl/verl/utils/profiler/empty_annotations.py b/Agent0/executor_train/verl/verl/utils/profiler/empty_annotations.py index ed18dd3..12eb04e 100644 --- a/Agent0/executor_train/verl/verl/utils/profiler/empty_annotations.py +++ b/Agent0/executor_train/verl/verl/utils/profiler/empty_annotations.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/profiler/mstx_profile.py b/Agent0/executor_train/verl/verl/utils/profiler/mstx_profile.py index c5c35ce..6145da8 100644 --- a/Agent0/executor_train/verl/verl/utils/profiler/mstx_profile.py +++ b/Agent0/executor_train/verl/verl/utils/profiler/mstx_profile.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Inspired from https://gitee.com/ascend/MindSpeed-RL/blob/master/mindspeed_rl/utils/utils.py +# Inspired from +# https://gitee.com/ascend/MindSpeed-RL/blob/master/mindspeed_rl/utils/utils.py import functools import os from contextlib import contextmanager @@ -81,7 +82,10 @@ def marked_timer(name: str, timing_raw: dict[str, float], **kwargs): mark_end_range(mark_range) -def get_npu_profiler(option: DictConfig, role: Optional[str] = None, profile_step: Optional[str] = None): +def get_npu_profiler( + option: DictConfig, + role: Optional[str] = None, + profile_step: Optional[str] = None): """Generate and return an NPU profiler object. Args: @@ -101,7 +105,9 @@ def get_npu_profiler(option: DictConfig, role: Optional[str] = None, profile_ste elif option.level == "level2": profile_level = torch_npu.profiler.ProfilerLevel.Level2 else: - raise ValueError(f"level only supports level0, 1, 2, and level_none, but gets {option.level}") + raise ValueError( + f"level only supports level0, 1, 2, and level_none, but gets { + option.level}") profile_save_path = option.save_path if profile_step: @@ -129,7 +135,9 @@ def get_npu_profiler(option: DictConfig, role: Optional[str] = None, profile_ste record_shapes=option.record_shapes, profile_memory=option.with_memory, activities=activites, - on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profile_save_path, analyse_flag=option.analysis), + on_trace_ready=torch_npu.profiler.tensorboard_trace_handler( + profile_save_path, analyse_flag=option.analysis + ), experimental_config=experimental_config, ) return prof @@ -162,12 +170,15 @@ def __init__(self, rank: int, config: ProfilerConfig, **kwargs): self.this_rank = rank in config.ranks def start(self, **kwargs): - role, profile_step = kwargs.get("role", None), kwargs.get("profile_step", None) + role, profile_step = kwargs.get( + "role", None), kwargs.get( + "profile_step", None) profile_step = str(profile_step) if profile_step is not None else None if self.this_rank and self.profile_option is not None: self.this_step = True if not self.discrete and NPUProfiler._define_count == 0: - self.profile_npu = get_npu_profiler(option=self.profile_option, role=role, profile_step=profile_step) + self.profile_npu = get_npu_profiler( + option=self.profile_option, role=role, profile_step=profile_step) self.profile_npu.start() NPUProfiler._define_count += 1 @@ -180,7 +191,9 @@ def stop(self): NPUProfiler._define_count -= 1 @staticmethod - def annotate(message: Optional[str] = None, role: Optional[str] = None, **kwargs) -> Callable: + def annotate( + message: Optional[str] = None, role: Optional[str] = None, **kwargs + ) -> Callable: """Decorate a Worker member function to profile the current rank in the current training step. Requires the target function to be a member function of a Worker, @@ -200,7 +213,9 @@ def wrapper(self, *args, **kwargs): if self.profiler.this_step and self.profile_option is not None: if self.profiler.discrete: - profile_npu = get_npu_profiler(option=self.profile_option, role=role) + profile_npu = get_npu_profiler( + option=self.profile_option, role=role + ) profile_npu.start() mark_range = mark_start_range(message=profile_name) diff --git a/Agent0/executor_train/verl/verl/utils/profiler/nvtx_profile.py b/Agent0/executor_train/verl/verl/utils/profiler/nvtx_profile.py index 9ebce37..25b4378 100644 --- a/Agent0/executor_train/verl/verl/utils/profiler/nvtx_profile.py +++ b/Agent0/executor_train/verl/verl/utils/profiler/nvtx_profile.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); @@ -41,7 +41,9 @@ def mark_start_range( category (str, optional): The category of the range. Defaults to None. """ - return nvtx.start_range(message=message, color=color, domain=domain, category=category) + return nvtx.start_range( + message=message, color=color, domain=domain, category=category + ) def mark_end_range(range_id: str) -> None: @@ -75,7 +77,9 @@ def mark_annotate( def decorator(func): profile_message = message or func.__name__ - return nvtx.annotate(profile_message, color=color, domain=domain, category=category)(func) + return nvtx.annotate( + profile_message, color=color, domain=domain, category=category + )(func) return decorator @@ -103,7 +107,9 @@ def marked_timer( Yields: None: This is a context manager that yields control back to the code block. """ - mark_range = mark_start_range(message=name, color=color, domain=domain, category=category) + mark_range = mark_start_range( + message=name, color=color, domain=domain, category=category + ) from .performance import _timer yield from _timer(name, timing_raw) @@ -120,7 +126,8 @@ def __init__(self, rank: int, config: Optional[ProfilerConfig], **kwargs): rank (int): The rank of the current process. config (Optional[ProfilerConfig]): Configuration for the profiler. If None, a default configuration is used. """ - # If no configuration is provided, create a default ProfilerConfig with an empty list of ranks + # If no configuration is provided, create a default ProfilerConfig with + # an empty list of ranks if not config: config = ProfilerConfig(ranks=[]) self.this_step: bool = False @@ -175,7 +182,12 @@ def wrapper(self, *args, **kwargs): if self.profiler.this_step: if self.profiler.discrete: torch.cuda.profiler.start() - mark_range = mark_start_range(message=profile_name, color=color, domain=domain, category=category) + mark_range = mark_start_range( + message=profile_name, + color=color, + domain=domain, + category=category, + ) result = func(self, *args, **kwargs) diff --git a/Agent0/executor_train/verl/verl/utils/profiler/performance.py b/Agent0/executor_train/verl/verl/utils/profiler/performance.py index 8991896..047c42c 100644 --- a/Agent0/executor_train/verl/verl/utils/profiler/performance.py +++ b/Agent0/executor_train/verl/verl/utils/profiler/performance.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -34,7 +34,8 @@ def _get_current_mem_info(unit: str = "GB", precision: int = 2) -> tuple[str]: mem_reserved = get_torch_device().memory_reserved() # use get_torch_device().mem_get_info to profile device memory # since vllm's sleep mode works below pytorch - # see https://github.com/vllm-project/vllm/pull/11743#issuecomment-2754338119 + # see + # https://github.com/vllm-project/vllm/pull/11743#issuecomment-2754338119 mem_free, mem_total = get_torch_device().mem_get_info() mem_used = mem_total - mem_free mem_allocated = f"{mem_allocated / divisor:.{precision}f}" @@ -44,7 +45,11 @@ def _get_current_mem_info(unit: str = "GB", precision: int = 2) -> tuple[str]: return mem_allocated, mem_reserved, mem_used, mem_total -def log_gpu_memory_usage(head: str, logger: logging.Logger = None, level=logging.DEBUG, rank: int = 0): +def log_gpu_memory_usage( + head: str, + logger: logging.Logger = None, + level=logging.DEBUG, + rank: int = 0): """Log GPU memory usage information. Args: @@ -53,7 +58,8 @@ def log_gpu_memory_usage(head: str, logger: logging.Logger = None, level=logging level: Logging level to use. Defaults to logging.DEBUG. rank (int): The rank of the process to log memory for. Defaults to 0. """ - if (not dist.is_initialized()) or (rank is None) or (dist.get_rank() == rank): + if (not dist.is_initialized()) or ( + rank is None) or (dist.get_rank() == rank): mem_allocated, mem_reserved, mem_used, mem_total = _get_current_mem_info() message = ( f"{head}, memory allocated (GB): {mem_allocated}, memory reserved (GB): {mem_reserved}, " @@ -77,7 +83,13 @@ class GPUMemoryLogger(DecoratorLoggerBase): ... return """ - def __init__(self, role: str, logger: logging.Logger = None, level=logging.DEBUG, log_only_rank_0: bool = True): + def __init__( + self, + role: str, + logger: logging.Logger = None, + level=logging.DEBUG, + log_only_rank_0: bool = True, + ): if dist.is_initialized() and dist.get_world_size() > 1: rank = dist.get_rank() else: @@ -198,8 +210,13 @@ def reduce_timing(timing_raw: dict[str, float]) -> dict[str, float]: for key in sorted(timing_raw.keys()): key_list.append(key) timing_list.append(timing_raw[key]) - timing_list = torch.tensor(timing_list, dtype=torch.float32, device=get_device_id()) - torch.distributed.all_reduce(timing_list, op=torch.distributed.ReduceOp.AVG) + timing_list = torch.tensor( + timing_list, + dtype=torch.float32, + device=get_device_id()) + torch.distributed.all_reduce( + timing_list, op=torch.distributed.ReduceOp.AVG) timing_list = [tensor.item() for tensor in timing_list.to("cpu")] - timing_generate = {key_list[i]: timing_list[i] for i in range(len(key_list))} + timing_generate = {key_list[i]: timing_list[i] + for i in range(len(key_list))} return timing_generate diff --git a/Agent0/executor_train/verl/verl/utils/profiler/profile.py b/Agent0/executor_train/verl/verl/utils/profiler/profile.py index 4e7ce4f..9d4913a 100644 --- a/Agent0/executor_train/verl/verl/utils/profiler/profile.py +++ b/Agent0/executor_train/verl/verl/utils/profiler/profile.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -39,7 +39,8 @@ class Profiler: """ def __init__(self, config): - # note : if we do not set use_profile, it will be set as None, so that all function will be skip + # note : if we do not set use_profile, it will be set as None, so that + # all function will be skip self.config = config self.skip_prof = False self.saved = False @@ -70,11 +71,15 @@ def _validate(self): if self.config.profile_ranks is None: print("[WARNING] Profile ranks is not set, default to rank 0") self.config.profile_ranks = [0] - assert self.config.step_start >= 0, "[ERROR] Profile step start must be greater than 0" - assert self.config.step_end >= 0, "[ERROR] Profile step end must be greater than 0" - assert self.config.step_start < self.config.step_end, ( - "[ERROR] Profile step start must be less than step end" - ) + assert ( + self.config.step_start >= 0 + ), "[ERROR] Profile step start must be greater than 0" + assert ( + self.config.step_end >= 0 + ), "[ERROR] Profile step end must be greater than 0" + assert ( + self.config.step_start < self.config.step_end + ), "[ERROR] Profile step start must be less than step end" def check(self): return self.prof is not None and not self.skip_prof @@ -97,9 +102,16 @@ def save(self): if self.prof is not None and not self.saved: if not os.path.exists(self.config.save_path): os.makedirs(self.config.save_path) - save_file_name = f"/prof_start_{self.config.step_start}_end_{self.config.step_end}_rank_{self.rank}.json" - print(f"[Profiler] Saving trace to {self.config.save_path + save_file_name}") - self.prof.export_chrome_trace(self.config.save_path + save_file_name) + save_file_name = f"/prof_start_{ + self.config.step_start}_end_{ + self.config.step_end}_rank_{ + self.rank}.json" + print( + f"[Profiler] Saving trace to { + self.config.save_path + + save_file_name}") + self.prof.export_chrome_trace( + self.config.save_path + save_file_name) self.skip_prof = True self.saved = True @@ -176,7 +188,11 @@ class DistProfiler: config (ProfilerConfig, optional): Configuration for the profiler. """ - def __init__(self, rank: int, config: Optional[ProfilerConfig] = None, **kwargs): + def __init__( + self, + rank: int, + config: Optional[ProfilerConfig] = None, + **kwargs): pass def start(self, **kwargs): diff --git a/Agent0/executor_train/verl/verl/utils/py_functional.py b/Agent0/executor_train/verl/verl/utils/py_functional.py index 1ea02ef..ff3ef41 100644 --- a/Agent0/executor_train/verl/verl/utils/py_functional.py +++ b/Agent0/executor_train/verl/verl/utils/py_functional.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,7 +27,12 @@ # --- Top-level helper for multiprocessing timeout --- # This function MUST be defined at the top level to be pickleable -def _mp_target_wrapper(target_func: Callable, mp_queue: multiprocessing.Queue, args: tuple, kwargs: dict[str, Any]): +def _mp_target_wrapper( + target_func: Callable, + mp_queue: multiprocessing.Queue, + args: tuple, + kwargs: dict[str, Any], +): """ Internal wrapper function executed in the child process. Calls the original target function and puts the result or exception into the queue. @@ -44,7 +49,10 @@ def _mp_target_wrapper(target_func: Callable, mp_queue: multiprocessing.Queue, a mp_queue.put((False, e)) # Indicate failure and put exception except (pickle.PicklingError, TypeError): # Fallback if the original exception cannot be pickled - mp_queue.put((False, RuntimeError(f"Original exception type {type(e).__name__} not pickleable: {e}"))) + mp_queue.put( + (False, RuntimeError( + f"Original exception type { + type(e).__name__} not pickleable: {e}"), )) # Renamed the function from timeout to timeout_limit @@ -75,18 +83,21 @@ def decorator(func): print( "WARN: The 'use_signals=True' option in the timeout decorator is deprecated. \ Signals are unreliable outside the main thread. \ - Please use the default multiprocessing-based timeout (use_signals=False)." - ) + Please use the default multiprocessing-based timeout (use_signals=False).") @wraps(func) def wrapper_signal(*args, **kwargs): def handler(signum, frame): - # Update function name in error message if needed (optional but good practice) - raise TimeoutError(f"Function {func.__name__} timed out after {seconds} seconds (signal)!") + # Update function name in error message if needed (optional + # but good practice) + raise TimeoutError( + f"Function { + func.__name__} timed out after {seconds} seconds (signal)!") old_handler = signal.getsignal(signal.SIGALRM) signal.signal(signal.SIGALRM, handler) - # Use setitimer for float seconds support, alarm only supports integers + # Use setitimer for float seconds support, alarm only supports + # integers signal.setitimer(signal.ITIMER_REAL, seconds) try: @@ -103,7 +114,9 @@ def handler(signum, frame): @wraps(func) def wrapper_mp(*args, **kwargs): q = multiprocessing.Queue(maxsize=1) - process = multiprocessing.Process(target=_mp_target_wrapper, args=(func, q, args, kwargs)) + process = multiprocessing.Process( + target=_mp_target_wrapper, args=(func, q, args, kwargs) + ) process.start() process.join(timeout=seconds) @@ -111,12 +124,19 @@ def wrapper_mp(*args, **kwargs): process.terminate() process.join(timeout=0.5) # Give it a moment to terminate if process.is_alive(): - print(f"Warning: Process {process.pid} did not terminate gracefully after timeout.") - # Update function name in error message if needed (optional but good practice) - raise TimeoutError(f"Function {func.__name__} timed out after {seconds} seconds (multiprocessing)!") + print( + f"Warning: Process { + process.pid} did not terminate gracefully after timeout.") + # Update function name in error message if needed (optional + # but good practice) + raise TimeoutError( + f"Function { + func.__name__} timed out after {seconds} seconds (multiprocessing)!") try: - success, result_or_exc = q.get(timeout=0.1) # Small timeout for queue read + success, result_or_exc = q.get( + timeout=0.1 + ) # Small timeout for queue read if success: return result_or_exc else: @@ -129,7 +149,8 @@ def wrapper_mp(*args, **kwargs): ) from err else: # Should have timed out if queue is empty after join unless process died unexpectedly - # Update function name in error message if needed (optional but good practice) + # Update function name in error message if needed + # (optional but good practice) raise TimeoutError( f"Operation timed out or process finished unexpectedly without result " f"(exitcode: {exitcode})." @@ -155,7 +176,9 @@ def union_two_dict(dict1: dict, dict2: dict): """ 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" + assert ( + dict2[key] == dict1[key] + ), f"{key} in meta_dict1 and meta_dict2 are not the same object" dict1[key] = val return dict1 @@ -277,7 +300,11 @@ def convert_to_regular_types(obj): from omegaconf import DictConfig, ListConfig 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) + 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): diff --git a/Agent0/executor_train/verl/verl/utils/ray_utils.py b/Agent0/executor_train/verl/verl/utils/ray_utils.py index a738c0f..2fed5ec 100644 --- a/Agent0/executor_train/verl/verl/utils/ray_utils.py +++ b/Agent0/executor_train/verl/verl/utils/ray_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -42,7 +42,8 @@ def ray_noset_visible_devices(env_vars=os.environ): "RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS", "RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR", ] - return any(env_vars.get(env_var) for env_var in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST) + return any(env_vars.get(env_var) + for env_var in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST) def parallel_put(data_list: list[Any], max_workers: Optional[int] = None): @@ -67,7 +68,12 @@ def put_data(index, data): 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)] + 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()) diff --git a/Agent0/executor_train/verl/verl/utils/rendezvous/__init__.py b/Agent0/executor_train/verl/verl/utils/rendezvous/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/utils/rendezvous/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/rendezvous/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/rendezvous/ray_backend.py b/Agent0/executor_train/verl/verl/utils/rendezvous/ray_backend.py index d991181..55a641e 100644 --- a/Agent0/executor_train/verl/verl/utils/rendezvous/ray_backend.py +++ b/Agent0/executor_train/verl/verl/utils/rendezvous/ray_backend.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,19 +31,27 @@ def get(self): 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] + 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("multiple actors with same name found: %s", matched_actors) + logging.warning( + "multiple actors with same name found: %s", + matched_actors) elif len(matched_actors) == 0: logging.info("failed to get any actor named %s", name) return None def create_nccl_communicator_in_ray( - rank: int, world_size: int, group_name: str, max_retries: int = 100, interval_s: int = 5 + rank: int, + world_size: int, + group_name: str, + max_retries: int = 100, + interval_s: int = 5, ): if rank == 0: nccl_id = get_unique_id() @@ -69,5 +77,9 @@ def create_nccl_communicator_in_ray( rank=rank, ) return communicator - logging.info("failed to get nccl_id for %d time, sleep for %d seconds", i + 1, interval_s) + logging.info( + "failed to get nccl_id for %d time, sleep for %d seconds", + i + 1, + interval_s, + ) time.sleep(interval_s) diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/__init__.py b/Agent0/executor_train/verl/verl/utils/reward_score/__init__.py index b298d41..84bf6d1 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -51,7 +51,8 @@ def default_compute_score( # [Optional] Math-Verify Integration # For enhanced accuracy, consider utilizing Math-Verify (https://github.com/huggingface/Math-Verify). # Note: Math-Verify needs to be manually installed via pip: `pip install math-verify`. - # To use it, override the `compute_score` function with the following implementation: + # To use it, override the `compute_score` function with the following + # implementation: # from . import math_verify # res = math_verify.compute_score(solution_str, ground_truth) @@ -75,16 +76,24 @@ def default_compute_score( if sandbox_fusion_url: from . import sandbox_fusion - # Pass the URL directly, ground_truth likely contains test cases here + # Pass the URL directly, ground_truth likely contains test cases + # here res = sandbox_fusion.compute_score( - sandbox_fusion_url, concurrent_semaphore, memory_limit_mb, solution_str, ground_truth, continuous=True + sandbox_fusion_url, + concurrent_semaphore, + memory_limit_mb, + solution_str, + ground_truth, + continuous=True, ) else: - # If no sandbox URL is provided, fall back to prime_code or raise error + # If no sandbox URL is provided, fall back to prime_code or raise + # error from . import prime_code # Assuming prime_code doesn't need the URL - res = prime_code.compute_score(solution_str, ground_truth, continuous=True) + res = prime_code.compute_score( + solution_str, ground_truth, continuous=True) elif data_source in ["hiyouga/geometry3k"]: from . import geo3k @@ -103,7 +112,9 @@ def default_compute_score( res = search_r1_like_qa_em.compute_score(solution_str, ground_truth) else: - raise NotImplementedError(f"Reward function is not implemented for {data_source=}") + raise NotImplementedError( + f"Reward function is not implemented for {data_source=}" + ) if isinstance(res, dict): return res @@ -127,7 +138,13 @@ def _default_compute_score( Legacy function API to be deprecated. Please use `default_compute_score` instead. """ return default_compute_score( - data_source, solution_str, ground_truth, extra_info, sandbox_fusion_url, concurrent_semaphore, memory_limit_mb + data_source, + solution_str, + ground_truth, + extra_info, + sandbox_fusion_url, + concurrent_semaphore, + memory_limit_mb, ) diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/geo3k.py b/Agent0/executor_train/verl/verl/utils/reward_score/geo3k.py index 8a85087..19b9e12 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/geo3k.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/geo3k.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -22,7 +22,10 @@ def format_reward(predict_str: str) -> float: return 1.0 if match_result else 0.0 -def acc_reward(predict_str: str, ground_truth: str, use_boxed: bool = True) -> float: +def acc_reward( + predict_str: str, + ground_truth: str, + use_boxed: bool = True) -> float: if use_boxed: answer = extract_boxed_content(predict_str) else: @@ -30,7 +33,12 @@ def acc_reward(predict_str: str, ground_truth: str, use_boxed: bool = True) -> f return 1.0 if grade_answer(answer, ground_truth) else 0.0 -def compute_score(predict_str: str, ground_truth: str, use_boxed: bool = True, format_score: float = 0.1) -> float: - return (1.0 - format_score) * acc_reward(predict_str, ground_truth, use_boxed) + format_score * format_reward( - predict_str - ) +def compute_score( + predict_str: str, + ground_truth: str, + use_boxed: bool = True, + format_score: float = 0.1, +) -> float: + return (1.0 - format_score) * acc_reward( + predict_str, ground_truth, use_boxed + ) + format_score * format_reward(predict_str) diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/gsm8k.py b/Agent0/executor_train/verl/verl/utils/reward_score/gsm8k.py index c2afafc..3399af3 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/gsm8k.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/gsm8k.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -41,7 +41,9 @@ def extract_solution(solution_str, method="strict"): return final_answer -def compute_score(solution_str, ground_truth, method="strict", format_score=0.0, score=1.0): +def compute_score( + solution_str, ground_truth, method="strict", format_score=0.0, score=1.0 +): """The scoring function for GSM8k. Reference: Trung, Luong, et al. "Reft: Reasoning with reinforced fine-tuning." Proceedings of the 62nd Annual diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/math.py b/Agent0/executor_train/verl/verl/utils/reward_score/math.py index 3fff7bc..83f52bf 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/math.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/math.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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. @@ -11,7 +11,8 @@ # WITHOUT 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 +# Adapted from +# https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/hendrycks_math/utils.py def compute_score(solution_str, ground_truth) -> float: @@ -28,7 +29,8 @@ def compute_score(solution_str, ground_truth) -> float: return retval -# string normalization from https://github.com/EleutherAI/lm-evaluation-harness/blob/master/lm_eval/tasks/hendrycks_math.py +# string normalization from +# https://github.com/EleutherAI/lm-evaluation-harness/blob/master/lm_eval/tasks/hendrycks_math.py def is_equiv(str1, str2, verbose=False): if str1 is None and str2 is None: print("WARNING: Both None") @@ -50,14 +52,14 @@ def remove_boxed(s): if "\\boxed " in s: left = "\\boxed " assert s[: len(left)] == left - return s[len(left) :] + return s[len(left):] left = "\\boxed{" assert s[: len(left)] == left assert s[-1] == "}" - return s[len(left) : -1] + return s[len(left): -1] def last_boxed_only_string(string): @@ -82,7 +84,7 @@ def last_boxed_only_string(string): break i += 1 - retval = None if right_brace_idx is None else string[idx : right_brace_idx + 1] + retval = None if right_brace_idx is None else string[idx: right_brace_idx + 1] return retval @@ -218,7 +220,8 @@ def strip_string(string): 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 + # 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 diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/math_batch.py b/Agent0/executor_train/verl/verl/utils/reward_score/math_batch.py index ed08086..a1fcd98 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/math_batch.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/math_batch.py @@ -1,4 +1,4 @@ -# Copyright 2025 Individual Contributor: Mert Unsal +# Copyright 2025-2026 Individual Contributor: Mert Unsal # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -15,12 +15,20 @@ from .math import compute_score -def compute_score_batched(data_sources, solution_strs, ground_truths, extra_infos): +def compute_score_batched( + data_sources, + solution_strs, + ground_truths, + extra_infos): """ This is a demonstration of how the batched reward function should look like. Typically, you want to use batched reward to speed up the process with parallelization """ return [ - compute_score(solution_str, ground_truth) - for solution_str, ground_truth in zip(solution_strs, ground_truths, strict=True) - ] + compute_score( + solution_str, + ground_truth) for solution_str, + ground_truth in zip( + solution_strs, + ground_truths, + strict=True)] diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/math_dapo.py b/Agent0/executor_train/verl/verl/utils/reward_score/math_dapo.py index 940500f..d3831d6 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/math_dapo.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/math_dapo.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 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. @@ -11,7 +11,8 @@ # WITHOUT 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 +# Adapted from +# https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/hendrycks_math/utils.py import re from typing import Optional @@ -44,7 +45,8 @@ def last_boxed_only_string(string: str) -> Optional[str]: break i += 1 - return string[idx : right_brace_idx + 1] if right_brace_idx is not None else None + return string[idx: right_brace_idx + + 1] if right_brace_idx is not None else None def remove_boxed(s: str) -> str: @@ -59,7 +61,7 @@ def remove_boxed(s: str) -> str: left = "\\boxed{" assert s[: len(left)] == left, f"box error: {s}" assert s[-1] == "}", f"box error: {s}" - return s[len(left) : -1] + return s[len(left): -1] # Constants for normalization @@ -163,7 +165,10 @@ def normalize_final_answer(final_answer: str) -> str: def is_correct_minerva( - solution_str: str, gt: str, gt_need_extract: bool = False, answer_pattern: str = r"(?i)Answer\s*:\s*([^\n]+)" + solution_str: str, + gt: str, + gt_need_extract: bool = False, + answer_pattern: str = r"(?i)Answer\s*:\s*([^\n]+)", ) -> tuple[bool, str]: """Check if the solution is correct according to Minerva criteria. @@ -206,19 +211,23 @@ def is_correct_strict_box( # Extract the relevant part of the prediction if pause_tokens_index is not None: assert len(pause_tokens_index) == 4 - pred = pred[pause_tokens_index[-1] - 100 :] + pred = pred[pause_tokens_index[-1] - 100:] else: pred = pred[-100:] # Extract and check the boxed answer boxed_pred = last_boxed_only_string(pred) - extracted_pred = remove_boxed(boxed_pred) if boxed_pred is not None else None + extracted_pred = remove_boxed( + boxed_pred) if boxed_pred is not None else None return 1 if (extracted_pred == gt) else -1, extracted_pred def verify( - solution_str: str, answer: str, strict_box_verify: bool = False, pause_tokens_index: Optional[list[int]] = None + solution_str: str, + answer: str, + strict_box_verify: bool = False, + pause_tokens_index: Optional[list[int]] = None, ) -> bool: """Verify if the solution is correct. @@ -232,7 +241,8 @@ def verify( True if the solution is correct, False otherwise """ if strict_box_verify: - correct, pred = is_correct_strict_box(solution_str, answer, pause_tokens_index) + correct, pred = is_correct_strict_box( + solution_str, answer, pause_tokens_index) return correct == 1, pred correct, pred = is_correct_minerva(solution_str, answer) @@ -257,10 +267,14 @@ def compute_score( Reward score (1.0 for correct, -1.0 for incorrect) """ # Limit solution length for efficiency - solution_str = solution_str[-300:] # The longest answer in MATH-500 has 159 characters + solution_str = solution_str[ + -300: + ] # The longest answer in MATH-500 has 159 characters # Verify the solution - correct, pred = verify(solution_str, ground_truth, strict_box_verify, pause_tokens_index) + correct, pred = verify( + solution_str, ground_truth, strict_box_verify, pause_tokens_index + ) reward = 1.0 if correct else -1.0 acc = correct diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/math_verify.py b/Agent0/executor_train/verl/verl/utils/reward_score/math_verify.py index c1ce7c1..0460041 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/math_verify.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/math_verify.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -17,14 +17,18 @@ from math_verify.metric import math_metric from math_verify.parser import ExprExtractionConfig, LatexExtractionConfig except ImportError: - print("To use Math-Verify, please install it first by running `pip install math-verify`.") + print( + "To use Math-Verify, please install it first by running `pip install math-verify`." + ) -def compute_score(model_output: str, ground_truth: str, timeout_score: float = 0) -> bool: +def compute_score( + model_output: str, ground_truth: str, timeout_score: float = 0 +) -> bool: verify_func = math_metric( - gold_extraction_target=(LatexExtractionConfig(),), - pred_extraction_target=(ExprExtractionConfig(), LatexExtractionConfig()), - ) + gold_extraction_target=( + LatexExtractionConfig(),), pred_extraction_target=( + ExprExtractionConfig(), LatexExtractionConfig()), ) ret_score = 0.0 # Wrap the ground truth in \boxed{} format for verification diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/__init__.py b/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/__init__.py index 214f99b..b536a4a 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/__init__.py @@ -19,7 +19,8 @@ def compute_score(completion, test_cases, continuous=False): - # try to get code solution from completion. if the completion is pure code, this will not take effect. + # try to get code solution from completion. if the completion is pure + # code, this will not take effect. solution = completion.split("```python")[-1].split("```")[0] try: try: @@ -28,9 +29,12 @@ def compute_score(completion, test_cases, continuous=False): except Exception as e: print(f"Error:{e}") - # Complete check on all in-out pairs first. If there is no failure, per-sample test can be skipped. + # Complete check on all in-out pairs first. If there is no failure, + # per-sample test can be skipped. try: - res, metadata = apps_check_correctness(in_outs=test_cases, generation=solution, timeout=5, debug=False) + res, metadata = apps_check_correctness( + in_outs=test_cases, generation=solution, timeout=5, debug=False + ) metadata = dict(enumerate(metadata))[0] success = all(map(lambda x: x is True, res)) if success: @@ -42,17 +46,22 @@ def compute_score(completion, test_cases, continuous=False): inputs = test_cases["inputs"] outputs = test_cases["outputs"] for i in range(len(inputs)): - test_cases_list.append({"inputs": [inputs[i]], "outputs": [outputs[i]]}) + test_cases_list.append( + {"inputs": [inputs[i]], "outputs": [outputs[i]]}) if continuous: # per sample test: if continuous score is needed, test first 10 samples regardless of failures - # do not test all samples cuz some problems have enormous test cases + # do not test all samples cuz some problems have enormous test + # cases metadata_list = [] res_list = [] for test_case_id, test_case in enumerate(test_cases_list): - res, metadata = apps_check_correctness(in_outs=test_case, generation=solution, timeout=10, debug=False) + res, metadata = apps_check_correctness( + in_outs=test_case, generation=solution, timeout=10, debug=False) try: - metadata = dict(enumerate(metadata))[0] # metadata can be empty occasionally + metadata = dict(enumerate(metadata))[ + 0 + ] # metadata can be empty occasionally except Exception: metadata = {} metadata["test_case"] = {} diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/testing_util.py b/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/testing_util.py index 2f22325..a6e1013 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/testing_util.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/testing_util.py @@ -41,7 +41,7 @@ def truncatefn(s, length=300): if len(s) <= length: return s - return s[: length // 2] + "...(truncated) ..." + s[-length // 2 :] + return s[: length // 2] + "...(truncated) ..." + s[-length // 2:] class CODE_TYPE(Enum): @@ -81,7 +81,9 @@ def combined_int_check(val): def clean_traceback(error_traceback): file_start = error_traceback.find('File ""') # print(file_start) - error_traceback = "Traceback (most recent call last):\n " + error_traceback[file_start:] + error_traceback = ( + "Traceback (most recent call last):\n " + error_traceback[file_start:] + ) return error_traceback @@ -146,8 +148,13 @@ def run_test(in_outs, test=None, debug=False, timeout=15): last_block = astree.body[-1] if isinstance(last_block, ast.If): condition = last_block.test - if ast.unparse(condition).strip() == "__name__ == '__main__'": - test = ast.unparse(astree.body[:-1]) + "\n" + ast.unparse(last_block.body) + if ast.unparse(condition).strip( + ) == "__name__ == '__main__'": + test = ( + ast.unparse(astree.body[:-1]) + + "\n" + + ast.unparse(last_block.body) + ) except Exception: pass @@ -155,7 +162,8 @@ def run_test(in_outs, test=None, debug=False, timeout=15): new_test = [] for x in tmp_test: - if (not x.startswith("from ")) and (not x.startswith("import ")): + if (not x.startswith("from ")) and ( + not x.startswith("import ")): new_test.append("\t" + x + "\n") else: new_test.append(x + "\n") @@ -201,7 +209,8 @@ def run_test(in_outs, test=None, debug=False, timeout=15): print(f"get method = {datetime.now().time()}") try: - method = getattr(tmp, method_name) # get_attr second arg must be str + # get_attr second arg must be str + method = getattr(tmp, method_name) except Exception: signal.alarm(0) error_traceback = traceback.format_exc() @@ -220,17 +229,22 @@ def run_test(in_outs, test=None, debug=False, timeout=15): raw_outputs = in_outs["outputs"][index] if which_type == CODE_TYPE.call_based: inputs = [json.loads(line) for line in inputs.split("\n")] - in_outs["outputs"][index] = json.loads(in_outs["outputs"][index]) + in_outs["outputs"][index] = json.loads( + in_outs["outputs"][index]) truncate_line_size = 300 // (raw_inputs.count("\n") + 1) raw_inputs = "\n".join( - [truncatefn(line, truncate_line_size) for line in raw_inputs.strip().split("\n")] + [ + truncatefn(line, truncate_line_size) + for line in raw_inputs.strip().split("\n") + ] ) raw_outputs = truncatefn(raw_outputs, 200) else: raw_inputs = truncatefn(raw_inputs) raw_outputs = truncatefn(raw_outputs, 200) - # JSON forces dictionaries to have string keys; this undoes this (assuming a singleton list) + # JSON forces dictionaries to have string keys; this undoes this + # (assuming a singleton list) try: if isinstance(inputs[0], dict): inputs = [{int(k): v for k, v in inputs[0].items()}] @@ -238,20 +252,24 @@ def run_test(in_outs, test=None, debug=False, timeout=15): pass try: if isinstance(in_outs["outputs"][index], dict): - in_outs["outputs"][index] = [{int(k): v for k, v in in_outs["outputs"][index].items()}] + in_outs["outputs"][index] = [ + {int(k): v for k, v in in_outs["outputs"][index].items()} + ] except Exception: pass try: if isinstance(in_outs["outputs"][index][0], dict): - in_outs["outputs"][index] = [{int(k): v for k, v in in_outs["outputs"][index][0].items()}] + in_outs["outputs"][index] = [ + {int(k): v for k, v in in_outs["outputs"][index][0].items()} + ] except Exception: pass if debug: print( - f"time: {datetime.now().time()} testing index = {index} inputs = {inputs}, {type(inputs)}. " - f"type = {which_type}" - ) + f"time: { + datetime.now().time()} testing index = {index} inputs = {inputs}, { + type(inputs)}. " f"type = {which_type}") if which_type == CODE_TYPE.call_based: # Call-based signal.alarm(timeout) faulthandler.enable() @@ -260,20 +278,29 @@ def run_test(in_outs, test=None, debug=False, timeout=15): raw_true_output = output raw_true_output_copy = json.dumps(output) - raw_true_output_copy = truncatefn(raw_true_output_copy, 200) + raw_true_output_copy = truncatefn( + raw_true_output_copy, 200) # ground truth sequences are not tuples if isinstance(output, tuple): output = list(output) tmp_result = output == in_outs["outputs"][index] - if isinstance(in_outs["outputs"][index], list) and in_outs["outputs"][index]: - tmp_result = tmp_result or (output == in_outs["outputs"][index][0]) + if ( + isinstance(in_outs["outputs"][index], list) + and in_outs["outputs"][index] + ): + tmp_result = tmp_result or ( + output == in_outs["outputs"][index][0] + ) # ground truth sequences are not tuples try: if isinstance(output[0], tuple): - tmp_result = tmp_result or ([list(x) for x in output] == in_outs["outputs"][index][0]) + tmp_result = tmp_result or ( + [list(x) for x in output] + == in_outs["outputs"][index][0] + ) except Exception: pass results.append(tmp_result) @@ -292,7 +319,8 @@ def run_test(in_outs, test=None, debug=False, timeout=15): error_traceback = traceback.format_exc() faulthandler.disable() if debug: - print(f"Standard input runtime error or time limit exceeded error = {e}") + print( + f"Standard input runtime error or time limit exceeded error = {e}") results.append(-1) return results, { "error": repr(e), @@ -302,9 +330,11 @@ def run_test(in_outs, test=None, debug=False, timeout=15): signal.alarm(0) if debug: print( - f"outputs = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, " - f"{type(inputs)}, {output == [in_outs['outputs'][index]]}" - ) + f"outputs = {output}, test outputs = { + in_outs['outputs'][index]}, inputs = {inputs}, " f"{ + type(inputs)}, { + output == [ + in_outs['outputs'][index]]}") elif which_type == CODE_TYPE.standard_input: # Standard input faulthandler.enable() passed = False @@ -312,7 +342,8 @@ def run_test(in_outs, test=None, debug=False, timeout=15): if isinstance(inputs, list): inputs = "\n".join(inputs) if isinstance(in_outs["outputs"][index], list): - in_outs["outputs"][index] = "\n".join(in_outs["outputs"][index]) + in_outs["outputs"][index] = "\n".join( + in_outs["outputs"][index]) signal.alarm(timeout) with Capturing() as output: @@ -325,7 +356,9 @@ def run_test(in_outs, test=None, debug=False, timeout=15): # runtime error or took too long signal.alarm(0) error_traceback = traceback.format_exc() - print(f"Call-based runtime error or time limit exceeded error = {repr(e)}{e}") + print( + f"Call-based runtime error or time limit exceeded error = {repr(e)}{e}" + ) results.append(-1) return results, { "error": repr(e), @@ -340,19 +373,26 @@ def run_test(in_outs, test=None, debug=False, timeout=15): nl = "\n" if not isinstance(inputs, list): print( - f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, " - f"inputs = {inputs.replace(nl, ' new-line ')}, {type(inputs)}, " - f"{output == [in_outs['outputs'][index]]}" - ) + f"not passed output = {output}, test outputs = { + in_outs['outputs'][index]}, " f"inputs = { + inputs.replace( + nl, ' new-line ')}, { + type(inputs)}, " f"{ + output == [ + in_outs['outputs'][index]]}") else: print( - f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, " - f"inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" - ) + f"not passed output = {output}, test outputs = { + in_outs['outputs'][index]}, " f"inputs = {inputs}, { + type(inputs)}, { + output == [ + in_outs['outputs'][index]]}") continue if passed and debug: - print(f"==> output = {output}, test outputs = {in_outs['outputs'][index]}") + print( + f"==> output = {output}, test outputs = { + in_outs['outputs'][index]}") if custom_compare_(output, in_outs["outputs"][index]): tmp_result = True @@ -367,9 +407,12 @@ def run_test(in_outs, test=None, debug=False, timeout=15): try: tmp_result = output == [in_outs["outputs"][index]] if isinstance(in_outs["outputs"][index], list): - tmp_result = tmp_result or (output == in_outs["outputs"][index]) + tmp_result = tmp_result or ( + output == in_outs["outputs"][index]) if isinstance(output[0], str): - tmp_result = tmp_result or ([e.strip() for e in output] == in_outs["outputs"][index]) + tmp_result = tmp_result or ( + [e.strip() for e in output] == in_outs["outputs"][index] + ) except Exception as e: if debug: print(f"Failed check1 exception = {e}") @@ -384,17 +427,22 @@ def run_test(in_outs, test=None, debug=False, timeout=15): for tmp_index, i in enumerate(in_outs["outputs"][index]): in_outs["outputs"][index][tmp_index] = i.split("\n") in_outs["outputs"][index][tmp_index] = [ - x.strip() for x in in_outs["outputs"][index][tmp_index] if x - ] + x.strip() for x in in_outs["outputs"][index][tmp_index] if x] else: - in_outs["outputs"][index] = in_outs["outputs"][index].split("\n") - in_outs["outputs"][index] = list(filter(len, in_outs["outputs"][index])) - in_outs["outputs"][index] = list(map(lambda x: x.strip(), in_outs["outputs"][index])) + in_outs["outputs"][index] = in_outs["outputs"][index].split( + "\n") + in_outs["outputs"][index] = list( + filter(len, in_outs["outputs"][index]) + ) + in_outs["outputs"][index] = list( + map(lambda x: x.strip(), in_outs["outputs"][index]) + ) try: tmp_result = output == [in_outs["outputs"][index]] if isinstance(in_outs["outputs"][index], list): - tmp_result = tmp_result or (output == in_outs["outputs"][index]) + tmp_result = tmp_result or ( + output == in_outs["outputs"][index]) except Exception as e: if debug: print(f"Failed check2 exception = {e}") @@ -418,9 +466,12 @@ def run_test(in_outs, test=None, debug=False, timeout=15): ) else: print( - f"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, " - f"{type(inputs)}, {output == [in_outs['outputs'][index]]} {tmp_result=}" - ) + f"@1 output = {output}, test outputs = { + in_outs['outputs'][index]}, inputs = {inputs}, " f"{ + type(inputs)}, { + output == [ + in_outs['outputs'][index]]} { + tmp_result=}") if debug: print(f"{tmp_result=} @a") @@ -428,7 +479,8 @@ def run_test(in_outs, test=None, debug=False, timeout=15): try: tmp_result = output == [in_outs["outputs"][index]] if isinstance(in_outs["outputs"][index], list): - tmp_result = tmp_result or (output == in_outs["outputs"][index]) + tmp_result = tmp_result or ( + output == in_outs["outputs"][index]) except Exception as e: if debug: print(f"Failed check3 exception = {e}") @@ -440,20 +492,26 @@ def run_test(in_outs, test=None, debug=False, timeout=15): try: all_ints = all( combined_int_check(e1) and combined_int_check(e2) - for e1, e2 in zip(output, in_outs["outputs"][index], strict=True) + for e1, e2 in zip( + output, in_outs["outputs"][index], strict=True + ) ) if not all_ints: if debug: print( [ combined_int_check(e1) and combined_int_check(e2) - for e1, e2 in zip(output, in_outs["outputs"][index], strict=True) + for e1, e2 in zip( + output, in_outs["outputs"][index], strict=True + ) ] ) output_float = [float(e) for e in output] - gt_float = [float(e) for e in in_outs["outputs"][index]] + gt_float = [float(e) + for e in in_outs["outputs"][index]] tmp_result = tmp_result or ( - (len(output_float) == len(gt_float)) and np.allclose(output_float, gt_float) + (len(output_float) == len(gt_float)) + and np.allclose(output_float, gt_float) ) except Exception: pass @@ -465,13 +523,17 @@ def run_test(in_outs, test=None, debug=False, timeout=15): if isinstance(output[0], list): all_ints = all( combined_int_check(e1) and combined_int_check(e2) - for e1, e2 in zip(output[0], in_outs["outputs"][index], strict=True) + for e1, e2 in zip( + output[0], in_outs["outputs"][index], strict=True + ) ) if not all_ints: output_float = [float(e) for e in output[0]] - gt_float = [float(e) for e in in_outs["outputs"][index][0]] + gt_float = [ + float(e) for e in in_outs["outputs"][index][0]] tmp_result = tmp_result or ( - (len(output_float) == len(gt_float)) and np.allclose(output_float, gt_float) + (len(output_float) == len(gt_float)) + and np.allclose(output_float, gt_float) ) except Exception: pass @@ -487,7 +549,8 @@ def run_test(in_outs, test=None, debug=False, timeout=15): for tmp_index, i in enumerate(in_outs["outputs"][index]): in_outs["outputs"][index][tmp_index] = set(i.split()) else: - in_outs["outputs"][index] = set(in_outs["outputs"][index].split()) + in_outs["outputs"][index] = set( + in_outs["outputs"][index].split()) if debug: print(f"{tmp_result=} @e") @@ -544,9 +607,11 @@ def run_test(in_outs, test=None, debug=False, timeout=15): ) else: print( - f"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, " - f"{type(inputs)}, {output == [in_outs['outputs'][index]]}" - ) + f"@2 output = {output}, test outputs = { + in_outs['outputs'][index]}, inputs = {inputs}, " f"{ + type(inputs)}, { + output == [ + in_outs['outputs'][index]]}") print(f"results = {results}") @@ -615,10 +680,17 @@ def reliability_guard(maximum_memory_bytes=None): if maximum_memory_bytes is not None: import resource - resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)) - resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)) + resource.setrlimit( + resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes) + ) + resource.setrlimit( + resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes) + ) if platform.uname().system != "Darwin": - resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)) + resource.setrlimit( + resource.RLIMIT_STACK, + (maximum_memory_bytes, + maximum_memory_bytes)) faulthandler.disable() diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/utils.py b/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/utils.py index 9123265..835eadc 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/utils.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/prime_code/utils.py @@ -12,7 +12,8 @@ # See the License for the specific language governing permissions and # limitations under the License. -# Borrowed from: https://huggingface.co/spaces/codeparrot/apps_metric/blob/main/utils.py +# Borrowed from: +# https://huggingface.co/spaces/codeparrot/apps_metric/blob/main/utils.py import multiprocessing import os @@ -28,7 +29,9 @@ def _temp_run(sample, generation, debug, result, metadata_list, timeout): sys.stdout = devnull sys.stderr = devnull try: - res, metadata = run_test(in_outs=sample, test=generation, debug=debug, timeout=timeout) + res, metadata = run_test( + in_outs=sample, test=generation, debug=debug, timeout=timeout + ) result.append(res) metadata_list.append(metadata) except Exception: @@ -38,7 +41,11 @@ def _temp_run(sample, generation, debug, result, metadata_list, timeout): metadata_list.append({}) -def check_correctness(in_outs: Optional[dict], generation, timeout=10, debug=True): +def check_correctness( + in_outs: Optional[dict], + generation, + timeout=10, + debug=True): """Check correctness of code generation with a global timeout. The global timeout is to catch some extreme/rare cases not handled by the timeouts inside `run_test`""" @@ -46,7 +53,10 @@ def check_correctness(in_outs: Optional[dict], generation, timeout=10, debug=Tru manager = multiprocessing.Manager() result = manager.list() metadata_list = manager.list() - p = multiprocessing.Process(target=_temp_run, args=(in_outs, generation, debug, result, metadata_list, timeout)) + p = multiprocessing.Process( + target=_temp_run, + args=(in_outs, generation, debug, result, metadata_list, timeout), + ) p.start() p.join(timeout=timeout + 1) if p.is_alive(): diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/__init__.py b/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/__init__.py index 04fd146..342edd9 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/__init__.py @@ -37,7 +37,7 @@ # sympy might hang -- we don't care about trying to be lenient in these cases BAD_SUBSTRINGS = ["^{", "^("] -BAD_REGEXES = ["\^[0-9]+\^", "\^[0-9][0-9]+"] +BAD_REGEXES = ["\\^[0-9]+\\^", "\\^[0-9][0-9]+"] TUPLE_CHARS = "()[]" @@ -46,7 +46,10 @@ def _sympy_parse(expr: str): py_expr = expr.replace("^", "**") return sympy_parser.parse_expr( py_expr, - transformations=(sympy_parser.standard_transformations + (sympy_parser.implicit_multiplication_application,)), + transformations=( + sympy_parser.standard_transformations + + (sympy_parser.implicit_multiplication_application,) + ), ) @@ -108,13 +111,13 @@ def _inject_implicit_mixed_number(step: str): e.g. 7 3/4 => 7+3/4 """ p1 = re.compile("([0-9]) +([0-9])") - step = p1.sub("\\1+\\2", step) ## implicit mults + step = p1.sub("\\1+\\2", step) # implicit mults return step def _strip_properly_formatted_commas(expr: str): # We want to be careful because we don't want to strip tuple commas - p1 = re.compile("(\d)(,)(\d\d\d)($|\D)") + p1 = re.compile("(\\d)(,)(\\d\\d\\d)($|\\D)") while True: next_expr = p1.sub("\\1\\3\\4", expr) if next_expr == expr: @@ -129,7 +132,7 @@ def _normalize(expr: str) -> str: return None # Remove enclosing `\text{}`. - m = re.search("^\\\\text\{(?P.+?)\}$", expr) + m = re.search("^\\\\text\\{(?P.+?)\\}$", expr) if m is not None: expr = m.group("text") @@ -163,8 +166,8 @@ def _normalize(expr: str) -> str: "yard", "liter", ]: - expr = re.sub(f"{unit}(es)?(s)? *(\^[0-9]+)?", "", expr) - expr = re.sub("\^ *\\\\circ", "", expr) + expr = re.sub(f"{unit}(es)?(s)? *(\\^[0-9]+)?", "", expr) + expr = re.sub("\\^ *\\\\circ", "", expr) if len(expr) > 0 and expr[0] == "{" and expr[-1] == "}": expr = expr[1:-1] @@ -198,7 +201,8 @@ def count_unknown_letters_in_expr(expr: str): def should_allow_eval(expr: str): - # we don't want to try parsing unknown text or functions of more than two variables + # we don't want to try parsing unknown text or functions of more than two + # variables if count_unknown_letters_in_expr(expr) > 2: return False @@ -253,8 +257,10 @@ def grade_answer(given_answer: str, ground_truth: str) -> bool: if given_answer is None: return False - ground_truth_normalized_mathd = math_normalize.normalize_answer(ground_truth) - given_answer_normalized_mathd = math_normalize.normalize_answer(given_answer) + ground_truth_normalized_mathd = math_normalize.normalize_answer( + ground_truth) + given_answer_normalized_mathd = math_normalize.normalize_answer( + given_answer) # be at least as lenient as mathd if ground_truth_normalized_mathd == given_answer_normalized_mathd: @@ -277,12 +283,17 @@ def grade_answer(given_answer: str, ground_truth: str) -> bool: if ( len(ground_truth_elems) > 1 - and (ground_truth_normalized[0] != given_normalized[0] or ground_truth_normalized[-1] != given_normalized[-1]) + and ( + ground_truth_normalized[0] != given_normalized[0] + or ground_truth_normalized[-1] != given_normalized[-1] + ) or len(ground_truth_elems) != len(given_elems) ): is_correct = False else: - for ground_truth_elem, given_elem in zip(ground_truth_elems, given_elems, strict=True): + for ground_truth_elem, given_elem in zip( + ground_truth_elems, given_elems, strict=True + ): if _is_frac(ground_truth_elem) and _is_frac(given_elem): # if fractions aren't reduced, then shouldn't be marked as correct # so, we don't want to allow sympy.simplify in this case @@ -293,11 +304,13 @@ def grade_answer(given_answer: str, ground_truth: str) -> bool: is_correct = False else: try: - is_correct = are_equal_under_sympy(ground_truth_elem, given_elem) + is_correct = are_equal_under_sympy( + ground_truth_elem, given_elem) except Exception as e: # if there's an error, we'll just say it's not correct is_correct = False - print(f"Error: {e} from are_equal_under_sympy, {ground_truth_elem}, {given_elem}") + print( + f"Error: {e} from are_equal_under_sympy, {ground_truth_elem}, {given_elem}") if not is_correct: break @@ -309,7 +322,7 @@ def remove_boxed(s): try: assert s[: len(left)] == left assert s[-1] == "}" - return s[len(left) : -1] + return s[len(left): -1] except Exception: return None @@ -341,7 +354,7 @@ def _last_boxed_only_string(string): if left_brace_idx is None or right_brace_idx is None: return None - return string[left_brace_idx + 1 : right_brace_idx].strip() + return string[left_brace_idx + 1: right_brace_idx].strip() def match_answer(response): @@ -350,11 +363,15 @@ def match_answer(response): ans_idx = response.lower().rfind(ans_marker) if ans_idx != -1: is_matched = True - response = response[ans_idx + len(ans_marker) :].strip() + response = response[ans_idx + len(ans_marker):].strip() if response.endswith("\n"): response = response[:-2] - for ans_marker in ["is answer", "is the answer", "are answers", "are the answers"]: + for ans_marker in [ + "is answer", + "is the answer", + "are answers", + "are the answers"]: ans_idx = response.lower().rfind(ans_marker) if ans_idx != -1: is_matched = True @@ -373,15 +390,29 @@ def match_answer(response): if dot_idx != -1: response = response[:dot_idx].strip() - for ans_marker in ["be ", "is ", "are ", "=", ": ", "get ", "be\n", "is\n", "are\n", ":\n", "get\n"]: + for ans_marker in [ + "be ", + "is ", + "are ", + "=", + ": ", + "get ", + "be\n", + "is\n", + "are\n", + ":\n", + "get\n", + ]: ans_idx = response.lower().rfind(ans_marker) if ans_idx != -1: is_matched = True - response = response[ans_idx + len(ans_marker) :].strip() + response = response[ans_idx + len(ans_marker):].strip() if response.endswith("\n"): response = response[:-2] - is_matched = is_matched if any([c.isdigit() for c in response]) else False # answer must have a digit + is_matched = ( + is_matched if any([c.isdigit() for c in response]) else False + ) # answer must have a digit # Grade return is_matched, response @@ -393,18 +424,27 @@ def compute_score(model_output: str, ground_truth: str) -> bool: is_matched, extracted_model_output = match_answer(model_output) format_correctness = "Step 2:" in model_output and "\\box" in model_output - # grade simple algebra questions. if succeeded, return; otherwise, proceed to more complex grading + # grade simple algebra questions. if succeeded, return; otherwise, proceed + # to more complex grading if grade_answer(extracted_model_output, ground_truth): return True, True, extracted_model_output try: - if "\pi" in extracted_model_output or "\pi" in ground_truth: + if "\\pi" in extracted_model_output or "\\pi" in ground_truth: equivs = [] for pi in [math.pi, 3.14]: - equivs.append(math_equal(extracted_model_output, ground_truth, timeout=True, pi=pi)) + equivs.append( + math_equal( + extracted_model_output, + ground_truth, + timeout=True, + pi=pi)) is_correct = any(equivs) else: - is_correct = math_equal(extracted_model_output, ground_truth, timeout=True) + is_correct = math_equal( + extracted_model_output, + ground_truth, + timeout=True) except Exception: is_correct = False diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/grader.py b/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/grader.py index d060584..058162b 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/grader.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/grader.py @@ -125,7 +125,8 @@ def normalize(answer, pi) -> str: # checking if answer is % or \\% and removing % if isinstance(answer, str) and ( - bool(re.match(r"^\d+(\.\d+)?%$", answer)) or bool(re.match(r"^\d+(\.\d+)?\\%$", answer)) + bool(re.match(r"^\d+(\.\d+)?%$", answer)) + or bool(re.match(r"^\d+(\.\d+)?\\%$", answer)) ): return answer.replace("\\%", "").replace("%", "") @@ -148,21 +149,24 @@ def handle_base(x) -> str: def handle_pi(string, pi): - if isinstance(string, str) and "\pi" in string: + if isinstance(string, str) and "\\pi" in string: # Find the first occurrence of "\pi" - idx = string.find("\pi") + idx = string.find("\\pi") - # Iterate over the string and find all occurrences of "\pi" with a valid previous character + # Iterate over the string and find all occurrences of "\pi" with a + # valid previous character while idx != -1: if idx > 0 and string[idx - 1].isdigit(): - # Replace "\pi" with "*math.pi" if the previous character is a digit - string = string[:idx] + f"*{pi}" + string[idx + 3 :] + # Replace "\pi" with "*math.pi" if the previous character is a + # digit + string = string[:idx] + f"*{pi}" + string[idx + 3:] else: - # Replace "\pi" with "1*math.pi" if the previous character is not a digit - string = string[:idx] + f"1*{pi}" + string[idx + 3 :] + # Replace "\pi" with "1*math.pi" if the previous character is + # not a digit + string = string[:idx] + f"1*{pi}" + string[idx + 3:] # Find the next occurrence of "\pi" - idx = string.find("\pi", idx + 1) + idx = string.find("\\pi", idx + 1) # Evaluate the expression using eval() function with contextlib.suppress(Exception): @@ -188,7 +192,9 @@ def math_equal( prediction = normalize(prediction, pi) reference = normalize(reference, pi) - if isinstance(prediction, str) and len(prediction) > 1000: # handling weird corner-cases + if ( + isinstance(prediction, str) and len(prediction) > 1000 + ): # handling weird corner-cases prediction = prediction[:1000] # 0. string comparison @@ -203,7 +209,11 @@ def math_equal( prediction = is_digit(prediction)[1] reference = is_digit(reference)[1] # number questions - gt_result = [reference / 100, reference, reference * 100] if include_percentage else [reference] + gt_result = ( + [reference / 100, reference, reference * 100] + if include_percentage + else [reference] + ) for item in gt_result: try: if isclose(item, prediction, rel_tol=tolerance): @@ -221,12 +231,18 @@ def math_equal( reference = str(reference).strip() prediction = str(prediction).strip() - ## deal with [], (), {} + # deal with [], (), {} prediction = format_intervals(prediction) pred_str, ref_str = prediction, reference - if (prediction.startswith("[") and prediction.endswith("]") and not reference.startswith("(")) or ( - prediction.startswith("(") and prediction.endswith(")") and not reference.startswith("[") + if ( + prediction.startswith("[") + and prediction.endswith("]") + and not reference.startswith("(") + ) or ( + prediction.startswith("(") + and prediction.endswith(")") + and not reference.startswith("[") ): pred_str = pred_str.strip("[]()") ref_str = ref_str.strip("[]()") @@ -236,7 +252,7 @@ def math_equal( if pred_str == ref_str: return True - ## [a, b] vs. [c, d], return a==c and b==d + # [a, b] vs. [c, d], return a==c and b==d if ( prediction and reference @@ -260,18 +276,15 @@ def math_equal( ref_parts = [item.strip() for item in reference.split(",")] if len(pred_parts) == len(ref_parts): - return bool( - all( - [ - math_equal(pred_parts[i], ref_parts[i], include_percentage, tolerance) - for i in range(len(pred_parts)) - ] - ) - ) + return bool(all([math_equal(pred_parts[i], + ref_parts[i], + include_percentage, + tolerance) for i in range(len(pred_parts))])) # if we have point == tuple of values - if prediction.startswith("Point") and reference[0] == "(" and reference[-1] == ")": - pred_parts = prediction[prediction.find("(") + 1 : -1].split(",") + if prediction.startswith( + "Point") and reference[0] == "(" and reference[-1] == ")": + pred_parts = prediction[prediction.find("(") + 1: -1].split(",") ref_parts = reference[1:-1].split(",") if len(pred_parts) == len(ref_parts) and all( [ @@ -295,7 +308,11 @@ def math_equal( return True except Exception: pass - elif "\begin{pmatrix}" in reference and prediction.startswith("[") and prediction.endswith("]"): + elif ( + "\begin{pmatrix}" in reference + and prediction.startswith("[") + and prediction.endswith("]") + ): if isinstance(eval(prediction), list): try: pred_matrix = eval(prediction) @@ -307,11 +324,14 @@ def math_equal( .rstrip("\end{pmatrix}") ) # noqa: B005 ref_matrix_items = ref_matrix_items.split("\\") - ref_matrix_items = [row.split("&") if "&" in row else row for row in ref_matrix_items] + ref_matrix_items = [ + row.split("&") if "&" in row else row for row in ref_matrix_items] if len(pred_matrix) == len(ref_matrix_items) and all( [ math_equal(pred, ref, include_percentage, tolerance) - for ref, pred in zip(ref_matrix_items, pred_matrix, strict=False) + for ref, pred in zip( + ref_matrix_items, pred_matrix, strict=False + ) ] ): return True diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/math_normalize.py b/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/math_normalize.py index 74d94cc..52a5ec7 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/math_normalize.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/prime_math/math_normalize.py @@ -47,7 +47,7 @@ def normalize_answer(answer: Optional[str]) -> Optional[str]: answer = answer.strip() try: # Remove enclosing `\text{}`. - m = re.search("^\\\\text\{(?P.+?)\}$", answer) + m = re.search("^\\\\text\\{(?P.+?)\\}$", answer) if m is not None: answer = m.group("text").strip() return _strip_string(answer) @@ -157,7 +157,7 @@ def _strip_string(string): # remove percentage string = string.replace("\\%", "") - string = string.replace("\%", "") + string = string.replace("\\%", "") # " 0." equivalent to " ." and "{0." equivalent to "{." Alternatively, add "0" if "." is the start of the string string = string.replace(" .", " 0.") @@ -186,7 +186,8 @@ def _strip_string(string): 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 + # 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 diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/sandbox_fusion/__init__.py b/Agent0/executor_train/verl/verl/utils/reward_score/sandbox_fusion/__init__.py index cd18498..14a159c 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/sandbox_fusion/__init__.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/sandbox_fusion/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,7 +26,13 @@ def compute_score( - sandbox_fusion_url, concurrent_semaphore, memory_limit_mb, completion, test_cases, continuous=False, timeout=10 + sandbox_fusion_url, + concurrent_semaphore, + memory_limit_mb, + completion, + test_cases, + continuous=False, + timeout=10, ): """ Computes the code score using the remote sandbox API. @@ -70,7 +76,9 @@ def compute_score( if not test_cases or "inputs" not in test_cases or "outputs" not in test_cases: logger.error("Invalid test_cases structure.") - return 0.0, [{"error": "Invalid test_cases structure (missing inputs/outputs)"}] + return 0.0, [ + {"error": "Invalid test_cases structure (missing inputs/outputs)"} + ] # Check all test cases # Note: The return value of check_correctness might need adaptation here @@ -95,7 +103,8 @@ def compute_score( if num_to_consider == 0: score = 0.0 else: - passed_count = sum(1 for r in res_list[:num_to_consider] if r is True) + passed_count = sum( + 1 for r in res_list[:num_to_consider] if r is True) score = passed_count / num_to_consider # Return all metadata, even if score is based on the first N final_metadata = metadata_list @@ -110,8 +119,14 @@ def compute_score( logger.error(f"Error during compute_score: {e}") traceback.print_exc() score = 0.0 - # Try to return partial metadata if available, otherwise return error info - final_metadata = metadata_list if "metadata_list" in locals() else [{"error": f"Unhandled exception: {e}"}] + # Try to return partial metadata if available, otherwise return error + # info + final_metadata = ( + metadata_list + if "metadata_list" in locals() + else [{"error": f"Unhandled exception: {e}"}] + ) # Ensure float and list are returned - return float(score), final_metadata if isinstance(final_metadata, list) else [final_metadata] + return float(score), (final_metadata if isinstance( + final_metadata, list) else [final_metadata]) diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/sandbox_fusion/utils.py b/Agent0/executor_train/verl/verl/utils/reward_score/sandbox_fusion/utils.py index d2154ca..ec853a9 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/sandbox_fusion/utils.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/sandbox_fusion/utils.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -110,8 +110,11 @@ def call_sandbox_api( "fetch_files": [], } ) - headers = {"Content-Type": "application/json", "Accept": "application/json"} - # Calculate a reasonable request timeout based on compile/run timeouts plus a buffer + headers = { + "Content-Type": "application/json", + "Accept": "application/json"} + # Calculate a reasonable request timeout based on compile/run timeouts + # plus a buffer request_timeout = compile_timeout + run_timeout + API_TIMEOUT last_error = None # Store the last error encountered @@ -131,16 +134,20 @@ def call_sandbox_api( # Check for Gateway Timeout (504) specifically for retrying if response.status_code == 504: last_error = ( - f"{log_prefix}API Request Error: Gateway Timeout (504) on attempt " - f"{attempt + 1}/{MAX_RETRIES}" - ) # <-- Use internal log_prefix + f"{log_prefix}API Request Error: Gateway Timeout (504) on attempt " f"{ + attempt + 1}/{MAX_RETRIES}") # <-- Use internal log_prefix logger.warning(last_error) if attempt < MAX_RETRIES - 1: # Don't sleep after the last attempt # Calculate increasing delay (e.g., 1s, 2s, 4s, ...) or (1s, 2s, 3s, ...) # Simple linear increase: delay = INITIAL_RETRY_DELAY * (attempt + 1) - # Exponential backoff: delay = INITIAL_RETRY_DELAY * (2 ** attempt) - delay = INITIAL_RETRY_DELAY * (attempt + 1) # Using linear increase for simplicity - logger.info(f"{log_prefix}Retrying after {delay} seconds...") # <-- Use internal log_prefix + # Exponential backoff: delay = INITIAL_RETRY_DELAY * (2 ** + # attempt) + delay = INITIAL_RETRY_DELAY * ( + attempt + 1 + ) # Using linear increase for simplicity + logger.info( + f"{log_prefix}Retrying after {delay} seconds..." + ) # <-- Use internal log_prefix time.sleep(delay) continue # Go to the next retry attempt @@ -149,26 +156,41 @@ def call_sandbox_api( # If successful (status code 2xx) logger.info( - f"{log_prefix}Sandbox API call successful on attempt {attempt + 1}" - ) # <-- Use internal log_prefix + f"{log_prefix}Sandbox API call successful on attempt { + attempt + 1}") # <-- Use internal log_prefix return response.json(), None except requests.exceptions.RequestException as e: - last_error = f"{log_prefix}API Request Error: {e}" # <-- Use internal log_prefix + last_error = ( + # <-- Use internal log_prefix + f"{log_prefix}API Request Error: {e}" + ) break # Exit retry loop on non-504 request errors except json.JSONDecodeError as e: raw_response_text = response.text if "response" in locals() else "N/A" - last_error = f"{log_prefix}API Response JSON Decode Error: {e}" # <-- Use internal log_prefix + # <-- Use internal log_prefix + last_error = f"{log_prefix}API Response JSON Decode Error: {e}" break # Exit retry loop on JSON decode errors except Exception as e: - last_error = f"{log_prefix}Unexpected Error: {e}" # <-- Use internal log_prefix + last_error = ( + # <-- Use internal log_prefix + f"{log_prefix}Unexpected Error: {e}" + ) break # Exit retry loop on other unexpected errors - # If loop finishes without returning success, return the last recorded error - logger.error(f"{log_prefix}Sandbox API call failed. Last error: {last_error}") # <-- Use internal log_prefix + # If loop finishes without returning success, return the last recorded + # error + logger.error( + f"{log_prefix}Sandbox API call failed. Last error: {last_error}" + ) # <-- Use internal log_prefix # Return the error message without the prefix, as the caller doesn't need the internal ID - # Ensure API call failure returns error message, leading to -1 in check_correctness - return None, last_error.replace(log_prefix, "API Call Failed: ") if last_error else "API Call Failed after retries" + # Ensure API call failure returns error message, leading to -1 in + # check_correctness + return None, ( + last_error.replace(log_prefix, "API Call Failed: ") + if last_error + else "API Call Failed after retries" + ) def _process_single_case( @@ -259,9 +281,9 @@ def _execute_user_function(): # Attempt to instantiate and get method. # Errors (e.g., Solution not a class, instantiation fails, method missing) # will be caught by the broad except block below. - _solution_instance = _Solution_class() + _solution_instance = _Solution_class() _target_callable = getattr(_solution_instance, _SANDBOX_FN_NAME) - + if not _target_callable: sys.stderr.write(f"WrapperError: Function or method '{{_SANDBOX_FN_NAME}}' not found.\\n") return None, True # result, error_occurred @@ -286,7 +308,7 @@ def _execute_user_function(): print(str(_result)) # Optional: To explicitly exit with an error code if the sandbox relies on it # else: - # sys.exit(1) + # sys.exit(1) """ current_generation_code = wrapper_code @@ -316,7 +338,8 @@ def _execute_user_function(): language=language, ) except Exception as e: - error_msg = f"API Request Exception during check_correctness for case {case_index + 1}: {e}" + error_msg = f"API Request Exception during check_correctness for case { + case_index + 1}: {e}" logger.error(f"Case {case_index + 1}: {error_msg}") traceback.print_exc() @@ -341,10 +364,13 @@ def _execute_user_function(): if error_msg: metadata["status"] = "api_error" - result_status = -1 # API request itself failed (includes timeout after retries) + # API request itself failed (includes timeout after retries) + result_status = -1 logger.error(f"Case {case_index}: API error occurred: {error_msg}") # Log code and input only on error for brevity - generation_to_log = generation[:200] + "..." if len(generation) > 200 else generation + generation_to_log = ( + generation[:200] + "..." if len(generation) > 200 else generation + ) logger.error(f"Case {case_index}: code: {generation_to_log}") logger.error(f"Case {case_index}: input: {str(stdin_data)}") elif api_response: @@ -365,7 +391,8 @@ def _execute_user_function(): if run_result: metadata["run_status"] = run_result.get("status") metadata["stdout"] = run_result.get("stdout") - metadata["stderr"] = run_result.get("stderr") # stderr during runtime + metadata["stderr"] = run_result.get( + "stderr") # stderr during runtime metadata["exit_code"] = run_result.get("return_code") metadata["duration"] = run_result.get("execution_time") @@ -377,17 +404,22 @@ def _execute_user_function(): result_status = -1 # Internal sandbox error elif api_status == "Failed": # --- Add debug logging --- - logger.debug(f"API returned Failed status. Response: {api_response}") + logger.debug( + f"API returned Failed status. Response: {api_response}") logger.debug(f"Compile Result: {compile_result}") logger.debug(f"Run Result: {run_result}") # --- Check the logic here --- # Compile failed or timed out is_compile_error = compile_result and ( metadata["compile_status"] in ["Error", "TimeLimitExceeded"] - or (metadata["compile_status"] == "Finished" and compile_result.get("return_code") != 0) + or ( + metadata["compile_status"] == "Finished" + and compile_result.get("return_code") != 0 + ) ) if is_compile_error: - # Differentiate between compile_error and compile_timeout based on specific status + # Differentiate between compile_error and compile_timeout based + # on specific status if metadata["compile_status"] == "TimeLimitExceeded": metadata["status"] = "compile_timeout" else: # Includes Error and Finished but return_code != 0 cases @@ -395,11 +427,15 @@ def _execute_user_function(): result_status = -4 # Run failed or timed out elif run_result: - # Modified condition: Check for TimeLimitExceeded OR (Finished with non-zero exit code) OR Error status + # Modified condition: Check for TimeLimitExceeded OR (Finished + # with non-zero exit code) OR Error status is_runtime_error = ( metadata["run_status"] == "TimeLimitExceeded" or metadata["run_status"] == "Error" - or (metadata["run_status"] == "Finished" and run_result.get("return_code") != 0) + or ( + metadata["run_status"] == "Finished" + and run_result.get("return_code") != 0 + ) ) if is_runtime_error: if metadata["run_status"] == "TimeLimitExceeded": @@ -409,28 +445,38 @@ def _execute_user_function(): metadata["status"] = "runtime_error" result_status = -2 else: - # Other Failed status with run_result, classify as unknown failure - logger.warning(f"Unknown run_status '{metadata['run_status']}' or state within Failed API status.") + # Other Failed status with run_result, classify as unknown + # failure + logger.warning( + f"Unknown run_status '{ + metadata['run_status']}' or state within Failed API status.") metadata["status"] = "unknown_failure" result_status = -1 # Default to -1 else: - # Status is Failed but neither a clear compile error nor run_result exists - logger.warning("API status Failed but cannot determine specific error type (compile/run).") + # Status is Failed but neither a clear compile error nor + # run_result exists + logger.warning( + "API status Failed but cannot determine specific error type (compile/run)." + ) metadata["status"] = "unknown_failure_state" result_status = -1 # Default to -1 elif api_status == "Success": # Run completed successfully, now check the answer if run_result and metadata["run_status"] == "Finished": - actual_output = metadata["stdout"] if metadata["stdout"] is not None else "" - # Note: Output might contain trailing newlines, need normalization - if str(actual_output).rstrip("\n") == str(expected_output).rstrip("\n"): + actual_output = ( + metadata["stdout"] if metadata["stdout"] is not None else "") + # Note: Output might contain trailing newlines, need + # normalization + if str(actual_output).rstrip("\n") == str( + expected_output).rstrip("\n"): result_status = True metadata["status"] = "success" else: result_status = False metadata["status"] = "wrong_answer" else: - # Status is Success but run_result status is not Finished, this is unexpected + # Status is Success but run_result status is not Finished, this + # is unexpected metadata["status"] = "unexpected_success_state" result_status = -1 # Classify as unknown error else: @@ -438,10 +484,13 @@ def _execute_user_function(): logger.warning(f"Unknown API status received: {api_status}") metadata["status"] = f"unknown_api_status_{api_status}" result_status = -1 # Default to -1 - else: # api_response is None and no error_msg (Should not happen with current call_sandbox_api logic) + # api_response is None and no error_msg (Should not happen with current + # call_sandbox_api logic) + else: metadata["status"] = "unknown_api_state" result_status = -1 - logger.error(f"Case {case_index}: Unknown API state (no response and no error message).") + logger.error( + f"Case {case_index}: Unknown API state (no response and no error message).") return result_status, metadata @@ -491,15 +540,25 @@ def check_correctness( return [], [] if len(inputs) != len(expected_outputs): - logger.warning(f"Mismatch between number of inputs ({len(inputs)}) and outputs ({len(expected_outputs)}).") + logger.warning( + f"Mismatch between number of inputs ({ + len(inputs)}) and outputs ({ + len(expected_outputs)}).") # Return error based on the number of inputs provided - return [-1] * num_cases, [{"error": "Input/output count mismatch", "case_index": i} for i in range(num_cases)] + return [-1] * num_cases, [ + {"error": "Input/output count mismatch", "case_index": i} + for i in range(num_cases) + ] first_compile_error_index = -1 - # max_workers is limited by sandbox_fusion_max_concurrent from concurrent_semaphore - with concurrent.futures.ThreadPoolExecutor(max_workers=max(32, os.cpu_count() * 5)) as executor: - # Submit all tasks, passing the concurrent_semaphore to _process_single_case + # max_workers is limited by sandbox_fusion_max_concurrent from + # concurrent_semaphore + with concurrent.futures.ThreadPoolExecutor( + max_workers=max(32, os.cpu_count() * 5) + ) as executor: + # Submit all tasks, passing the concurrent_semaphore to + # _process_single_case future_to_index = { executor.submit( _process_single_case, @@ -527,13 +586,18 @@ def check_correctness( # Check for compile error (-4) if result_status == -4: - if first_compile_error_index == -1 or index < first_compile_error_index: + if ( + first_compile_error_index == -1 + or index < first_compile_error_index + ): first_compile_error_index = index # Optimization: could potentially cancel futures for index > first_compile_error_index - # However, cancellation is not guaranteed. Post-processing is safer. + # However, cancellation is not guaranteed. Post-processing + # is safer. except Exception as exc: - logger.error(f"Test case {index} generated an exception: {exc}") + logger.error( + f"Test case {index} generated an exception: {exc}") traceback.print_exc() results[index] = -1 # Mark as API/internal error metadata_list[index] = { @@ -550,11 +614,16 @@ def check_correctness( f"Compile error detected in case {first_compile_error_index}. Marking subsequent cases as compile errors." ) for i in range(first_compile_error_index + 1, num_cases): - # Only update if not already processed (though it should be None or have a result) - if results[i] != -4: # Avoid overwriting if it somehow already got -4 + # Only update if not already processed (though it should be None or + # have a result) + if results[i] != - \ + 4: # Avoid overwriting if it somehow already got -4 results[i] = -4 - # Update or create metadata for skipped cases due to compile error - if metadata_list[i] is None: # If future failed before returning metadata + # Update or create metadata for skipped cases due to compile + # error + if ( + metadata_list[i] is None + ): # If future failed before returning metadata metadata_list[i] = { "case_index": i, "input": str(inputs[i]), diff --git a/Agent0/executor_train/verl/verl/utils/reward_score/search_r1_like_qa_em.py b/Agent0/executor_train/verl/verl/utils/reward_score/search_r1_like_qa_em.py index 56782fc..644f0ab 100644 --- a/Agent0/executor_train/verl/verl/utils/reward_score/search_r1_like_qa_em.py +++ b/Agent0/executor_train/verl/verl/utils/reward_score/search_r1_like_qa_em.py @@ -1,156 +1,161 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 Search-R1 Contributors -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT 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/PeterGriffinJin/Search-R1/blob/main/verl/utils/reward_score/qa_em.py - -import random -import re -import string - - -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)))) - - -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 - - -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 - - -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"(.*?)" - match = re.finditer(answer_pattern, solution_str, re.DOTALL) - matches = list(match) - - # If there are 0 matches, return None - if len(matches) < 1: - return None - - # If there are 2 or more matches, return the last one - return matches[-1].group(1).strip() - - -def count_answer_tags(text): - opening_tags = text.count("") - closing_tags = text.count("") - - return opening_tags, closing_tags - - -def compute_score(solution_str, ground_truth, method="strict", format_score=0.0, score=1.0): - """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) - open_count, close_count = count_answer_tags(solution_str) - do_print = random.randint(1, 64) == 1 - - if do_print: - print("--------------------------------") - print(f"Golden answers: {ground_truth['target']}") - if answer is not None: - print(f"Extracted answer is not None: {answer}") - else: - print("Extracted answer: None!") - print(f"Solution string: {solution_str}") - - if answer is None: - return 0 - else: - if em_check(answer, ground_truth["target"]): - if open_count > 10 or close_count > 10: # prevent output a lot of - score = score / 4 - return score - return score - else: - return format_score - - -def compute_score_subem(solution_str, ground_truth, method="strict", format_score=0.0, score=1.0): - """The scoring function for substring exact match (EM). - - Args: - solution_str: the solution text - ground_truth: the ground truth - method: the method to extract the solution, choices are 'strict' and 'flexible' - format_score: the score for the format - score: the score for the correct answer - """ - answer = extract_solution(solution_str=solution_str) - do_print = random.randint(1, 64) == 1 - - if do_print: - print("--------------------------------") - print(f"Golden answers: {ground_truth['target']}") - print(f"Extracted answer: {answer}") - print(f"Solution string: {solution_str}") - - if answer is None: - return 0 - else: - if subem_check(answer, ground_truth["target"]): - return score - else: - return format_score +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025 Search-R1 Contributors +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT 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/PeterGriffinJin/Search-R1/blob/main/verl/utils/reward_score/qa_em.py + +import random +import re +import string + + +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)))) + + +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 + + +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 + + +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"(.*?)" + match = re.finditer(answer_pattern, solution_str, re.DOTALL) + matches = list(match) + + # If there are 0 matches, return None + if len(matches) < 1: + return None + + # If there are 2 or more matches, return the last one + return matches[-1].group(1).strip() + + +def count_answer_tags(text): + opening_tags = text.count("") + closing_tags = text.count("") + + return opening_tags, closing_tags + + +def compute_score( + solution_str, ground_truth, method="strict", format_score=0.0, score=1.0 +): + """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) + open_count, close_count = count_answer_tags(solution_str) + do_print = random.randint(1, 64) == 1 + + if do_print: + print("--------------------------------") + print(f"Golden answers: {ground_truth['target']}") + if answer is not None: + print(f"Extracted answer is not None: {answer}") + else: + print("Extracted answer: None!") + print(f"Solution string: {solution_str}") + + if answer is None: + return 0 + else: + if em_check(answer, ground_truth["target"]): + if open_count > 10 or close_count > 10: # prevent output a lot of + score = score / 4 + return score + return score + else: + return format_score + + +def compute_score_subem( + solution_str, ground_truth, method="strict", format_score=0.0, score=1.0 +): + """The scoring function for substring exact match (EM). + + Args: + solution_str: the solution text + ground_truth: the ground truth + method: the method to extract the solution, choices are 'strict' and 'flexible' + format_score: the score for the format + score: the score for the correct answer + """ + answer = extract_solution(solution_str=solution_str) + do_print = random.randint(1, 64) == 1 + + if do_print: + print("--------------------------------") + print(f"Golden answers: {ground_truth['target']}") + print(f"Extracted answer: {answer}") + print(f"Solution string: {solution_str}") + + if answer is None: + return 0 + else: + if subem_check(answer, ground_truth["target"]): + return score + else: + return format_score diff --git a/Agent0/executor_train/verl/verl/utils/rollout_trace.py b/Agent0/executor_train/verl/verl/utils/rollout_trace.py index 114006d..021efc7 100644 --- a/Agent0/executor_train/verl/verl/utils/rollout_trace.py +++ b/Agent0/executor_train/verl/verl/utils/rollout_trace.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -37,7 +37,13 @@ def get_instance(cls) -> "RolloutTraceConfig": return cls._instance @classmethod - def init(cls, project_name: str, experiment_name: str, backend: str, token2text: bool = False): + def init( + cls, + project_name: str, + experiment_name: str, + backend: str, + token2text: bool = False, + ): config = cls.get_instance() config.backend = backend config.token2text = token2text @@ -82,7 +88,8 @@ def rollout_trace_attr(sample_index=None, step=None, rollout_n=None): attributes["step"] = step if rollout_n is not None: attributes["rollout_n"] = rollout_n - attributes["experiment_name"] = RolloutTraceConfig.get_instance().experiment_name + attributes["experiment_name"] = RolloutTraceConfig.get_instance( + ).experiment_name if not attributes or backend is None: yield @@ -123,15 +130,23 @@ async def async_wrapper(self, *args, **kwargs): del inputs["self"] async def add_token2text(self, result): - if hasattr(result, "prompt_ids") and hasattr(self, "tokenizer") and hasattr(self.tokenizer, "decode"): + if ( + hasattr(result, "prompt_ids") + and hasattr(self, "tokenizer") + and hasattr(self.tokenizer, "decode") + ): _result = [result] loop = asyncio.get_running_loop() if hasattr(result, "prompt_ids"): - prompt_text = await loop.run_in_executor(None, self.tokenizer.decode, result.prompt_ids) + prompt_text = await loop.run_in_executor( + None, self.tokenizer.decode, result.prompt_ids + ) _result.append(prompt_text) if hasattr(result, "response_ids"): - response_text = await loop.run_in_executor(None, self.tokenizer.decode, result.response_ids) + response_text = await loop.run_in_executor( + None, self.tokenizer.decode, result.response_ids + ) _result.append(response_text) return _result return result @@ -141,7 +156,9 @@ async def add_token2text(self, result): from weave.trace.context import call_context cur_attributes = {**call_context.call_attributes.get()} - call = tracer.create_call(op=func.__qualname__, inputs=inputs, attributes=cur_attributes) + call = tracer.create_call( + op=func.__qualname__, inputs=inputs, attributes=cur_attributes + ) try: result = await func(self, *args, **kwargs) @@ -177,7 +194,9 @@ def wrapper(self, *args, **kwargs): from weave.trace.context import call_context cur_attributes = {**call_context.call_attributes.get()} - call = tracer.create_call(op=func.__qualname__, inputs=inputs, attributes=cur_attributes) + call = tracer.create_call( + op=func.__qualname__, inputs=inputs, attributes=cur_attributes + ) try: result = func(self, *args, **kwargs) tracer.finish_call(call, output=result) diff --git a/Agent0/executor_train/verl/verl/utils/seqlen_balancing.py b/Agent0/executor_train/verl/verl/utils/seqlen_balancing.py index 4938e8f..f3ab264 100644 --- a/Agent0/executor_train/verl/verl/utils/seqlen_balancing.py +++ b/Agent0/executor_train/verl/verl/utils/seqlen_balancing.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -21,7 +21,10 @@ from verl.utils.device import get_device_name -def karmarkar_karp(seqlen_list: list[int], k_partitions: int, equal_size: bool): +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: @@ -94,10 +97,13 @@ def __repr__(self) -> str: repr_str += "]" return repr_str - sorted_seqlen_list = sorted([(seqlen, i) for i, seqlen in enumerate(seqlen_list)]) + 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" + 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): @@ -106,7 +112,10 @@ def __repr__(self) -> str: 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)) + heapq.heappush( + states_pq, State( + items=[ + (idx, seqlen)], k=k_partitions)) while len(states_pq) > 1: state0 = heapq.heappop(states_pq) @@ -119,15 +128,19 @@ def __repr__(self) -> str: 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)}" - ) + assert len(partition) * k_partitions == len( + seqlen_list + ), f"{len(partition)} * {k_partitions} != {len(seqlen_list)}" return partitions -def greedy_partition(seqlen_list: list[int], k_partitions: int, equal_size: bool): +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)] + 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: @@ -139,13 +152,15 @@ def greedy_partition(seqlen_list: list[int], k_partitions: int, equal_size: bool 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)}" - ) + assert len(partition) * k_partitions == len( + seqlen_list + ), f"{len(partition)} * {k_partitions} != {len(seqlen_list)}" return partitions -def get_seqlen_balanced_partitions(seqlen_list: list[int], k_partitions: int, equal_size: bool): +def get_seqlen_balanced_partitions( + seqlen_list: list[int], k_partitions: int, equal_size: bool +): """ Calculates partitions of indices from seqlen_list such that the sum of sequence lengths in each partition is balanced. Uses the Karmarkar-Karp differencing method. @@ -171,10 +186,13 @@ def get_seqlen_balanced_partitions(seqlen_list: list[int], k_partitions: int, eq AssertionError: If equal_size is True and len(seqlen_list) is not divisible by k_partitions. AssertionError: If any resulting partition is empty. """ - assert len(seqlen_list) >= k_partitions, f"number of items:[{len(seqlen_list)}] < k_partitions:[{k_partitions}]" + 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}" + assert len(partitions) == k_partitions, f"{ + len(partitions)} != {k_partitions}" seen_idx = set() sorted_partitions = [None] * k_partitions for i, partition in enumerate(partitions): @@ -185,11 +203,15 @@ def _check_and_sort_partitions(partitions): 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) + partitions = karmarkar_karp( + seqlen_list=seqlen_list, + k_partitions=k_partitions, + equal_size=equal_size) return _check_and_sort_partitions(partitions) -def log_seqlen_unbalance(seqlen_list: list[int], partitions: list[list[int]], prefix): +def log_seqlen_unbalance( + seqlen_list: list[int], partitions: list[list[int]], prefix): """ Calculate and log metrics related to sequence length imbalance before and after partitioning. @@ -212,7 +234,7 @@ def log_seqlen_unbalance(seqlen_list: list[int], partitions: list[list[int]], pr # Iterate over each batch of sequence lengths for offset in range(0, len(seqlen_list), batch_size): - cur_sum_seqlen = sum(seqlen_list[offset : offset + batch_size]) + 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: @@ -270,34 +292,44 @@ def rearrange_micro_batches( """ # 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=}" - ) + 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() # NOTE: num_microbatches <= batch_size, so take the min of this two. - num_micro_batches = min(len(seq_len_effective), ceildiv(total_seqlen, max_token_len)) + num_micro_batches = min( + len(seq_len_effective), ceildiv(total_seqlen, max_token_len) + ) if min_num_micro_batch is not None: # used to support pp num_micro_batches = max(min_num_micro_batch, num_micro_batches) if dist.is_initialized() and same_micro_num_in_dp: - num_micro_batches = torch.tensor([num_micro_batches], device=get_device_name()) - dist.all_reduce(num_micro_batches, op=dist.ReduceOp.MAX, group=dp_group) + num_micro_batches = torch.tensor( + [num_micro_batches], device=get_device_name()) + dist.all_reduce( + num_micro_batches, + op=dist.ReduceOp.MAX, + group=dp_group) num_micro_batches = num_micro_batches.cpu().item() if num_batches_divided_by is not None: - num_micro_batches = roundup_divisible(num_micro_batches, num_batches_divided_by) + num_micro_batches = roundup_divisible( + num_micro_batches, num_batches_divided_by) 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_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.append(batch[idx: idx + 1]) curr_micro_batch = torch.cat(curr_micro_batch) micro_batches.append(curr_micro_batch) diff --git a/Agent0/executor_train/verl/verl/utils/tokenizer.py b/Agent0/executor_train/verl/verl/utils/tokenizer.py index 668ea3e..039099b 100644 --- a/Agent0/executor_train/verl/verl/utils/tokenizer.py +++ b/Agent0/executor_train/verl/verl/utils/tokenizer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,13 +27,22 @@ def set_pad_token_id(tokenizer): """ if tokenizer.pad_token_id is None: tokenizer.pad_token_id = tokenizer.eos_token_id - warnings.warn(f"tokenizer.pad_token_id is None. Now set to {tokenizer.eos_token_id}", stacklevel=1) + warnings.warn( + f"tokenizer.pad_token_id is None. Now set to { + tokenizer.eos_token_id}", stacklevel=1, ) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token - warnings.warn(f"tokenizer.pad_token is None. Now set to {tokenizer.eos_token}", stacklevel=1) + warnings.warn( + f"tokenizer.pad_token is None. Now set to {tokenizer.eos_token}", + stacklevel=1, + ) -def hf_tokenizer(name_or_path, correct_pad_token=True, correct_gemma2=True, **kwargs): +def hf_tokenizer( + name_or_path, + correct_pad_token=True, + correct_gemma2=True, + **kwargs): """Create a huggingface pretrained tokenizer which correctness handles eos and pad tokens. Args: @@ -49,11 +58,16 @@ def hf_tokenizer(name_or_path, correct_pad_token=True, correct_gemma2=True, **kw """ from transformers import AutoTokenizer - if correct_gemma2 and isinstance(name_or_path, str) and "gemma-2-2b-it" in name_or_path: + if ( + correct_gemma2 + and isinstance(name_or_path, str) + and "gemma-2-2b-it" in name_or_path + ): # the EOS token in gemma2 is ambiguious, which may worsen RL performance. # https://huggingface.co/google/gemma-2-2b-it/commit/17a01657f5c87135bcdd0ec7abb4b2dece04408a warnings.warn( - "Found gemma-2-2b-it tokenizer. Set eos_token and eos_token_id to and 107.", stacklevel=1 + "Found gemma-2-2b-it tokenizer. Set eos_token and eos_token_id to and 107.", + stacklevel=1, ) kwargs["eos_token"] = "" kwargs["eos_token_id"] = 107 @@ -80,7 +94,10 @@ def hf_processor(name_or_path, **kwargs): processor = None # TODO(haibin.lin): try-catch should be removed after adding transformer version req to setup.py to avoid # silent failure - warnings.warn(f"Failed to create processor: {e}. This may affect multimodal processing", stacklevel=1) + warnings.warn( + f"Failed to create processor: {e}. This may affect multimodal processing", + stacklevel=1, + ) # Avoid load tokenizer, see: # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/auto/processing_auto.py#L344 if processor is not None and "Processor" not in processor.__class__.__name__: diff --git a/Agent0/executor_train/verl/verl/utils/torch_dtypes.py b/Agent0/executor_train/verl/verl/utils/torch_dtypes.py index f2f445c..0f7e870 100644 --- a/Agent0/executor_train/verl/verl/utils/torch_dtypes.py +++ b/Agent0/executor_train/verl/verl/utils/torch_dtypes.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/utils/torch_functional.py b/Agent0/executor_train/verl/verl/utils/torch_functional.py index df91ad7..f3aa96e 100644 --- a/Agent0/executor_train/verl/verl/utils/torch_functional.py +++ b/Agent0/executor_train/verl/verl/utils/torch_functional.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -41,7 +41,8 @@ try: import torch_npu - NPU_CROSS_ENTROPY_LOSS_AVAILABLE = hasattr(torch_npu, "npu_cross_entropy_loss") + NPU_CROSS_ENTROPY_LOSS_AVAILABLE = hasattr( + torch_npu, "npu_cross_entropy_loss") except ImportError: NPU_CROSS_ENTROPY_LOSS_AVAILABLE = False @@ -83,7 +84,9 @@ def logprobs_from_logits(logits, labels, inplace_backward=True): last_dim = logits.shape[-1] logits = logits.reshape(-1, last_dim) labels = labels.reshape(-1) - output = logprobs_from_logits_flash_attn(logits, labels, inplace_backward=inplace_backward) + output = logprobs_from_logits_flash_attn( + logits, labels, inplace_backward=inplace_backward + ) output = output.view(*batch_dim) elif NPU_CROSS_ENTROPY_LOSS_AVAILABLE: output = logprobs_from_logits_torch_npu(logits, labels) @@ -93,17 +96,20 @@ def logprobs_from_logits(logits, labels, inplace_backward=True): def logprobs_from_logits_flash_attn(logits, labels, inplace_backward=True): - output = cross_entropy_loss(logits, labels, inplace_backward=inplace_backward) - assert isinstance(output, tuple), ( - "please make sure flash-attn>=2.4.3 where cross_entropy_loss returns Tuple[losses, z_losses]." - ) + output = cross_entropy_loss( + logits, labels, inplace_backward=inplace_backward) + assert isinstance( + output, tuple + ), "please make sure flash-attn>=2.4.3 where cross_entropy_loss returns Tuple[losses, z_losses]." return -output[0] def logprobs_from_logits_torch_npu(logits, labels): batch_dim = logits.shape[:-1] logits = logits.reshape(-1, logits.shape[-1]) - loss, _, _, _ = torch_npu.npu_cross_entropy_loss(logits, labels.reshape(-1), reduction="none") + loss, _, _, _ = torch_npu.npu_cross_entropy_loss( + logits, labels.reshape(-1), reduction="none" + ) return -loss.view(*batch_dim) @@ -118,16 +124,27 @@ def logprobs_from_logits_v2(logits: torch.FloatTensor, labels): A memory efficient implementation of logprobs_from_logits """ if logits.dtype in [torch.float32, torch.float64]: - logits_labels = torch.gather(logits, dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) + logits_labels = torch.gather( + logits, dim=-1, index=labels.unsqueeze(-1) + ).squeeze(-1) # loop to reduce peak mem consumption - logsumexp_values = torch.stack([torch.logsumexp(logit, dim=-1) for logit in logits]) - logprobs_labels = logits_labels - logsumexp_values # log_softmax(x_i) = x_i - logsumexp(x) + logsumexp_values = torch.stack( + [torch.logsumexp(logit, dim=-1) for logit in logits] + ) + logprobs_labels = ( + logits_labels - logsumexp_values + ) # log_softmax(x_i) = x_i - logsumexp(x) else: - # logsumexp approach is unstable with bfloat16, fall back to slightly less efficent approach + # logsumexp approach is unstable with bfloat16, fall back to slightly + # less efficent approach logprobs_labels = [] - for row_logits, row_labels in zip(logits, labels, strict=True): # loop to reduce peak mem consumption + for row_logits, row_labels in zip( + logits, labels, strict=True + ): # loop to reduce peak mem consumption row_logprobs = F.log_softmax(row_logits, dim=-1) - row_logprobs_labels = row_logprobs.gather(dim=-1, index=row_labels.unsqueeze(-1)).squeeze(-1) + row_logprobs_labels = row_logprobs.gather( + dim=-1, index=row_labels.unsqueeze(-1) + ).squeeze(-1) logprobs_labels.append(row_logprobs_labels) logprobs_labels = torch.stack(logprobs_labels) return logprobs_labels @@ -149,14 +166,18 @@ def entropy_from_logits(logits: torch.Tensor): return entropy -def entropy_from_logits_with_chunking(logits: torch.Tensor, chunk_size: int = 2048): +def entropy_from_logits_with_chunking( + logits: torch.Tensor, + chunk_size: int = 2048): """Memory-efficient entropy calculation with chunking.""" entropy = torch.zeros(logits.shape[0], device=logits.device) for i in range(0, logits.shape[0], chunk_size): - logits_chunk = logits[i : i + chunk_size].float() + logits_chunk = logits[i: i + chunk_size].float() pd_chunk = torch.nn.functional.softmax(logits_chunk, dim=-1) - entropy_chunk = torch.logsumexp(logits_chunk, dim=-1) - torch.sum(pd_chunk * logits_chunk, dim=-1) - entropy[i : i + chunk_size] = entropy_chunk + entropy_chunk = torch.logsumexp(logits_chunk, dim=-1) - torch.sum( + pd_chunk * logits_chunk, dim=-1 + ) + entropy[i: i + chunk_size] = entropy_chunk return entropy @@ -197,7 +218,9 @@ def masked_var(values, mask, unbiased=True): # 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.") + 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 @@ -223,7 +246,10 @@ def masked_whiten(values, mask, shift_mean=True): return whitened -def get_response_mask(response_id: torch.Tensor, eos_token: int | list[int] = 2, dtype=torch.int64): +def get_response_mask( + response_id: torch.Tensor, + eos_token: int | list[int] = 2, + dtype=torch.int64): """ end of sentence token can be int or list: 1 or [1, 2] e.g. @@ -242,7 +268,9 @@ def get_response_mask(response_id: torch.Tensor, eos_token: int | list[int] = 2, [1, 1, 1, 0, 0, 0, 0], [1, 1, 1, 1, 1, 0, 0]]) """ - eos_mask = torch.isin(response_id, torch.tensor(eos_token, device=response_id.device)).int() + eos_mask = torch.isin( + response_id, torch.tensor(eos_token, device=response_id.device) + ).int() return (eos_mask.cumsum(dim=1) - eos_mask).eq(0).to(dtype) @@ -250,20 +278,25 @@ def compute_grad_norm(model: nn.Module): total_grad_square = 0 for param in model.parameters(): if param.grad is not None: - total_grad_square += torch.sum(torch.square(param.grad.detach())).item() + total_grad_square += torch.sum( + torch.square(param.grad.detach())).item() return total_grad_square -def broadcast_dict_tensor(tensors: dict[str, torch.Tensor] | TensorDict, src, group): +def broadcast_dict_tensor( + tensors: 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) + torch.distributed.broadcast( + tensors[key], src=src, group=group, async_op=False) -def allgather_dict_tensors(tensors: dict[str, torch.Tensor] | TensorDict, size, group, dim=0): +def allgather_dict_tensors( + tensors: dict[str, torch.Tensor] | TensorDict, size, group, dim=0 +): """ TODO: optimize this. - We can use async ops @@ -288,19 +321,24 @@ def allgather_dict_tensors(tensors: dict[str, torch.Tensor] | TensorDict, size, 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) + 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) + output = TensorDict(source=output, + batch_size=tensors.batch_size[0] * size) return output -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}" - ) +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) @@ -309,8 +347,15 @@ def pad_2d_list_to_length(response, pad_token_id, max_length=None): pad a 2D list (e.g. responses, logprobs) to a 2D tensor. """ response_length = max(len(sub_list) for sub_list in response) - target_length = max_length if max_length is not None and max_length > response_length else response_length - padded_response = [tuple(sub_list) + (pad_token_id,) * (target_length - len(sub_list)) for sub_list in response] + target_length = ( + max_length + if max_length is not None and max_length > response_length + else response_length + ) + padded_response = [ + tuple(sub_list) + (pad_token_id,) * (target_length - len(sub_list)) + for sub_list in response + ] tensor = torch.tensor(padded_response) return tensor @@ -324,7 +369,11 @@ def pad_sequence_to_length(tensors, max_seq_len, pad_token_id, left_pad=False): if tensors.shape[-1] >= max_seq_len: return tensors # (0, max_seq_len - tensors.shape[-1]) means right pad to max_seq_length and no left pad - pad_tuple = (max_seq_len - tensors.shape[-1], 0) if left_pad else (0, max_seq_len - tensors.shape[-1]) + 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) @@ -355,11 +404,16 @@ def postprocess_data( 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 + 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 - ) + 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 @@ -371,18 +425,30 @@ def postprocess_data( elif truncation == "middle": left_half = max_length // 2 right_half = max_length - left_half - input_ids = torch.cat([input_ids[:, :left_half], input_ids[:, -right_half:]], dim=-1) - attention_mask = torch.cat([attention_mask[:, :left_half], attention_mask[:, -right_half:]], dim=-1) + input_ids = torch.cat( + [input_ids[:, :left_half], input_ids[:, -right_half:]], dim=-1 + ) + attention_mask = torch.cat( + [attention_mask[:, :left_half], attention_mask[:, -right_half:]], dim=-1 + ) elif truncation == "error": - raise NotImplementedError(f"{sequence_length=} is larger than {max_length=}") + raise NotImplementedError( + f"{sequence_length=} is larger than {max_length=}" + ) else: - raise NotImplementedError(f"Unknown truncation method {truncation}") + raise NotImplementedError( + f"Unknown truncation method {truncation}") return input_ids, attention_mask def tokenize_and_postprocess_data( - prompt: str, tokenizer: PreTrainedTokenizer, max_length: int, pad_token_id: int, left_pad=True, truncation="error" + prompt: str, + tokenizer: PreTrainedTokenizer, + max_length: int, + pad_token_id: int, + left_pad=True, + truncation="error", ): """Tokenize text and process outputs to consistent tensor shapes. @@ -397,11 +463,20 @@ def tokenize_and_postprocess_data( Returns: Tuple of (input_ids, attention_mask) from postprocess_data """ - input_data = tokenizer(prompt, return_tensors="pt", add_special_tokens=False) + input_data = tokenizer( + prompt, + return_tensors="pt", + add_special_tokens=False) input_ids = input_data["input_ids"] attention_mask = input_data["attention_mask"] - return postprocess_data(input_ids, attention_mask, max_length, pad_token_id, left_pad, truncation) + return postprocess_data( + input_ids, + attention_mask, + max_length, + pad_token_id, + left_pad, + truncation) def remove_pad_token(input_ids: torch.Tensor, attention_mask: torch.Tensor): @@ -415,7 +490,8 @@ def remove_pad_token(input_ids: torch.Tensor, attention_mask: torch.Tensor): """ no_padding_batch = [] for ids, mask in zip(input_ids, attention_mask, strict=True): - no_padding_batch.append((ids[len(ids) - mask.sum() :]).cpu().numpy().tolist()) + no_padding_batch.append( + (ids[len(ids) - mask.sum():]).cpu().numpy().tolist()) return no_padding_batch @@ -429,13 +505,16 @@ def log_probs_from_logits_response(input_ids, logits, response_length): Returns: response_log_prob: """ - response_logits = logits[:, -response_length - 1 : -1] + response_logits = logits[:, -response_length - 1: -1] response = input_ids[:, -response_length:] - response_log_prob = logprobs_from_logits(logits=response_logits, labels=response) + response_log_prob = logprobs_from_logits( + logits=response_logits, labels=response) return response_log_prob -def log_probs_from_logits_response_rmpad(input_ids, attention_mask, logits_rmpad, response_length): +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. @@ -451,18 +530,29 @@ def log_probs_from_logits_response_rmpad(input_ids, attention_mask, logits_rmpad 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, 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_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 + 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] + output = full_output.squeeze(-1)[ + :, -response_length - 1: -1 + ] # [batch_size, response_length] return output -def log_probs_from_logits_all_rmpad(input_ids_rmpad, logits_rmpad, indices, batch_size, seqlen, response_length): +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. @@ -479,14 +569,23 @@ def log_probs_from_logits_all_rmpad(input_ids_rmpad, logits_rmpad, indices, batc """ 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.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_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 + 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] + output = full_output.squeeze(-1)[ + :, -response_length - 1: -1 + ] # [batch_size, response_length] return output @@ -542,8 +641,12 @@ def get_cosine_schedule_with_warmup( def lr_lambda(current_step): if current_step < num_warmup_steps: - return min_lr_ratio + (1.0 - min_lr_ratio) * (float(current_step) / float(max(1, num_warmup_steps))) - progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps)) + return min_lr_ratio + (1.0 - min_lr_ratio) * ( + float(current_step) / float(max(1, num_warmup_steps)) + ) + progress = float(current_step - num_warmup_steps) / float( + max(1, num_training_steps - num_warmup_steps) + ) x = math.cos(math.pi * float(num_cycles) * 2.0 * progress) return max(min_lr_ratio, x * coef + intercept) @@ -588,23 +691,31 @@ def prepare_decoder_attention_mask(attention_mask, input_shape, inputs_embeds): 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 - ) + 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 + expanded_attn_mask + if combined_attention_mask is None + else expanded_attn_mask + combined_attention_mask ) return combined_attention_mask # Copied from transformers.models.bart.modeling_bart._make_causal_mask -def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device): +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 = 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) @@ -612,25 +723,37 @@ def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: t # Copied from transformers.models.bart.modeling_bart._expand_mask -def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None): +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) + 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) + return inverted_mask.masked_fill( + inverted_mask.to(torch.bool), torch.finfo(dtype).min + ) def get_unpad_data(attention_mask): seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() max_seqlen_in_batch = seqlens_in_batch.max().item() - cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + cu_seqlens = F.pad( + torch.cumsum( + seqlens_in_batch, + dim=0, + dtype=torch.int32), + (1, + 0)) return ( indices, cu_seqlens, @@ -685,8 +808,13 @@ def lr_lambda(current_step): if current_step < num_warmup_steps + num_stable_steps: return 1.0 if current_step < num_training_steps: - progress = float(current_step - num_warmup_steps - num_stable_steps) / float(max(1, num_decay_steps)) - value = max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress))) + progress = float( + current_step - num_warmup_steps - num_stable_steps + ) / float(max(1, num_decay_steps)) + value = max( + 0.0, + 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)), + ) return (1.0 - min_lr_ratio) * value + min_lr_ratio return min_lr_ratio @@ -701,12 +829,16 @@ def check_device_is_available(): This context manager checks if CUDA is available and raises an error if it is not. """ if not get_torch_device().is_available(): - raise RuntimeError("Device {} must be initialized before importing this module.".format(get_device_name())) + raise RuntimeError( + "Device {} must be initialized before importing this module.".format( + get_device_name())) yield -def distributed_mean_max_min_std(local_tensor, compute_max=True, compute_min=True, compute_std=True): +def distributed_mean_max_min_std( + local_tensor, compute_max=True, compute_min=True, compute_std=True +): """Compute distributed statistics across all processes. Args: @@ -720,7 +852,9 @@ def distributed_mean_max_min_std(local_tensor, compute_max=True, compute_min=Tru """ # Sum the local tensor across all processes local_sum = torch.sum(local_tensor) - local_num = torch.tensor(torch.numel(local_tensor), device=get_device_name()) + local_num = torch.tensor( + torch.numel(local_tensor), + device=get_device_name()) torch.distributed.all_reduce(local_sum, op=torch.distributed.ReduceOp.SUM) torch.distributed.all_reduce(local_num, op=torch.distributed.ReduceOp.SUM) @@ -729,19 +863,22 @@ def distributed_mean_max_min_std(local_tensor, compute_max=True, compute_min=Tru if compute_max: local_max = torch.max(local_tensor) - torch.distributed.all_reduce(local_max, op=torch.distributed.ReduceOp.MAX) + torch.distributed.all_reduce( + local_max, op=torch.distributed.ReduceOp.MAX) else: local_max = None if compute_min: local_min = torch.min(local_tensor) - torch.distributed.all_reduce(local_min, op=torch.distributed.ReduceOp.MIN) + torch.distributed.all_reduce( + local_min, op=torch.distributed.ReduceOp.MIN) else: local_min = None if compute_std: square_diff = torch.sum(torch.pow(local_tensor - global_mean, 2)) - torch.distributed.all_reduce(square_diff, op=torch.distributed.ReduceOp.SUM) + torch.distributed.all_reduce( + square_diff, op=torch.distributed.ReduceOp.SUM) global_std = torch.sqrt(square_diff / (local_num - 1)) else: global_std = None diff --git a/Agent0/executor_train/verl/verl/utils/tracking.py b/Agent0/executor_train/verl/verl/utils/tracking.py index 07f45a3..76f3570 100644 --- a/Agent0/executor_train/verl/verl/utils/tracking.py +++ b/Agent0/executor_train/verl/verl/utils/tracking.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -34,16 +34,34 @@ class Tracking: logger: Dictionary of initialized logger instances for each backend. """ - supported_backend = ["wandb", "mlflow", "swanlab", "vemlp_wandb", "tensorboard", "console", "clearml"] - - def __init__(self, project_name, experiment_name, default_backend: str | list[str] = "console", config=None): + supported_backend = [ + "wandb", + "mlflow", + "swanlab", + "vemlp_wandb", + "tensorboard", + "console", + "clearml", + ] + + def __init__( + self, + project_name, + experiment_name, + default_backend: str | list[str] = "console", + config=None, + ): if isinstance(default_backend, str): default_backend = [default_backend] for backend in default_backend: if backend == "tracking": import warnings - warnings.warn("`tracking` logger is deprecated. use `wandb` instead.", DeprecationWarning, stacklevel=2) + warnings.warn( + "`tracking` logger is deprecated. use `wandb` instead.", + DeprecationWarning, + stacklevel=2, + ) else: assert backend in self.supported_backend, f"{backend} is not supported" @@ -54,8 +72,14 @@ def __init__(self, project_name, experiment_name, default_backend: str | list[st settings = None if config and config["trainer"].get("wandb_proxy", None): - settings = wandb.Settings(https_proxy=config["trainer"]["wandb_proxy"]) - wandb.init(project=project_name, name=experiment_name, config=config, settings=settings) + settings = wandb.Settings( + https_proxy=config["trainer"]["wandb_proxy"]) + wandb.init( + project=project_name, + name=experiment_name, + config=config, + settings=settings, + ) self.logger["wandb"] = wandb if "mlflow" in default_backend: @@ -70,7 +94,9 @@ def __init__(self, project_name, experiment_name, default_backend: str | list[st # Project_name is actually experiment_name in MLFlow # If experiment does not exist, will create a new experiment experiment = mlflow.set_experiment(project_name) - mlflow.start_run(experiment_id=experiment.experiment_id, run_name=experiment_name) + mlflow.start_run( + experiment_id=experiment.experiment_id, + run_name=experiment_name) mlflow.log_params(_compute_mlflow_params_from_objects(config)) self.logger["mlflow"] = _MlflowLoggingAdapter() @@ -83,10 +109,13 @@ def __init__(self, project_name, experiment_name, default_backend: str | list[st SWANLAB_LOG_DIR = os.environ.get("SWANLAB_LOG_DIR", "swanlog") SWANLAB_MODE = os.environ.get("SWANLAB_MODE", "cloud") if SWANLAB_API_KEY: - swanlab.login(SWANLAB_API_KEY) # NOTE: previous login information will be overwritten + swanlab.login( + SWANLAB_API_KEY + ) # NOTE: previous login information will be overwritten if config is None: - config = {} # make sure config is not None, otherwise **config will raise error + config = ( + {}) # make sure config is not None, otherwise **config will raise error swanlab.init( project=project_name, experiment_name=experiment_name, @@ -117,7 +146,9 @@ def __init__(self, project_name, experiment_name, default_backend: str | list[st self.logger["vemlp_wandb"] = vemlp_wandb if "tensorboard" in default_backend: - self.logger["tensorboard"] = _TensorboardAdapter(project_name, experiment_name) + self.logger["tensorboard"] = _TensorboardAdapter( + project_name, experiment_name + ) if "console" in default_backend: from verl.utils.logger import LocalLogger @@ -126,7 +157,9 @@ def __init__(self, project_name, experiment_name, default_backend: str | list[st self.logger["console"] = self.console_logger if "clearml" in default_backend: - self.logger["clearml"] = ClearMLLogger(project_name, experiment_name, config) + self.logger["clearml"] = ClearMLLogger( + project_name, experiment_name, config + ) def log(self, data, step, backend=None): for default_backend, logger_instance in self.logger.items(): @@ -205,7 +238,9 @@ def __init__(self, project_name, experiment_name): from torch.utils.tensorboard import SummaryWriter - tensorboard_dir = os.environ.get("TENSORBOARD_DIR", f"tensorboard_log/{project_name}/{experiment_name}") + tensorboard_dir = os.environ.get( + "TENSORBOARD_DIR", + f"tensorboard_log/{project_name}/{experiment_name}") os.makedirs(tensorboard_dir, exist_ok=True) print(f"Saving tensorboard log to {tensorboard_dir}.") self.writer = SummaryWriter(tensorboard_dir) @@ -230,11 +265,16 @@ 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="/") + return _flatten_dict( + _transform_params_to_json_serializable( + params, convert_list_to_dict=True), sep="/", ) def _transform_params_to_json_serializable(x, convert_list_to_dict: bool): - _transform = partial(_transform_params_to_json_serializable, convert_list_to_dict=convert_list_to_dict) + _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)) @@ -242,7 +282,9 @@ def _transform_params_to_json_serializable(x, convert_list_to_dict: bool): 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)} + 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): @@ -294,7 +336,11 @@ def _log_generations_to_wandb(self, samples, step, wandb): # Create column names for all samples columns = ["step"] + sum( - [[f"input_{i + 1}", f"output_{i + 1}", f"score_{i + 1}"] for i in range(len(samples))], [] + [ + [f"input_{i + 1}", f"output_{i + 1}", f"score_{i + 1}"] + for i in range(len(samples)) + ], + [], ) if not hasattr(self, "validation_table"): @@ -302,8 +348,11 @@ def _log_generations_to_wandb(self, samples, step, wandb): self.validation_table = wandb.Table(columns=columns) # Create a new table with same columns and existing data - # Workaround for https://github.com/wandb/wandb/issues/2981#issuecomment-1997445737 - new_table = wandb.Table(columns=columns, data=self.validation_table.data) + # Workaround for + # https://github.com/wandb/wandb/issues/2981#issuecomment-1997445737 + new_table = wandb.Table( + columns=columns, + data=self.validation_table.data) # Add new row with all data row_data = [] @@ -343,16 +392,21 @@ def log_generations_to_mlflow(self, samples, step): try: with tempfile.TemporaryDirectory() as tmp_dir: - validation_gen_step_file = Path(tmp_dir, f"val_step{step}.json") + validation_gen_step_file = Path( + tmp_dir, f"val_step{step}.json") row_data = [] for sample in samples: - data = {"input": sample[0], "output": sample[1], "score": sample[2]} + data = { + "input": sample[0], + "output": sample[1], + "score": sample[2]} row_data.append(data) with open(validation_gen_step_file, "w") as file: json.dump(row_data, file) mlflow.log_artifact(validation_gen_step_file) except Exception as e: - print(f"WARNING: save validation generation file to mlflow failed with error {e}") + print( + f"WARNING: save validation generation file to mlflow failed with error {e}") def log_generations_to_clearml(self, samples, step): """Log validation generation to clearml as table""" @@ -388,7 +442,8 @@ def log_generations_to_tensorboard(self, samples, step): if not hasattr(self, "writer"): from torch.utils.tensorboard import SummaryWriter - tensorboard_dir = os.environ.get("TENSORBOARD_DIR", "tensorboard_log") + tensorboard_dir = os.environ.get( + "TENSORBOARD_DIR", "tensorboard_log") os.makedirs(tensorboard_dir, exist_ok=True) self.writer = SummaryWriter(log_dir=tensorboard_dir) diff --git a/Agent0/executor_train/verl/verl/utils/ulysses.py b/Agent0/executor_train/verl/verl/utils/ulysses.py index b37c691..85ed586 100644 --- a/Agent0/executor_train/verl/verl/utils/ulysses.py +++ b/Agent0/executor_train/verl/verl/utils/ulysses.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -43,7 +43,8 @@ def get_ulysses_sequence_parallel_group() -> Optional[dist.ProcessGroup]: return _ULYSSES_SEQUENCE_PARALLEL_GROUP -def get_ulysses_sequence_parallel_world_size(group: ProcessGroup = None) -> int: +def get_ulysses_sequence_parallel_world_size( + group: ProcessGroup = None) -> int: """ Get ulysses sequence parallel world size. """ @@ -83,7 +84,9 @@ def gather_seq_scatter_heads( return x -def gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int, group: ProcessGroup = None) -> Tensor: +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: @@ -114,7 +117,9 @@ def _unpad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor: return x[slc] -def slice_input_tensor(x: Tensor, dim: int, padding: bool = True, group: ProcessGroup = None) -> Tensor: +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() @@ -139,9 +144,17 @@ def all_to_all_tensor( ): 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) + 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(): @@ -152,13 +165,23 @@ def wait(): return torch.cat(output_list, dim=gather_dim).contiguous() -def all_gather_tensor(local_tensor: Tensor, group: Optional[dist.ProcessGroup] = None, async_op: bool = False): +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) + 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 @@ -176,14 +199,26 @@ def forward( 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) + return all_to_all_tensor( + local_input, + scatter_dim, + gather_dim, + group, + async_op) @staticmethod - def backward(ctx: Any, *grad_output: Tensor) -> tuple[None, Tensor, None, None]: - input_t = torch.cat(grad_output[1:], dim=ctx.gather_dim).contiguous() if ctx.async_op else grad_output[0] + def backward(ctx: Any, * + grad_output: Tensor) -> tuple[None, Tensor, None, None]: + input_t = ( + torch.cat(grad_output[1:], dim=ctx.gather_dim).contiguous() + if ctx.async_op + else grad_output[0] + ) return ( None, - all_to_all_tensor(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group, False), + all_to_all_tensor( + input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group, False + ), None, None, None, @@ -226,7 +261,9 @@ def backward(ctx: Any, grad_output: Tensor) -> Any: grad_output = grad_output * ctx.sp_world_size return ( None, - grad_output.split(ctx.part_size, dim=ctx.gather_dim)[ctx.sp_rank].contiguous(), + grad_output.split(ctx.part_size, dim=ctx.gather_dim)[ + ctx.sp_rank + ].contiguous(), None, None, None, @@ -262,14 +299,20 @@ def gather_outpus_and_unpad( 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" + assert isinstance( + padding_size, int + ), "padding size is not given or is not an integer" if padding_size == 0: return x x = _unpad_tensor(x, unpad_dim, padding_size) return x -def ulysses_pad(input_ids_rmpad: torch.Tensor, position_ids_rmpad: Optional[torch.Tensor] = None, sp_size: int = 1): +def ulysses_pad( + input_ids_rmpad: torch.Tensor, + position_ids_rmpad: Optional[torch.Tensor] = None, + sp_size: int = 1, +): if position_ids_rmpad is not None: assert position_ids_rmpad.size(-2) == 1 assert input_ids_rmpad.size(-1) == position_ids_rmpad.size(-1) @@ -278,17 +321,24 @@ def ulysses_pad(input_ids_rmpad: torch.Tensor, position_ids_rmpad: Optional[torc _, 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) + 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) + pad_pos_ids = torch.arange( + pad_size, device=position_ids_rmpad.device + ).unsqueeze(0) if position_ids_rmpad.dim() == 3: pad_pos_ids = pad_pos_ids.unsqueeze(0).repeat(3, 1, 1) - position_ids_rmpad = torch.cat((position_ids_rmpad, pad_pos_ids), dim=-1) + position_ids_rmpad = torch.cat( + (position_ids_rmpad, pad_pos_ids), dim=-1) return input_ids_rmpad, position_ids_rmpad, pad_size def ulysses_pad_and_slice_inputs( - input_ids_rmpad: torch.Tensor, position_ids_rmpad: Optional[torch.Tensor] = None, sp_size: int = 1 + 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 @@ -308,15 +358,19 @@ def ulysses_pad_and_slice_inputs( torch.Tensor: padded and sliced position_ids int: pad size """ - input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad(input_ids_rmpad, position_ids_rmpad, sp_size) + input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad( + input_ids_rmpad, position_ids_rmpad, sp_size + ) input_ids_rmpad = slice_input_tensor(input_ids_rmpad, dim=1, padding=False) if position_ids_rmpad is not None: - position_ids_rmpad = slice_input_tensor(position_ids_rmpad, dim=1, padding=False) + position_ids_rmpad = slice_input_tensor( + position_ids_rmpad, dim=1, padding=False + ) return input_ids_rmpad, position_ids_rmpad, pad_size def validate_ulysses_config(num_heads, ulysses_sequence_size): if ulysses_sequence_size > 1: - assert num_heads % ulysses_sequence_size == 0, ( - f"num_heads ({num_heads}) must be divisible by ulysses sequence size({ulysses_sequence_size})" - ) + assert ( + num_heads % ulysses_sequence_size == 0 + ), f"num_heads ({num_heads}) must be divisible by ulysses sequence size({ulysses_sequence_size})" diff --git a/Agent0/executor_train/verl/verl/utils/vllm_utils.py b/Agent0/executor_train/verl/verl/utils/vllm_utils.py index 25ee665..8b10b3e 100644 --- a/Agent0/executor_train/verl/verl/utils/vllm_utils.py +++ b/Agent0/executor_train/verl/verl/utils/vllm_utils.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,7 +27,10 @@ SUPPORTED_MOE_MODELS = [] try: - from vllm.model_executor.models.deepseek_v2 import DeepseekV2ForCausalLM, DeepseekV3ForCausalLM + from vllm.model_executor.models.deepseek_v2 import ( + DeepseekV2ForCausalLM, + DeepseekV3ForCausalLM, + ) SUPPORTED_MOE_MODELS.append(DeepseekV2ForCausalLM) SUPPORTED_MOE_MODELS.append(DeepseekV3ForCausalLM) @@ -90,9 +93,17 @@ def patch_vllm_moe_model_weight_loader(model): if not isinstance(model, tuple(SUPPORTED_MOE_MODELS)): return - model = getattr(model, "model", None) or getattr(model, "language_model", None) + model = getattr( + model, + "model", + None) or getattr( + model, + "language_model", + None) if model is None: - raise ValueError("The provided model does not have a valid 'model' or 'language_model' attribute.") + raise ValueError( + "The provided model does not have a valid 'model' or 'language_model' attribute." + ) for layer in model.layers: mlp_attr = MLP_ATTR_MAPPING.get(type(model), DEFAULT_MLP_ATTR) @@ -112,7 +123,8 @@ class TensorLoRARequest(LoRARequest): class VLLMHijack: @staticmethod def hijack(): - def hijack__load_adapter(self, lora_request: TensorLoRARequest) -> LoRAModel: + def hijack__load_adapter( + self, lora_request: TensorLoRARequest) -> LoRAModel: """ based on vllm.lora.worker_manager.WorkerLoRAManager._load_adapter, support load adapter with lora tensors @@ -127,7 +139,8 @@ def hijack__load_adapter(self, lora_request: TensorLoRARequest) -> LoRAModel: expected_lora_modules: list[str] = [] for module in supported_lora_modules: if module in packed_modules_mapping: - expected_lora_modules.extend(packed_modules_mapping[module]) + expected_lora_modules.extend( + packed_modules_mapping[module]) else: expected_lora_modules.append(module) @@ -141,9 +154,12 @@ def hijack__load_adapter(self, lora_request: TensorLoRARequest) -> LoRAModel: lora_tensors = lora_request.lora_tensors peft_helper = PEFTHelper.from_dict(peft_config) else: - lora_path = get_adapter_absolute_path(lora_request.lora_path) + lora_path = get_adapter_absolute_path( + lora_request.lora_path) - peft_helper = PEFTHelper.from_local_dir(lora_path, self.max_position_embeddings) + peft_helper = PEFTHelper.from_local_dir( + lora_path, self.max_position_embeddings + ) # Validates the LoRA configuration against requirements before # loading weights, throwing an exception if validation fails. @@ -153,7 +169,10 @@ def hijack__load_adapter(self, lora_request: TensorLoRARequest) -> LoRAModel: # to ensure correct loading of lora weights. model = self._adapter_manager.model hf_to_vllm_mapper = None - if hasattr(model, "hf_to_vllm_mapper") and model.hf_to_vllm_mapper is not None: + if ( + hasattr(model, "hf_to_vllm_mapper") + and model.hf_to_vllm_mapper is not None + ): hf_to_vllm_mapper = model.hf_to_vllm_mapper if isinstance(lora_request, TensorLoRARequest): @@ -164,7 +183,8 @@ def hijack__load_adapter(self, lora_request: TensorLoRARequest) -> LoRAModel: device="cpu", dtype=self.lora_config.lora_dtype, embeddings=None, - target_embedding_padding=self.vocab_size + self.lora_config.lora_extra_vocab_size, + target_embedding_padding=self.vocab_size + + self.lora_config.lora_extra_vocab_size, embedding_modules=self.embedding_modules, embedding_padding_modules=self.embedding_padding_modules, weights_mapper=hf_to_vllm_mapper, @@ -177,7 +197,8 @@ def hijack__load_adapter(self, lora_request: TensorLoRARequest) -> LoRAModel: lora_model_id=lora_request.lora_int_id, device="cpu", dtype=self.lora_config.lora_dtype, - target_embedding_padding=self.vocab_size + self.lora_config.lora_extra_vocab_size, + target_embedding_padding=self.vocab_size + + self.lora_config.lora_extra_vocab_size, embedding_modules=self.embedding_modules, embedding_padding_modules=self.embedding_padding_modules, weights_mapper=hf_to_vllm_mapper, @@ -187,15 +208,18 @@ def hijack__load_adapter(self, lora_request: TensorLoRARequest) -> LoRAModel: if lora.extra_vocab_size > self.lora_config.lora_extra_vocab_size: raise ValueError( - f"LoRA added vocab size {lora.extra_vocab_size} is greater than lora_extra_vocab_size " - f"{self.lora_config.lora_extra_vocab_size}." - ) + f"LoRA added vocab size { + lora.extra_vocab_size} is greater than lora_extra_vocab_size " f"{ + self.lora_config.lora_extra_vocab_size}.") return lora def do_hijack(target_cls, target_method_name, hooking_method): setattr(target_cls, target_method_name, hooking_method) - do_hijack(LRUCacheWorkerLoRAManager, "_load_adapter", hijack__load_adapter) + do_hijack( + LRUCacheWorkerLoRAManager, + "_load_adapter", + hijack__load_adapter) def is_version_ge(pkg: str = "vllm", minver: str = "0.7.3"): diff --git a/Agent0/executor_train/verl/verl/workers/__init__.py b/Agent0/executor_train/verl/verl/workers/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/workers/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/actor/__init__.py b/Agent0/executor_train/verl/verl/workers/actor/__init__.py index 7a1404e..f71ffa7 100644 --- a/Agent0/executor_train/verl/verl/workers/actor/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/actor/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/actor/base.py b/Agent0/executor_train/verl/verl/workers/actor/base.py index 2d1ba29..e6399a7 100644 --- a/Agent0/executor_train/verl/verl/workers/actor/base.py +++ b/Agent0/executor_train/verl/verl/workers/actor/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/actor/dp_actor.py b/Agent0/executor_train/verl/verl/workers/actor/dp_actor.py index f18bf6b..a26a807 100644 --- a/Agent0/executor_train/verl/verl/workers/actor/dp_actor.py +++ b/Agent0/executor_train/verl/verl/workers/actor/dp_actor.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -27,20 +27,44 @@ import verl.utils.torch_functional as verl_F from verl import DataProto -from verl.trainer.ppo.core_algos import agg_loss, compute_policy_loss, get_policy_loss_fn, kl_penalty -from verl.utils.device import get_device_id, get_device_name, is_cuda_available, is_npu_available +from verl.trainer.ppo.core_algos import ( + agg_loss, + compute_policy_loss, + get_policy_loss_fn, + kl_penalty, +) +from verl.utils.device import ( + get_device_id, + get_device_name, + is_cuda_available, + is_npu_available, +) from verl.utils.fsdp_utils import FSDPModule, fsdp2_clip_grad_norm_ from verl.utils.profiler import GPUMemoryLogger from verl.utils.py_functional import append_to_dict from verl.utils.seqlen_balancing import get_reverse_idx, rearrange_micro_batches from verl.utils.torch_functional import logprobs_from_logits -from verl.utils.ulysses import gather_outpus_and_unpad, ulysses_pad, ulysses_pad_and_slice_inputs +from verl.utils.ulysses import ( + gather_outpus_and_unpad, + ulysses_pad, + ulysses_pad_and_slice_inputs, +) from verl.workers.actor import BasePPOActor if is_cuda_available: - from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input + from flash_attn.bert_padding import ( + index_first_axis, + pad_input, + rearrange, + unpad_input, + ) elif is_npu_available: - from transformers.integrations.npu_flash_attention import index_first_axis, pad_input, rearrange, unpad_input + from transformers.integrations.npu_flash_attention import ( + index_first_axis, + pad_input, + rearrange, + unpad_input, + ) __all__ = ["DataParallelPPOActor"] @@ -50,7 +74,12 @@ class DataParallelPPOActor(BasePPOActor): - def __init__(self, config, actor_module: nn.Module, actor_optimizer: torch.optim.Optimizer = None): + 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 @@ -73,7 +102,9 @@ def __init__(self, config, actor_module: nn.Module, actor_optimizer: torch.optim self.compute_entropy_from_logits = ( torch.compile(entropy_from_logits, dynamic=True) - if self.config.get("use_torch_compile", True) # use torch compile by default + if self.config.get( + "use_torch_compile", True + ) # use torch compile by default else entropy_from_logits ) self.device_name = get_device_name() @@ -89,13 +120,16 @@ def _forward_micro_batch( response_length = micro_batch["responses"].size(-1) multi_modal_inputs = {} if "multi_modal_inputs" in micro_batch.keys(): - if "image_bound" in micro_batch["multi_modal_inputs"][0]: # minicpm-o logic + # minicpm-o logic + if "image_bound" in micro_batch["multi_modal_inputs"][0]: for key in micro_batch["multi_modal_inputs"][0].keys(): - multi_modal_inputs[key] = [inputs[key] for inputs in micro_batch["multi_modal_inputs"]] + multi_modal_inputs[key] = [inputs[key] + for inputs in micro_batch["multi_modal_inputs"]] else: for key in micro_batch["multi_modal_inputs"][0].keys(): multi_modal_inputs[key] = torch.cat( - [inputs[key] for inputs in micro_batch["multi_modal_inputs"]], dim=0 + [inputs[key] for inputs in micro_batch["multi_modal_inputs"]], + dim=0, ) with torch.autocast(device_type=self.device_name, dtype=torch.bfloat16): @@ -105,35 +139,45 @@ def _forward_micro_batch( position_ids = micro_batch["position_ids"] entropy = None if position_ids.dim() == 3: # qwen2vl mrope - position_ids = position_ids.transpose(0, 1) # (bsz, 3, seqlen) -> (3, bsz, seqlen) + position_ids = position_ids.transpose( + 0, 1 + ) # (bsz, 3, seqlen) -> (3, bsz, seqlen) if self.use_remove_padding: input_ids_rmpad, indices, cu_seqlens, *_ = 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) + input_ids_rmpad = input_ids_rmpad.transpose( + 0, 1) # (1, total_nnz) # unpad the position_ids to align the rotary if position_ids.dim() == 3: position_ids_rmpad = ( - index_first_axis(rearrange(position_ids, "c b s ... -> (b s) c ..."), indices) - .transpose(0, 1) - .unsqueeze(1) - ) # (3, bsz, seqlen) -> (3, 1, bsz * seqlen) + index_first_axis( + rearrange( + position_ids, + "c b s ... -> (b s) c ..."), + indices) .transpose( + 0, + 1) .unsqueeze(1)) # (3, bsz, seqlen) -> (3, 1, bsz * seqlen) else: position_ids_rmpad = index_first_axis( - rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices + rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), + indices, ).transpose(0, 1) if "image_bound" in multi_modal_inputs: - from verl.utils.dataset.vision_utils import process_multi_modal_inputs_for_minicpmo + from verl.utils.dataset.vision_utils import ( + process_multi_modal_inputs_for_minicpmo, + ) multi_modal_inputs = process_multi_modal_inputs_for_minicpmo( - input_ids, attention_mask, position_ids, cu_seqlens, multi_modal_inputs - ) + input_ids, attention_mask, position_ids, cu_seqlens, multi_modal_inputs, ) # for compute the log_prob - input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=1) # (1, total_nnz) + 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: @@ -146,10 +190,12 @@ def _forward_micro_batch( sp_size=self.ulysses_sequence_parallel_size, ) else: - input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs( - input_ids_rmpad, - position_ids_rmpad=position_ids_rmpad, - sp_size=self.ulysses_sequence_parallel_size, + input_ids_rmpad, position_ids_rmpad, pad_size = ( + ulysses_pad_and_slice_inputs( + input_ids_rmpad, + position_ids_rmpad=position_ids_rmpad, + sp_size=self.ulysses_sequence_parallel_size, + ) ) input_ids_rmpad_rolled, _, _ = ulysses_pad_and_slice_inputs( input_ids_rmpad_rolled, @@ -157,9 +203,12 @@ def _forward_micro_batch( sp_size=self.ulysses_sequence_parallel_size, ) - input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0) # ((total_nnz / sp) + pad) + 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 + # only pass input_ids and position_ids to enable + # flash_attn_varlen extra_args = {} if self.use_fused_kernels: extra_args["temperature"] = temperature @@ -179,10 +228,12 @@ def _forward_micro_batch( entropy_rmpad = output.entropy.squeeze(0) # (total_nnz,) else: - logits_rmpad = output.logits.squeeze(0) # (total_nnz, vocab_size) + logits_rmpad = output.logits.squeeze( + 0) # (total_nnz, vocab_size) logits_rmpad.div_(temperature) - # if use_sp: ((total_nnz / sp) + pad) ; if not use_sp: (batch, seqlen) + # if use_sp: ((total_nnz / sp) + pad) ; if not use_sp: + # (batch, seqlen) inplace_backward = True if calculate_entropy: inplace_backward = False @@ -195,7 +246,9 @@ def _forward_micro_batch( # compute entropy if calculate_entropy: if not self.config.entropy_checkpointing: - entropy_rmpad = self.compute_entropy_from_logits(logits_rmpad) # ((total_nnz / sp) + pad) + entropy_rmpad = self.compute_entropy_from_logits( + logits_rmpad + ) # ((total_nnz / sp) + pad) else: entropy_rmpad = torch.utils.checkpoint.checkpoint( self.compute_entropy_from_logits, logits_rmpad @@ -234,8 +287,12 @@ def _forward_micro_batch( # only return response part: if calculate_entropy: - 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) + 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 extra_args = {} @@ -253,20 +310,29 @@ def _forward_micro_batch( ) # prevent model thinks we are generating if self.use_fused_kernels: - log_probs = output.log_probs[:, -response_length - 1 : -1] - entropy = output.entropy[:, -response_length - 1 : -1] # (bsz, response_length) + log_probs = output.log_probs[:, -response_length - 1: -1] + entropy = output.entropy[ + :, -response_length - 1: -1 + ] # (bsz, response_length) else: logits = output.logits logits.div_(temperature) - logits = logits[:, -response_length - 1 : -1, :] # (bsz, response_length, vocab_size) - log_probs = logprobs_from_logits(logits, micro_batch["responses"]) + logits = logits[ + :, -response_length - 1: -1, : + ] # (bsz, response_length, vocab_size) + log_probs = logprobs_from_logits( + logits, micro_batch["responses"]) if calculate_entropy: if not self.config.entropy_checkpointing: - entropy = verl_F.entropy_from_logits(logits) # (bsz, response_length) + entropy = verl_F.entropy_from_logits( + logits + ) # (bsz, response_length) else: - entropy = torch.utils.checkpoint.checkpoint(verl_F.entropy_from_logits, logits) + entropy = torch.utils.checkpoint.checkpoint( + verl_F.entropy_from_logits, logits + ) return entropy, log_probs @@ -274,22 +340,32 @@ 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) + grad_norm = self.actor_module.clip_grad_norm_( + max_norm=self.config.grad_clip + ) elif isinstance(self.actor_module, FSDPModule): - grad_norm = fsdp2_clip_grad_norm_(self.actor_module.parameters(), max_norm=self.config.grad_clip) + grad_norm = fsdp2_clip_grad_norm_( + self.actor_module.parameters(), 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) + grad_norm = torch.nn.utils.clip_grad_norm_( + self.actor_module.parameters(), max_norm=self.config.grad_clip + ) # if grad_norm is not finite, skip the update if not torch.isfinite(grad_norm): - print(f"WARN: rank {torch.distributed.get_rank()} grad_norm is not finite: {grad_norm}") + print( + f"WARN: rank { + torch.distributed.get_rank()} grad_norm is not finite: {grad_norm}") self.actor_optimizer.zero_grad() else: self.actor_optimizer.step() return grad_norm @GPUMemoryLogger(role="dp actor", logger=logger) - def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Tensor: + def compute_log_prob( + self, data: DataProto, calculate_entropy=False + ) -> torch.Tensor: """Compute the log probability of the responses given input_ids, attention_mask and position_ids Args: @@ -311,26 +387,43 @@ def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Te 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 silent error + temperature = data.meta_info[ + "temperature" + ] # temperature must be in the data.meta_info to avoid silent error use_dynamic_bsz = data.meta_info["use_dynamic_bsz"] def _get_micro_batches(data: DataProto) -> tuple[list, list | None]: - select_keys = ["responses", "input_ids", "attention_mask", "position_ids"] + select_keys = [ + "responses", + "input_ids", + "attention_mask", + "position_ids"] batch = data.select(batch_keys=select_keys).batch has_multi_modal_inputs = "multi_modal_inputs" in data.non_tensor_batch if has_multi_modal_inputs: - all_multi_modal_inputs_list = data.non_tensor_batch["multi_modal_inputs"] + all_multi_modal_inputs_list = data.non_tensor_batch[ + "multi_modal_inputs" + ] if use_dynamic_bsz: - max_token_len = data.meta_info["max_token_len"] * self.ulysses_sequence_parallel_size - rearranged_text_micro_batches, textual_indices = rearrange_micro_batches( - batch=batch, max_token_len=max_token_len + max_token_len = ( + data.meta_info["max_token_len"] + * self.ulysses_sequence_parallel_size + ) + rearranged_text_micro_batches, textual_indices = ( + rearrange_micro_batches( + batch=batch, max_token_len=max_token_len + ) ) final_micro_batches_list = [] - for i, text_mb_td in enumerate(rearranged_text_micro_batches): + for i, text_mb_td in enumerate( + rearranged_text_micro_batches): current_original_indices = textual_indices[i] - current_mm_inputs_list = [all_multi_modal_inputs_list[idx] for idx in current_original_indices] + current_mm_inputs_list = [ + all_multi_modal_inputs_list[idx] + for idx in current_original_indices + ] mb_dict = {k: v for k, v in text_mb_td.items()} mb_dict["multi_modal_inputs"] = current_mm_inputs_list @@ -341,8 +434,13 @@ def _get_micro_batches(data: DataProto) -> tuple[list, list | None]: micro_batches_dp = data.chunk(num_micro_batches) return micro_batches_dp, None elif use_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) + 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 + ) return micro_batches, indices else: micro_batches = batch.split(micro_batch_size) @@ -354,10 +452,14 @@ def _get_micro_batches(data: DataProto) -> tuple[list, list | None]: entropy_lst = [] for micro_batch in micro_batches: if isinstance(micro_batch, DataProto): - micro_batch = {**micro_batch.batch, **micro_batch.non_tensor_batch} + micro_batch = { + **micro_batch.batch, + **micro_batch.non_tensor_batch} with torch.no_grad(): entropy, log_probs = self._forward_micro_batch( - micro_batch, temperature=temperature, calculate_entropy=calculate_entropy + micro_batch, + temperature=temperature, + calculate_entropy=calculate_entropy, ) log_probs_lst.append(log_probs) if calculate_entropy: @@ -369,8 +471,11 @@ def _get_micro_batches(data: DataProto) -> tuple[list, list | None]: entropys = torch.concat(entropy_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) + 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] if calculate_entropy: entropys = entropys[revert_indices] @@ -382,7 +487,9 @@ def update_policy(self, data: DataProto): # make sure we are in training mode self.actor_module.train() - temperature = data.meta_info["temperature"] # temperature must be in the data.meta_info to avoid silent error + temperature = data.meta_info[ + "temperature" + ] # temperature must be in the data.meta_info to avoid silent error select_keys = [ "responses", @@ -401,9 +508,12 @@ def update_policy(self, data: DataProto): # Split to make minibatch iterator for updating the actor # See PPO paper for details. https://arxiv.org/abs/1707.06347 if has_multi_modal_inputs: - num_mini_batches = data.batch.batch_size[0] // self.config.ppo_mini_batch_size + num_mini_batches = ( + data.batch.batch_size[0] // self.config.ppo_mini_batch_size + ) non_tensor_select_keys = ["multi_modal_inputs"] - dataloader = data.select(select_keys, non_tensor_select_keys).chunk(num_mini_batches) + dataloader = data.select( + select_keys, non_tensor_select_keys).chunk(num_mini_batches) else: dataloader = batch.split(self.config.ppo_mini_batch_size) @@ -415,38 +525,63 @@ def update_policy(self, data: DataProto): if has_multi_modal_inputs: micro_batches = [] if self.config.use_dynamic_bsz: - all_multi_modal_inputs_list = data.non_tensor_batch["multi_modal_inputs"] + all_multi_modal_inputs_list = data.non_tensor_batch[ + "multi_modal_inputs" + ] batch_tensordict_for_rearrange = data.batch - max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size - rearranged_text_micro_batches_tds, textual_indices = rearrange_micro_batches( - batch=batch_tensordict_for_rearrange, max_token_len=max_token_len + max_token_len = ( + self.config.ppo_max_token_len_per_gpu + * self.ulysses_sequence_parallel_size + ) + rearranged_text_micro_batches_tds, textual_indices = ( + rearrange_micro_batches( + batch=batch_tensordict_for_rearrange, + max_token_len=max_token_len, + ) ) for current_original_indices, text_mb_td in zip( - textual_indices, rearranged_text_micro_batches_tds, strict=True + textual_indices, + rearranged_text_micro_batches_tds, + strict=True, ): current_mm_inputs_list = [ - all_multi_modal_inputs_list[idx] for idx in current_original_indices + all_multi_modal_inputs_list[idx] + for idx in current_original_indices ] mb_dict = {k: v for k, v in text_mb_td.items()} mb_dict["multi_modal_inputs"] = current_mm_inputs_list micro_batches.append(mb_dict) else: self.gradient_accumulation = ( - self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu + self.config.ppo_mini_batch_size + // self.config.ppo_micro_batch_size_per_gpu + ) + num_micro_batches = ( + mini_batch.batch.batch_size[0] + // self.config.ppo_micro_batch_size_per_gpu ) - num_micro_batches = mini_batch.batch.batch_size[0] // self.config.ppo_micro_batch_size_per_gpu - micro_batches = data.select(select_keys, non_tensor_select_keys).chunk(num_micro_batches) + micro_batches = data.select( + select_keys, non_tensor_select_keys + ).chunk(num_micro_batches) elif 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) + 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: self.gradient_accumulation = ( - self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu + self.config.ppo_mini_batch_size + // self.config.ppo_micro_batch_size_per_gpu ) # split batch into micro_batches - micro_batches = mini_batch.split(self.config.ppo_micro_batch_size_per_gpu) + micro_batches = mini_batch.split( + self.config.ppo_micro_batch_size_per_gpu + ) self.actor_optimizer.zero_grad() @@ -455,29 +590,42 @@ def update_policy(self, data: DataProto): # Support all hardwares if isinstance(data, DataProto): - data = {**data.batch.to(get_device_id()), **data.non_tensor_batch} + data = { + **data.batch.to(get_device_id()), + **data.non_tensor_batch, + } elif isinstance(data, dict): for k, v in data.items(): if isinstance(v, torch.Tensor): data[k] = v.to(get_device_id()) elif k == "multi_modal_inputs" and v is not None: data[k] = [ - {kk: vv.to(get_device_id()) for kk, vv in item_dict.items()} for item_dict in v + { + kk: vv.to(get_device_id()) + for kk, vv in item_dict.items() + } + for item_dict in v ] else: data[k] = v else: - data = data.to(get_device_id()) # actor device is cpu when using offload + data = data.to( + get_device_id() + ) # actor device is cpu when using offload response_mask = data["response_mask"] old_log_prob = data["old_log_probs"] advantages = data["advantages"] clip_ratio = self.config.clip_ratio clip_ratio_low = ( - self.config.clip_ratio_low if self.config.clip_ratio_low is not None else clip_ratio + self.config.clip_ratio_low + if self.config.clip_ratio_low is not None + else clip_ratio ) clip_ratio_high = ( - self.config.clip_ratio_high if self.config.clip_ratio_high is not None else clip_ratio + self.config.clip_ratio_high + if self.config.clip_ratio_high is not None + else clip_ratio ) clip_ratio_c = self.config.get("clip_ratio_c", 3.0) entropy_coeff = self.config.entropy_coeff @@ -488,37 +636,48 @@ def update_policy(self, data: DataProto): if entropy_coeff != 0: calculate_entropy = True entropy, log_prob = self._forward_micro_batch( - micro_batch=data, temperature=temperature, calculate_entropy=calculate_entropy + micro_batch=data, + temperature=temperature, + calculate_entropy=calculate_entropy, ) - loss_mode = self.config.policy_loss.get("loss_mode", "vanilla") + loss_mode = self.config.policy_loss.get( + "loss_mode", "vanilla") if self.config.policy_loss.loss_mode == "vanilla": - pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = compute_policy_loss( - old_log_prob=old_log_prob, - log_prob=log_prob, - advantages=advantages, - response_mask=response_mask, - cliprange=clip_ratio, - cliprange_low=clip_ratio_low, - cliprange_high=clip_ratio_high, - clip_ratio_c=clip_ratio_c, - loss_agg_mode=loss_agg_mode, + pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = ( + compute_policy_loss( + old_log_prob=old_log_prob, + log_prob=log_prob, + advantages=advantages, + response_mask=response_mask, + cliprange=clip_ratio, + cliprange_low=clip_ratio_low, + cliprange_high=clip_ratio_high, + clip_ratio_c=clip_ratio_c, + loss_agg_mode=loss_agg_mode, + ) ) else: policy_loss_fn = get_policy_loss_fn(loss_mode) - pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = policy_loss_fn( - old_log_prob=old_log_prob, - log_prob=log_prob, - advantages=advantages, - response_mask=advantages, - loss_agg_mode=loss_agg_mode, - config=self.config, + pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = ( + policy_loss_fn( + old_log_prob=old_log_prob, + log_prob=log_prob, + advantages=advantages, + response_mask=advantages, + loss_agg_mode=loss_agg_mode, + config=self.config, + ) ) if entropy_coeff != 0: - entropy_loss = agg_loss(loss_mat=entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode) + entropy_loss = agg_loss( + loss_mat=entropy, + loss_mask=response_mask, + loss_agg_mode=loss_agg_mode, + ) # compute policy loss policy_loss = pg_loss - entropy_loss * entropy_coeff @@ -529,9 +688,15 @@ def update_policy(self, data: DataProto): ref_log_prob = data["ref_log_prob"] # compute kl loss kld = kl_penalty( - logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type + logprob=log_prob, + ref_logprob=ref_log_prob, + kl_penalty=self.config.kl_loss_type, + ) + kl_loss = agg_loss( + loss_mat=kld, + loss_mask=response_mask, + loss_agg_mode=loss_agg_mode, ) - kl_loss = agg_loss(loss_mat=kld, loss_mask=response_mask, loss_agg_mode=loss_agg_mode) policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef micro_batch_metrics["actor/kl_loss"] = kl_loss.detach().item() @@ -539,7 +704,9 @@ def update_policy(self, data: DataProto): if self.config.use_dynamic_bsz: # relative to the dynamic bsz - loss = policy_loss * (len(data) / self.config.ppo_mini_batch_size) + loss = policy_loss * ( + len(data) / self.config.ppo_mini_batch_size + ) else: loss = policy_loss / self.gradient_accumulation loss.backward() @@ -550,12 +717,12 @@ def update_policy(self, data: DataProto): "actor/pg_clipfrac": pg_clipfrac.detach().item(), "actor/ppo_kl": ppo_kl.detach().item(), "actor/pg_clipfrac_lower": pg_clipfrac_lower.detach().item(), - } - ) + }) append_to_dict(metrics, micro_batch_metrics) grad_norm = self._optimizer_step() - mini_batch_metrics = {"actor/grad_norm": grad_norm.detach().item()} + mini_batch_metrics = { + "actor/grad_norm": grad_norm.detach().item()} append_to_dict(metrics, mini_batch_metrics) self.actor_optimizer.zero_grad() return metrics diff --git a/Agent0/executor_train/verl/verl/workers/actor/megatron_actor.py b/Agent0/executor_train/verl/verl/workers/actor/megatron_actor.py index 08238d4..ca97e50 100644 --- a/Agent0/executor_train/verl/verl/workers/actor/megatron_actor.py +++ b/Agent0/executor_train/verl/verl/workers/actor/megatron_actor.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -37,10 +37,18 @@ from torch import nn from verl import DataProto -from verl.trainer.ppo.core_algos import agg_loss, compute_policy_loss, get_policy_loss_fn, kl_penalty +from verl.trainer.ppo.core_algos import ( + agg_loss, + compute_policy_loss, + get_policy_loss_fn, + kl_penalty, +) from verl.utils.device import get_device_id, get_torch_device from verl.utils.megatron.pipeline_parallel import make_batch_generator -from verl.utils.megatron.tensor_parallel import vocab_parallel_entropy, vocab_parallel_log_probs_from_logits +from verl.utils.megatron.tensor_parallel import ( + vocab_parallel_entropy, + vocab_parallel_log_probs_from_logits, +) from verl.utils.megatron_utils import get_model_config from verl.utils.profiler import GPUMemoryLogger from verl.utils.profiler.profile import Profiler @@ -152,14 +160,18 @@ def _validate_config(self, config) -> None: """Validate config options not implemented for Megatron backend""" assert config.get("ulysses_sequence_parallel_size", 1) == 1 if config.get("shuffle", False): - assert config.data_loader_seed is not None, "If shuffle dataloader, seed must be manually set" + assert ( + config.data_loader_seed is not None + ), "If shuffle dataloader, seed must be manually set" if config.megatron.tensor_model_parallel_size == 1: print("[Warining] Because actor tp size == 1, set sp to False") config.megatron.sequence_parallel = False self.config = config @GPUMemoryLogger(role="megatron actor", logger=logger) - def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Tensor: + def compute_log_prob( + self, data: DataProto, calculate_entropy=False + ) -> torch.Tensor: """Compute the log probability of the responses given input_ids, attention_mask and position_ids Args: @@ -182,25 +194,39 @@ def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Te use_dynamic_bsz = data.meta_info.get("use_dynamic_bsz", False) micro_batch_size = data.meta_info.get("micro_batch_size", None) max_token_len = data.meta_info.get("max_token_len", None) - assert micro_batch_size is not None, "micro batch size is needed for forward compute" + assert ( + micro_batch_size is not None + ), "micro batch size is needed for forward compute" if use_dynamic_bsz: - assert max_token_len is not None, "max_token_len must be set when use_dynamic_bsz is True" + assert ( + max_token_len is not None + ), "max_token_len must be set when use_dynamic_bsz is True" max_token_len = max_token_len * self.config.megatron.context_parallel_size - def compute_logprobs_fn(output, data, use_dynamic_bsz=False, indices=None): + def compute_logprobs_fn( + output, + data, + use_dynamic_bsz=False, + indices=None): response = data["responses"] response_length = response.size(1) - log_probs = output["log_probs"][:, -response_length - 1 : -1].contiguous() + log_probs = output["log_probs"][:, - + response_length - 1: -1].contiguous() 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) + recompute_old_log_prob = self.config.get( + "recompute_old_log_prob", True) entropys = torch.Tensor() if recompute_old_log_prob: - select_keys = ["responses", "input_ids", "attention_mask", "position_ids"] + 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) @@ -219,19 +245,29 @@ def compute_logprobs_fn(output, data, use_dynamic_bsz=False, indices=None): if mpu.is_pipeline_last_stage(ignore_virtual=True): # only on last rank. It should be on every tp rank if calculate_entropy: - log_probs = [o[0]["log_probs"] for o in output["output"]] # (bs, seq_size) + log_probs = [ + o[0]["log_probs"] for o in output["output"] + ] # (bs, seq_size) else: - log_probs = [o["log_probs"] for o in output["output"]] # (bs, seq_size) + log_probs = [ + o["log_probs"] for o in output["output"] + ] # (bs, seq_size) log_probs = torch.cat(log_probs, dim=0).to(torch.float32) if use_dynamic_bsz: indices = output["indices"] 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) + 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] else: log_probs = torch.empty( - size=(batch_size, response_length), dtype=torch.float32, device=input_ids.device + size=(batch_size, response_length), + dtype=torch.float32, + device=input_ids.device, ) # broadcast across pp ranks @@ -244,17 +280,25 @@ def compute_logprobs_fn(output, data, use_dynamic_bsz=False, indices=None): if calculate_entropy: # Note that o[0] is metrics, o[1] is entropy if mpu.is_pipeline_last_stage(ignore_virtual=True): - entropys = torch.cat([o[1] for o in output["output"]], dim=0) + entropys = torch.cat( + [o[1] for o in output["output"]], dim=0) entropys = entropys.to(torch.float32) if use_dynamic_bsz: indices = output["indices"] - indices = list(itertools.chain.from_iterable(indices)) - assert len(indices) == entropys.size(0), f"{len(indices)} vs. {entropys.size()}" - revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long) + indices = list( + itertools.chain.from_iterable(indices)) + assert len(indices) == entropys.size( + 0 + ), f"{len(indices)} vs. {entropys.size()}" + revert_indices = torch.tensor( + get_reverse_idx(indices), dtype=torch.long + ) entropys = entropys[revert_indices] else: entropys = torch.empty( - size=(batch_size, response_length), dtype=torch.float32, device=input_ids.device + size=(batch_size, response_length), + dtype=torch.float32, + device=input_ids.device, ) # broadcast across pp ranks torch.distributed.broadcast( @@ -295,10 +339,19 @@ def make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]: Returns: """ - select_keys = ["responses", "input_ids", "attention_mask", "position_ids", "old_log_probs", "advantages"] + select_keys = [ + "responses", + "input_ids", + "attention_mask", + "position_ids", + "old_log_probs", + "advantages", + ] if self.config.use_kl_loss: select_keys.append("ref_log_prob") - self.has_multi_modal_inputs = "multi_modal_inputs" in data.non_tensor_batch.keys() + self.has_multi_modal_inputs = ( + "multi_modal_inputs" in data.non_tensor_batch.keys() + ) if self.has_multi_modal_inputs: data = data.select(select_keys, ["multi_modal_inputs"]) else: @@ -335,41 +388,58 @@ def forward_backward_batch( group=mpu.get_pipeline_model_parallel_group(), ) # split into micro-batches - mini_batch.batch["attention_mask"] = mini_batch.batch["attention_mask"].to(bool) - self.has_multi_modal_inputs = "multi_modal_inputs" in mini_batch.non_tensor_batch.keys() + mini_batch.batch["attention_mask"] = mini_batch.batch["attention_mask"].to( + bool) + self.has_multi_modal_inputs = ( + "multi_modal_inputs" in mini_batch.non_tensor_batch.keys() + ) if self.has_multi_modal_inputs: - mini_batch.batch["multi_modal_inputs"] = mini_batch.non_tensor_batch["multi_modal_inputs"] + mini_batch.batch["multi_modal_inputs"] = mini_batch.non_tensor_batch[ + "multi_modal_inputs" + ] mini_batch.batch["multi_modal_inputs_idx"] = torch.Tensor( list(range(len(mini_batch.non_tensor_batch["multi_modal_inputs"]))) ).to(torch.int64) - if mini_batch.batch["position_ids"].dim() == 3: # qwen2vl mrope [bs, 3, seq_len] + if ( + mini_batch.batch["position_ids"].dim() == 3 + ): # qwen2vl mrope [bs, 3, seq_len] mini_batch.batch["position_ids"] = mini_batch.batch["position_ids"][ :, 0 ] # mcore patch recompute qwen2vl's pos ids during forward indices = None if use_dynamic_bsz: - assert max_token_len is not None, "max_token_len must be set when use_dynamic_bsz is True" + assert ( + max_token_len is not None + ), "max_token_len must be set when use_dynamic_bsz is True" vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() if vpp_size is not None and vpp_size > 1: - microbatch_group_size_per_vp_stage = self.tf_config.microbatch_group_size_per_vp_stage + microbatch_group_size_per_vp_stage = ( + self.tf_config.microbatch_group_size_per_vp_stage + ) micro_batches, indices = rearrange_micro_batches( batch=mini_batch.batch, num_batches_divided_by=microbatch_group_size_per_vp_stage, max_token_len=max_token_len, ) - assert len(micro_batches) % self.tf_config.microbatch_group_size_per_vp_stage == 0, ( + assert ( + len(micro_batches) + % self.tf_config.microbatch_group_size_per_vp_stage + == 0 + ), ( f"micro_batches {micro_batches} must be divisible by microbatch_group_size_per_vp_stage " f"{microbatch_group_size_per_vp_stage} for megatron backend" ) else: - micro_batches, indices = rearrange_micro_batches(batch=mini_batch.batch, max_token_len=max_token_len) + micro_batches, indices = rearrange_micro_batches( + batch=mini_batch.batch, max_token_len=max_token_len + ) total_seqlen = max_token_len else: - assert micro_batch_size is not None, ( - "micro_batch_size is needed to be passed in when not using dynamic batch size" - ) + assert ( + micro_batch_size is not None + ), "micro_batch_size is needed to be passed in when not using dynamic batch size" micro_batches = mini_batch.batch.split(micro_batch_size) seq_len = micro_batches[0]["input_ids"].shape[1] total_seqlen = micro_batch_size * seq_len @@ -380,7 +450,8 @@ def forward_backward_batch( def loss_func(output, data, meta_info): # For memory efficiency - # We move calculation of entropy to compute_log_probs, forward_only == True + # We move calculation of entropy to compute_log_probs, forward_only + # == True device = output["log_probs"].device metrics = {} if forward_only: @@ -400,7 +471,8 @@ def loss_func(output, data, meta_info): loss_agg_mode = self.config.loss_agg_mode # compute policy loss - log_prob = output["log_probs"][:, -response_length - 1 : -1].contiguous() + log_prob = output["log_probs"][:, - + response_length - 1: -1].contiguous() ret_entropy = None stats = {} if not forward_only: @@ -408,8 +480,16 @@ def loss_func(output, data, meta_info): advantages = data["advantages"] clip_ratio = self.config.clip_ratio - clip_ratio_low = self.config.clip_ratio_low if self.config.clip_ratio_low is not None else clip_ratio - clip_ratio_high = self.config.clip_ratio_high if self.config.clip_ratio_high is not None else clip_ratio + clip_ratio_low = ( + self.config.clip_ratio_low + if self.config.clip_ratio_low is not None + else clip_ratio + ) + clip_ratio_high = ( + self.config.clip_ratio_high + if self.config.clip_ratio_high is not None + else clip_ratio + ) clip_ratio_c = self.config.get("clip_ratio_c", 3.0) entropy_coeff = self.config.entropy_coeff @@ -418,16 +498,18 @@ def loss_func(output, data, meta_info): loss_mode = self.config.policy_loss.get("loss_mode", "vanilla") if self.config.policy_loss.loss_mode == "vanilla": - pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = compute_policy_loss( - old_log_prob=old_log_prob, - log_prob=log_prob, - advantages=advantages, - response_mask=response_mask, - cliprange=clip_ratio, - cliprange_low=clip_ratio_low, - cliprange_high=clip_ratio_high, - clip_ratio_c=clip_ratio_c, - loss_agg_mode=loss_agg_mode, + pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = ( + compute_policy_loss( + old_log_prob=old_log_prob, + log_prob=log_prob, + advantages=advantages, + response_mask=response_mask, + cliprange=clip_ratio, + cliprange_low=clip_ratio_low, + cliprange_high=clip_ratio_high, + clip_ratio_c=clip_ratio_c, + loss_agg_mode=loss_agg_mode, + ) ) else: @@ -441,20 +523,22 @@ def loss_func(output, data, meta_info): config=self.config, ) - stats.update( - { - "actor/pg_loss": pg_loss.detach().item(), - "actor/pg_clipfrac": pg_clipfrac.detach().item(), - "actor/ppo_kl": ppo_kl.detach().item(), - "actor/pg_clipfrac_lower": pg_clipfrac_lower.detach().item(), - } - ) + stats.update({"actor/pg_loss": pg_loss.detach().item(), + "actor/pg_clipfrac": pg_clipfrac.detach().item(), + "actor/ppo_kl": ppo_kl.detach().item(), + "actor/pg_clipfrac_lower": pg_clipfrac_lower.detach().item(), + }) policy_loss = pg_loss if calculate_entropy: - entropy = output["entropy"][:, -response_length - 1 : -1].contiguous() + entropy = output["entropy"][:, - + response_length - 1: -1].contiguous() if not forward_only: - entropy_loss = agg_loss(loss_mat=entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode) + entropy_loss = agg_loss( + loss_mat=entropy, + loss_mask=response_mask, + loss_agg_mode=loss_agg_mode, + ) entropy_coeff = meta_info["entropy_coeff"] policy_loss = pg_loss - entropy_coeff * entropy_loss else: @@ -466,8 +550,16 @@ def loss_func(output, data, meta_info): if self.config.use_kl_loss: ref_log_prob = data["ref_log_prob"] # compute kl loss - kld = kl_penalty(logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type) - kl_loss = agg_loss(loss_mat=kld, loss_mask=response_mask, loss_agg_mode=self.config.loss_agg_mode) + kld = kl_penalty( + logprob=log_prob, + ref_logprob=ref_log_prob, + kl_penalty=self.config.kl_loss_type, + ) + kl_loss = agg_loss( + loss_mat=kld, + loss_mask=response_mask, + loss_agg_mode=self.config.loss_agg_mode, + ) policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef metrics["actor/kl_loss"] = kl_loss.detach().item() @@ -490,17 +582,25 @@ def forward_step(batch_iter, model): idxs = batch["multi_modal_inputs_idx"] mmi = batch["multi_modal_inputs"] multi_modal_inputs[key] = torch.cat( - [mmi[idx].get(key) for idx in idxs if mmi[idx].get(key) is not None], dim=0 + [ + mmi[idx].get(key) + for idx in idxs + if mmi[idx].get(key) is not None + ], + dim=0, ) responses = batch["responses"] response_length = responses.size(1) label = position_ids.clone() - label[:, -response_length - 1 : -1] = responses + label[:, -response_length - 1: -1] = responses label_mask = attention_mask.clone() label_mask[:, : -response_length - 1] = False label_mask[:, -1] = False - from verl.models.mcore import get_mcore_forward_fn, get_mcore_forward_fused_fn + from verl.models.mcore import ( + get_mcore_forward_fn, + get_mcore_forward_fused_fn, + ) if self.use_fused_kernels: forward_fn = get_mcore_forward_fused_fn(self.hf_config) @@ -525,12 +625,14 @@ def logits_processor(logits, label, label_mask): if calculate_entropy: entropy = vocab_parallel_entropy(logits) ret["entropy"] = entropy - log_probs = vocab_parallel_log_probs_from_logits(logits, label) + log_probs = vocab_parallel_log_probs_from_logits( + logits, label) log_probs = log_probs.masked_fill(~label_mask, 0.0) ret["log_probs"] = log_probs return ret - logits_processor_args = {"label": label, "label_mask": label_mask} + logits_processor_args = { + "label": label, "label_mask": label_mask} output = forward_fn( model, input_ids, @@ -554,10 +656,13 @@ def logits_processor(logits, label, label_mask): 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(micro_batches, vpp_size=len(self.actor_module)) + batch_generator = make_batch_generator( + micro_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) + # 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, @@ -608,9 +713,11 @@ def update_policy(self, dataloader: Iterable[DataProto]) -> dict: for data in dataloader: data.to(get_device_id()) self.actor_optimizer.zero_grad() - # use use_contiguous_buffers_in_local_ddp and no overlap_dp_param_comm + # 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 + # if use distributed optimizer, zero grad buffer will be + # handled by optimizer chunk.zero_grad_buffer() calculate_entropy = self.config.entropy_coeff != 0 @@ -620,7 +727,10 @@ def update_policy(self, dataloader: Iterable[DataProto]) -> dict: micro_batch_size = self.config.ppo_micro_batch_size_per_gpu max_token_len = None if self.config.use_dynamic_bsz: - max_token_len = self.config.ppo_max_token_len_per_gpu * self.config.megatron.context_parallel_size + max_token_len = ( + self.config.ppo_max_token_len_per_gpu + * self.config.megatron.context_parallel_size + ) metric_micro_batch = self.forward_backward_batch( data, calculate_entropy=calculate_entropy, @@ -631,10 +741,15 @@ def update_policy(self, dataloader: Iterable[DataProto]) -> dict: ) metric_micro_batch = metric_micro_batch["output"] for metric in metric_micro_batch: - # Note that o[0] is metrics, o[1] is entropy, o[2] is response_mask - append_to_dict(metrics, metric[0]) # append the metric from this micro-batch to global metrics. - - update_successful, grad_norm, num_zeros_in_grad = self.actor_optimizer.step() + # Note that o[0] is metrics, o[1] is entropy, o[2] is + # response_mask + append_to_dict( + metrics, metric[0] + ) # append the metric from this micro-batch to global metrics. + + update_successful, grad_norm, num_zeros_in_grad = ( + self.actor_optimizer.step() + ) data = {"actor/grad_norm": grad_norm} append_to_dict(metrics, data) diff --git a/Agent0/executor_train/verl/verl/workers/critic/__init__.py b/Agent0/executor_train/verl/verl/workers/critic/__init__.py index 80808f1..282166f 100644 --- a/Agent0/executor_train/verl/verl/workers/critic/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/critic/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/critic/base.py b/Agent0/executor_train/verl/verl/workers/critic/base.py index 8201758..07c4f60 100644 --- a/Agent0/executor_train/verl/verl/workers/critic/base.py +++ b/Agent0/executor_train/verl/verl/workers/critic/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/critic/dp_critic.py b/Agent0/executor_train/verl/verl/workers/critic/dp_critic.py index ac77758..1bee538 100644 --- a/Agent0/executor_train/verl/verl/workers/critic/dp_critic.py +++ b/Agent0/executor_train/verl/verl/workers/critic/dp_critic.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -26,7 +26,12 @@ from verl import DataProto from verl.trainer.ppo import core_algos -from verl.utils.device import get_device_id, get_device_name, is_cuda_available, is_npu_available +from verl.utils.device import ( + get_device_id, + get_device_name, + is_cuda_available, + is_npu_available, +) from verl.utils.fsdp_utils import FSDPModule, fsdp2_clip_grad_norm_ from verl.utils.profiler import GPUMemoryLogger from verl.utils.py_functional import append_to_dict @@ -36,23 +41,37 @@ from verl.workers.critic import BasePPOCritic if is_cuda_available: - from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input + from flash_attn.bert_padding import ( + index_first_axis, + pad_input, + rearrange, + unpad_input, + ) elif is_npu_available: - from transformers.integrations.npu_flash_attention import index_first_axis, pad_input, rearrange, unpad_input + from transformers.integrations.npu_flash_attention import ( + index_first_axis, + pad_input, + rearrange, + unpad_input, + ) logger = logging.getLogger(__file__) logger.setLevel(os.getenv("VERL_LOGGING_LEVEL", "WARN")) class DataParallelPPOCritic(BasePPOCritic): - def __init__(self, config, critic_module: nn.Module, critic_optimizer: optim.Optimizer): + 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) + self.use_remove_padding = self.config.model.get( + "use_remove_padding", False) print(f"Critic use_remove_padding={self.use_remove_padding}") - self.ulysses_sequence_parallel_size = self.config.get("ulysses_sequence_parallel_size", 1) + self.ulysses_sequence_parallel_size = self.config.get( + "ulysses_sequence_parallel_size", 1 + ) self.device_name = get_device_name() def _forward_micro_batch(self, micro_batch): @@ -76,27 +95,37 @@ def _forward_micro_batch(self, micro_batch): 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) + input_ids_rmpad = input_ids_rmpad.transpose( + 0, 1) # (1, total_nnz) # unpad the position_ids to align the rotary if position_ids.dim() == 3: position_ids_rmpad = ( - index_first_axis(rearrange(position_ids, "c b s ... -> (b s) c ..."), indices) - .transpose(0, 1) - .unsqueeze(1) - ) # (3, bsz, seqlen) -> (3, 1, bsz * seqlen) + index_first_axis( + rearrange( + position_ids, + "c b s ... -> (b s) c ..."), + indices) .transpose( + 0, + 1) .unsqueeze(1)) # (3, bsz, seqlen) -> (3, 1, bsz * seqlen) else: position_ids_rmpad = index_first_axis( - rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices + 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 + 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 + # only pass input_ids and position_ids to enable + # flash_attn_varlen output = self.critic_module( input_ids=input_ids_rmpad, attention_mask=None, @@ -115,12 +144,13 @@ def _forward_micro_batch(self, micro_batch): # 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 - ) + 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] + 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, @@ -134,18 +164,23 @@ def _forward_micro_batch(self, micro_batch): values = output[2] else: values = output.logits - values = values[:, -response_length - 1 : -1].squeeze(-1) + 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) + grad_norm = self.critic_module.clip_grad_norm_( + self.config.grad_clip) elif isinstance(self.critic_module, FSDPModule): - grad_norm = fsdp2_clip_grad_norm_(self.critic_module.parameters(), max_norm=self.config.grad_clip) + grad_norm = fsdp2_clip_grad_norm_( + self.critic_module.parameters(), max_norm=self.config.grad_clip + ) else: - grad_norm = torch.nn.utils.clip_grad_norm_(self.critic_module.parameters(), max_norm=self.config.grad_clip) + grad_norm = torch.nn.utils.clip_grad_norm_( + self.critic_module.parameters(), max_norm=self.config.grad_clip + ) # if grad_norm is not finite, skip the update if not torch.isfinite(grad_norm): @@ -159,7 +194,11 @@ def _optimizer_step(self): 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"] + 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"] has_multi_modal_inputs = "multi_modal_inputs" in data.non_tensor_batch.keys() @@ -167,18 +206,25 @@ def compute_values(self, data: DataProto) -> torch.Tensor: if has_multi_modal_inputs: num_micro_batches = data.batch.batch_size[0] // micro_batch_size non_tensor_select_keys = ["multi_modal_inputs"] - micro_batches = data.select(select_keys, non_tensor_select_keys).chunk(num_micro_batches) + micro_batches = data.select( + select_keys, non_tensor_select_keys).chunk(num_micro_batches) elif 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) + 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: if isinstance(micro_batch, DataProto): - micro_batch = {**micro_batch.batch, **micro_batch.non_tensor_batch} + micro_batch = { + **micro_batch.batch, + **micro_batch.non_tensor_batch} with torch.no_grad(): values = self._forward_micro_batch(micro_batch) @@ -187,8 +233,10 @@ def compute_values(self, data: DataProto) -> torch.Tensor: 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) + 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] response_mask = data.batch["response_mask"] @@ -201,16 +249,27 @@ def update_critic(self, data: DataProto): self.critic_module.train() metrics = {} - select_keys = ["input_ids", "responses", "response_mask", "attention_mask", "position_ids", "values", "returns"] + select_keys = [ + "input_ids", + "responses", + "response_mask", + "attention_mask", + "position_ids", + "values", + "returns", + ] batch = data.select(batch_keys=select_keys).batch has_multi_modal_inputs = "multi_modal_inputs" in data.non_tensor_batch.keys() # Split to make minibatch iterator for updating the actor # See PPO paper for details. https://arxiv.org/abs/1707.06347 if has_multi_modal_inputs: - num_mini_batches = data.batch.batch_size[0] // self.config.ppo_mini_batch_size + num_mini_batches = ( + data.batch.batch_size[0] // self.config.ppo_mini_batch_size + ) non_tensor_select_keys = ["multi_modal_inputs"] - dataloader = data.select(select_keys, non_tensor_select_keys).chunk(num_mini_batches) + dataloader = data.select( + select_keys, non_tensor_select_keys).chunk(num_mini_batches) else: dataloader = batch.split(self.config.ppo_mini_batch_size) @@ -219,18 +278,32 @@ def update_critic(self, data: DataProto): # split batch into micro_batches mini_batch = data if has_multi_modal_inputs: - num_micro_batches = mini_batch.batch.batch_size[0] // self.config.ppo_micro_batch_size_per_gpu - micro_batches = data.select(select_keys, non_tensor_select_keys).chunk(num_micro_batches) + num_micro_batches = ( + mini_batch.batch.batch_size[0] + // self.config.ppo_micro_batch_size_per_gpu + ) + micro_batches = data.select( + select_keys, non_tensor_select_keys + ).chunk(num_micro_batches) self.gradient_accumulation = ( - self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu + self.config.ppo_mini_batch_size + // self.config.ppo_micro_batch_size_per_gpu ) elif 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) + 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_per_gpu) + micro_batches = mini_batch.split( + self.config.ppo_micro_batch_size_per_gpu + ) self.gradient_accumulation = ( - self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu + self.config.ppo_mini_batch_size + // self.config.ppo_micro_batch_size_per_gpu ) self.critic_optimizer.zero_grad() @@ -240,9 +313,14 @@ def update_critic(self, data: DataProto): # Support all devices if isinstance(data, DataProto): - data = {**data.batch.to(get_device_id()), **data.non_tensor_batch} + data = { + **data.batch.to(get_device_id()), + **data.non_tensor_batch, + } else: - data = data.to(get_device_id()) # critic device is cpu when using offload + data = data.to( + get_device_id() + ) # critic device is cpu when using offload response_mask = data["response_mask"] values = data["values"] returns = data["returns"] @@ -261,7 +339,8 @@ def update_critic(self, data: DataProto): ) if self.config.use_dynamic_bsz: # relative to the dynamic bsz - loss = vf_loss * (len(data) / self.config.ppo_mini_batch_size) + loss = vf_loss * \ + (len(data) / self.config.ppo_mini_batch_size) else: loss = vf_loss / self.gradient_accumulation @@ -271,14 +350,16 @@ def update_critic(self, data: DataProto): { "critic/vf_loss": vf_loss.detach().item(), "critic/vf_clipfrac": vf_clipfrac.detach().item(), - "critic/vpred_mean": masked_mean(vpreds, response_mask).detach().item(), - } - ) + "critic/vpred_mean": masked_mean( + vpreds, + response_mask) .detach() .item(), + }) append_to_dict(metrics, micro_batch_metrics) grad_norm = self._optimizer_step() - mini_batch_metrics = {"critic/grad_norm": grad_norm.detach().item()} + mini_batch_metrics = { + "critic/grad_norm": grad_norm.detach().item()} append_to_dict(metrics, mini_batch_metrics) self.critic_optimizer.zero_grad() return metrics diff --git a/Agent0/executor_train/verl/verl/workers/critic/megatron_critic.py b/Agent0/executor_train/verl/verl/workers/critic/megatron_critic.py index 1d44a88..3473909 100644 --- a/Agent0/executor_train/verl/verl/workers/critic/megatron_critic.py +++ b/Agent0/executor_train/verl/verl/workers/critic/megatron_critic.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -64,7 +64,8 @@ def __init__( 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. + # 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, @@ -83,7 +84,9 @@ def _validate_config(self, config) -> None: """Validate config options not implemented for Megatron backend""" assert config.get("ulysses_sequence_parallel_size", 1) == 1 if config.shuffle: - assert config.data_loader_seed is not None, "If shuffle dataloader, seed must be manually set" + assert ( + config.data_loader_seed is not None + ), "If shuffle dataloader, seed must be manually set" if config.megatron.tensor_model_parallel_size == 1: print("[Warining] Because critic tp size == 1, set sp to False") config.megatron.sequence_parallel = False @@ -97,9 +100,13 @@ def compute_values(self, data: DataProto) -> DataProto: use_dynamic_bsz = data.meta_info.get("use_dynamic_bsz", False) micro_batch_size = data.meta_info.get("micro_batch_size", None) max_token_len = data.meta_info.get("max_token_len", None) - assert micro_batch_size is not None, "micro batch size is needed for forward compute" + assert ( + micro_batch_size is not None + ), "micro batch size is needed for forward compute" if use_dynamic_bsz: - assert max_token_len is not None, "max_token_len must be set when use_dynamic_bsz is True" + assert ( + max_token_len is not None + ), "max_token_len must be set when use_dynamic_bsz is True" max_token_len = max_token_len * self.config.megatron.context_parallel_size response_length = responses.size(1) with torch.no_grad(): @@ -113,21 +120,27 @@ def compute_values(self, data: DataProto) -> DataProto: ) if mpu.is_pipeline_last_stage(ignore_virtual=True): # only on last rank. It should be on every tp rank - values = [o["vpreds"] for o in output["output"]] # (bs, seq_size, vocal_size) + values = [ + o["vpreds"] for o in output["output"] + ] # (bs, seq_size, vocal_size) values = torch.cat(values, dim=0).to(torch.float32) if use_dynamic_bsz: indices = output["indices"] 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) + 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] else: values = torch.empty_like(attention_mask, dtype=torch.float32) # each tp ranks should contain the same value - values = values[ - :, -response_length - 1 : -1 - ] # Values are predicted at the ends of prefixes, e.g., the last prompt token + # Values are predicted at the ends of prefixes, e.g., the last + # prompt token + values = values[:, -response_length - 1: -1] response_mask = attention_mask[:, -response_length:] values = values * response_mask # Only action tokens have values values = values.contiguous() @@ -145,7 +158,14 @@ def compute_values(self, data: DataProto) -> DataProto: return values def make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]: - select_keys = ["input_ids", "responses", "attention_mask", "position_ids", "values", "returns"] + 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, @@ -173,30 +193,41 @@ def forward_backward_batch( group=mpu.get_pipeline_model_parallel_group(), ) # split into micro-batches - mini_batch.batch["attention_mask"] = mini_batch.batch["attention_mask"].to(bool) + mini_batch.batch["attention_mask"] = mini_batch.batch["attention_mask"].to( + bool) indices = None if use_dynamic_bsz: - assert max_token_len is not None, "max_token_len must be set when use_dynamic_bsz is True" + assert ( + max_token_len is not None + ), "max_token_len must be set when use_dynamic_bsz is True" vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() if vpp_size is not None and vpp_size > 1: - microbatch_group_size_per_vp_stage = self.tf_config.microbatch_group_size_per_vp_stage + microbatch_group_size_per_vp_stage = ( + self.tf_config.microbatch_group_size_per_vp_stage + ) micro_batches, indices = rearrange_micro_batches( batch=mini_batch.batch, num_batches_divided_by=microbatch_group_size_per_vp_stage, max_token_len=max_token_len, ) - assert len(micro_batches) % self.tf_config.microbatch_group_size_per_vp_stage == 0, ( + assert ( + len(micro_batches) + % self.tf_config.microbatch_group_size_per_vp_stage + == 0 + ), ( f"micro_batches {micro_batches} must be divisible by microbatch_group_size_per_vp_stage " f"{microbatch_group_size_per_vp_stage} for megatron backend" ) else: - micro_batches, indices = rearrange_micro_batches(batch=mini_batch.batch, max_token_len=max_token_len) + micro_batches, indices = rearrange_micro_batches( + batch=mini_batch.batch, max_token_len=max_token_len + ) total_seqlen = max_token_len else: - assert micro_batch_size is not None, ( - "micro_batch_size is needed to be passed in when not using dynamic batch size" - ) + assert ( + micro_batch_size is not None + ), "micro_batch_size is needed to be passed in when not using dynamic batch size" micro_batches = mini_batch.batch.split(micro_batch_size) seq_len = micro_batches[0]["input_ids"].shape[1] total_seqlen = micro_batch_size * seq_len @@ -208,7 +239,9 @@ def loss_func(output, data, meta_info): nonlocal use_dynamic_bsz if forward_only: - return torch.tensor(1.0, device=output.device), {"vpreds": output} + return torch.tensor( + 1.0, device=output.device), { + "vpreds": output} responses = data["responses"] attention_mask = data["attention_mask"] @@ -221,7 +254,7 @@ def loss_func(output, data, meta_info): cliprange_value = self.config.cliprange_value vpreds = output # (bs, sequence_length) - vpreds = vpreds[:, -response_length - 1 : -1] + vpreds = vpreds[:, -response_length - 1: -1] vf_loss, vf_clipfrac = core_algos.compute_value_loss( vpreds=vpreds, @@ -261,10 +294,13 @@ def forward_step(batch_iter, model): return output, partial(loss_func, data=batch, meta_info={}) # batch should be a list of batches inside micro-batches - batch_generator = make_batch_generator(micro_batches, vpp_size=len(self.critic_module)) + batch_generator = make_batch_generator( + micro_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) + # 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, @@ -298,14 +334,18 @@ def update_critic(self, dataloader: Iterable[DataProto]): 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 + # use use_contiguous_buffers_in_local_ddp and no + # overlap_dp_param_comm for chunk in self.critic_module: chunk.zero_grad_buffer() micro_batch_size = self.config.ppo_micro_batch_size_per_gpu max_token_len = None if self.config.use_dynamic_bsz: - max_token_len = self.config.ppo_max_token_len_per_gpu * self.config.megatron.context_parallel_size + max_token_len = ( + self.config.ppo_max_token_len_per_gpu + * self.config.megatron.context_parallel_size + ) metric_micro_batch = self.forward_backward_batch( data, forward_only=False, @@ -315,7 +355,9 @@ def update_critic(self, dataloader: Iterable[DataProto]): mini_batch_size=self.config.ppo_mini_batch_size, ) metric_micro_batch = metric_micro_batch["output"] - update_successful, grad_norm, num_zeros_in_grad = self.critic_optimizer.step() + update_successful, grad_norm, num_zeros_in_grad = ( + self.critic_optimizer.step() + ) learning_rate = self.critic_optimizer.param_groups[-1]["lr"] data = {"critic/grad_norm": grad_norm, "critic/lr": learning_rate} append_to_dict(metrics, data) @@ -327,7 +369,9 @@ def update_critic(self, dataloader: Iterable[DataProto]): raise NotImplementedError for metric in metric_micro_batch: - append_to_dict(metrics, metric) # append the metric from this micro-batch to global metrics. + append_to_dict( + metrics, metric + ) # append the metric from this micro-batch to global metrics. # add empty cache after each compute get_torch_device().empty_cache() diff --git a/Agent0/executor_train/verl/verl/workers/fsdp_workers.py b/Agent0/executor_train/verl/verl/workers/fsdp_workers.py index f9bb475..90c3723 100644 --- a/Agent0/executor_train/verl/verl/workers/fsdp_workers.py +++ b/Agent0/executor_train/verl/verl/workers/fsdp_workers.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -69,7 +69,12 @@ ) from verl.utils.import_utils import import_external_libs from verl.utils.model import compute_position_id_with_mask -from verl.utils.profiler import DistProfiler, DistProfilerExtension, log_gpu_memory_usage, simple_timer +from verl.utils.profiler import ( + DistProfiler, + DistProfilerExtension, + log_gpu_memory_usage, + simple_timer, +) from verl.utils.profiler.performance import reduce_timing from verl.utils.py_functional import convert_to_regular_types from verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager @@ -82,10 +87,14 @@ def create_device_mesh(world_size, fsdp_size): if fsdp_size < 0 or fsdp_size >= world_size: - device_mesh = init_device_mesh(device_name, mesh_shape=(world_size,), mesh_dim_names=["fsdp"]) + device_mesh = init_device_mesh( + device_name, mesh_shape=(world_size,), mesh_dim_names=["fsdp"] + ) else: device_mesh = init_device_mesh( - device_name, mesh_shape=(world_size // fsdp_size, fsdp_size), mesh_dim_names=["ddp", "fsdp"] + device_name, + mesh_shape=(world_size // fsdp_size, fsdp_size), + mesh_dim_names=["ddp", "fsdp"], ) return device_mesh @@ -98,7 +107,9 @@ def get_sharding_strategy(device_mesh): elif device_mesh.ndim == 2: sharding_strategy = ShardingStrategy.HYBRID_SHARD else: - raise NotImplementedError(f"Get device mesh ndim={device_mesh.ndim}, but only support 1 or 2") + raise NotImplementedError( + f"Get device mesh ndim={device_mesh.ndim}, but only support 1 or 2" + ) return sharding_strategy @@ -128,26 +139,45 @@ def __init__(self, config: DictConfig, role: str, **kwargs): # build device mesh for FSDP world_size = torch.distributed.get_world_size() # TODO(sgm): support FSDP hybrid shard for larger model - self.device_mesh = create_device_mesh(world_size=world_size, fsdp_size=self.config.actor.fsdp_config.fsdp_size) + self.device_mesh = create_device_mesh( + world_size=world_size, + fsdp_size=self.config.actor.fsdp_config.fsdp_size) # 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) + 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( - device_name, mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=["dp", "sp"] + device_name, + mesh_shape=(dp, self.ulysses_sequence_parallel_size), + mesh_dim_names=["dp", "sp"], ) - self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) + self.ulysses_sharding_manager = FSDPUlyssesShardingManager( + self.ulysses_device_mesh + ) self._lora_rank = self.config.model.get("lora_rank", 0) self._is_lora = self._lora_rank > 0 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"] + 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(haibin.lin): @@ -155,56 +185,82 @@ def __init__(self, config: DictConfig, role: str, **kwargs): # it will actually convert the ProfilerConfig dataclass back to a DictConfig. # We can still use ProfilerConfig for testing purpose (tests/utils/test_nvtx_profile.py) # as they provides DictConfig-like interface - # The benefit of creating the dataclass config is to perform validation during __post_init__ + # The benefit of creating the dataclass config is to perform validation + # during __post_init__ profiler_config = omega_conf_to_dataclass(config.get("profiler")) DistProfilerExtension.__init__( - self, DistProfiler(rank=self.rank, config=profiler_config, option=self.profile_option) + self, + DistProfiler( + rank=self.rank, + config=profiler_config, + option=self.profile_option), ) self._is_offload_param = 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_optimizer = self.config.actor.fsdp_config.get("optimizer_offload", False) + self._is_offload_param = self.config.actor.fsdp_config.get( + "param_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) + 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.config.rollout.n - self.config.actor.ppo_mini_batch_size //= self.device_mesh.size() // self.ulysses_sequence_parallel_size - assert self.config.actor.ppo_mini_batch_size > 0, ( - f"ppo_mini_batch_size {self.config.actor.ppo_mini_batch_size} should be larger than 0 after " - f"normalization" + self.config.actor.ppo_mini_batch_size //= ( + self.device_mesh.size() // self.ulysses_sequence_parallel_size ) + assert self.config.actor.ppo_mini_batch_size > 0, (f"ppo_mini_batch_size { + self.config.actor.ppo_mini_batch_size} should be larger than 0 after " f"normalization") # micro bsz if self.config.actor.ppo_micro_batch_size is not None: self.config.actor.ppo_micro_batch_size //= ( - self.device_mesh.size() // self.ulysses_sequence_parallel_size + self.device_mesh.size() // self.ulysses_sequence_parallel_size) + self.config.actor.ppo_micro_batch_size_per_gpu = ( + self.config.actor.ppo_micro_batch_size ) - self.config.actor.ppo_micro_batch_size_per_gpu = self.config.actor.ppo_micro_batch_size if self.config.actor.ppo_micro_batch_size_per_gpu is not None: - assert self.config.actor.ppo_mini_batch_size % self.config.actor.ppo_micro_batch_size_per_gpu == 0, ( - f"normalized ppo_mini_batch_size {self.config.actor.ppo_mini_batch_size} should be divisible by " - f"ppo_micro_batch_size_per_gpu {self.config.actor.ppo_micro_batch_size_per_gpu}" - ) - assert self.config.actor.ppo_mini_batch_size // self.config.actor.ppo_micro_batch_size_per_gpu > 0, ( + assert ( + self.config.actor.ppo_mini_batch_size % + self.config.actor.ppo_micro_batch_size_per_gpu == 0), (f"normalized ppo_mini_batch_size { + self.config.actor.ppo_mini_batch_size} should be divisible by " f"ppo_micro_batch_size_per_gpu { + self.config.actor.ppo_micro_batch_size_per_gpu}") + assert ( + self.config.actor.ppo_mini_batch_size + // self.config.actor.ppo_micro_batch_size_per_gpu + > 0 + ), ( f"normalized ppo_mini_batch_size {self.config.actor.ppo_mini_batch_size} should be larger than " f"ppo_micro_batch_size_per_gpu {self.config.actor.ppo_micro_batch_size_per_gpu}" ) # normalize rollout config - if self._is_rollout and self.config.rollout.log_prob_micro_batch_size is not None: + if ( + self._is_rollout + and self.config.rollout.log_prob_micro_batch_size is not None + ): self.config.rollout.log_prob_micro_batch_size //= ( self.device_mesh.size() // self.ulysses_sequence_parallel_size ) - self.config.rollout.log_prob_micro_batch_size_per_gpu = self.config.rollout.log_prob_micro_batch_size + self.config.rollout.log_prob_micro_batch_size_per_gpu = ( + self.config.rollout.log_prob_micro_batch_size + ) # normalize ref config if self._is_ref and self.config.ref.log_prob_micro_batch_size is not None: - self.config.ref.log_prob_micro_batch_size //= self.device_mesh.size() // self.ulysses_sequence_parallel_size - self.config.ref.log_prob_micro_batch_size_per_gpu = self.config.ref.log_prob_micro_batch_size + self.config.ref.log_prob_micro_batch_size //= ( + self.device_mesh.size() // self.ulysses_sequence_parallel_size + ) + self.config.ref.log_prob_micro_batch_size_per_gpu = ( + self.config.ref.log_prob_micro_batch_size + ) def _build_model_optimizer( self, @@ -222,20 +278,33 @@ def _build_model_optimizer( ): from torch import optim from torch.distributed.fsdp import CPUOffload, MixedPrecision - from transformers import AutoConfig, AutoModelForCausalLM, AutoModelForVision2Seq + from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoModelForVision2Seq, + ) - from verl.utils.model import get_generation_config, print_model_size, update_model_config + from verl.utils.model import ( + get_generation_config, + print_model_size, + update_model_config, + ) from verl.utils.torch_dtypes import PrecisionType assert role in ["actor", "ref"] - log_gpu_memory_usage(f"Before init {role} from HF AutoModel", logger=logger) + log_gpu_memory_usage( + f"Before init {role} from HF AutoModel", + logger=logger) local_path = 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) - self.processor = hf_processor(local_path, trust_remote_code=trust_remote_code) + # 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) + self.processor = hf_processor( + local_path, trust_remote_code=trust_remote_code) if self.config.model.get("custom_chat_template", None) is not None: if self.processor is not None: @@ -251,14 +320,18 @@ def _build_model_optimizer( # override model kwargs actor_model_config = AutoConfig.from_pretrained( - local_path, trust_remote_code=trust_remote_code, attn_implementation="flash_attention_2" + local_path, + trust_remote_code=trust_remote_code, + attn_implementation="flash_attention_2", ) # patch for kimi-vl if getattr(actor_model_config, "model_type", None) == "kimi_vl": actor_model_config.text_config.topk_method = "greedy" - self.generation_config = get_generation_config(local_path, trust_remote_code=trust_remote_code) + self.generation_config = get_generation_config( + local_path, trust_remote_code=trust_remote_code + ) override_config_kwargs = { "bos_token_id": self.tokenizer.bos_token_id, @@ -266,18 +339,22 @@ def _build_model_optimizer( "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) + 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, mesh=self.device_mesh + use_meta_tensor=not actor_model_config.tie_word_embeddings, + mesh=self.device_mesh, ) with init_context(), warnings.catch_warnings(): warnings.simplefilter("ignore") - if type(actor_model_config) in AutoModelForVision2Seq._model_mapping.keys(): + if type( + actor_model_config) in AutoModelForVision2Seq._model_mapping.keys(): actor_module_class = AutoModelForVision2Seq else: actor_module_class = AutoModelForCausalLM @@ -291,13 +368,18 @@ def _build_model_optimizer( # Apply Liger kernel to the model if use_liger is set to True if use_liger: - from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance + from liger_kernel.transformers.monkey_patch import ( + _apply_liger_kernel_to_instance, + ) _apply_liger_kernel_to_instance(model=actor_module) - fused_kernel_options = self.config.model.get("fused_kernel_options", None) + fused_kernel_options = self.config.model.get( + "fused_kernel_options", None) fused_kernels_backend = ( - fused_kernel_options.get("impl_backend", None) if fused_kernel_options is not None else None + fused_kernel_options.get("impl_backend", None) + if fused_kernel_options is not None + else None ) apply_monkey_patch( @@ -308,43 +390,64 @@ def _build_model_optimizer( fused_kernels_backend=fused_kernels_backend, ) - # some parameters may not in torch_dtype. TODO(zhangchi.usc1992) remove this after we switch to fsdp2 + # 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}) + actor_module.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) if self._is_lora: print("Applying LoRA to actor module") actor_module.enable_input_require_grads() - # Convert config to regular Python types before creating PEFT model + # 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), - "exclude_modules": convert_to_regular_types(self.config.model.exclude_modules), + "target_modules": convert_to_regular_types( + self.config.model.target_modules + ), + "exclude_modules": convert_to_regular_types( + self.config.model.exclude_modules + ), "bias": "none", } - actor_module = get_peft_model(actor_module, LoraConfig(**lora_config)) + actor_module = get_peft_model( + actor_module, LoraConfig(**lora_config)) torch.distributed.barrier() if self.rank == 0: print_model_size(actor_module) - log_gpu_memory_usage(f"After init {role} from HF AutoModel", logger=logger) + log_gpu_memory_usage( + f"After init {role} 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")) + 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) + mixed_precision = MixedPrecision( + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + buffer_dtype=buffer_dtype, + ) auto_wrap_policy = get_fsdp_wrap_policy( module=actor_module, @@ -353,7 +456,8 @@ def _build_model_optimizer( ) 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 + # TODO(zhangchi.usc1992, shengguangming) fix me. Current, + # auto_wrap_policy causes HFRollout to hang in Gemma auto_wrap_policy = None if self.rank == 0: @@ -364,8 +468,10 @@ def _build_model_optimizer( # TODO: add transformer policy # We force reference policy to use CPUOffload to save memory. - # We force turn off CPUOffload for actor because it causes incorrect results when using grad accumulation - cpu_offload = None if role == "actor" else CPUOffload(offload_params=True) + # We force turn off CPUOffload for actor because it causes incorrect + # results when using grad accumulation + cpu_offload = None if role == "actor" else CPUOffload( + offload_params=True) fsdp_strategy = self.config.actor.strategy if fsdp_strategy == "fsdp": actor_module_fsdp = FSDP( @@ -378,20 +484,30 @@ def _build_model_optimizer( mixed_precision=mixed_precision, sync_module_states=True, device_mesh=self.device_mesh, - use_orig_params=self.config.actor.fsdp_config.get("use_orig_params", False), - forward_prefetch=self.config.actor.fsdp_config.get("forward_prefetch", False), + use_orig_params=self.config.actor.fsdp_config.get( + "use_orig_params", False + ), + forward_prefetch=self.config.actor.fsdp_config.get( + "forward_prefetch", False + ), ) elif fsdp_strategy == "fsdp2": - assert CPUOffloadPolicy is not None, "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" + assert ( + CPUOffloadPolicy is not None + ), "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" mp_policy = MixedPrecisionPolicy( - param_dtype=param_dtype, reduce_dtype=reduce_dtype, cast_forward_inputs=True + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + cast_forward_inputs=True, ) if role == "actor" and fsdp_config.offload_policy: cpu_offload = CPUOffloadPolicy(pin_memory=True) self._is_offload_param = False self._is_offload_optimizer = False else: - cpu_offload = None if role == "actor" else CPUOffloadPolicy(pin_memory=True) + cpu_offload = ( + None if role == "actor" else CPUOffloadPolicy( + pin_memory=True)) fsdp_kwargs = { "mesh": fsdp_mesh, @@ -401,19 +517,25 @@ def _build_model_optimizer( } full_state = actor_module.state_dict() apply_fsdp2(actor_module, fsdp_kwargs, fsdp_config) - fsdp2_load_full_state_dict(actor_module, full_state, fsdp_mesh, cpu_offload) + fsdp2_load_full_state_dict( + actor_module, full_state, fsdp_mesh, cpu_offload) actor_module_fsdp = actor_module else: raise NotImplementedError(f"not implement {fsdp_strategy}") if enable_activation_offload: - enable_activation_offloading(actor_module_fsdp, fsdp_strategy, enable_gradient_checkpointing) + enable_activation_offloading( + actor_module_fsdp, fsdp_strategy, enable_gradient_checkpointing + ) log_gpu_memory_usage(f"After {role} FSDP init", logger=logger) # TODO: add more optimizer args into config if role == "actor" and optim_config is not None: - from verl.utils.torch_functional import get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup + from verl.utils.torch_functional import ( + get_constant_schedule_with_warmup, + get_cosine_schedule_with_warmup, + ) actor_optimizer = optim.AdamW( actor_module_fsdp.parameters(), @@ -428,16 +550,17 @@ def _build_model_optimizer( min_lr_ratio = optim_config.get("min_lr_ratio", 0.0) num_cycles = optim_config.get("num_cycles", 0.5) if num_warmup_steps < 0: - num_warmup_steps_ratio = optim_config.get("lr_warmup_steps_ratio", 0.0) + num_warmup_steps_ratio = optim_config.get( + "lr_warmup_steps_ratio", 0.0) num_warmup_steps = int(num_warmup_steps_ratio * total_steps) if self.rank == 0: - print(f"Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}") + print( + f"Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}") if warmup_style == "constant": actor_lr_scheduler = get_constant_schedule_with_warmup( - optimizer=actor_optimizer, num_warmup_steps=num_warmup_steps - ) + optimizer=actor_optimizer, num_warmup_steps=num_warmup_steps) elif warmup_style == "cosine": actor_lr_scheduler = get_cosine_schedule_with_warmup( optimizer=actor_optimizer, @@ -447,14 +570,21 @@ def _build_model_optimizer( num_cycles=num_cycles, ) else: - raise NotImplementedError(f"Warmup style {warmup_style} is not supported") + raise NotImplementedError( + f"Warmup style {warmup_style} is not supported" + ) log_gpu_memory_usage(f"After {role} optimizer init", logger=logger) else: actor_optimizer = None actor_lr_scheduler = None - return actor_module_fsdp, actor_optimizer, actor_lr_scheduler, actor_model_config + return ( + actor_module_fsdp, + actor_optimizer, + actor_lr_scheduler, + actor_model_config, + ) def _build_rollout(self, trust_remote_code=False): from torch.distributed.device_mesh import init_device_mesh @@ -462,18 +592,22 @@ def _build_rollout(self, trust_remote_code=False): # 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}" - ) + 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( - device_name, mesh_shape=(dp, infer_tp), mesh_dim_names=["dp", "infer_tp"] - ) + device_name, mesh_shape=( + dp, infer_tp), mesh_dim_names=[ + "dp", "infer_tp"]) rollout_name = self.config.rollout.name if rollout_name == "hf": from verl.workers.rollout import HFRollout from verl.workers.sharding_manager.base import BaseShardingManager - rollout = HFRollout(module=self.actor_module_fsdp, config=self.config.rollout) + rollout = HFRollout( + module=self.actor_module_fsdp, config=self.config.rollout + ) rollout_sharding_manager = BaseShardingManager() # TODO: a sharding manager that do nothing? @@ -481,17 +615,30 @@ def _build_rollout(self, trust_remote_code=False): from verl.workers.rollout.vllm_rollout import vLLMRollout from verl.workers.sharding_manager.fsdp_vllm import FSDPVLLMShardingManager - log_gpu_memory_usage(f"Before building {rollout_name} rollout", logger=logger) - local_path = copy_to_local(self.config.model.path, use_shm=self.config.model.get("use_shm", False)) + log_gpu_memory_usage( + f"Before building {rollout_name} rollout", logger=logger + ) + local_path = copy_to_local( + self.config.model.path, + use_shm=self.config.model.get( + "use_shm", + False)) lora_kwargs = ( - {"lora_kwargs": {"enable_lora": True, "max_loras": 1, "max_lora_rank": self._lora_rank}} + { + "lora_kwargs": { + "enable_lora": True, + "max_loras": 1, + "max_lora_rank": self._lora_rank, + } + } if self._is_lora else {} ) # lora_kwargs = {} from verl.workers.rollout.vllm_rollout import vLLMAsyncRollout - vllm_rollout_cls = vLLMRollout if self.config.rollout.mode == "sync" else vLLMAsyncRollout + vllm_rollout_cls = ( + vLLMRollout if self.config.rollout.mode == "sync" else vLLMAsyncRollout) rollout = vllm_rollout_cls( model_path=local_path, config=self.config.rollout, @@ -502,7 +649,9 @@ def _build_rollout(self, trust_remote_code=False): **lora_kwargs, ) - log_gpu_memory_usage(f"After building {rollout_name} rollout", logger=logger) + log_gpu_memory_usage( + f"After building {rollout_name} rollout", logger=logger + ) full_params = torch.distributed.get_world_size() == 1 rollout_sharding_manager = FSDPVLLMShardingManager( module=self.actor_module_fsdp, @@ -513,9 +662,13 @@ def _build_rollout(self, trust_remote_code=False): device_mesh=rollout_device_mesh, offload_param=self._is_offload_param, load_format=self.config.rollout.load_format, - layered_summon=self.config.rollout.get("layered_summon", False), + layered_summon=self.config.rollout.get( + "layered_summon", + False), ) - log_gpu_memory_usage("After building sharding manager", logger=logger) + log_gpu_memory_usage( + "After building sharding manager", + logger=logger) elif rollout_name == "sglang": from verl.workers.rollout.sglang_rollout import SGLangRollout @@ -526,19 +679,27 @@ def _build_rollout(self, trust_remote_code=False): # "RuntimeError: No CUDA GPUs are available". # For this reason, sharding_manager.__init__ should not import FSDPSGLangShardingManager and # we import it here use the abs path. - # check: https://github.com/sgl-project/sglang/blob/00f42707eaddfc2c0528e5b1e0094025c640b7a0/python/sglang/srt/layers/quantization/fp8_utils.py#L76 - from verl.workers.sharding_manager.fsdp_sglang import FSDPSGLangShardingManager + # check: + # https://github.com/sgl-project/sglang/blob/00f42707eaddfc2c0528e5b1e0094025c640b7a0/python/sglang/srt/layers/quantization/fp8_utils.py#L76 + from verl.workers.sharding_manager.fsdp_sglang import ( + FSDPSGLangShardingManager, + ) local_path = copy_to_local(self.config.model.path) - log_gpu_memory_usage(f"Before building {rollout_name} rollout", logger=logger) + log_gpu_memory_usage( + f"Before building {rollout_name} rollout", logger=logger + ) rollout = SGLangRollout( actor_module=local_path, config=self.config.rollout, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), model_hf_config=self.actor_model_config, trust_remote_code=trust_remote_code, ) - log_gpu_memory_usage(f"After building {rollout_name} rollout", logger=logger) + log_gpu_memory_usage( + f"After building {rollout_name} rollout", logger=logger + ) if torch.distributed.get_world_size() == 1: self.config.rollout.load_format = "dummy_hf" @@ -552,10 +713,14 @@ def _build_rollout(self, trust_remote_code=False): offload_param=self._is_offload_param, multi_stage_wake_up=self.config.rollout.multi_stage_wake_up, ) - log_gpu_memory_usage("After building sharding manager", logger=logger) + log_gpu_memory_usage( + "After building sharding manager", + logger=logger) else: - raise NotImplementedError(f"Rollout name: {self.config.rollout.name} is not supported") + raise NotImplementedError( + f"Rollout name: {self.config.rollout.name} is not supported" + ) return rollout, rollout_sharding_manager @@ -566,7 +731,9 @@ 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)) - override_model_config = OmegaConf.to_container(self.config.model.get("override_config", OmegaConf.create())) + 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) use_shm = self.config.model.get("use_shm", False) @@ -594,11 +761,19 @@ def init_model(self): override_model_config=override_model_config, use_remove_padding=use_remove_padding, use_fused_kernels=use_fused_kernels, - enable_gradient_checkpointing=self.config.model.get("enable_gradient_checkpointing", False), - trust_remote_code=self.config.model.get("trust_remote_code", False), - use_liger=self.config.model.get("use_liger", False), + enable_gradient_checkpointing=self.config.model.get( + "enable_gradient_checkpointing", + False), + trust_remote_code=self.config.model.get( + "trust_remote_code", + False), + use_liger=self.config.model.get( + "use_liger", + False), role="actor", - enable_activation_offload=self.config.model.get("enable_activation_offload", False), + enable_activation_offload=self.config.model.get( + "enable_activation_offload", + False), ) # get the original unwrapped module @@ -607,11 +782,15 @@ def init_model(self): if self._is_offload_param: offload_fsdp_model_to_cpu(self.actor_module_fsdp) - log_gpu_memory_usage("After offload actor model during init", logger=logger) + log_gpu_memory_usage( + "After offload actor model 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) + log_gpu_memory_usage( + "After offload actor optimizer during init", logger=logger + ) if self._is_actor: OmegaConf.set_struct(self.config.actor, True) @@ -619,7 +798,9 @@ def init_model(self): self.config.actor.use_remove_padding = use_remove_padding self.config.actor.use_fused_kernels = use_fused_kernels self.actor = DataParallelPPOActor( - config=self.config.actor, actor_module=self.actor_module_fsdp, actor_optimizer=self.actor_optimizer + config=self.config.actor, + actor_module=self.actor_module_fsdp, + actor_optimizer=self.actor_optimizer, ) if self._is_rollout: @@ -636,15 +817,21 @@ def init_model(self): override_model_config=override_model_config, use_remove_padding=use_remove_padding, use_fused_kernels=use_fused_kernels, - trust_remote_code=self.config.model.get("trust_remote_code", False), - use_liger=self.config.model.get("use_liger", False), + trust_remote_code=self.config.model.get( + "trust_remote_code", + False), + use_liger=self.config.model.get( + "use_liger", + False), role="ref", )[0] OmegaConf.set_struct(self.config.ref, True) with open_dict(self.config.ref): self.config.ref.use_remove_padding = use_remove_padding self.config.ref.use_fused_kernels = use_fused_kernels - self.ref_policy = DataParallelPPOActor(config=self.config.ref, actor_module=self.ref_module_fsdp) + 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) @@ -652,20 +839,25 @@ def init_model(self): model=self.actor_module_fsdp, optimizer=self.actor.actor_optimizer, lr_scheduler=self.actor_lr_scheduler, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), checkpoint_config=self.config.actor.checkpoint, ) if not self._is_actor and self._is_rollout: # If ActorRolloutRefWorker is initialized as a standalone rollout, - # create a checkpoint manager for FSDP model to allow loading FSDP checkpoints for rollout. + # create a checkpoint manager for FSDP model to allow loading FSDP + # checkpoints for rollout. - checkpoint_contents = OmegaConf.create({"load_contents": ["model"], "save_contents": []}) + checkpoint_contents = OmegaConf.create( + {"load_contents": ["model"], "save_contents": []} + ) self.checkpoint_manager = FSDPCheckpointManager( model=self.actor_module_fsdp, optimizer=None, lr_scheduler=None, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), checkpoint_config=checkpoint_contents, ) @@ -679,7 +871,9 @@ def update_actor(self, data: DataProto): if self._is_offload_param: load_fsdp_model_to_gpu(self.actor_module_fsdp) if self._is_offload_optimizer: - load_fsdp_optimizer(optimizer=self.actor_optimizer, device_id=get_device_id()) + load_fsdp_optimizer( + optimizer=self.actor_optimizer, device_id=get_device_id() + ) with self.ulysses_sharding_manager: data = self.ulysses_sharding_manager.preprocess_data(data=data) @@ -688,13 +882,23 @@ def update_actor(self, data: DataProto): 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) + estimated_flops, promised_flops = self.flops_counter.estimate_flops( + global_num_tokens, delta_time) metrics["perf/mfu/actor"] = ( - estimated_flops * self.config.actor.ppo_epochs / promised_flops / self.world_size + estimated_flops + * self.config.actor.ppo_epochs + / promised_flops + / self.world_size + ) + metrics["perf/max_memory_allocated_gb"] = ( + get_torch_device().max_memory_allocated() / (1024**3) + ) + metrics["perf/max_memory_reserved_gb"] = ( + get_torch_device().max_memory_reserved() / (1024**3) + ) + metrics["perf/cpu_memory_used_gb"] = psutil.virtual_memory().used / ( + 1024**3 ) - metrics["perf/max_memory_allocated_gb"] = get_torch_device().max_memory_allocated() / (1024**3) - metrics["perf/max_memory_reserved_gb"] = get_torch_device().max_memory_reserved() / (1024**3) - metrics["perf/cpu_memory_used_gb"] = psutil.virtual_memory().used / (1024**3) lr = self.actor_lr_scheduler.get_last_lr()[0] metrics["actor/lr"] = lr @@ -703,15 +907,20 @@ def update_actor(self, data: DataProto): # TODO: here, we should return all metrics output = DataProto(meta_info={"metrics": metrics}) - output = self.ulysses_sharding_manager.postprocess_data(data=output) + output = self.ulysses_sharding_manager.postprocess_data( + data=output) output = output.to("cpu") if self._is_offload_param: offload_fsdp_model_to_cpu(self.actor_module_fsdp) - log_gpu_memory_usage("After offload actor model during update_actor", logger=logger) + log_gpu_memory_usage( + "After offload actor model during update_actor", logger=logger + ) if self._is_offload_optimizer: offload_fsdp_optimizer(optimizer=self.actor_optimizer) - log_gpu_memory_usage("After offload actor optimizer during update_actor", logger=logger) + log_gpu_memory_usage( + "After offload actor optimizer during update_actor", + logger=logger) return output @@ -724,17 +933,23 @@ def generate_sequences(self, prompts: DataProto): assert self._is_rollout meta_info = { - "eos_token_id": self.generation_config.eos_token_id - if self.generation_config is not None - else self.tokenizer.eos_token_id, - "pad_token_id": self.generation_config.pad_token_id - if self.generation_config is not None - else self.tokenizer.pad_token_id, + "eos_token_id": ( + self.generation_config.eos_token_id + if self.generation_config is not None + else self.tokenizer.eos_token_id + ), + "pad_token_id": ( + self.generation_config.pad_token_id + if self.generation_config is not None + else self.tokenizer.pad_token_id + ), } prompts.meta_info.update(meta_info) timing_generate = {} with self.rollout_sharding_manager: - log_gpu_memory_usage("After entering rollout sharding manager", logger=logger) + log_gpu_memory_usage( + "After entering rollout sharding manager", logger=logger + ) prompts = self.rollout_sharding_manager.preprocess_data(prompts) with simple_timer("generate_sequences", timing_generate): @@ -768,18 +983,25 @@ def compute_log_prob(self, data: DataProto): from contextlib import nullcontext is_lora = data.meta_info.pop("is_lora", False) - adapter_ctx = self.actor.actor_module.disable_adapter() if is_lora else nullcontext() + adapter_ctx = (self.actor.actor_module.disable_adapter() + if is_lora else nullcontext()) data = data.to(get_device_id()) # we should always recompute old_log_probs when it is HybridEngine - data.meta_info["micro_batch_size"] = self.config.rollout.log_prob_micro_batch_size_per_gpu - data.meta_info["max_token_len"] = self.config.rollout.log_prob_max_token_len_per_gpu + data.meta_info["micro_batch_size"] = ( + self.config.rollout.log_prob_micro_batch_size_per_gpu + ) + data.meta_info["max_token_len"] = ( + self.config.rollout.log_prob_max_token_len_per_gpu + ) data.meta_info["use_dynamic_bsz"] = self.config.rollout.log_prob_use_dynamic_bsz data.meta_info["temperature"] = self.config.rollout.temperature # perform recompute log_prob with self.ulysses_sharding_manager: data = self.ulysses_sharding_manager.preprocess_data(data) with adapter_ctx: - output, entropys = self.actor.compute_log_prob(data=data, calculate_entropy=True) + output, entropys = self.actor.compute_log_prob( + data=data, calculate_entropy=True + ) output = DataProto.from_dict( tensors={"old_log_probs": output, "entropys": entropys}, meta_info={"temperature": self.config.rollout.temperature}, @@ -795,7 +1017,9 @@ def compute_log_prob(self, data: DataProto): if self._is_offload_param: offload_fsdp_model_to_cpu(self.actor_module_fsdp) - log_gpu_memory_usage("After offload actor model during compute_log_prob", logger=logger) + log_gpu_memory_usage( + "After offload actor model during compute_log_prob", + logger=logger) return output @@ -807,7 +1031,9 @@ def compute_ref_log_prob(self, data: DataProto): data.meta_info["is_lora"] = True data = self.compute_log_prob(data) # this old_log_probs is in fact ref_log_prob - data = DataProto.from_dict(tensors={"ref_log_prob": data.batch["old_log_probs"]}) + data = DataProto.from_dict( + tensors={"ref_log_prob": data.batch["old_log_probs"]} + ) return data assert self._is_ref # else: @@ -822,7 +1048,9 @@ def compute_ref_log_prob(self, data: DataProto): 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, calculate_entropy=False) + output, _ = self.ref_policy.compute_log_prob( + data=data, calculate_entropy=False + ) output = DataProto.from_dict(tensors={"ref_log_prob": output}) output = self.ulysses_sharding_manager.postprocess_data(output) @@ -830,13 +1058,16 @@ def compute_ref_log_prob(self, data: DataProto): # https://pytorch.org/docs/stable/notes/fsdp.html#fsdp-notes # unshard the root FSDP module - if self.world_size > 1 and fsdp_version(self.ref_policy.actor_module) == 1: + if self.world_size > 1 and fsdp_version( + self.ref_policy.actor_module) == 1: self.ref_policy.actor_module._handle.reshard(True) return output @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None): + def save_checkpoint( + self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None + ): from verl.utils.logger import log_with_rank # only support save and load ckpt for actor @@ -846,11 +1077,19 @@ def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to load_fsdp_model_to_gpu(self.actor_module_fsdp) self.checkpoint_manager.save_checkpoint( - local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep + local_path=local_path, + hdfs_path=hdfs_path, + global_step=global_step, + max_ckpt_to_keep=max_ckpt_to_keep, ) dist.barrier() - if self._is_lora and hasattr(getattr(self, "actor_module", self.actor_module_fsdp), "peft_config"): + if self._is_lora and hasattr( + getattr( + self, + "actor_module", + self.actor_module_fsdp), + "peft_config"): lora_save_path = os.path.join(local_path, "lora_adapter") peft_model = getattr(self, "actor_module", self.actor_module_fsdp) peft_config = {} @@ -859,18 +1098,32 @@ def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to peft_config = asdict(peft_model.peft_config.get("default", {})) peft_config["task_type"] = peft_config["task_type"].value peft_config["peft_type"] = peft_config["peft_type"].value - peft_config["target_modules"] = list(peft_config["target_modules"]) + peft_config["target_modules"] = list( + peft_config["target_modules"]) try: if fsdp_version(self.actor_module_fsdp) > 0: - self.actor_module_fsdp = self.actor_module_fsdp.to(get_device_name()) - lora_params = layered_summon_lora_params(self.actor_module_fsdp) + self.actor_module_fsdp = self.actor_module_fsdp.to( + get_device_name() + ) + lora_params = layered_summon_lora_params( + self.actor_module_fsdp) if dist.get_rank() == 0: - save_file(lora_params, os.path.join(lora_save_path, "adapter_model.safetensors")) - with open(os.path.join(lora_save_path, "adapter_config.json"), "w", encoding="utf-8") as f: - json.dump(peft_config, f, ensure_ascii=False, indent=4) + save_file( + lora_params, os.path.join( + lora_save_path, "adapter_model.safetensors"), ) + with open( + os.path.join(lora_save_path, "adapter_config.json"), + "w", + encoding="utf-8", + ) as f: + json.dump( + peft_config, f, ensure_ascii=False, indent=4) except Exception as e: log_with_rank( - f"Save LoRA Adapter Error ({e})", rank=dist.get_rank(), logger=logger, log_only_rank_0=True + f"Save LoRA Adapter Error ({e})", + rank=dist.get_rank(), + logger=logger, + log_only_rank_0=True, ) dist.barrier() @@ -885,7 +1138,11 @@ def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to offload_fsdp_model_to_cpu(self.actor_module_fsdp) @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=False): + def load_checkpoint( + self, + local_path, + hdfs_path=None, + del_local_after_load=False): assert self._is_actor or (not self._is_actor and self._is_rollout), ( f"Checkpoint loading is only supported for Actor or standalone Rollout Workers, but got " f"{self._is_actor} and {self._is_rollout}" @@ -895,7 +1152,9 @@ def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=False load_fsdp_model_to_gpu(self.actor_module_fsdp) self.checkpoint_manager.load_checkpoint( - local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load + local_path=local_path, + hdfs_path=hdfs_path, + del_local_after_load=del_local_after_load, ) if self._is_offload_param: @@ -919,13 +1178,15 @@ class CriticWorker(Worker, DistProfilerExtension): def __init__(self, config): Worker.__init__(self) DistProfilerExtension.__init__( - self, DistProfiler(rank=self.rank, config=omega_conf_to_dataclass(config.get("profiler"))) - ) + self, DistProfiler( + rank=self.rank, config=omega_conf_to_dataclass( + config.get("profiler"))), ) import torch.distributed if not torch.distributed.is_initialized(): torch.distributed.init_process_group( - backend=get_nccl_backend(), init_method=os.environ.get("DIST_INIT_METHOD", None) + backend=get_nccl_backend(), + init_method=os.environ.get("DIST_INIT_METHOD", None), ) self.config = config @@ -934,17 +1195,25 @@ def __init__(self, config): from torch.distributed.device_mesh import init_device_mesh fsdp_size = self.config.model.fsdp_config.fsdp_size - self.device_mesh = create_device_mesh(world_size=world_size, fsdp_size=fsdp_size) + self.device_mesh = create_device_mesh( + world_size=world_size, fsdp_size=fsdp_size + ) self.ulysses_device_mesh = None - self.ulysses_sequence_parallel_size = self.config.get("ulysses_sequence_parallel_size", 1) + 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( - device_name, mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=["dp", "sp"] + device_name, + mesh_shape=(dp, self.ulysses_sequence_parallel_size), + mesh_dim_names=["dp", "sp"], ) - self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) + self.ulysses_sharding_manager = FSDPUlyssesShardingManager( + self.ulysses_device_mesh + ) # set FSDP offload params self._is_offload_param = self.config.model.fsdp_config.param_offload @@ -952,23 +1221,34 @@ def __init__(self, config): # normalize config self.config.ppo_mini_batch_size *= self.config.rollout_n - self.config.ppo_mini_batch_size //= torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size + self.config.ppo_mini_batch_size //= ( + torch.distributed.get_world_size() // + self.ulysses_sequence_parallel_size) if self.config.ppo_micro_batch_size is not None: self.config.ppo_micro_batch_size //= ( - torch.distributed.get_world_size() // self.ulysses_sequence_parallel_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 + torch.distributed.get_world_size() + // self.ulysses_sequence_parallel_size ) self.config.ppo_micro_batch_size_per_gpu = self.config.ppo_micro_batch_size - self.config.forward_micro_batch_size_per_gpu = self.config.forward_micro_batch_size + self.config.forward_micro_batch_size_per_gpu = ( + self.config.forward_micro_batch_size + ) if self.config.ppo_micro_batch_size_per_gpu is not None: - assert self.config.ppo_mini_batch_size % self.config.ppo_micro_batch_size_per_gpu == 0, ( - f"normalized ppo_mini_batch_size {self.config.ppo_mini_batch_size} should be divisible by " - f"ppo_micro_batch_size_per_gpu {self.config.ppo_micro_batch_size_per_gpu}" - ) - assert self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu > 0, ( + assert ( + self.config.ppo_mini_batch_size % + self.config.ppo_micro_batch_size_per_gpu == 0), (f"normalized ppo_mini_batch_size { + self.config.ppo_mini_batch_size} should be divisible by " f"ppo_micro_batch_size_per_gpu { + self.config.ppo_micro_batch_size_per_gpu}") + assert ( + self.config.ppo_mini_batch_size + // self.config.ppo_micro_batch_size_per_gpu + > 0 + ), ( f"normalized ppo_mini_batch_size {self.config.ppo_mini_batch_size} should be larger than " f"ppo_micro_batch_size_per_gpu {self.config.ppo_micro_batch_size_per_gpu}" ) @@ -985,11 +1265,19 @@ def _build_critic_model_optimizer(self, config): use_shm = config.model.get("use_shm", False) local_path = copy_to_local(config.model.path, use_shm=use_shm) # 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. + # using random initialized model from any architecture. May not be the + # same as Actor. - tokenizer_path = copy_to_local(config.model.tokenizer_path, use_shm=use_shm) - self.tokenizer = hf_tokenizer(tokenizer_path, trust_remote_code=config.model.get("trust_remote_code", False)) - self.processor = hf_processor(tokenizer_path, trust_remote_code=config.model.get("trust_remote_code", False)) + tokenizer_path = copy_to_local( + config.model.tokenizer_path, use_shm=use_shm) + self.tokenizer = hf_tokenizer( + tokenizer_path, + trust_remote_code=config.model.get("trust_remote_code", False), + ) + self.processor = hf_processor( + tokenizer_path, + trust_remote_code=config.model.get("trust_remote_code", False), + ) if self.config.model.get("custom_chat_template", None) is not None: if self.processor is not None: @@ -997,7 +1285,9 @@ def _build_critic_model_optimizer(self, config): else: self.tokenizer.chat_template = self.config.model.custom_chat_template - override_config = OmegaConf.to_container(self.config.model.get("override_config", OmegaConf.create())) + 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, @@ -1023,7 +1313,8 @@ def _build_critic_model_optimizer(self, config): critic_model_config.text_config.topk_method = "greedy" init_context = get_init_weight_context_manager( - use_meta_tensor=not critic_model_config.tie_word_embeddings, mesh=self.device_mesh + use_meta_tensor=not critic_model_config.tie_word_embeddings, + mesh=self.device_mesh, ) with init_context(), warnings.catch_warnings(): @@ -1051,7 +1342,9 @@ def _build_critic_model_optimizer(self, config): critic_module.to(torch_dtype) if config.model.get("enable_gradient_checkpointing", False): - critic_module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) + critic_module.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) if self._is_lora: print("Applying LoRA to critic module") @@ -1061,10 +1354,13 @@ def _build_critic_model_optimizer(self, 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), + "target_modules": convert_to_regular_types( + self.config.model.target_modules + ), "bias": "none", } - critic_module = get_peft_model(critic_module, LoraConfig(**lora_config)) + critic_module = get_peft_model( + critic_module, LoraConfig(**lora_config)) if self.rank == 0: print_model_size(critic_module) @@ -1074,15 +1370,25 @@ def _build_critic_model_optimizer(self, 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")) + 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) + 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, @@ -1095,7 +1401,8 @@ def _build_critic_model_optimizer(self, config): fsdp_mesh = self.device_mesh sharding_strategy = get_sharding_strategy(fsdp_mesh) - # Note: We force turn off CPUOffload for critic because it causes incorrect results when using grad accumulation + # Note: We force turn off CPUOffload for critic because it causes + # incorrect results when using grad accumulation if config.strategy == "fsdp": critic_module = FSDP( critic_module, @@ -1111,9 +1418,13 @@ def _build_critic_model_optimizer(self, config): cpu_offload=None, ) elif config.strategy == "fsdp2": - assert CPUOffloadPolicy is not None, "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" + assert ( + CPUOffloadPolicy is not None + ), "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" mp_policy = MixedPrecisionPolicy( - param_dtype=param_dtype, reduce_dtype=reduce_dtype, cast_forward_inputs=True + param_dtype=param_dtype, + reduce_dtype=reduce_dtype, + cast_forward_inputs=True, ) offload_policy = None if fsdp_config.offload_policy: @@ -1129,13 +1440,19 @@ def _build_critic_model_optimizer(self, config): } full_state = critic_module.state_dict() apply_fsdp2(critic_module, fsdp_kwargs, fsdp_config) - fsdp2_load_full_state_dict(critic_module, full_state, fsdp_mesh, offload_policy) + fsdp2_load_full_state_dict( + critic_module, full_state, fsdp_mesh, offload_policy + ) else: raise NotImplementedError(f"Unknown strategy {config.strategy}") if config.model.get("enable_activation_offload", False): - enable_gradient_checkpointing = config.model.get("enable_gradient_checkpointing", False) - enable_activation_offloading(critic_module, config.strategy, enable_gradient_checkpointing) + enable_gradient_checkpointing = config.model.get( + "enable_gradient_checkpointing", False + ) + enable_activation_offloading( + critic_module, config.strategy, enable_gradient_checkpointing + ) log_gpu_memory_usage("After critic FSDP", logger=None) @@ -1150,13 +1467,18 @@ def _build_critic_model_optimizer(self, config): num_warmup_steps = int(config.optim.get("lr_warmup_steps", -1)) warmup_style = config.optim.get("warmup_style", "constant") if num_warmup_steps < 0: - num_warmup_steps_ratio = config.optim.get("lr_warmup_steps_ratio", 0.0) + num_warmup_steps_ratio = config.optim.get( + "lr_warmup_steps_ratio", 0.0) num_warmup_steps = int(num_warmup_steps_ratio * total_steps) if self.rank == 0: - print(f"Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}") + print( + f"Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}") - from verl.utils.torch_functional import get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup + from verl.utils.torch_functional import ( + get_constant_schedule_with_warmup, + get_cosine_schedule_with_warmup, + ) if warmup_style == "constant": critic_lr_scheduler = get_constant_schedule_with_warmup( @@ -1164,10 +1486,13 @@ def _build_critic_model_optimizer(self, config): ) elif warmup_style == "cosine": critic_lr_scheduler = get_cosine_schedule_with_warmup( - optimizer=critic_optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=total_steps + optimizer=critic_optimizer, + num_warmup_steps=num_warmup_steps, + num_training_steps=total_steps, ) else: - raise NotImplementedError(f"Warmup style {warmup_style} is not supported") + raise NotImplementedError( + f"Warmup style {warmup_style} is not supported") return critic_module, critic_optimizer, critic_lr_scheduler @@ -1178,19 +1503,25 @@ def init_model(self): from verl.workers.critic import DataParallelPPOCritic - self.critic_module, self.critic_optimizer, self.critic_lr_scheduler = self._build_critic_model_optimizer( - self.config + self.critic_module, self.critic_optimizer, self.critic_lr_scheduler = ( + self._build_critic_model_optimizer(self.config) ) if self._is_offload_param: offload_fsdp_model_to_cpu(self.critic_module) - log_gpu_memory_usage("After offload critic model during init", logger=logger) + log_gpu_memory_usage( + "After offload critic model during init", logger=logger + ) if self._is_offload_optimizer: offload_fsdp_optimizer(optimizer=self.critic_optimizer) - log_gpu_memory_usage("After offload critic optimizer during init", logger=logger) + log_gpu_memory_usage( + "After offload critic optimizer during init", logger=logger + ) self.critic = DataParallelPPOCritic( - config=self.config, critic_module=self.critic_module, critic_optimizer=self.critic_optimizer + config=self.config, + critic_module=self.critic_module, + critic_optimizer=self.critic_optimizer, ) self.flops_counter = FlopsCounter(self.critic_model_config) @@ -1198,7 +1529,8 @@ def init_model(self): model=self.critic_module, optimizer=self.critic_optimizer, lr_scheduler=self.critic_lr_scheduler, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), checkpoint_config=self.config.checkpoint, ) @@ -1219,7 +1551,8 @@ def compute_values(self, data: DataProto): 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 = self.ulysses_sharding_manager.postprocess_data( + data=output) output = output.to("cpu") if self._is_offload_param: @@ -1234,7 +1567,9 @@ def update_critic(self, data: DataProto): if self._is_offload_param: load_fsdp_model_to_gpu(self.critic_module) if self._is_offload_optimizer: - load_fsdp_optimizer(optimizer=self.critic_optimizer, device_id=get_device_id()) + load_fsdp_optimizer( + optimizer=self.critic_optimizer, device_id=get_device_id() + ) # perform forward computation with self.ulysses_sharding_manager: @@ -1245,15 +1580,22 @@ def update_critic(self, data: DataProto): 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["perf/mfu/critic"] = estimated_flops * self.config.ppo_epochs / promised_flops / self.world_size + estimated_flops, promised_flops = self.flops_counter.estimate_flops( + global_num_tokens, delta_time) + metrics["perf/mfu/critic"] = ( + estimated_flops + * self.config.ppo_epochs + / promised_flops + / self.world_size + ) lr = self.critic_lr_scheduler.get_last_lr()[0] metrics["critic/lr"] = lr self.critic_lr_scheduler.step() output = DataProto(batch=None, meta_info={"metrics": metrics}) - output = self.ulysses_sharding_manager.postprocess_data(data=output) + output = self.ulysses_sharding_manager.postprocess_data( + data=output) if self._is_offload_param: offload_fsdp_model_to_cpu(self.critic_module) @@ -1264,14 +1606,19 @@ def update_critic(self, data: DataProto): return output @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None): + def save_checkpoint( + self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None + ): import torch if self._is_offload_param: load_fsdp_model_to_gpu(self.critic_module) self.checkpoint_manager.save_checkpoint( - local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep + local_path=local_path, + hdfs_path=hdfs_path, + global_step=global_step, + max_ckpt_to_keep=max_ckpt_to_keep, ) torch.distributed.barrier() @@ -1279,14 +1626,20 @@ def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to offload_fsdp_model_to_cpu(self.critic_module) @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=True): + def load_checkpoint( + self, + local_path, + hdfs_path=None, + del_local_after_load=True): import torch if self._is_offload_param: load_fsdp_model_to_gpu(self.critic_module) self.checkpoint_manager.load_checkpoint( - local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load + local_path=local_path, + hdfs_path=hdfs_path, + del_local_after_load=del_local_after_load, ) torch.distributed.barrier() @@ -1306,14 +1659,16 @@ class RewardModelWorker(Worker, DistProfilerExtension): def __init__(self, config): Worker.__init__(self) DistProfilerExtension.__init__( - self, DistProfiler(rank=self.rank, config=omega_conf_to_dataclass(config.get("profiler"))) - ) + self, DistProfiler( + rank=self.rank, config=omega_conf_to_dataclass( + config.get("profiler"))), ) import torch.distributed if not torch.distributed.is_initialized(): torch.distributed.init_process_group( - backend=get_nccl_backend(), init_method=os.environ.get("DIST_INIT_METHOD", None) + backend=get_nccl_backend(), + init_method=os.environ.get("DIST_INIT_METHOD", None), ) self.config = config @@ -1322,19 +1677,28 @@ def __init__(self, config): from torch.distributed.device_mesh import init_device_mesh fsdp_size = self.config.model.fsdp_config.fsdp_size - self.device_mesh = create_device_mesh(world_size=world_size, fsdp_size=fsdp_size) + self.device_mesh = create_device_mesh( + world_size=world_size, fsdp_size=fsdp_size + ) self.ulysses_device_mesh = None - self.ulysses_sequence_parallel_size = self.config.get("ulysses_sequence_parallel_size", 1) + 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( - device_name, mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=["dp", "sp"] + device_name, + mesh_shape=(dp, self.ulysses_sequence_parallel_size), + mesh_dim_names=["dp", "sp"], ) - self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh) + self.ulysses_sharding_manager = FSDPUlyssesShardingManager( + self.ulysses_device_mesh + ) - self.use_remove_padding = self.config.model.get("use_remove_padding", False) + self.use_remove_padding = self.config.model.get( + "use_remove_padding", False) # normalize config if self.config.micro_batch_size is not None: @@ -1354,20 +1718,28 @@ def _build_model(self, config): self._do_switch_chat_template = False else: self._do_switch_chat_template = True - input_tokenizer_local_path = copy_to_local(config.model.input_tokenizer, use_shm=use_shm) + input_tokenizer_local_path = copy_to_local( + config.model.input_tokenizer, use_shm=use_shm + ) self.input_tokenizer = hf_tokenizer( - input_tokenizer_local_path, trust_remote_code=config.model.get("trust_remote_code", False) + 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), ) - 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 = AutoConfig.from_pretrained( + local_path, trust_remote_code=trust_remote_code + ) model_config.num_labels = 1 - # note that we have to create model in fp32. Otherwise, the optimizer is in bf16, which is incorrect + # 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, mesh=self.device_mesh - ) + use_meta_tensor=not model_config.tie_word_embeddings, mesh=self.device_mesh) with init_context(), warnings.catch_warnings(): warnings.simplefilter("ignore") @@ -1382,13 +1754,17 @@ def _build_model(self, config): apply_monkey_patch( model=reward_module, - use_remove_padding=config.model.get("use_remove_padding", False), + use_remove_padding=config.model.get( + "use_remove_padding", + False), ulysses_sp_size=self.ulysses_sequence_parallel_size, ) reward_module.to(torch.bfloat16) - auto_wrap_policy = get_fsdp_wrap_policy(module=reward_module, config=self.config.model.fsdp_config) + auto_wrap_policy = get_fsdp_wrap_policy( + module=reward_module, config=self.config.model.fsdp_config + ) fsdp_mesh = self.device_mesh sharding_strategy = get_sharding_strategy(fsdp_mesh) @@ -1407,7 +1783,9 @@ def _build_model(self, config): device_mesh=self.device_mesh, ) elif config.strategy == "fsdp2": - assert CPUOffloadPolicy is not None, "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" + assert ( + CPUOffloadPolicy is not None + ), "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" cpu_offload = CPUOffloadPolicy(pin_memory=True) fsdp_kwargs = { "mesh": fsdp_mesh, @@ -1416,7 +1794,9 @@ def _build_model(self, config): } full_state = reward_module.state_dict() apply_fsdp2(reward_module, fsdp_kwargs, config.model.fsdp_config) - fsdp2_load_full_state_dict(reward_module, full_state, fsdp_mesh, cpu_offload) + fsdp2_load_full_state_dict( + reward_module, full_state, fsdp_mesh, cpu_offload + ) else: raise NotImplementedError(f"Unknown strategy: {config.strategy}") return reward_module @@ -1429,7 +1809,12 @@ def init_model(self): def _forward_micro_batch(self, micro_batch): if is_cuda_available: - from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input + from flash_attn.bert_padding import ( + index_first_axis, + pad_input, + rearrange, + unpad_input, + ) elif is_npu_available: from transformers.integrations.npu_flash_attention import ( index_first_axis, @@ -1438,43 +1823,63 @@ def _forward_micro_batch(self, micro_batch): unpad_input, ) - from verl.utils.ulysses import gather_outpus_and_unpad, ulysses_pad_and_slice_inputs + from verl.utils.ulysses import ( + gather_outpus_and_unpad, + ulysses_pad_and_slice_inputs, + ) - with torch.no_grad(), torch.autocast(device_type=device_name, dtype=torch.bfloat16): + with torch.no_grad(), torch.autocast( + device_type=device_name, 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 position_ids.dim() == 3: # qwen2vl mrope - position_ids = position_ids.transpose(0, 1) # (bsz, 3, seqlen) -> (3, bsz, seqlen) + position_ids = position_ids.transpose( + 0, 1 + ) # (bsz, 3, seqlen) -> (3, bsz, seqlen) 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) + input_ids_rmpad = input_ids_rmpad.transpose( + 0, 1) # (1, total_nnz) # unpad the position_ids to align the rotary if position_ids.dim() == 3: position_ids_rmpad = ( - index_first_axis(rearrange(position_ids, "c b s ... -> (b s) c ..."), indices) - .transpose(0, 1) - .unsqueeze(1) - ) # (3, bsz, seqlen) -> (3, 1, bsz * seqlen) + index_first_axis( + rearrange( + position_ids, + "c b s ... -> (b s) c ..."), + indices) .transpose( + 0, + 1) .unsqueeze(1)) # (3, bsz, seqlen) -> (3, 1, bsz * seqlen) else: position_ids_rmpad = index_first_axis( - rearrange(position_ids.unsqueeze(-1), "b s ... -> (b s) ..."), indices + 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 + 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 + # 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 + input_ids=input_ids_rmpad, + attention_mask=None, + position_ids=position_ids_rmpad, + use_cache=False, ) reward_rmpad = output.logits reward_rmpad = reward_rmpad.squeeze(0) # (total_nnz) @@ -1482,20 +1887,26 @@ def _forward_micro_batch(self, micro_batch): # 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 - ) + 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) + 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, use_cache=False + input_ids=input_ids, + attention_mask=attention_mask, + position_ids=position_ids, + use_cache=False, ) 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,) + 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 @@ -1507,8 +1918,12 @@ def _expand_to_token_level(self, data: DataProto, scores: torch.Tensor): response_length = data.batch["responses"].shape[-1] if position_ids.dim() == 3: # qwen2vl mrope [bs, 3, seq_len] position_ids = position_ids[:, 0, :] - 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) + 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 @@ -1535,7 +1950,9 @@ def _switch_chat_template(self, data: DataProto): # 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_length = data.batch["attention_mask"][i][ + -response_length: + ].sum() valid_response_ids = response_ids[:valid_response_length] # decode @@ -1552,12 +1969,16 @@ def _switch_chat_template(self, data: DataProto): # for debugging purpose print(f"Switch template. chat: {prompt_with_chat_template}") - # the maximum length is actually determined by the reward model itself + # 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 - model_inputs = target_tokenizer(prompt_with_chat_template, return_tensors="pt", add_special_tokens=False) + model_inputs = target_tokenizer( + prompt_with_chat_template, + return_tensors="pt", + add_special_tokens=False) input_ids, attention_mask = verl_F.postprocess_data( input_ids=model_inputs["input_ids"], attention_mask=model_inputs["attention_mask"], @@ -1575,7 +1996,11 @@ def _switch_chat_template(self, data: DataProto): 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} + rm_inputs = { + "input_ids": rm_input_ids, + "attention_mask": rm_attention_mask, + "position_ids": rm_position_ids, + } return DataProto.from_dict(rm_inputs) @@ -1606,15 +2031,23 @@ def compute_rm_score(self, data: DataProto): # perform forward computation with self.ulysses_sharding_manager: - rm_data = self.ulysses_sharding_manager.preprocess_data(data=rm_data) + 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) + 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_per_gpu) + micro_batches = rm_data.batch.split( + self.config.micro_batch_size_per_gpu + ) output = [] for micro_batch in micro_batches: rm_score = self._forward_micro_batch(micro_batch) @@ -1623,14 +2056,21 @@ def compute_rm_score(self, data: DataProto): 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) + 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) + # 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) # https://pytorch.org/docs/stable/notes/fsdp.html#fsdp-notes # unshard the root FSDP module @@ -1641,7 +2081,7 @@ def compute_rm_score(self, data: DataProto): return output -# ================================= Async related workers ================================= +# ================================= Async related workers ================ class AsyncActorRolloutRefWorker(ActorRolloutRefWorker): def _build_rollout(self, trust_remote_code=False): rollout, rollout_sharding_manager = super()._build_rollout(trust_remote_code) @@ -1660,7 +2100,9 @@ def _build_rollout(self, trust_remote_code=False): @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO) def generate_sequences(self, prompts: DataProto): - raise NotImplementedError("AsyncActorRolloutRefWorker does not support generate_sequences") + raise NotImplementedError( + "AsyncActorRolloutRefWorker does not support generate_sequences" + ) # ============================ vLLM related ============================ @@ -1681,7 +2123,11 @@ async def chat_completion(self, json_request): return ret @register(dispatch_mode=Dispatch.DIRECT_ROLLOUT_METHOD, blocking=False) - async def generate(self, prompt_ids: list[int], sampling_params: dict[str, Any], request_id: str) -> list[int]: + async def generate(self, + prompt_ids: list[int], + sampling_params: dict[str, + Any], + request_id: str) -> list[int]: ret = await self.rollout.generate(prompt_ids, sampling_params, request_id) return ret diff --git a/Agent0/executor_train/verl/verl/workers/megatron_workers.py b/Agent0/executor_train/verl/verl/workers/megatron_workers.py index e761f0e..0b262c7 100644 --- a/Agent0/executor_train/verl/verl/workers/megatron_workers.py +++ b/Agent0/executor_train/verl/verl/workers/megatron_workers.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -34,7 +34,12 @@ from verl.utils import hf_tokenizer from verl.utils.checkpoint.megatron_checkpoint_manager import MegatronCheckpointManager from verl.utils.config import omega_conf_to_dataclass -from verl.utils.device import get_device_id, get_device_name, get_nccl_backend, get_torch_device +from verl.utils.device import ( + get_device_id, + get_device_name, + get_nccl_backend, + get_torch_device, +) from verl.utils.flops_counter import FlopsCounter from verl.utils.fs import copy_to_local from verl.utils.megatron_utils import ( @@ -43,7 +48,11 @@ offload_megatron_model_to_cpu, offload_megatron_optimizer, ) -from verl.utils.model import get_hf_model_path, load_mcore_dist_weights, load_megatron_gptmodel_weights +from verl.utils.model import ( + get_hf_model_path, + load_mcore_dist_weights, + load_megatron_gptmodel_weights, +) from verl.utils.profiler import ( DistProfiler, DistProfilerExtension, @@ -94,12 +103,15 @@ def __init__(self, config: DictConfig, role: str, **kwargs): # Therefore, we only require one distribute initialization. # To utilize different parallel strategy 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 + # 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=get_nccl_backend(), - timeout=datetime.timedelta(seconds=self.config.get("nccl_timeout", 600)), + timeout=datetime.timedelta( + seconds=self.config.get("nccl_timeout", 600) + ), init_method=os.environ.get("DIST_INIT_METHOD", None), ) get_torch_device().set_device(rank) @@ -121,14 +133,27 @@ def __init__(self, config: DictConfig, role: str, **kwargs): 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"] + 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"] profiler_config = omega_conf_to_dataclass(config.get("profiler")) - DistProfilerExtension.__init__(self, DistProfiler(rank=self.rank, config=profiler_config)) + DistProfilerExtension.__init__( + self, DistProfiler(rank=self.rank, config=profiler_config) + ) # TODO(sgm): Currently, we only support reference model param offload # will support other offload later @@ -141,27 +166,59 @@ def __init__(self, config: DictConfig, role: str, **kwargs): self.config.actor.ppo_mini_batch_size *= self.config.rollout.n self.config.actor.ppo_mini_batch_size //= mpu.get_data_parallel_world_size() if self.config.actor.get("ppo_micro_batch_size", None): - 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.config.actor.ppo_micro_batch_size_per_gpu = self.config.actor.ppo_micro_batch_size - self.config.rollout.log_prob_micro_batch_size_per_gpu = self.config.rollout.log_prob_micro_batch_size - - self._is_offload_param = self.config.actor.megatron.get("param_offload", False) - self._is_offload_grad = self.config.actor.megatron.get("grad_offload", False) - self._is_offload_optimizer = self.config.actor.megatron.get("optimizer_offload", False) + 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.config.actor.ppo_micro_batch_size_per_gpu = ( + self.config.actor.ppo_micro_batch_size + ) + self.config.rollout.log_prob_micro_batch_size_per_gpu = ( + self.config.rollout.log_prob_micro_batch_size + ) + + self._is_offload_param = self.config.actor.megatron.get( + "param_offload", False + ) + self._is_offload_grad = self.config.actor.megatron.get( + "grad_offload", False + ) + self._is_offload_optimizer = self.config.actor.megatron.get( + "optimizer_offload", False + ) elif self._is_ref: if self.config.ref.get("log_prob_micro_batch_size", None): - self.config.ref.log_prob_micro_batch_size //= mpu.get_data_parallel_world_size() - self.config.ref.log_prob_micro_batch_size_per_gpu = self.config.ref.log_prob_micro_batch_size + self.config.ref.log_prob_micro_batch_size //= ( + mpu.get_data_parallel_world_size() + ) + self.config.ref.log_prob_micro_batch_size_per_gpu = ( + self.config.ref.log_prob_micro_batch_size + ) else: - assert self.config.ref.get("log_prob_micro_batch_size_per_gpu", None) is not None, ( + assert ( + self.config.ref.get("log_prob_micro_batch_size_per_gpu", None) + is not None + ), ( "Please note that in the ref policy configuration, `log_prob_micro_batch_size_per_gpu` and " "`log_prob_micro_batch_size` should not be None at the same time." ) - self._ref_is_offload_param = self.config.ref.megatron.get("param_offload", False) + self._ref_is_offload_param = self.config.ref.megatron.get( + "param_offload", False + ) - def _build_model_optimizer(self, model_path, optim_config, override_model_config, override_transformer_config): - from verl.utils.megatron.optimizer import get_megatron_optimizer, get_megatron_optimizer_param_scheduler + def _build_model_optimizer( + self, + model_path, + optim_config, + override_model_config, + override_transformer_config, + ): + from verl.utils.megatron.optimizer import ( + get_megatron_optimizer, + get_megatron_optimizer_param_scheduler, + ) from verl.utils.megatron_utils import get_model, init_megatron_optim_config from verl.utils.model import get_generation_config, print_model_size @@ -181,10 +238,13 @@ def make_model(wrap_with_ddp=False): from verl.models.mcore.mbridge import freeze_moe_router post_model_creation_callbacks = [] - if override_model_config.get("moe_config", {}).get("freeze_moe_router", False): + if override_model_config.get("moe_config", {}).get( + "freeze_moe_router", False + ): post_model_creation_callbacks.append(freeze_moe_router) return self.bridge.get_model( - post_model_creation_callbacks=post_model_creation_callbacks, wrap_with_ddp=wrap_with_ddp + post_model_creation_callbacks=post_model_creation_callbacks, + wrap_with_ddp=wrap_with_ddp, ) else: @@ -198,7 +258,11 @@ def megatron_actor_model_provider(pre_process, post_process): post_process, share_embeddings_and_output_weights=self.share_embeddings_and_output_weights, value=False, - freeze_moe_router=override_model_config.get("moe_config", {}).get("freeze_moe_router", False), + freeze_moe_router=override_model_config.get( + "moe_config", + {}).get( + "freeze_moe_router", + False), ) parallel_model.to(get_device_name()) return parallel_model @@ -215,29 +279,40 @@ def megatron_actor_model_provider(pre_process, post_process): if self.config.actor.load_weight: if self.config.actor.megatron.use_dist_checkpointing: load_mcore_dist_weights( - actor_module, self.config.actor.megatron.dist_checkpointing_path, is_value_model=False + actor_module, + self.config.actor.megatron.dist_checkpointing_path, + is_value_model=False, ) else: if self.bridge is not None: local_model_path = get_hf_model_path(self.config) - self.bridge.load_weights(actor_module, local_model_path) + self.bridge.load_weights( + actor_module, local_model_path) else: load_megatron_gptmodel_weights( - self.config, self.hf_config, actor_module, params_dtype=self.dtype, is_value_model=False + self.config, + self.hf_config, + actor_module, + params_dtype=self.dtype, + is_value_model=False, ) if self.rank == 0: print_model_size(actor_module[0]) log_gpu_memory_usage("After MegatronPPOActor init", logger=logger) elif self._is_ref: - print(f"self.config.ref.load_weight: {self.config.ref.load_weight}") + print( + f"self.config.ref.load_weight: { + self.config.ref.load_weight}") ref_module = make_model(wrap_with_ddp=False) if self.config.ref.load_weight: # should align with the actor: assert self.config.actor.load_weight == self.config.ref.load_weight print("load ref weight start") if self.config.ref.megatron.use_dist_checkpointing: load_mcore_dist_weights( - ref_module, self.config.ref.megatron.dist_checkpointing_path, is_value_model=False + ref_module, + self.config.ref.megatron.dist_checkpointing_path, + is_value_model=False, ) else: if self.bridge is not None: @@ -245,7 +320,11 @@ def megatron_actor_model_provider(pre_process, post_process): self.bridge.load_weights(ref_module, local_model_path) else: load_megatron_gptmodel_weights( - self.config, self.hf_config, ref_module, params_dtype=self.dtype, is_value_model=False + self.config, + self.hf_config, + ref_module, + params_dtype=self.dtype, + is_value_model=False, ) log_gpu_memory_usage("After ref module init", logger=logger) return ref_module, self.hf_config @@ -253,7 +332,9 @@ def megatron_actor_model_provider(pre_process, post_process): # TODO: add more optimizer args into config if self._is_actor: optim_config_megatron = init_megatron_optim_config(optim_config) - actor_optimizer = get_megatron_optimizer(model=actor_module, config=optim_config_megatron) + actor_optimizer = get_megatron_optimizer( + model=actor_module, config=optim_config_megatron + ) actor_optimizer_scheduler = get_megatron_optimizer_param_scheduler( optimizer=actor_optimizer, config=optim_config ) @@ -264,7 +345,13 @@ def megatron_actor_model_provider(pre_process, post_process): log_gpu_memory_usage("After actor optimizer init", logger=logger) - return actor_module, actor_optimizer, actor_optimizer_scheduler, self.hf_config, optim_config + return ( + actor_module, + actor_optimizer, + actor_optimizer_scheduler, + self.hf_config, + optim_config, + ) def _build_rollout(self, trust_remote_code=False): from torch.distributed.device_mesh import init_device_mesh @@ -277,25 +364,36 @@ def _build_rollout(self, trust_remote_code=False): from torch.distributed.device_mesh import init_device_mesh from verl.workers.rollout.vllm_rollout import vLLMRollout - from verl.workers.sharding_manager.megatron_vllm import MegatronVLLMShardingManager + from verl.workers.sharding_manager.megatron_vllm import ( + MegatronVLLMShardingManager, + ) # 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. + # we will reorganize their weight format when resharding from actor + # to rollout. 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}" - ) + 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( - get_device_name(), mesh_shape=(dp, infer_tp), mesh_dim_names=["dp", "infer_tp"] + get_device_name(), + mesh_shape=(dp, infer_tp), + mesh_dim_names=["dp", "infer_tp"], ) log_gpu_memory_usage("Before building vllm rollout", logger=None) - local_path = copy_to_local(self.config.model.path, use_shm=self.config.model.get("use_shm", False)) + local_path = copy_to_local( + self.config.model.path, + use_shm=self.config.model.get( + "use_shm", + False)) from verl.workers.rollout.vllm_rollout import vLLMAsyncRollout - vllm_rollout_cls = vLLMRollout if self.config.rollout.mode == "sync" else vLLMAsyncRollout + vllm_rollout_cls = ( + vLLMRollout if self.config.rollout.mode == "sync" else vLLMAsyncRollout) rollout = vllm_rollout_cls( model_path=local_path, config=self.config.rollout, @@ -309,7 +407,9 @@ def _build_rollout(self, trust_remote_code=False): # perform weight resharding between actor and rollout from verl.models.mcore import get_mcore_weight_converter - weight_converter = get_mcore_weight_converter(self.actor_model_config, self.dtype) + weight_converter = get_mcore_weight_converter( + self.actor_model_config, self.dtype + ) sharding_manager = MegatronVLLMShardingManager( inference_engine=rollout.inference_engine, model_config=self.actor_model_config, @@ -322,7 +422,9 @@ def _build_rollout(self, trust_remote_code=False): offload_param=self._is_offload_param, bridge=self.bridge, ) - log_gpu_memory_usage("After building sharding manager", logger=logger) + log_gpu_memory_usage( + "After building sharding manager", + logger=logger) elif self.config.rollout.name == "sglang": from verl.workers.rollout.sglang_rollout import SGLangRollout @@ -333,33 +435,47 @@ def _build_rollout(self, trust_remote_code=False): # potentially lead to: "RuntimeError: No CUDA GPUs are available". # For this reason, sharding_manager.__init__ should not import FSDPSGLangShardingManager and we import it # here use the abs path. - # check: https://github.com/sgl-project/sglang/blob/00f42707eaddfc2c0528e5b1e0094025c640b7a0/python/sglang/srt/layers/quantization/fp8_utils.py#L76 - from verl.workers.sharding_manager.megatron_sglang import MegatronSGLangShardingManager + # check: + # https://github.com/sgl-project/sglang/blob/00f42707eaddfc2c0528e5b1e0094025c640b7a0/python/sglang/srt/layers/quantization/fp8_utils.py#L76 + from verl.workers.sharding_manager.megatron_sglang import ( + MegatronSGLangShardingManager, + ) 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}" - ) + 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( - "cpu", mesh_shape=(dp, infer_tp, 1), mesh_dim_names=("dp", "tp", "pp") - ) + "cpu", mesh_shape=( + dp, infer_tp, 1), mesh_dim_names=( + "dp", "tp", "pp")) local_path = copy_to_local(self.config.model.path) - log_gpu_memory_usage(f"Before building {self.config.rollout.name} rollout", logger=None) + log_gpu_memory_usage( + f"Before building { + self.config.rollout.name} rollout", + logger=None) rollout = SGLangRollout( actor_module=local_path, config=self.config.rollout, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), model_hf_config=self.actor_model_config, trust_remote_code=trust_remote_code, device_mesh=rollout_device_mesh, ) - log_gpu_memory_usage(f"After building {self.config.rollout.name} rollout", logger=None) + log_gpu_memory_usage( + f"After building { + self.config.rollout.name} rollout", + logger=None) from verl.models.mcore import get_mcore_weight_converter - weight_converter = get_mcore_weight_converter(self.actor_model_config, self.dtype) + weight_converter = get_mcore_weight_converter( + self.actor_model_config, self.dtype + ) sharding_manager = MegatronSGLangShardingManager( actor_module=self.actor.actor_module, inference_engine=rollout._engine, @@ -372,10 +488,14 @@ def _build_rollout(self, trust_remote_code=False): device_mesh=rollout_device_mesh, offload_param=self._is_offload_param, ) - log_gpu_memory_usage("After building sharding manager", logger=logger) + log_gpu_memory_usage( + "After building sharding manager", + logger=logger) else: - raise NotImplementedError("Only vllmRollout is supported with Megatron now") - print(f"rollout and sharding manager init done sharding_manager: {sharding_manager}") + raise NotImplementedError( + "Only vllmRollout is supported with Megatron now") + print( + f"rollout and sharding manager init done sharding_manager: {sharding_manager}") return rollout, sharding_manager @register(dispatch_mode=Dispatch.ONE_TO_ALL) @@ -388,19 +508,29 @@ def init_model(self): from verl.utils.torch_dtypes import PrecisionType - override_model_config = OmegaConf.to_container(self.config.model.get("override_config", OmegaConf.create())) + override_model_config = OmegaConf.to_container( + self.config.model.get("override_config", OmegaConf.create()) + ) if self._is_actor: override_transformer_config = OmegaConf.to_container( - self.config.actor.megatron.get("override_transformer_config", OmegaConf.create()), resolve=True + self.config.actor.megatron.get( + "override_transformer_config", OmegaConf.create() + ), + resolve=True, ) elif self._is_ref: override_transformer_config = OmegaConf.to_container( - self.config.ref.megatron.get("override_transformer_config", OmegaConf.create()), resolve=True + self.config.ref.megatron.get( + "override_transformer_config", OmegaConf.create() + ), + resolve=True, ) else: override_transformer_config = None self.param_dtype = torch.bfloat16 - log_gpu_memory_usage("Before init actor model and optimizer", logger=logger) + log_gpu_memory_usage( + "Before init actor model and optimizer", + logger=logger) self.dtype = PrecisionType.to_dtype(self.param_dtype) if self._is_actor or self._is_rollout: # we need the model for actor and rollout @@ -419,15 +549,20 @@ def init_model(self): ) if self._is_offload_param: offload_megatron_model_to_cpu(self.actor_module) - log_gpu_memory_usage("After offload actor params and grad during init", logger=logger) + log_gpu_memory_usage( + "After offload actor params and grad during init", + logger=logger) if self._is_offload_optimizer: offload_megatron_optimizer(self.actor_optimizer) - log_gpu_memory_usage("After offload actor optimizer during init", logger=logger) + log_gpu_memory_usage( + "After offload actor optimizer during init", logger=logger + ) if self._is_actor: OmegaConf.set_struct(self.config.actor, True) with open_dict(self.config.actor): - use_fused_kernels = self.config.model.get("use_fused_kernels", False) + use_fused_kernels = self.config.model.get( + "use_fused_kernels", False) self.config.actor.use_fused_kernels = use_fused_kernels self.actor = MegatronPPOActor( config=self.config.actor, @@ -465,7 +600,9 @@ def init_model(self): ) if self._ref_is_offload_param: offload_megatron_model_to_cpu(self.ref_module) - log_gpu_memory_usage("After offload ref params during init", logger=logger) + log_gpu_memory_usage( + "After offload ref params during init", logger=logger + ) if self._is_actor: self.flops_counter = FlopsCounter(self.actor_model_config) @@ -480,7 +617,8 @@ def init_model(self): hf_config=self.hf_config, param_dtype=self.param_dtype, share_embeddings_and_output_weights=self.share_embeddings_and_output_weights, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), optimizer=self.actor_optimizer, optimizer_scheduler=self.actor_optimizer_scheduler, use_distributed_optimizer=self.config.actor.megatron.use_distributed_optimizer, @@ -498,10 +636,14 @@ def update_actor(self, data: DataProto): assert self._is_actor if self._is_offload_param: load_megatron_model_to_gpu(self.actor_module) - log_gpu_memory_usage("After load actor params and grad during update_actor", logger=logger) + log_gpu_memory_usage( + "After load actor params and grad during update_actor", + logger=logger) if self._is_offload_optimizer: load_megatron_optimizer(self.actor_optimizer) - log_gpu_memory_usage("After load actor optimizer during update_actor", logger=logger) + log_gpu_memory_usage( + "After load actor optimizer during update_actor", logger=logger + ) data.batch = data.batch.to(get_device_name()) micro_batch_size = self.config.actor.ppo_micro_batch_size_per_gpu @@ -511,11 +653,23 @@ def update_actor(self, data: DataProto): metrics = self.actor.update_policy(dataloader=dataloader) 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["perf/mfu/actor"] = estimated_flops * self.config.actor.ppo_epochs / promised_flops / self.world_size - metrics["perf/max_memory_allocated_gb"] = get_torch_device().max_memory_allocated() / (1024**3) - metrics["perf/max_memory_reserved_gb"] = get_torch_device().max_memory_reserved() / (1024**3) - metrics["perf/cpu_memory_used_gb"] = psutil.virtual_memory().used / (1024**3) + estimated_flops, promised_flops = self.flops_counter.estimate_flops( + global_num_tokens, delta_time + ) + metrics["perf/mfu/actor"] = ( + estimated_flops + * self.config.actor.ppo_epochs + / promised_flops + / self.world_size + ) + metrics["perf/max_memory_allocated_gb"] = ( + get_torch_device().max_memory_allocated() / (1024**3) + ) + metrics["perf/max_memory_reserved_gb"] = ( + get_torch_device().max_memory_reserved() / (1024**3) + ) + metrics["perf/cpu_memory_used_gb"] = psutil.virtual_memory().used / \ + (1024**3) from verl.utils.megatron.optimizer import get_megatron_last_lr metrics["actor/lr"] = get_megatron_last_lr(self.actor_optimizer) @@ -527,10 +681,14 @@ def update_actor(self, data: DataProto): if self._is_offload_param: offload_megatron_model_to_cpu(self.actor_module) - log_gpu_memory_usage("After offload actor params and grad during update_actor", logger=logger) + log_gpu_memory_usage( + "After offload actor params and grad during update_actor", + logger=logger) if self._is_offload_optimizer: offload_megatron_optimizer(self.actor_optimizer) - log_gpu_memory_usage("After offload actor optimizer during update_actor", logger=logger) + log_gpu_memory_usage( + "After offload actor optimizer during update_actor", + logger=logger) get_torch_device().empty_cache() return output @@ -542,12 +700,16 @@ def generate_sequences(self, prompts: DataProto): assert self._is_rollout prompts.batch = prompts.batch.to(get_device_name()) meta_info = { - "eos_token_id": self.generation_config.eos_token_id - if self.generation_config is not None - else self.tokenizer.eos_token_id, - "pad_token_id": self.generation_config.pad_token_id - if self.generation_config is not None - else self.tokenizer.pad_token_id, + "eos_token_id": ( + self.generation_config.eos_token_id + if self.generation_config is not None + else self.tokenizer.eos_token_id + ), + "pad_token_id": ( + self.generation_config.pad_token_id + if self.generation_config is not None + else self.tokenizer.pad_token_id + ), } prompts.meta_info.update(meta_info) if self._is_offload_optimizer: @@ -555,7 +717,9 @@ def generate_sequences(self, prompts: DataProto): timing_generate = {} with self.sharding_manager: - log_gpu_memory_usage("After entering sharding manager", logger=logger) + log_gpu_memory_usage( + "After entering sharding manager", + logger=logger) prompts = self.sharding_manager.preprocess_data(prompts) with simple_timer("generate_sequences", timing_generate): output = self.rollout.generate_sequences(prompts=prompts) @@ -579,19 +743,26 @@ def compute_ref_log_prob(self, data: DataProto): assert self._is_ref if self._ref_is_offload_param: load_megatron_model_to_gpu(self.ref_module, load_grad=False) - log_gpu_memory_usage("After load ref params and grad during compute_ref_log_prob", logger=logger) + log_gpu_memory_usage( + "After load ref params and grad during compute_ref_log_prob", + logger=logger, + ) micro_batch_size = self.config.ref.log_prob_micro_batch_size_per_gpu data.meta_info["micro_batch_size"] = micro_batch_size 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 data.meta_info["temperature"] = self.config.rollout.temperature data = data.to(get_device_id()) - output, _ = self.ref_policy.compute_log_prob(data=data, calculate_entropy=False) + output, _ = self.ref_policy.compute_log_prob( + data=data, calculate_entropy=False) output = DataProto.from_dict(tensors={"ref_log_prob": output}) output = output.to("cpu") if self._ref_is_offload_param: offload_megatron_model_to_cpu(self.ref_module) - log_gpu_memory_usage("After offload ref params and grad during compute_ref_log_prob", logger=logger) + log_gpu_memory_usage( + "After offload ref params and grad during compute_ref_log_prob", + logger=logger, + ) get_torch_device().empty_cache() return output @@ -602,14 +773,23 @@ def compute_log_prob(self, data: DataProto): assert self._is_actor if self._is_offload_param: load_megatron_model_to_gpu(self.actor_module, load_grad=False) - log_gpu_memory_usage("After load actor params and grad during compute_log_prob", logger=logger) + log_gpu_memory_usage( + "After load actor params and grad during compute_log_prob", + logger=logger, + ) # we should always recompute old_log_probs when it is HybridEngine - data.meta_info["micro_batch_size"] = self.config.rollout.log_prob_micro_batch_size_per_gpu - data.meta_info["max_token_len"] = self.config.rollout.log_prob_max_token_len_per_gpu + data.meta_info["micro_batch_size"] = ( + self.config.rollout.log_prob_micro_batch_size_per_gpu + ) + data.meta_info["max_token_len"] = ( + self.config.rollout.log_prob_max_token_len_per_gpu + ) data.meta_info["use_dynamic_bsz"] = self.config.rollout.log_prob_use_dynamic_bsz data.meta_info["temperature"] = self.config.rollout.temperature data = data.to(get_device_id()) - output, entropys = self.actor.compute_log_prob(data=data, calculate_entropy=True) + output, entropys = self.actor.compute_log_prob( + data=data, calculate_entropy=True + ) output = DataProto.from_dict( tensors={"old_log_probs": output, "entropys": entropys}, meta_info={"temperature": self.config.rollout.temperature}, @@ -618,16 +798,23 @@ def compute_log_prob(self, data: DataProto): # clear kv cache if self._is_offload_param: offload_megatron_model_to_cpu(self.actor_module) - log_gpu_memory_usage("After offload actor params and grad during compute_log_prob", logger=logger) + log_gpu_memory_usage( + "After offload actor params and grad during compute_log_prob", + logger=logger, + ) get_torch_device().empty_cache() return output @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def load_checkpoint(self, checkpoint_path, hdfs_path=None, del_local_after_load=True): + def load_checkpoint( + self, checkpoint_path, hdfs_path=None, del_local_after_load=True + ): if self._is_offload_param: load_megatron_model_to_gpu(self.actor_module) self.checkpoint_mananager.load_checkpoint( - local_path=checkpoint_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load + local_path=checkpoint_path, + hdfs_path=hdfs_path, + del_local_after_load=del_local_after_load, ) if self._is_offload_param: offload_megatron_model_to_cpu(self.actor_module) @@ -635,15 +822,26 @@ def load_checkpoint(self, checkpoint_path, hdfs_path=None, del_local_after_load= offload_megatron_optimizer(self.actor_optimizer) @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def load_pretrained_model(self, checkpoint_path, del_local_after_load=True): + def load_pretrained_model( + self, + checkpoint_path, + del_local_after_load=True): pass @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def save_checkpoint(self, checkpoint_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None): + def save_checkpoint( + self, + checkpoint_path, + hdfs_path=None, + global_step=0, + max_ckpt_to_keep=None): if self._is_offload_param: load_megatron_model_to_gpu(self.actor_module) self.checkpoint_mananager.save_checkpoint( - local_path=checkpoint_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep + local_path=checkpoint_path, + hdfs_path=hdfs_path, + global_step=global_step, + max_ckpt_to_keep=max_ckpt_to_keep, ) torch.distributed.barrier() if self._is_offload_param: @@ -673,9 +871,12 @@ def execute_method(self, method: str | bytes, *args, **kwargs): """Called by ExternalRayDistributedExecutor collective_rpc.""" if self.vllm_tp_rank == 0 and method != "execute_model": print( - f"[DP={self.vllm_dp_rank},TP={self.vllm_tp_rank}] execute_method: " - f"{method if isinstance(method, str) else 'Callable'}" - ) + f"[DP={ + self.vllm_dp_rank},TP={ + self.vllm_tp_rank}] execute_method: " f"{ + method if isinstance( + method, + str) else 'Callable'}") return self.rollout.execute_method(method, *args, **kwargs) @register(dispatch_mode=Dispatch.DIRECT_ROLLOUT_METHOD) @@ -690,7 +891,11 @@ async def chat_completion(self, json_request): return ret @register(dispatch_mode=Dispatch.DIRECT_ROLLOUT_METHOD, blocking=False) - async def generate(self, prompt_ids: list[int], sampling_params: dict[str, Any], request_id: str) -> list[int]: + async def generate(self, + prompt_ids: list[int], + sampling_params: dict[str, + Any], + request_id: str) -> list[int]: ret = await self.rollout.generate(prompt_ids, sampling_params, request_id) return ret @@ -713,8 +918,9 @@ class CriticWorker(MegatronWorker, DistProfilerExtension): def __init__(self, config): MegatronWorker.__init__(self) DistProfilerExtension.__init__( - self, DistProfiler(rank=self.rank, config=omega_conf_to_dataclass(config.get("profiler"))) - ) + self, DistProfiler( + rank=self.rank, config=omega_conf_to_dataclass( + config.get("profiler"))), ) self.config = config # NOTE(sgm): We utilize colocate WorkerGroup by default. @@ -722,12 +928,15 @@ def __init__(self, config): # Therefore, we only require one distribute initialization. # To utilize different parallel strategy 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 + # 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=get_nccl_backend(), - timeout=datetime.timedelta(seconds=self.config.get("nccl_timeout", 600)), + timeout=datetime.timedelta( + seconds=self.config.get("nccl_timeout", 600) + ), init_method=os.environ.get("DIST_INIT_METHOD", None), ) get_torch_device().set_device(rank) @@ -762,11 +971,18 @@ def __init__(self, config): # TODO(sgm): support critic model offload def _build_critic_model_optimizer( - self, model_path, optim_config, override_model_config, override_transformer_config + self, + model_path, + optim_config, + override_model_config, + override_transformer_config, ): from megatron.core.models.gpt.gpt_model import ModelType - from verl.utils.megatron.optimizer import get_megatron_optimizer, get_megatron_optimizer_param_scheduler + from verl.utils.megatron.optimizer import ( + get_megatron_optimizer, + get_megatron_optimizer_param_scheduler, + ) from verl.utils.megatron_utils import get_model, init_megatron_optim_config from verl.utils.model import print_model_size @@ -784,10 +1000,13 @@ def _build_critic_model_optimizer( from verl.models.mcore.mbridge import freeze_moe_router, make_value_model post_model_creation_callbacks = [make_value_model] - if override_model_config.get("moe_config", {}).get("freeze_moe_router", False): + if override_model_config.get("moe_config", {}).get( + "freeze_moe_router", False + ): post_model_creation_callbacks.append(freeze_moe_router) critic_module = self.bridge.get_model( - post_model_creation_callbacks=post_model_creation_callbacks, wrap_with_ddp=True + post_model_creation_callbacks=post_model_creation_callbacks, + wrap_with_ddp=True, ) else: @@ -801,7 +1020,11 @@ def megatron_critic_model_provider(pre_process, post_process): post_process, share_embeddings_and_output_weights=False, value=True, - freeze_moe_router=override_model_config.get("moe_config", {}).get("freeze_moe_router", False), + freeze_moe_router=override_model_config.get( + "moe_config", + {}).get( + "freeze_moe_router", + False), ) parallel_model.to(get_device_name()) return parallel_model @@ -821,7 +1044,9 @@ def megatron_critic_model_provider(pre_process, post_process): t0 = time.time() if self.config.megatron.use_dist_checkpointing: load_mcore_dist_weights( - critic_module, self.config.megatron.dist_checkpointing_path, is_value_model=True + critic_module, + self.config.megatron.dist_checkpointing_path, + is_value_model=True, ) else: if self.bridge is not None: @@ -829,7 +1054,11 @@ def megatron_critic_model_provider(pre_process, post_process): self.bridge.load_weights(critic_module, local_model_path) else: load_megatron_gptmodel_weights( - self.config, self.hf_config, critic_module, params_dtype=self.dtype, is_value_model=True + self.config, + self.hf_config, + critic_module, + params_dtype=self.dtype, + is_value_model=True, ) t1 = time.time() if torch.distributed.get_rank() == 0: @@ -839,12 +1068,20 @@ def megatron_critic_model_provider(pre_process, post_process): # TODO: add more optimizer args into config optim_config_megatron = init_megatron_optim_config(optim_config) - critic_optimizer = get_megatron_optimizer(model=critic_module, config=optim_config_megatron) + critic_optimizer = get_megatron_optimizer( + model=critic_module, config=optim_config_megatron + ) critic_optimizer_scheduler = get_megatron_optimizer_param_scheduler( optimizer=critic_optimizer, config=optim_config ) get_torch_device().empty_cache() - return critic_module, critic_optimizer, critic_optimizer_scheduler, self.hf_config, optim_config + return ( + critic_module, + critic_optimizer, + critic_optimizer_scheduler, + self.hf_config, + optim_config, + ) @register(dispatch_mode=Dispatch.ONE_TO_ALL) def init_model(self): @@ -857,10 +1094,11 @@ def init_model(self): import importlib importlib.import_module(self.config.model.external_lib) - override_model_config = OmegaConf.to_container(self.config.model.get("override_config", OmegaConf.create())) - override_transformer_config = OmegaConf.to_container( - self.config.megatron.get("override_transformer_config", OmegaConf.create()), resolve=True + override_model_config = OmegaConf.to_container( + self.config.model.get("override_config", OmegaConf.create()) ) + override_transformer_config = OmegaConf.to_container(self.config.megatron.get( + "override_transformer_config", OmegaConf.create()), resolve=True, ) self.param_dtype = torch.bfloat16 self.dtype = PrecisionType.to_dtype(self.param_dtype) ( @@ -901,7 +1139,8 @@ def init_model(self): hf_config=self.hf_config, param_dtype=self.param_dtype, share_embeddings_and_output_weights=False, - processing_class=self.processor if self.processor is not None else self.tokenizer, + processing_class=( + self.processor if self.processor is not None else self.tokenizer), optimizer=self.critic_optimizer, optimizer_scheduler=self.critic_optimizer_scheduler, use_distributed_optimizer=self.config.megatron.use_distributed_optimizer, @@ -942,8 +1181,13 @@ def update_critic(self, data: DataProto): metrics = self.critic.update_critic(dataloader=dataloader) 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["perf/mfu/critic"] = estimated_flops * self.config.ppo_epochs / promised_flops / self.world_size + estimated_flops, promised_flops = self.flops_counter.estimate_flops( + global_num_tokens, delta_time + ) + metrics["perf/mfu/critic"] = (estimated_flops * + self.config.ppo_epochs / + promised_flops / + self.world_size) from verl.utils.megatron.optimizer import get_megatron_last_lr metrics["critic/lr"] = get_megatron_last_lr(self.critic_optimizer) @@ -959,11 +1203,15 @@ def update_critic(self, data: DataProto): return output @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def load_checkpoint(self, checkpoint_path, hdfs_path=None, del_local_after_load=True): + def load_checkpoint( + self, checkpoint_path, hdfs_path=None, del_local_after_load=True + ): if self._is_offload_param: load_megatron_model_to_gpu(self.critic_module) self.checkpoint_mananager.load_checkpoint( - local_path=checkpoint_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load + local_path=checkpoint_path, + hdfs_path=hdfs_path, + del_local_after_load=del_local_after_load, ) if self._is_offload_param: offload_megatron_model_to_cpu(self.critic_module) @@ -971,11 +1219,19 @@ def load_checkpoint(self, checkpoint_path, hdfs_path=None, del_local_after_load= offload_megatron_optimizer(self.critic_optimizer) @register(dispatch_mode=Dispatch.ONE_TO_ALL) - def save_checkpoint(self, checkpoint_path, hdfs_path=None, global_steps=0, max_ckpt_to_keep=None): + def save_checkpoint( + self, + checkpoint_path, + hdfs_path=None, + global_steps=0, + max_ckpt_to_keep=None): if self._is_offload_param: load_megatron_model_to_gpu(self.critic_module) self.checkpoint_mananager.save_checkpoint( - local_path=checkpoint_path, hdfs_path=hdfs_path, global_step=global_steps, max_ckpt_to_keep=max_ckpt_to_keep + local_path=checkpoint_path, + hdfs_path=hdfs_path, + global_step=global_steps, + max_ckpt_to_keep=max_ckpt_to_keep, ) if self._is_offload_param: offload_megatron_model_to_cpu(self.critic_module) @@ -989,8 +1245,9 @@ class RewardModelWorker(MegatronWorker, DistProfilerExtension): def __init__(self, config): MegatronWorker.__init__(self) DistProfilerExtension.__init__( - self, DistProfiler(rank=self.rank, config=omega_conf_to_dataclass(config.get("profiler"))) - ) + self, DistProfiler( + rank=self.rank, config=omega_conf_to_dataclass( + config.get("profiler"))), ) self.config = config # NOTE(sgm): We utilize colocate WorkerGroup by default. @@ -998,12 +1255,15 @@ def __init__(self, config): # Therefore, we only require one distribute initialization. # To utilize different parallel strategy 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 + # 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=get_nccl_backend(), - timeout=datetime.timedelta(seconds=self.config.get("nccl_timeout", 600)), + timeout=datetime.timedelta( + seconds=self.config.get("nccl_timeout", 600) + ), init_method=os.environ.get("DIST_INIT_METHOD", None), ) get_torch_device().set_device(rank) @@ -1029,7 +1289,12 @@ def __init__(self, config): self.config.micro_batch_size //= mpu.get_data_parallel_world_size() self.config.micro_batch_size_per_gpu = self.config.micro_batch_size - def _build_rm_model(self, model_path, tokenizer, override_model_config, override_transformer_config): + def _build_rm_model( + self, + model_path, + tokenizer, + override_model_config, + override_transformer_config): from megatron.core.models.gpt.gpt_model import ModelType from verl.utils.megatron_utils import get_model @@ -1047,10 +1312,13 @@ def _build_rm_model(self, model_path, tokenizer, override_model_config, override from verl.models.mcore.mbridge import freeze_moe_router, make_value_model post_model_creation_callbacks = [make_value_model] - if override_model_config.get("moe_config", {}).get("freeze_moe_router", False): + if override_model_config.get("moe_config", {}).get( + "freeze_moe_router", False + ): post_model_creation_callbacks.append(freeze_moe_router) reward_model = self.bridge.get_model( - post_model_creation_callbacks=post_model_creation_callbacks, wrap_with_ddp=False + post_model_creation_callbacks=post_model_creation_callbacks, + wrap_with_ddp=False, ) else: @@ -1081,14 +1349,22 @@ def megatron_rm_model_provider(pre_process, post_process): if self.config.load_weight: if self.config.megatron.use_dist_checkpointing: - load_mcore_dist_weights(reward_model, self.config.megatron.dist_checkpointing_path, is_value_model=True) + load_mcore_dist_weights( + reward_model, + self.config.megatron.dist_checkpointing_path, + is_value_model=True, + ) else: if self.bridge is not None: local_model_path = get_hf_model_path(self.config) self.bridge.load_weights(reward_model, local_model_path) else: load_megatron_gptmodel_weights( - self.config, self.hf_config, reward_model, params_dtype=self.dtype, is_value_model=True + self.config, + self.hf_config, + reward_model, + params_dtype=self.dtype, + is_value_model=True, ) # TODO: add more optimizer args into config @@ -1106,20 +1382,27 @@ def init_model(self): import importlib importlib.import_module(self.config.model.external_lib) - override_model_config = OmegaConf.to_container(self.config.model.get("override_config", OmegaConf.create())) - override_transformer_config = OmegaConf.to_container( - self.config.megatron.get("override_transformer_config", OmegaConf.create()), resolve=True + override_model_config = OmegaConf.to_container( + self.config.model.get("override_config", OmegaConf.create()) ) + override_transformer_config = OmegaConf.to_container(self.config.megatron.get( + "override_transformer_config", OmegaConf.create()), resolve=True, ) use_shm = self.config.model.get("use_shm", False) - sft_tokenizer_local_path = copy_to_local(self.config.model.input_tokenizer, use_shm=use_shm) + sft_tokenizer_local_path = copy_to_local( + self.config.model.input_tokenizer, use_shm=use_shm + ) 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_to_local(rm_tokenizer_path, use_shm=use_shm) + rm_tokenizer_local_path = copy_to_local( + rm_tokenizer_path, use_shm=use_shm) rm_tokenizer = hf_tokenizer( - rm_tokenizer_local_path, trust_remote_code=self.config.model.get("trust_remote_code", False) + rm_tokenizer_local_path, + trust_remote_code=self.config.model.get( + "trust_remote_code", + False), ) self.param_dtype = torch.bfloat16 @@ -1144,7 +1427,8 @@ def init_model(self): ) # TODO: reward model use itself tokenizer instead of sft tokenizer - # the input_ids, responses, attention_mask and position_ids may be different! + # the input_ids, responses, attention_mask and position_ids may be + # different! @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO) @DistProfiler.annotate(color="brown") def compute_rm_score(self, data: DataProto): diff --git a/Agent0/executor_train/verl/verl/workers/reward_manager/__init__.py b/Agent0/executor_train/verl/verl/workers/reward_manager/__init__.py index 566631b..5c2bf1b 100644 --- a/Agent0/executor_train/verl/verl/workers/reward_manager/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/reward_manager/__init__.py @@ -18,7 +18,8 @@ from .naive import NaiveRewardManager from .prime import PrimeRewardManager -# Note(haibin.lin): no need to include all reward managers here in case of complicated dependencies +# Note(haibin.lin): no need to include all reward managers here in case of +# complicated dependencies __all__ = [ "BatchRewardManager", "DAPORewardManager", diff --git a/Agent0/executor_train/verl/verl/workers/reward_manager/batch.py b/Agent0/executor_train/verl/verl/workers/reward_manager/batch.py index 8d1b112..1a09699 100644 --- a/Agent0/executor_train/verl/verl/workers/reward_manager/batch.py +++ b/Agent0/executor_train/verl/verl/workers/reward_manager/batch.py @@ -1,4 +1,4 @@ -# Copyright 2025 Individual Contributor: Mert Unsal +# Copyright 2025-2026 Individual Contributor: Mert Unsal # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -33,7 +33,14 @@ class BatchRewardManager: reward_kwargs (dict): The keyword arguments to pass to the reward function. """ - def __init__(self, tokenizer, num_examine, compute_score, reward_fn_key="data_source", **reward_kwargs): + def __init__( + self, + tokenizer, + num_examine, + compute_score, + reward_fn_key="data_source", + **reward_kwargs + ): self.tokenizer = tokenizer self.num_examine = num_examine self.compute_score = compute_score @@ -52,10 +59,15 @@ def verify(self, data): for i in range(len(data)): valid_len = valid_response_lengths[i] valid_response_ids = response_ids[i][:valid_len] - response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True) + response_str = self.tokenizer.decode( + valid_response_ids, skip_special_tokens=True + ) responses_str.append(response_str) - ground_truths = [item.non_tensor_batch["reward_model"].get("ground_truth", None) for item in data] + ground_truths = [ + item.non_tensor_batch["reward_model"].get("ground_truth", None) + for item in data + ] data_sources = data.non_tensor_batch[self.reward_fn_key] extras = data.non_tensor_batch.get("extra_info", [None] * len(data)) @@ -70,14 +82,16 @@ def verify(self, data): return scores def __call__(self, data: DataProto, return_dict=False): - # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn + # If there is rm score, we directly return rm score. Otherwise, we + # compute via rm_score_fn if "rm_scores" in data.batch.keys(): if return_dict: return {"reward_tensor": data.batch["rm_scores"]} else: return data.batch["rm_scores"] - reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) + reward_tensor = torch.zeros_like( + data.batch["responses"], dtype=torch.float32) reward_extra_info = defaultdict(list) prompt_ids = data.batch["prompts"] prompt_len = prompt_ids.shape[-1] @@ -105,18 +119,30 @@ def __call__(self, data: DataProto, return_dict=False): data_source = data_sources[i] if already_printed.get(data_source, 0) < self.num_examine: - response_str = self.tokenizer.decode(data.batch["responses"][i][:length], skip_special_tokens=True) - prompt_str = self.tokenizer.decode(data.batch["prompts"][i], skip_special_tokens=True) - ground_truth = data[i].non_tensor_batch["reward_model"].get("ground_truth", None) + response_str = self.tokenizer.decode( + data.batch["responses"][i][:length], skip_special_tokens=True + ) + prompt_str = self.tokenizer.decode( + data.batch["prompts"][i], skip_special_tokens=True + ) + ground_truth = ( + data[i].non_tensor_batch["reward_model"].get( + "ground_truth", None)) print("[prompt]", prompt_str) print("[response]", response_str) print("[ground_truth]", ground_truth) print("[score]", scores[i]) - already_printed[data_source] = already_printed.get(data_source, 0) + 1 + already_printed[data_source] = already_printed.get( + data_source, 0) + 1 - data.batch["acc"] = torch.tensor(rewards, dtype=torch.float32, device=prompt_ids.device) + data.batch["acc"] = torch.tensor( + rewards, dtype=torch.float32, device=prompt_ids.device + ) if return_dict: - return {"reward_tensor": reward_tensor, "reward_extra_info": reward_extra_info} + return { + "reward_tensor": reward_tensor, + "reward_extra_info": reward_extra_info, + } else: return reward_tensor diff --git a/Agent0/executor_train/verl/verl/workers/reward_manager/dapo.py b/Agent0/executor_train/verl/verl/workers/reward_manager/dapo.py index 3ba9afe..306d7f2 100644 --- a/Agent0/executor_train/verl/verl/workers/reward_manager/dapo.py +++ b/Agent0/executor_train/verl/verl/workers/reward_manager/dapo.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -35,31 +35,34 @@ def __init__( overlong_buffer_cfg=None, ) -> None: self.tokenizer = tokenizer - self.num_examine = num_examine # the number of batches of decoded responses to print to the console + # the number of batches of decoded responses to print to the console + self.num_examine = num_examine self.compute_score = compute_score or default_compute_score self.reward_fn_key = reward_fn_key self.overlong_buffer_cfg = overlong_buffer_cfg self.max_resp_len = max_resp_len if self.overlong_buffer_cfg is not None: - assert self.max_resp_len is not None, ( - f"max_resp_len must be provided if {overlong_buffer_cfg=}, but got None" - ) - assert self.max_resp_len >= self.overlong_buffer_cfg.len, ( - "max_resp_len must be larger than overlong_buffer.len" - ) + assert ( + self.max_resp_len is not None), f"max_resp_len must be provided if { + overlong_buffer_cfg=}, but got None" + assert ( + self.max_resp_len >= self.overlong_buffer_cfg.len + ), "max_resp_len must be larger than overlong_buffer.len" def __call__(self, data: DataProto, return_dict: bool = False): """We will expand this function gradually based on the available datasets""" - # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn + # If there is rm score, we directly return rm score. Otherwise, we + # compute via rm_score_fn if "rm_scores" in data.batch.keys(): if return_dict: return {"reward_tensor": data.batch["rm_scores"]} else: return data.batch["rm_scores"] - reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) + reward_tensor = torch.zeros_like( + data.batch["responses"], dtype=torch.float32) reward_extra_info = defaultdict(list) already_print_data_sources = {} @@ -71,16 +74,24 @@ def __call__(self, data: DataProto, return_dict: bool = False): prompt_length = prompt_ids.shape[-1] - valid_prompt_length = data_item.batch["attention_mask"][:prompt_length].sum() + 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_length = data_item.batch["attention_mask"][ + prompt_length: + ].sum() valid_response_ids = response_ids[:valid_response_length] # decode - prompt_str = self.tokenizer.decode(valid_prompt_ids, skip_special_tokens=True) - response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True) + prompt_str = self.tokenizer.decode( + valid_prompt_ids, skip_special_tokens=True + ) + response_str = self.tokenizer.decode( + valid_response_ids, skip_special_tokens=True + ) eos_token = self.tokenizer.eos_token if response_str.endswith(eos_token): response_str = response_str[: -len(eos_token)] @@ -114,10 +125,13 @@ def __call__(self, data: DataProto, return_dict: bool = False): expected_len = self.max_resp_len - overlong_buffer_len exceed_len = valid_response_length - expected_len overlong_penalty_factor = self.overlong_buffer_cfg.penalty_factor - overlong_reward = min(-exceed_len / overlong_buffer_len * overlong_penalty_factor, 0) + overlong_reward = min( + -exceed_len / overlong_buffer_len * overlong_penalty_factor, 0 + ) reward += overlong_reward if self.overlong_buffer_cfg.log: - reward_extra_info["overlong_reward"].append(overlong_reward) + reward_extra_info["overlong_reward"].append( + overlong_reward) reward_extra_info["overlong"].append(overlong_reward < 0) reward_tensor[i, valid_response_length - 1] = reward diff --git a/Agent0/executor_train/verl/verl/workers/reward_manager/naive.py b/Agent0/executor_train/verl/verl/workers/reward_manager/naive.py index f6f979e..6cf61d5 100644 --- a/Agent0/executor_train/verl/verl/workers/reward_manager/naive.py +++ b/Agent0/executor_train/verl/verl/workers/reward_manager/naive.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -25,7 +25,12 @@ class NaiveRewardManager: """The reward manager.""" - def __init__(self, tokenizer, num_examine, compute_score=None, reward_fn_key="data_source") -> None: + def __init__( + self, + tokenizer, + num_examine, + compute_score=None, + reward_fn_key="data_source") -> None: """ Initialize the NaiveRewardManager instance. @@ -37,21 +42,26 @@ def __init__(self, tokenizer, num_examine, compute_score=None, reward_fn_key="da "data_source". """ self.tokenizer = tokenizer # Store the tokenizer for decoding token IDs - self.num_examine = num_examine # the number of batches of decoded responses to print to the console + # the number of batches of decoded responses to print to the console + self.num_examine = num_examine self.compute_score = compute_score or default_compute_score - self.reward_fn_key = reward_fn_key # Store the key for accessing the data source + self.reward_fn_key = ( + reward_fn_key # Store the key for accessing the data source + ) def __call__(self, data: DataProto, return_dict=False): """We will expand this function gradually based on the available datasets""" - # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn + # If there is rm score, we directly return rm score. Otherwise, we + # compute via rm_score_fn if "rm_scores" in data.batch.keys(): if return_dict: return {"reward_tensor": data.batch["rm_scores"]} else: return data.batch["rm_scores"] - reward_tensor = torch.zeros_like(data.batch["responses"], dtype=torch.float32) + reward_tensor = torch.zeros_like( + data.batch["responses"], dtype=torch.float32) reward_extra_info = defaultdict(list) already_print_data_sources = {} @@ -63,16 +73,24 @@ def __call__(self, data: DataProto, return_dict=False): prompt_length = prompt_ids.shape[-1] - valid_prompt_length = data_item.batch["attention_mask"][:prompt_length].sum() + 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_length = data_item.batch["attention_mask"][ + prompt_length: + ].sum() valid_response_ids = response_ids[:valid_response_length] # decode - prompt_str = self.tokenizer.decode(valid_prompt_ids, skip_special_tokens=True) - response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True) + prompt_str = self.tokenizer.decode( + valid_prompt_ids, skip_special_tokens=True + ) + response_str = self.tokenizer.decode( + valid_response_ids, skip_special_tokens=True + ) ground_truth = data_item.non_tensor_batch["reward_model"]["ground_truth"] data_source = data_item.non_tensor_batch[self.reward_fn_key] diff --git a/Agent0/executor_train/verl/verl/workers/reward_manager/prime.py b/Agent0/executor_train/verl/verl/workers/reward_manager/prime.py index f2c526b..3865869 100644 --- a/Agent0/executor_train/verl/verl/workers/reward_manager/prime.py +++ b/Agent0/executor_train/verl/verl/workers/reward_manager/prime.py @@ -26,11 +26,27 @@ from verl.workers.reward_manager import register -async def single_compute_score(evaluation_func, completion, reference, task, task_extra_info, executor, timeout=300.0): +async def single_compute_score( + evaluation_func, + completion, + reference, + task, + task_extra_info, + executor, + timeout=300.0, +): loop = asyncio.get_running_loop() try: # Ensure process_completion is called properly - future = loop.run_in_executor(executor, partial(evaluation_func, task, completion, reference, task_extra_info)) + future = loop.run_in_executor( + executor, + partial( + evaluation_func, + task, + completion, + reference, + task_extra_info), + ) return await asyncio.wait_for(future, timeout=timeout) except asyncio.TimeoutError: print(f"[Timeout] Task timeout: {completion}") @@ -41,19 +57,28 @@ async def single_compute_score(evaluation_func, completion, reference, task, tas async def parallel_compute_score_async( - evaluation_func, completions, references, tasks, extra_info=None, num_processes=64 -): + evaluation_func, + completions, + references, + tasks, + extra_info=None, + num_processes=64): if extra_info is None: extra_info = [None] * len(tasks) scores = [] with ProcessPoolExecutor(max_workers=num_processes) as executor: # to prevent very occasional starvation caused by some anomalous programs ( like infinite loop ), the - # exceptions in async programs will instantly halt the evaluation, and all summoned processes will be killed. + # exceptions in async programs will instantly halt the evaluation, and + # all summoned processes will be killed. try: # Create tasks for all rows tasks_async = [ - single_compute_score(evaluation_func, c, r, t, ei, executor, timeout=300.0) - for c, r, t, ei in zip(completions, references, tasks, extra_info, strict=True) + single_compute_score( + evaluation_func, c, r, t, ei, executor, timeout=300.0 + ) + for c, r, t, ei in zip( + completions, references, tasks, extra_info, strict=True + ) ] results = await asyncio.gather(*tasks_async, return_exceptions=False) except Exception as e: @@ -75,7 +100,9 @@ async def parallel_compute_score_async( print(f"[Shutdown] {terminated_count} subprocess(es) terminated.") # Process results - for result, completion, reference, task in zip(results, completions, references, tasks, strict=True): + for result, completion, reference, task in zip( + results, completions, references, tasks, strict=True + ): if isinstance(result, Exception) or result is None: # Handle failed or timed-out tasks scores.append(0.0) @@ -86,12 +113,25 @@ async def parallel_compute_score_async( return scores -def run_reward_scoring(evaluation_func, completions, references, tasks, extra_info=None, num_processes=64): +def run_reward_scoring( + evaluation_func, + completions, + references, + tasks, + extra_info=None, + num_processes=64): loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: return loop.run_until_complete( - parallel_compute_score_async(evaluation_func, completions, references, tasks, extra_info, num_processes) + parallel_compute_score_async( + evaluation_func, + completions, + references, + tasks, + extra_info, + num_processes, + ) ) finally: loop.close() @@ -111,7 +151,8 @@ def __init__( reward_fn_key: str = "data_source", ) -> None: self.tokenizer = tokenizer - self.num_examine = num_examine # the number of batches of decoded responses to print to the console + # the number of batches of decoded responses to print to the console + self.num_examine = num_examine self.compute_score = compute_score or default_compute_score self.reward_fn_key = reward_fn_key @@ -123,8 +164,13 @@ def verify(self, data): prompt_ids = data.batch["prompts"] response_ids = data.batch["responses"] - sequences_str = self.tokenizer.batch_decode(response_ids, skip_special_tokens=True) - ground_truth = [data_item.non_tensor_batch["reward_model"]["ground_truth"] for data_item in data] + sequences_str = self.tokenizer.batch_decode( + response_ids, skip_special_tokens=True + ) + ground_truth = [ + data_item.non_tensor_batch["reward_model"]["ground_truth"] + for data_item in data + ] data_sources = data.non_tensor_batch[self.reward_fn_key] extra_info = data.non_tensor_batch.get("extra_info", None) @@ -142,19 +188,24 @@ def verify(self, data): print("[Timeout] Global reward scoring timed out. Setting all as 0.") scores = [0.0 for _ in range(len(sequences_str))] except Exception as e: - print(f"[Error] Unexpected error during scoring. Setting all as 0. {e}") + print( + f"[Error] Unexpected error during scoring. Setting all as 0. {e}") scores = [0.0 for _ in range(len(sequences_str))] - data.batch["acc"] = torch.tensor(scores, dtype=torch.float32, device=prompt_ids.device) + data.batch["acc"] = torch.tensor( + scores, dtype=torch.float32, device=prompt_ids.device + ) return scores def __call__(self, data: DataProto, return_dict: bool = False): """We will expand this function gradually based on the available datasets""" - # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn + # 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) + reward_tensor = torch.zeros_like( + data.batch["responses"], dtype=torch.float32) already_print_data_sources = {} @@ -163,8 +214,11 @@ def __call__(self, data: DataProto, return_dict: bool = False): prompt_length = prompt_ids.shape[-1] response_ids = data.batch["responses"] - valid_response_length = data.batch["attention_mask"][:, prompt_length:].sum(dim=-1) - sequences_str = self.tokenizer.batch_decode(response_ids, skip_special_tokens=True) + valid_response_length = data.batch["attention_mask"][:, prompt_length:].sum( + dim=-1) + sequences_str = self.tokenizer.batch_decode( + response_ids, skip_special_tokens=True + ) data_sources = data.non_tensor_batch["data_source"] scores = self.verify(data) diff --git a/Agent0/executor_train/verl/verl/workers/reward_manager/registry.py b/Agent0/executor_train/verl/verl/workers/reward_manager/registry.py index 3fc34ef..cb55356 100644 --- a/Agent0/executor_train/verl/verl/workers/reward_manager/registry.py +++ b/Agent0/executor_train/verl/verl/workers/reward_manager/registry.py @@ -1,4 +1,4 @@ -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -28,8 +28,8 @@ def register(name): def decorator(cls): if name in REWARD_MANAGER_REGISTRY and REWARD_MANAGER_REGISTRY[name] != cls: raise ValueError( - f"Reward manager {name} has already been registered: {REWARD_MANAGER_REGISTRY[name]} vs {cls}" - ) + f"Reward manager {name} has already been registered: { + REWARD_MANAGER_REGISTRY[name]} vs {cls}") REWARD_MANAGER_REGISTRY[name] = cls return cls diff --git a/Agent0/executor_train/verl/verl/workers/reward_model/__init__.py b/Agent0/executor_train/verl/verl/workers/reward_model/__init__.py index db412bd..4d900d6 100644 --- a/Agent0/executor_train/verl/verl/workers/reward_model/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/reward_model/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/reward_model/base.py b/Agent0/executor_train/verl/verl/workers/reward_model/base.py index cb719bd..8d413be 100644 --- a/Agent0/executor_train/verl/verl/workers/reward_model/base.py +++ b/Agent0/executor_train/verl/verl/workers/reward_model/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/reward_model/megatron/__init__.py b/Agent0/executor_train/verl/verl/workers/reward_model/megatron/__init__.py index 5bd4da2..eed2a2d 100644 --- a/Agent0/executor_train/verl/verl/workers/reward_model/megatron/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/reward_model/megatron/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/reward_model/megatron/reward_model.py b/Agent0/executor_train/verl/verl/workers/reward_model/megatron/reward_model.py index 01b1324..9b20f3d 100644 --- a/Agent0/executor_train/verl/verl/workers/reward_model/megatron/reward_model.py +++ b/Agent0/executor_train/verl/verl/workers/reward_model/megatron/reward_model.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -67,7 +67,11 @@ def re_encode_by_rm_tokenizer(self, data: DataProto) -> DataProto: 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_values = { + "input_ids": input_ids, + "attention_mask": attention_mask, + "position_ids": position_ids, + } _, ori_seqlen = input_ids.size(0), input_ids.size(1) input_ids_for_rm = [] attention_mask_for_rm = [] @@ -77,8 +81,9 @@ def re_encode_by_rm_tokenizer(self, data: DataProto) -> DataProto: for id, mask in zip(input_ids, attention_mask, strict=True): # 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] + 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 @@ -91,30 +96,41 @@ def re_encode_by_rm_tokenizer(self, data: DataProto) -> DataProto: if print_decode and torch.distributed.get_rank() == 0: # only print first decode result print( - f"device {get_device_id()}: sft decode result:\n{decode_result}\n \ - \ndevice {get_device_id()}: sft decode result with \ - rm chat template:\n{decode_with_rm_chat}\n\n" - ) + f"device { + get_device_id()}: sft decode result:\n{decode_result}\n \ + \ndevice { + get_device_id()}: 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 - ) + 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) + 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) + # 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}") + 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) + 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)) + 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) @@ -142,9 +158,13 @@ def compute_reward(self, data: DataProto) -> DataProto: use_dynamic_bsz = data.meta_info.get("use_dynamic_bsz", False) micro_batch_size = data.meta_info.get("micro_batch_size", None) max_token_len = data.meta_info.get("max_token_len", None) - assert micro_batch_size is not None, "micro batch size is needed for forward compute" + assert ( + micro_batch_size is not None + ), "micro batch size is needed for forward compute" if use_dynamic_bsz: - assert max_token_len is not None, "use_dynamic_bsz is True, but max_token_len is None!" + assert ( + max_token_len is not None + ), "use_dynamic_bsz is True, but max_token_len is None!" max_token_len = max_token_len * self.config.megatron.context_parallel_size responses = data.batch["responses"] @@ -153,15 +173,22 @@ def compute_reward(self, data: DataProto) -> DataProto: with torch.no_grad(): output = self.forward_batch( - data, use_dynamic_bsz=use_dynamic_bsz, micro_batch_size=micro_batch_size, max_token_len=max_token_len + data, + use_dynamic_bsz=use_dynamic_bsz, + micro_batch_size=micro_batch_size, + max_token_len=max_token_len, ) if mpu.is_pipeline_last_stage(ignore_virtual=True): logits = torch.cat(output["output"], dim=0) if use_dynamic_bsz: indices = output["indices"] indices = list(itertools.chain.from_iterable(indices)) - assert len(indices) == logits.size(0), f"{len(indices)} vs. {logits.size()}" - revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long) + assert len(indices) == logits.size( + 0 + ), f"{len(indices)} vs. {logits.size()}" + revert_indices = torch.tensor( + get_reverse_idx(indices), dtype=torch.long + ) logits = logits[revert_indices] else: logits = torch.empty( @@ -181,8 +208,12 @@ def compute_reward(self, data: DataProto) -> DataProto: # (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) + 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) @@ -190,12 +221,16 @@ def compute_reward(self, data: DataProto) -> DataProto: 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) + token_level_rewards = rewards.expand( + attention_mask.shape[0], attention_mask.shape[1] + ) # (bs, ori_seqlen) # assign last valid token reward to ori position if position_ids.dim() == 3: # qwen2vl mrope [bs, 3, seq_len] position_ids = position_ids[:, 0, :] - eos_mask_idx = torch.argmax(position_ids * attention_mask, dim=-1) # (bs,) + 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.0 @@ -208,11 +243,19 @@ def compute_reward(self, data: DataProto) -> DataProto: # add empty cache after each compute get_torch_device().empty_cache() - batch = TensorDict({"rm_scores": token_level_rewards}, batch_size=input_ids.shape[0]) + batch = TensorDict( + {"rm_scores": token_level_rewards}, batch_size=input_ids.shape[0] + ) return DataProto(batch=batch) - def forward_batch(self, data: DataProto, use_dynamic_bsz=False, micro_batch_size=None, max_token_len=None): + def forward_batch( + self, + data: DataProto, + use_dynamic_bsz=False, + micro_batch_size=None, + max_token_len=None, + ): """ We assume: - The model takes input: (input_ids, attention_mask, position_ids). No rmpad for the input @@ -228,37 +271,52 @@ def forward_batch(self, data: DataProto, use_dynamic_bsz=False, micro_batch_size group=mpu.get_pipeline_model_parallel_group(), ) - mini_batch.batch["attention_mask"] = mini_batch.batch["attention_mask"].to(bool) + mini_batch.batch["attention_mask"] = mini_batch.batch["attention_mask"].to( + bool) - self.has_multi_modal_inputs = "multi_modal_inputs" in mini_batch.non_tensor_batch.keys() + self.has_multi_modal_inputs = ( + "multi_modal_inputs" in mini_batch.non_tensor_batch.keys() + ) if self.has_multi_modal_inputs: - mini_batch.batch["multi_modal_inputs"] = mini_batch.non_tensor_batch["multi_modal_inputs"] + mini_batch.batch["multi_modal_inputs"] = mini_batch.non_tensor_batch[ + "multi_modal_inputs" + ] mini_batch.batch["multi_modal_inputs_idx"] = torch.Tensor( list(range(len(mini_batch.non_tensor_batch["multi_modal_inputs"]))) ).to(torch.int64) indices = None if use_dynamic_bsz: - assert max_token_len is not None, "max_token_len must be set when use_dynamic_bsz is True" + assert ( + max_token_len is not None + ), "max_token_len must be set when use_dynamic_bsz is True" vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size() if vpp_size is not None and vpp_size > 1: - microbatch_group_size_per_vp_stage = self.tf_config.microbatch_group_size_per_vp_stage + microbatch_group_size_per_vp_stage = ( + self.tf_config.microbatch_group_size_per_vp_stage + ) micro_batches, indices = rearrange_micro_batches( batch=mini_batch.batch, num_batches_divided_by=microbatch_group_size_per_vp_stage, max_token_len=max_token_len, ) - assert len(micro_batches) % self.tf_config.microbatch_group_size_per_vp_stage == 0, ( + assert ( + len(micro_batches) + % self.tf_config.microbatch_group_size_per_vp_stage + == 0 + ), ( f"micro_batches {micro_batches} must be divisible by microbatch_group_size_per_vp_stage " f"{microbatch_group_size_per_vp_stage} for megatron backend" ) else: - micro_batches, indices = rearrange_micro_batches(batch=mini_batch.batch, max_token_len=max_token_len) + micro_batches, indices = rearrange_micro_batches( + batch=mini_batch.batch, max_token_len=max_token_len + ) total_seqlen = max_token_len else: - assert micro_batch_size is not None, ( - "micro_batch_size is needed to be passed in when not using dynamic batch size" - ) + assert ( + micro_batch_size is not None + ), "micro_batch_size is needed to be passed in when not using dynamic batch size" micro_batches = mini_batch.batch.split(micro_batch_size) seq_len = micro_batches[0]["input_ids"].shape[1] total_seqlen = micro_batch_size * seq_len @@ -283,7 +341,11 @@ def forward_step(batch_iter, model): if "multi_modal_inputs" in batch: for key in batch["multi_modal_inputs"][0].keys(): multi_modal_inputs[key] = torch.cat( - [batch["multi_modal_inputs"][i][key] for i in batch["multi_modal_inputs_idx"]], dim=0 + [ + batch["multi_modal_inputs"][i][key] + for i in batch["multi_modal_inputs_idx"] + ], + dim=0, ) output = forward_fn( @@ -299,10 +361,13 @@ def forward_step(batch_iter, model): return output, loss_func # batch should be a list of batches inside micro-batches - batch_generator = make_batch_generator(micro_batches, vpp_size=len(self.reward_model_module)) + batch_generator = make_batch_generator( + micro_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) + # 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, @@ -346,5 +411,6 @@ 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(get_device_id(), non_blocking=True) + param.data = param.data.to( + get_device_id(), non_blocking=True) self.device = get_device_name() diff --git a/Agent0/executor_train/verl/verl/workers/rollout/__init__.py b/Agent0/executor_train/verl/verl/workers/rollout/__init__.py index 5efcd33..1e3c7d1 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/rollout/async_server.py b/Agent0/executor_train/verl/verl/workers/rollout/async_server.py index da59c37..75c0965 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/async_server.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/async_server.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -56,15 +56,22 @@ async def lifespan(app: fastapi.FastAPI): yield # There's no way to gracefully restart uvicorn server if port is already in use, - # so we exit the process directly and let AsyncLLMServerManager restart it. - print("FastAPI shutdown, maybe address already in use, exit process immediately.") + # so we exit the process directly and let AsyncLLMServerManager + # restart it. + print( + "FastAPI shutdown, maybe address already in use, exit process immediately." + ) os._exit(-1) app = fastapi.FastAPI(lifespan=lifespan) - app.router.add_api_route("/v1/chat/completions", self.chat_completion, methods=["POST"]) + app.router.add_api_route( + "/v1/chat/completions", self.chat_completion, methods=["POST"] + ) self.port = _get_free_port() - config = uvicorn.Config(app, host=["::", "0.0.0.0"], port=self.port, log_level="warning") + config = uvicorn.Config( + app, host=["::", "0.0.0.0"], port=self.port, log_level="warning" + ) server = uvicorn.Server(config) await server.serve() @@ -82,7 +89,11 @@ async def chat_completion(self, raw_request: Request): raise NotImplementedError @abstractmethod - async def generate(self, prompt_ids: list[int], sampling_params: dict[str, Any], request_id: str) -> list[int]: + async def generate(self, + prompt_ids: list[int], + sampling_params: dict[str, + Any], + request_id: str) -> list[int]: """Generate response ids given prompt ids. Args: @@ -128,7 +139,9 @@ def __init__(self, config: DictConfig, worker_group: RayWorkerGroup): self.rollout_tp_size = self.config.rollout.tensor_model_parallel_size self.rollout_dp_size = self.worker_group.world_size // self.rollout_tp_size - register_center = ray.get_actor(f"{self.worker_group.name_prefix}_register_center") + register_center = ray.get_actor( + f"{self.worker_group.name_prefix}_register_center" + ) workers_info = ray.get(register_center.get_worker_info.remote()) assert len(workers_info) == self.worker_group.world_size @@ -142,20 +155,28 @@ def __init__(self, config: DictConfig, worker_group: RayWorkerGroup): rollout_backend_class=self.config.rollout.agent.custom_async_server.name, ) else: - server_class = async_server_class(rollout_backend=self.config.rollout.name) + server_class = async_server_class( + rollout_backend=self.config.rollout.name) # Start all server instances, restart if address already in use. unready_dp_ranks = set(range(self.rollout_dp_size)) while len(unready_dp_ranks) > 0: servers = { rollout_dp_rank: server_class.options( - # make sure AsyncvLLMServer colocates with its corresponding workers + # make sure AsyncvLLMServer colocates with its + # corresponding workers scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy( - node_id=workers_info[rollout_dp_rank * self.rollout_tp_size], + node_id=workers_info[rollout_dp_rank * + self.rollout_tp_size], soft=False, ), name=f"async_llm_server_{rollout_dp_rank}", - ).remote(config, self.rollout_dp_size, rollout_dp_rank, self.worker_group.name_prefix) + ).remote( + config, + self.rollout_dp_size, + rollout_dp_rank, + self.worker_group.name_prefix, + ) for rollout_dp_rank in unready_dp_ranks } @@ -167,17 +188,21 @@ def __init__(self, config: DictConfig, worker_group: RayWorkerGroup): unready_dp_ranks.remove(rollout_dp_rank) except Exception: ray.kill(server) - print(f"rollout server {rollout_dp_rank} failed, maybe address already in use, restarting...") + print( + f"rollout server {rollout_dp_rank} failed, maybe address already in use, restarting...") # All server instances are ready, init AsyncLLM engine. - ray.get([server.init_engine.remote() for server in self.async_llm_servers]) + ray.get([server.init_engine.remote() + for server in self.async_llm_servers]) # Init user provided chat scheduler in sperate thread. self.chat_scheduler: ChatCompletionScheduler = None self.chat_scheduler_exception: Exception = None self.chat_scheduler_loop = None self.chat_scheduler_ready = threading.Event() - self.chat_scheduler_thread = threading.Thread(target=self._init_chat_scheduler, daemon=True) + self.chat_scheduler_thread = threading.Thread( + target=self._init_chat_scheduler, daemon=True + ) self.chat_scheduler_thread.start() self.chat_scheduler_ready.wait() @@ -200,12 +225,14 @@ def _init_chat_scheduler(self): def wake_up(self): """Wake up all vllm instances.""" if self.config.rollout.free_cache_engine: - ray.get([server.wake_up.remote() for server in self.async_llm_servers]) + ray.get([server.wake_up.remote() + for server in self.async_llm_servers]) def sleep(self): """Sleep all vllm instances.""" if self.config.rollout.free_cache_engine: - ray.get([server.sleep.remote() for server in self.async_llm_servers]) + ray.get([server.sleep.remote() + for server in self.async_llm_servers]) def submit_chat_completions( self, @@ -228,18 +255,22 @@ def submit_chat_completions( ) future.result() - def generate_sequences(self, prompts: DataProto, **sampling_params) -> DataProto: + def generate_sequences(self, prompts: DataProto, ** + sampling_params) -> DataProto: """Generate multiple sequences in parallel via chat scheduler.""" assert self.chat_scheduler is not None, "chat scheduler is not initialized." future = asyncio.run_coroutine_threadsafe( - self.chat_scheduler.generate_sequences(prompts, **sampling_params), self.chat_scheduler_loop + self.chat_scheduler.generate_sequences(prompts, **sampling_params), + self.chat_scheduler_loop, ) return future.result() def async_server_class( - rollout_backend: str, rollout_backend_module: Optional[str] = None, rollout_backend_class: Optional[str] = None + rollout_backend: str, + rollout_backend_module: Optional[str] = None, + rollout_backend_class: Optional[str] = None, ) -> type[AsyncServerBase]: """Get async server class. @@ -254,21 +285,29 @@ def async_server_class( if rollout_backend_class is None and rollout_backend_module is None: # If both are None, use the default backend class # Do not change the original import behavior - # importlib.import_module and from ... import ... have subtle differences in ray + # importlib.import_module and from ... import ... have subtle + # differences in ray if rollout_backend == "vllm": - from verl.workers.rollout.vllm_rollout.vllm_async_server import AsyncvLLMServer + from verl.workers.rollout.vllm_rollout.vllm_async_server import ( + AsyncvLLMServer, + ) return AsyncvLLMServer elif rollout_backend == "sglang": - from verl.workers.rollout.sglang_rollout.async_sglang_server import AsyncSglangServer + from verl.workers.rollout.sglang_rollout.async_sglang_server import ( + AsyncSglangServer, ) return AsyncSglangServer else: - raise NotImplementedError(f"rollout backend {rollout_backend} is not supported") + raise NotImplementedError( + f"rollout backend {rollout_backend} is not supported" + ) if rollout_backend_module is None or rollout_backend_class is None: - raise ValueError("rollout_backend_module and rollout_backend_class must be both provided for customization") + raise ValueError( + "rollout_backend_module and rollout_backend_class must be both provided for customization" + ) from verl.utils.import_utils import load_extern_type diff --git a/Agent0/executor_train/verl/verl/workers/rollout/base.py b/Agent0/executor_train/verl/verl/workers/rollout/base.py index 0319824..d96d5f2 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/base.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/rollout/chat_scheduler.py b/Agent0/executor_train/verl/verl/workers/rollout/chat_scheduler.py index 268c82d..9ad1b0f 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/chat_scheduler.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/chat_scheduler.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -41,16 +41,25 @@ class CompletionCallback(ABC): - def __init__(self, config: DictConfig, scheduler: "ChatCompletionScheduler"): + def __init__( + self, + config: DictConfig, + scheduler: "ChatCompletionScheduler"): self.config = config self.scheduler = scheduler # Initialize tools from config file - self.max_assistant_turns = config.actor_rollout_ref.rollout.multi_turn.max_assistant_turns + self.max_assistant_turns = ( + config.actor_rollout_ref.rollout.multi_turn.max_assistant_turns + ) tool_config_path = config.actor_rollout_ref.rollout.multi_turn.tool_config_path - tool_list = initialize_tools_from_config(tool_config_path) if tool_config_path else [] + tool_list = (initialize_tools_from_config( + tool_config_path) if tool_config_path else []) self.tools = {tool.name: tool for tool in tool_list} - self._tool_schemas = [tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) for tool in tool_list] + self._tool_schemas = [ + tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) + for tool in tool_list + ] print(f"Initialized tools: {self.tools}", flush=True) local_path = copy_to_local(config.actor_rollout_ref.model.path) @@ -67,7 +76,12 @@ def extra_body(self) -> dict[str, Any]: return None @abstractmethod - async def __call__(self, messages: list[dict[str, str]], completions: ChatCompletion, info: dict[str, Any]): + async def __call__( + self, + messages: list[dict[str, str]], + completions: ChatCompletion, + info: dict[str, Any], + ): """Call back function to process completions. Args: @@ -78,7 +92,9 @@ async def __call__(self, messages: list[dict[str, str]], completions: ChatComple raise NotImplementedError @abstractmethod - def postprocess(self, batch: DataProto, batch_conversations: list[list[dict[str, str]]], n: int) -> DataProto: + def postprocess( + self, batch: DataProto, batch_conversations: list[list[dict[str, str]]], n: int + ) -> DataProto: """Post process batch data. Args: @@ -96,45 +112,69 @@ def postprocess(self, batch: DataProto, batch_conversations: list[list[dict[str, class ToolCompletionCallback(CompletionCallback): - def __init__(self, config: DictConfig, scheduler: "ChatCompletionScheduler"): + def __init__( + self, + config: DictConfig, + scheduler: "ChatCompletionScheduler"): super().__init__(config, scheduler) - # TODO: add reward manager to calculate reward score once a sample finish + # TODO: add reward manager to calculate reward score once a sample + # finish - async def __call__(self, messages: list[dict[str, str]], completions: ChatCompletion, info: dict[str, Any]): - message = completions.choices[0].message.model_dump(exclude_unset=True, exclude_none=True) + async def __call__( + self, + messages: list[dict[str, str]], + completions: ChatCompletion, + info: dict[str, Any], + ): + message = completions.choices[0].message.model_dump( + exclude_unset=True, exclude_none=True + ) if "content" not in message: message["content"] = "" messages.append(message) finish_reason = completions.choices[0].finish_reason # STEP 0: check if we reach max turns - if self.max_assistant_turns and len(messages) >= self.max_assistant_turns: - print(f"[id={completions.id},turn={len(messages)},finish_reason={finish_reason}] Reach max turns, done!") + if self.max_assistant_turns and len( + messages) >= self.max_assistant_turns: + print( + f"[id={ + completions.id},turn={ + len(messages)},finish_reason={finish_reason}] Reach max turns, done!") return # STEP 1: check if the model called tools if finish_reason != "tool_calls": - print(f"[id={completions.id},turn={len(messages)},finish_reason={finish_reason}] No tool called, done!") + print( + f"[id={ + completions.id},turn={ + len(messages)},finish_reason={finish_reason}] No tool called, done!") return # STEP 2: call tools tool_calls = completions.choices[0].message.tool_calls - print(f"[id={completions.id},turn={len(messages)},finish_reason={finish_reason}] Call {len(tool_calls)} tools") + print( + f"[id={ + completions.id},turn={ + len(messages)},finish_reason={finish_reason}] Call { + len(tool_calls)} tools") tasks = [] for tool_call in tool_calls: tasks.append(self._call_tool(tool_call)) tool_responses = await asyncio.gather(*tasks) if any(isinstance(item, Exception) for item in tool_responses): print( - f"[id={completions.id},turn={len(messages)},finish_reason={finish_reason}] Error when calling tools, " - f"done!" - ) + f"[id={ + completions.id},turn={ + len(messages)},finish_reason={finish_reason}] Error when calling tools, " f"done!") return messages.extend(tool_responses) # STEP 3: resubmit completion request with tool responses - self.scheduler.submit_chat_completions(messages=messages, request_id=completions.id, info=info) + self.scheduler.submit_chat_completions( + messages=messages, request_id=completions.id, info=info + ) async def _call_tool(self, tool_call) -> dict[str, str]: """Call tool and return tool response.""" @@ -144,7 +184,9 @@ async def _call_tool(self, tool_call) -> dict[str, str]: instance_id = await tool.create() try: - tool_response, tool_reward_score, tool_metrics = await tool.execute(instance_id, tool_args) + tool_response, tool_reward_score, tool_metrics = await tool.execute( + instance_id, tool_args + ) except Exception as e: logger.exception(f"Error when executing tool: {e}") return e @@ -157,7 +199,9 @@ async def _call_tool(self, tool_call) -> dict[str, str]: "tool_call_id": tool_call.id, } - def postprocess(self, batch: DataProto, batch_conversations: list[list[dict[str, str]]], n: int) -> DataProto: + def postprocess( + self, batch: DataProto, batch_conversations: list[list[dict[str, str]]], n: int + ) -> DataProto: # NOTE: consistent with batch version of generate_sequences in vllm_rollout_spmd.py # prompts: left pad # responses: right pad @@ -168,7 +212,10 @@ def postprocess(self, batch: DataProto, batch_conversations: list[list[dict[str, # prompts: [prompt] from input dataset prompts = [ self.tokenizer.apply_chat_template( - prompt, tools=self.tool_schemas, add_generation_prompt=True, tokenize=False + prompt, + tools=self.tool_schemas, + add_generation_prompt=True, + tokenize=False, ) for prompt in batch.non_tensor_batch["raw_prompt"] ] @@ -177,19 +224,33 @@ def postprocess(self, batch: DataProto, batch_conversations: list[list[dict[str, # sequences: [prompt + response] sequences = [ self.tokenizer.apply_chat_template( - conversation, tools=self.tool_schemas, add_generation_prompt=False, tokenize=False + conversation, + tools=self.tool_schemas, + add_generation_prompt=False, + tokenize=False, ) for conversation in batch_conversations ] # responses: [response] - responses = [sequence[len(prompts[i // n]) :] for i, sequence in enumerate(sequences)] - - prompts = self.tokenizer(prompts, return_tensors="pt", padding="longest", padding_side="left") - responses = self.tokenizer(responses, return_tensors="pt", padding="longest", padding_side="right") + responses = [sequence[len(prompts[i // n]):] + for i, sequence in enumerate(sequences)] + + prompts = self.tokenizer( + prompts, + return_tensors="pt", + padding="longest", + padding_side="left") + responses = self.tokenizer( + responses, + return_tensors="pt", + padding="longest", + padding_side="right") if n > 1: - prompts["input_ids"] = prompts["input_ids"].repeat_interleave(n, dim=0) - prompts["attention_mask"] = prompts["attention_mask"].repeat_interleave(n, dim=0) + prompts["input_ids"] = prompts["input_ids"].repeat_interleave( + n, dim=0) + prompts["attention_mask"] = prompts["attention_mask"].repeat_interleave( + n, dim=0) # response_mask: response mask with tools calling masked out response_mask = self._mask_out_tools_calling_tokens( @@ -199,8 +260,11 @@ def postprocess(self, batch: DataProto, batch_conversations: list[list[dict[str, responses["attention_mask"], ) - input_ids = torch.cat([prompts["input_ids"], responses["input_ids"]], dim=1) - attention_mask = torch.cat([prompts["attention_mask"], responses["attention_mask"]], dim=1) + input_ids = torch.cat( + [prompts["input_ids"], responses["input_ids"]], dim=1) + attention_mask = torch.cat( + [prompts["attention_mask"], responses["attention_mask"]], dim=1 + ) position_ids = (attention_mask.cumsum(dim=1) - 1) * attention_mask batch = TensorDict( @@ -208,15 +272,22 @@ def postprocess(self, batch: DataProto, batch_conversations: list[list[dict[str, "prompts": prompts["input_ids"], # [bsz, prompt_length] "responses": responses["input_ids"], # [bsz, response_length] "response_mask": response_mask, # [bsz, response_length] - "input_ids": input_ids, # [bsz, prompt_length + response_length] - "attention_mask": attention_mask, # [bsz, prompt_length + response_length] - "position_ids": position_ids, # [bsz, prompt_length + response_length] + # [bsz, prompt_length + response_length] + "input_ids": input_ids, + # [bsz, prompt_length + response_length] + "attention_mask": attention_mask, + # [bsz, prompt_length + response_length] + "position_ids": position_ids, }, batch_size=len(input_ids), ) - num_turns = np.array([len(conversation) for conversation in batch_conversations], dtype=np.int32) - return DataProto(batch=batch, non_tensor_batch={"__num_turns__": num_turns}) + num_turns = np.array( + [len(conversation) for conversation in batch_conversations], dtype=np.int32 + ) + return DataProto( + batch=batch, non_tensor_batch={ + "__num_turns__": num_turns}) def _mask_out_tools_calling_tokens( self, @@ -237,12 +308,16 @@ def _mask_out_tools_calling_tokens( mask: (batch_size, response_length) """ batch_size = input_ids.size(0) - assert len(raw_prompts) == batch_size, f"{len(raw_prompts)} != {batch_size}" - assert len(batch_conversations) == batch_size, f"{len(batch_conversations)} != {batch_size}" + assert len(raw_prompts) == batch_size, f"{ + len(raw_prompts)} != {batch_size}" + assert ( + len(batch_conversations) == batch_size + ), f"{len(batch_conversations)} != {batch_size}" # Deduplicate adjacent tool calls, since they're merged into one turn. # [user, assistant, tool, tool, assistant] -> [user, assistant, tool, assistant] - # TODO: it's chat_template specific, find a more generic way to do this. + # TODO: it's chat_template specific, find a more generic way to do + # this. def deduplicate_adjacent_tool_calls(roles): result = [] for role, group in itertools.groupby(roles): @@ -254,17 +329,24 @@ def deduplicate_adjacent_tool_calls(roles): loss_mask = attention_mask.clone() for i in range(batch_size): - responses = batch_conversations[i][len(raw_prompts[i]) :] + responses = batch_conversations[i][len(raw_prompts[i]):] assert len(responses) > 0, f"responses is empty: {responses}" - roles = deduplicate_adjacent_tool_calls([response["role"] for response in responses]) + roles = deduplicate_adjacent_tool_calls( + [response["role"] for response in responses] + ) # Each turn should be: [BOS]...[EOS] - eos_indices = input_ids[i].eq(self.tokenizer.eos_token_id).nonzero().squeeze(1)[: len(roles)] + eos_indices = ( + input_ids[i] + .eq(self.tokenizer.eos_token_id) + .nonzero() + .squeeze(1)[: len(roles)] + ) for j in range(len(roles)): if roles[j] == "tool": bos = eos_indices[j - 1] + 1 if j > 0 else 0 eos = eos_indices[j] - loss_mask[i, bos : eos + 1] = 0 + loss_mask[i, bos: eos + 1] = 0 return loss_mask @@ -288,7 +370,8 @@ def __init__( self.model_name = "/".join(model_path.split("/")[-2:]) # Least requests load balancing - self.weighted_addresses = [[0, address] for address in server_addresses] + self.weighted_addresses = [[0, address] + for address in server_addresses] heapq.heapify(self.weighted_addresses) # LRU cache to map request_id to address @@ -297,13 +380,18 @@ def __init__( self.background_tasks = set() if self.config.multi_turn.completion_callback is None: self.completion_callback = ToolCompletionCallback(config, self) - logger.warning("completion_callback is None, use ToolCompletionCallback") + logger.warning( + "completion_callback is None, use ToolCompletionCallback") else: - module_path, class_name = self.config.multi_turn.completion_callback.rsplit(".", 1) + module_path, class_name = self.config.multi_turn.completion_callback.rsplit( + ".", 1) module = importlib.import_module(module_path) - self.completion_callback = getattr(module, class_name)(config, self) + self.completion_callback = getattr( + module, class_name)(config, self) - def submit_chat_completions(self, *, messages: list[dict[str, str]], request_id: str, info: dict[str, Any]): + def submit_chat_completions( + self, *, messages: list[dict[str, str]], request_id: str, info: dict[str, Any] + ): """Submit chat completion request without wait, completion_callback will be called when the request is done. Args: @@ -312,7 +400,9 @@ def submit_chat_completions(self, *, messages: list[dict[str, str]], request_id: info: Any other auxiliary information pass across multi-turn. """ info["__depth__"] += 1 - task = asyncio.create_task(self._submit_chat_completions_and_callback(messages, request_id, info)) + task = asyncio.create_task( + self._submit_chat_completions_and_callback( + messages, request_id, info)) # โ€œfire-and-forgetโ€ background tasks self.background_tasks.add(task) @@ -332,7 +422,9 @@ async def _submit_chat_completions_and_callback( else: address = self.weighted_addresses[0][1] self.weighted_addresses[0][0] += 1 - heapq.heapreplace(self.weighted_addresses, self.weighted_addresses[0]) + heapq.heapreplace( + self.weighted_addresses, + self.weighted_addresses[0]) # use new request_id to avoid duplicate request_id problem request_id = uuid4().hex @@ -340,7 +432,8 @@ async def _submit_chat_completions_and_callback( completions, exception = None, None try: - # NOTE: OpenAI client uses httpx, seems to have performance issue in high concurrency requests. + # NOTE: OpenAI client uses httpx, seems to have performance issue + # in high concurrency requests. completions = await self._chat_completions_aiohttp( address, messages=messages, @@ -356,22 +449,33 @@ async def _submit_chat_completions_and_callback( info["__depth__"] -= 1 if exception is not None: - logger.exception(f"chat completion failed with exception: {exception}") + logger.exception( + f"chat completion failed with exception: {exception}") else: try: await self.completion_callback(messages, completions, info) except Exception as e: - logger.exception(f"completion callback failed with exception: {e}") + logger.exception( + f"completion callback failed with exception: {e}") # No more ongoing completion requests if info["__depth__"] == 0: info["__done__"].set() - async def _chat_completions_openai(self, address: str, **chat_complete_request) -> ChatCompletion: - client = AsyncOpenAI(base_url=f"http://{address}/v1", api_key="token-abc123", timeout=None, max_retries=0) + async def _chat_completions_openai( + self, address: str, **chat_complete_request + ) -> ChatCompletion: + client = AsyncOpenAI( + base_url=f"http://{address}/v1", + api_key="token-abc123", + timeout=None, + max_retries=0, + ) return await client.chat.completions.create(**chat_complete_request) - async def _chat_completions_aiohttp(self, address: str, **chat_complete_request) -> ChatCompletion: + async def _chat_completions_aiohttp( + self, address: str, **chat_complete_request + ) -> ChatCompletion: try: extra_body = chat_complete_request.pop("extra_body", {}) chat_complete_request.update(extra_body or {}) @@ -401,14 +505,19 @@ async def generate_sequences(self, batch: DataProto) -> DataProto: kwargs["top_p"] = self.config.val_kwargs.top_p kwargs["temperature"] = self.config.val_kwargs.temperature - print(f"[ChatCompletionScheduler] generate_sequences sampling params: {kwargs}") + print( + f"[ChatCompletionScheduler] generate_sequences sampling params: {kwargs}") # NOTE: For multi-turn rollout, repeat raw_prompt n times and process each prompt independently, - # validation dataset has already been repeated in `PPOTrainer._validate`. + # validation dataset has already been repeated in + # `PPOTrainer._validate`. n = 1 if batch.meta_info.get("validate", False) else self.config.n tasks, batch_conversations = [], [None] * len(batch) * n - for batch_index, conversation in enumerate(batch.non_tensor_batch["raw_prompt"].repeat(n, axis=0)): - # raw_prompt: [{"role": "user", "content": ""}, ["role": "assistant", "content"], ...] + for batch_index, conversation in enumerate( + batch.non_tensor_batch["raw_prompt"].repeat(n, axis=0) + ): + # raw_prompt: [{"role": "user", "content": ""}, ["role": + # "assistant", "content"], ...] batch_conversations[batch_index] = conversation.tolist() tasks.append( @@ -422,13 +531,19 @@ async def generate_sequences(self, batch: DataProto) -> DataProto: ) await asyncio.gather(*tasks) - output_batch = self.completion_callback.postprocess(batch, batch_conversations, n=n) - output_batch.meta_info["timing"] = {"generate_sequences": time.time() - t_start} + output_batch = self.completion_callback.postprocess( + batch, batch_conversations, n=n + ) + output_batch.meta_info["timing"] = { + "generate_sequences": time.time() - t_start} print("[ChatCompletionScheduler] generate_sequences done") return output_batch async def _submit_chat_completions_semaphore( - self, messages: list[dict[str, str]], request_id: str, sampling_params: dict[str, Any] + self, + messages: list[dict[str, str]], + request_id: str, + sampling_params: dict[str, Any], ): done = asyncio.Event() @@ -438,7 +553,9 @@ async def _submit_chat_completions_semaphore( "__sampling_params__": sampling_params, } - self.submit_chat_completions(messages=messages, request_id=request_id, info=info) + self.submit_chat_completions( + messages=messages, request_id=request_id, info=info + ) # Wait until all completion requests are done await done.wait() diff --git a/Agent0/executor_train/verl/verl/workers/rollout/hf_rollout.py b/Agent0/executor_train/verl/verl/workers/rollout/hf_rollout.py index 32d0bc8..86764e7 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/hf_rollout.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/hf_rollout.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -44,7 +44,9 @@ def __init__(self, module: nn.Module, config): 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) + 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) @@ -56,10 +58,15 @@ def _generate_minibatch(self, prompts: DataProto) -> DataProto: do_sample = prompts.meta_info.get("do_sample", self.config.do_sample) is_validate = prompts.meta_info.get("validate", False) - temperature = prompts.meta_info.get("temperature", self.config.temperature) - response_length = prompts.meta_info.get("response_length", self.config.response_length) + temperature = prompts.meta_info.get( + "temperature", self.config.temperature) + response_length = prompts.meta_info.get( + "response_length", self.config.response_length + ) top_p = prompts.meta_info.get("top_p", self.config.get("top_p", 1.0)) - top_k = max(0, prompts.meta_info.get("top_k", self.config.get("top_k", 0))) # to be compatible with vllm + top_k = max( + 0, prompts.meta_info.get("top_k", self.config.get("top_k", 0)) + ) # to be compatible with vllm if not do_sample: # do_sample==False -> greedy decoding @@ -72,7 +79,9 @@ def _generate_minibatch(self, prompts: DataProto) -> DataProto: kwargs = { "do_sample": True, "num_beams": 1, - "top_k": max(0, self.config.val_kwargs.top_k), # to be compatible with vllm + "top_k": max( + 0, self.config.val_kwargs.top_k + ), # to be compatible with vllm "top_p": self.config.val_kwargs.top_p, "temperature": self.config.val_kwargs.temperature, "num_return_sequences": 1, # if validate, already repeat in ray_trainer @@ -93,7 +102,8 @@ def _generate_minibatch(self, prompts: DataProto) -> DataProto: idx = prompts.batch["input_ids"] # (bs, prompt_length) prompt_length = idx.size(1) - attention_mask = prompts.batch["attention_mask"] # left-padded attention_mask + # left-padded attention_mask + attention_mask = prompts.batch["attention_mask"] position_ids = prompts.batch["position_ids"] # used to construct attention_mask @@ -104,9 +114,14 @@ def _generate_minibatch(self, prompts: DataProto) -> DataProto: param_ctx = contextlib.nullcontext() if isinstance(self.module, FSDP): - # recurse need to set to False according to https://github.com/pytorch/pytorch/issues/100069 - param_ctx = FSDP.summon_full_params(self.module, writeback=False, recurse=False) - with param_ctx, torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16): + # recurse need to set to False according to + # https://github.com/pytorch/pytorch/issues/100069 + param_ctx = FSDP.summon_full_params( + self.module, writeback=False, recurse=False + ) + with param_ctx, torch.autocast( + device_type=get_device_name(), dtype=torch.bfloat16 + ): output = self.module.generate( input_ids=idx, attention_mask=attention_mask, @@ -131,7 +146,11 @@ def _generate_minibatch(self, prompts: DataProto) -> DataProto: delta_length = sequence_length - seq.shape[1] if delta_length > 0: - delta_tokens = torch.ones(size=(generated_batch_size, delta_length), device=seq.device, dtype=seq.dtype) + delta_tokens = torch.ones( + size=(generated_batch_size, delta_length), + device=seq.device, + dtype=seq.dtype, + ) delta_tokens = pad_token_id * delta_tokens seq = torch.cat((seq, delta_tokens), dim=1) assert seq.shape[1] == sequence_length @@ -139,23 +158,34 @@ def _generate_minibatch(self, prompts: DataProto) -> DataProto: # make necessary reputations if num_return_sequences > 1 num_return_sequences = kwargs.get("num_return_sequences", 1) if num_return_sequences > 1: - position_ids = position_ids.repeat_interleave(num_return_sequences, dim=0) - attention_mask = attention_mask.repeat_interleave(num_return_sequences, dim=0) + position_ids = position_ids.repeat_interleave( + num_return_sequences, dim=0) + attention_mask = attention_mask.repeat_interleave( + num_return_sequences, dim=0 + ) - prompt = seq[:, :prompt_length] # (generated_batch_size, prompt_length) - response = seq[:, prompt_length:] # (generated_batch_size, response_length) + # (generated_batch_size, prompt_length) + prompt = seq[:, :prompt_length] + # (generated_batch_size, response_length) + response = seq[:, prompt_length:] response_length = response.size(1) - delta_position_id = torch.arange(1, response_length + 1, device=position_ids.device) - delta_position_id = delta_position_id.unsqueeze(0).repeat(generated_batch_size, 1) + delta_position_id = torch.arange( + 1, response_length + 1, device=position_ids.device + ) + delta_position_id = delta_position_id.unsqueeze(0).repeat( + generated_batch_size, 1 + ) response_position_ids = position_ids[:, -1:] + delta_position_id position_ids = torch.cat([position_ids, response_position_ids], dim=-1) response_attention_mask = get_response_mask( - response_id=response, eos_token=eos_token_id, dtype=attention_mask.dtype - ) - attention_mask = torch.cat((attention_mask, response_attention_mask), dim=-1) + 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( { diff --git a/Agent0/executor_train/verl/verl/workers/rollout/naive/__init__.py b/Agent0/executor_train/verl/verl/workers/rollout/naive/__init__.py index cb6c23b..9f75b1c 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/naive/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/naive/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/rollout/naive/naive_rollout.py b/Agent0/executor_train/verl/verl/workers/rollout/naive/naive_rollout.py index fe56dc4..dc123c9 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/naive/naive_rollout.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/naive/naive_rollout.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -51,7 +51,8 @@ def __init__(self, module: nn.Module, config): 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 + # left-padded attention_mask + attention_mask = prompts.batch["attention_mask"] position_ids = prompts.batch["position_ids"] # used to construct attention_mask @@ -62,7 +63,11 @@ def generate_sequences(self, prompts: DataProto) -> DataProto: self.module.eval() - prev_attention_mask = torch.ones(size=(batch_size, 1), dtype=attention_mask.dtype, device=attention_mask.device) + 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): @@ -71,13 +76,20 @@ def generate_sequences(self, prompts: DataProto) -> DataProto: 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) + 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) + # 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))) + 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) @@ -87,19 +99,24 @@ def generate_sequences(self, prompts: DataProto) -> DataProto: else: idx_next = torch.argmax(probs, dim=-1, keepdim=True) - attention_mask = torch.cat((attention_mask, prev_attention_mask), dim=-1) + attention_mask = torch.cat( + (attention_mask, prev_attention_mask), dim=-1) for token_id in eos_token_id: - prev_attention_mask = torch.logical_and(idx_next != token_id, prev_attention_mask.bool()) + prev_attention_mask = torch.logical_and( + idx_next != token_id, prev_attention_mask.bool() + ) prev_attention_mask.to(attention_mask.dtype) - position_ids = torch.cat((position_ids, position_ids[:, -1:] + 1), dim=-1) + 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) + # (bs, response_length, vocab_size) + logits = torch.stack(logits_lst, dim=1) prompts = idx[:, :prompt_length] # (bs, prompt_length) response = idx[:, prompt_length:] # (bs, response_length) log_probs = logprobs_from_logits(logits=logits, labels=response) diff --git a/Agent0/executor_train/verl/verl/workers/rollout/schemas.py b/Agent0/executor_train/verl/verl/workers/rollout/schemas.py index 99f860a..3ba9e38 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/schemas.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/schemas.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -122,19 +122,27 @@ class AsyncRolloutRequest(BaseModel): @classmethod def initialize_request(cls, values): if not (messages := values.get("messages")): - raise ValueError("messages is required for AsyncRolloutRequest initialization") + raise ValueError( + "messages is required for AsyncRolloutRequest initialization" + ) if not (max_prompt_len := values.get("max_prompt_len")): - raise ValueError("max_prompt_len is required for AsyncRolloutRequest initialization") + raise ValueError( + "max_prompt_len is required for AsyncRolloutRequest initialization" + ) if not (processing_class := values.pop("processing_class", None)): - raise ValueError("processing_class is required for AsyncRolloutRequest initialization") + raise ValueError( + "processing_class is required for AsyncRolloutRequest initialization" + ) values["messages"] = [Message.model_validate(msg) for msg in messages] - # If there is no multi_modal_keys, we assume the multi-modal data is image and video. + # If there is no multi_modal_keys, we assume the multi-modal data is + # image and video. if not values.get("multi_modal_keys"): values["multi_modal_keys"] = ["image", "video"] if not values.get("multi_modal_data"): - values["multi_modal_data"] = {key: [] for key in values["multi_modal_keys"]} + values["multi_modal_data"] = {key: [] + for key in values["multi_modal_keys"]} else: # check if all multi_modal_keys are in multi_modal_data for key in values["multi_modal_keys"]: @@ -144,7 +152,9 @@ def initialize_request(cls, values): values["multi_modal_inputs"] = {} tools = ( - [tool.model_dump() for tool in tool_schemas] if (tool_schemas := values.get("tool_schemas", [])) else None + [tool.model_dump() for tool in tool_schemas] + if (tool_schemas := values.get("tool_schemas", [])) + else None ) multi_modal_data = values["multi_modal_data"] @@ -189,13 +199,25 @@ def initialize_request(cls, values): multi_modal_inputs.pop("attention_mask", None) values["multi_modal_inputs"] = multi_modal_inputs - values["position_ids"] = values["prompt_position_ids"] = cls._get_position_ids( - processing_class, values["input_ids"], values["attention_mask"], multi_modal_inputs + values["position_ids"] = values["prompt_position_ids"] = ( + cls._get_position_ids( + processing_class, + values["input_ids"], + values["attention_mask"], + multi_modal_inputs, + ) ) - values["prompt_ids"], values["prompt_attention_mask"] = values["input_ids"], values["attention_mask"] - values["loss_mask"] = values["prompt_loss_mask"] = torch.zeros_like(values["input_ids"], dtype=torch.bool) - values["generation_prompt_ids"] = values["input_ids"][..., tokens_without_prompt.shape[-1] :] + values["prompt_ids"], values["prompt_attention_mask"] = ( + values["input_ids"], + values["attention_mask"], + ) + values["loss_mask"] = values["prompt_loss_mask"] = torch.zeros_like( + values["input_ids"], dtype=torch.bool + ) + values["generation_prompt_ids"] = values["input_ids"][ + ..., tokens_without_prompt.shape[-1]: + ] values["base_conv_wo_gen_prompt_end_pos"] = cls._handle_apply_chat_template( processing_class, BASE_CHAT_HISTORY, @@ -218,7 +240,9 @@ def initialize_request(cls, values): @staticmethod def _handle_apply_chat_template( - processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), messages: list[Message], multi_modal_data: dict[str, Any], tools: Optional[list[OpenAIFunctionToolSchema]] = None, @@ -227,24 +251,45 @@ def _handle_apply_chat_template( return_dict: bool = False, ): raw_prompt = processing_class.apply_chat_template( - messages, tools=tools, add_generation_prompt=add_generation_prompt, tokenize=False + messages, + tools=tools, + add_generation_prompt=add_generation_prompt, + tokenize=False, ) if not tokenize: return raw_prompt - if isinstance(processing_class, PreTrainedTokenizer) or isinstance(processing_class, PreTrainedTokenizerFast): + if isinstance(processing_class, PreTrainedTokenizer) or isinstance( + processing_class, PreTrainedTokenizerFast + ): if any(len(values) > 0 for values in multi_modal_data.values()): logger.warning( "There is multi_modal_data but you are not using a processor. Multi-modal data will be ignored." ) - model_inputs = processing_class(text=[raw_prompt], return_tensors="pt") + model_inputs = processing_class( + text=[raw_prompt], return_tensors="pt") elif isinstance(processing_class, ProcessorMixin): - # When we update multi_model_keys, we also need to update this logic - images = images if len(images := multi_modal_data.get("image", [])) > 0 else None - videos = videos if len(videos := multi_modal_data.get("video", [])) > 0 else None - model_inputs = processing_class(text=[raw_prompt], images=images, videos=videos, return_tensors="pt") + # When we update multi_model_keys, we also need to update this + # logic + images = ( + images if len( + images := multi_modal_data.get( + "image", + [])) > 0 else None) + videos = ( + videos if len( + videos := multi_modal_data.get( + "video", + [])) > 0 else None) + model_inputs = processing_class( + text=[raw_prompt], + images=images, + videos=videos, + return_tensors="pt") else: - raise ValueError(f"Unsupported processing class type: {type(processing_class)}") + raise ValueError( + f"Unsupported processing class type: {type(processing_class)}" + ) model_inputs = dict(model_inputs) if return_dict: @@ -254,7 +299,9 @@ def _handle_apply_chat_template( @staticmethod def _get_position_ids( - processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), input_ids: torch.Tensor, attention_mask: torch.Tensor, multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None, @@ -262,7 +309,8 @@ def _get_position_ids( # special case for qwen2vl is_qwen2vl = ( hasattr(processing_class, "image_processor") - and "Qwen2VLImageProcessor" in processing_class.image_processor.__class__.__name__ + and "Qwen2VLImageProcessor" + in processing_class.image_processor.__class__.__name__ ) if is_qwen2vl: from verl.models.transformers.qwen2_vl import get_rope_index @@ -271,14 +319,15 @@ def _get_position_ids( if multi_modal_inputs: image_grid_thw = multi_modal_inputs.get("image_grid_thw") video_grid_thw = multi_modal_inputs.get("video_grid_thw") - second_per_grid_ts = multi_modal_inputs.get("second_per_grid_ts") - - assert input_ids.dim() == 2 and input_ids.shape[0] == 1, ( - f"input_ids should be 2D with batch size 1, but got shape {input_ids.shape}" - ) - assert attention_mask.dim() == 2 and attention_mask.shape[0] == 1, ( - f"attention_mask should be 2D with batch size 1, but got shape {attention_mask.shape}" - ) + second_per_grid_ts = multi_modal_inputs.get( + "second_per_grid_ts") + + assert ( + input_ids.dim() == 2 and input_ids.shape[0] == 1 + ), f"input_ids should be 2D with batch size 1, but got shape {input_ids.shape}" + assert ( + attention_mask.dim() == 2 and attention_mask.shape[0] == 1 + ), f"attention_mask should be 2D with batch size 1, but got shape {attention_mask.shape}" new_position_ids = get_rope_index( processing_class, input_ids=input_ids.squeeze(0), @@ -289,11 +338,14 @@ def _get_position_ids( ) return new_position_ids # (3, seq_len) else: - return compute_position_id_with_mask(attention_mask) # (1, seq_len) + return compute_position_id_with_mask( + attention_mask) # (1, seq_len) def _update_input_ids( self, - processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), new_input_ids: torch.Tensor, attention_mask: bool, loss_mask: bool, @@ -304,7 +356,8 @@ def _update_input_ids( """ self.input_ids = torch.cat([self.input_ids, new_input_ids], dim=-1) attention_mask = torch.ones_like(new_input_ids) * int(attention_mask) - self.attention_mask = torch.cat([self.attention_mask, attention_mask], dim=-1) + self.attention_mask = torch.cat( + [self.attention_mask, attention_mask], dim=-1) loss_mask = torch.ones_like(new_input_ids) * int(loss_mask) self.loss_mask = torch.cat([self.loss_mask, loss_mask], dim=-1) @@ -312,23 +365,28 @@ def _update_input_ids( self._update_multi_modal_inputs(new_multi_modal_inputs) new_position_ids = self._get_position_ids( - processing_class, new_input_ids, attention_mask, new_multi_modal_inputs - ) + processing_class, + new_input_ids, + attention_mask, + new_multi_modal_inputs) last_pos = self.position_ids[..., -1:] new_position_ids = new_position_ids + (last_pos + 1) - self.position_ids = torch.cat([self.position_ids, new_position_ids], dim=-1) + self.position_ids = torch.cat( + [self.position_ids, new_position_ids], dim=-1) assert ( self.input_ids.shape[-1] == self.attention_mask.shape[-1] == self.position_ids.shape[-1] == self.loss_mask.shape[-1] - ), f"""Request {self.request_id} has different length of {self.input_ids.shape[-1]=}, + ), f"""Request {self.request_id} has different length of {self.input_ids.shape[-1]=}, {self.attention_mask.shape[-1]=}, {self.position_ids.shape[-1]=}, {self.loss_mask.shape[-1]=}""" - def _update_multi_modal_inputs(self, new_multi_modal_inputs: dict[str, torch.Tensor]) -> None: + def _update_multi_modal_inputs( + self, new_multi_modal_inputs: dict[str, torch.Tensor] + ) -> None: """ Update the multi_modal_inputs of the request in additive manner. """ @@ -341,7 +399,10 @@ def _update_multi_modal_inputs(self, new_multi_modal_inputs: dict[str, torch.Ten ) def get_generation_prompt_ids( - self, processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + self, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), ) -> list[int]: """ Get the generation prompt ids for rollout engine. @@ -350,15 +411,26 @@ def get_generation_prompt_ids( """ generation_prompt_ids = ( None - if self.input_ids[..., -self.generation_prompt_ids.shape[-1] :].eq(self.generation_prompt_ids).all() + if self.input_ids[..., -self.generation_prompt_ids.shape[-1]:] + .eq(self.generation_prompt_ids) + .all() else self.generation_prompt_ids ) if generation_prompt_ids is not None: - self._update_input_ids(processing_class, generation_prompt_ids, attention_mask=True, loss_mask=False) + self._update_input_ids( + processing_class, + generation_prompt_ids, + attention_mask=True, + loss_mask=False, + ) if self.use_inference_chat_template: messages = [msg.model_dump() for msg in self.messages] - tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None + tools = ( + [tool.model_dump() for tool in self.tool_schemas] + if self.tool_schemas + else None + ) generation_prompt_ids = self._handle_apply_chat_template( processing_class, messages, @@ -373,52 +445,84 @@ def get_generation_prompt_ids( def add_user_message( self, - processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), content: str, ) -> None: self.messages.append(Message(role="user", content=content)) messages = [*BASE_CHAT_HISTORY, self.messages[-1]] - tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None + tools = ( + [tool.model_dump() for tool in self.tool_schemas] + if self.tool_schemas + else None + ) # We don't need to pass multi_modal_data here because we don't have any multi-modal data from Engine # Inference, it is pure text. content_ids = self._handle_apply_chat_template( - processing_class, messages, multi_modal_data={}, tools=tools, add_generation_prompt=False, tokenize=True - )[..., self.base_conv_wo_gen_prompt_end_pos :] - self._update_input_ids(processing_class, content_ids, attention_mask=True, loss_mask=False) + processing_class, + messages, + multi_modal_data={}, + tools=tools, + add_generation_prompt=False, + tokenize=True, + )[..., self.base_conv_wo_gen_prompt_end_pos:] + self._update_input_ids( + processing_class, content_ids, attention_mask=True, loss_mask=False + ) def add_assistant_message( self, - processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), content: str, tool_calls: Optional[list[OpenAIFunctionToolCall]] = None, ) -> None: - self.messages.append(Message(role="assistant", content=content, tool_calls=tool_calls)) + self.messages.append( + Message(role="assistant", content=content, tool_calls=tool_calls) + ) messages = [*BASE_CHAT_HISTORY, self.messages[-1]] - tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None + tools = ( + [tool.model_dump() for tool in self.tool_schemas] + if self.tool_schemas + else None + ) # We don't need to pass multi_modal_data here because we don't have any multi-modal data from Engine # Inference, it is pure text. content_ids = self._handle_apply_chat_template( - processing_class, messages, multi_modal_data={}, tools=tools, add_generation_prompt=False, tokenize=True - )[..., self.base_conv_with_gen_prompt_end_pos :] - self._update_input_ids(processing_class, content_ids, attention_mask=True, loss_mask=True) + processing_class, + messages, + multi_modal_data={}, + tools=tools, + add_generation_prompt=False, + tokenize=True, + )[..., self.base_conv_with_gen_prompt_end_pos:] + self._update_input_ids( + processing_class, content_ids, attention_mask=True, loss_mask=True + ) def add_tool_response_messages( self, - processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), contents: list[str | dict[str, Any]], ) -> None: if not contents: return # We also handle the case when tool returns image - # We require the processing of the image and video to be done at tool.execute() level + # We require the processing of the image and video to be done at + # tool.execute() level delta_multi_modal_data = {key: [] for key in self.multi_modal_keys} for content in contents: if isinstance(content, dict): content_list = [] - # When we update multi_model_keys, we also need to update this logic + # When we update multi_model_keys, we also need to update this + # logic if "image" in content: if not isinstance(content["image"], list): raise ValueError( @@ -427,7 +531,8 @@ def add_tool_response_messages( f"Example: {{'image': [img1]}} or {{'image': [img1, img2, ...]}}." ) - content_list.extend([{"type": "image"} for _ in content["image"]]) + content_list.extend([{"type": "image"} + for _ in content["image"]]) delta_multi_modal_data["image"].extend(content["image"]) if "video" in content: if not isinstance(content["video"], list): @@ -437,28 +542,38 @@ def add_tool_response_messages( f"Example: {{'video': [video1]}} or {{'video': [video1, video2, ...]}}." ) - content_list.extend([{"type": "video"} for _ in content["video"]]) + content_list.extend([{"type": "video"} + for _ in content["video"]]) delta_multi_modal_data["video"].extend(content["video"]) if "text" in content: - content_list.append({"type": "text", "text": content["text"]}) + content_list.append( + {"type": "text", "text": content["text"]}) for key in content: if key not in ["image", "video", "text"]: logger.warning( f"Tool response message contains unexpected key: {key} " f"while we only support `image`, `video`, and `text`." ) - self.messages.append(Message(role="tool", content=content_list)) + self.messages.append( + Message( + role="tool", + content=content_list)) else: self.messages.append(Message(role="tool", content=content)) - messages = [*BASE_CHAT_HISTORY, *self.messages[-len(contents) :]] - tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None + messages = [*BASE_CHAT_HISTORY, *self.messages[-len(contents):]] + tools = ( + [tool.model_dump() for tool in self.tool_schemas] + if self.tool_schemas + else None + ) for key in self.multi_modal_keys: if len(delta_multi_modal_data[key]) > 0: self.multi_modal_data[key].extend(delta_multi_modal_data[key]) - # We just passed the new multi-modal data to the chat template to update the input_ids. + # We just passed the new multi-modal data to the chat template to + # update the input_ids. content_info = self._handle_apply_chat_template( processing_class, messages, @@ -468,7 +583,9 @@ def add_tool_response_messages( tokenize=True, return_dict=True, ) - content_ids = content_info["input_ids"][..., self.base_conv_wo_gen_prompt_end_pos :] + content_ids = content_info["input_ids"][ + ..., self.base_conv_wo_gen_prompt_end_pos: + ] # process multi_modal_inputs multi_modal_inputs = content_info.copy() @@ -492,7 +609,9 @@ def update_metrics(self, metrics: Any, tool_id: str) -> None: def _get_prompt_diffs( self, - processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), full_prompt_ids: torch.Tensor, current_prompt_ids: torch.Tensor, diff_surrounding_chars: int = 10, @@ -524,9 +643,14 @@ def _get_prompt_diffs( """ full_prompt_ids = full_prompt_ids.squeeze(0) current_prompt_ids = current_prompt_ids.squeeze(0) - full_prompt = processing_class.decode(full_prompt_ids, skip_special_tokens=False) - current_prompt = processing_class.decode(current_prompt_ids, skip_special_tokens=False) - s = difflib.SequenceMatcher(None, full_prompt, current_prompt, autojunk=False) + full_prompt = processing_class.decode( + full_prompt_ids, skip_special_tokens=False + ) + current_prompt = processing_class.decode( + current_prompt_ids, skip_special_tokens=False + ) + s = difflib.SequenceMatcher( + None, full_prompt, current_prompt, autojunk=False) diffs = [] for tag, i1, i2, j1, j2 in s.get_opcodes(): if tag == "equal": @@ -549,7 +673,9 @@ def _get_prompt_diffs( def finalize( self, - processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), reward_scores: dict[str, list[float]], finish_reason_type: FinishReasonTypeEnum = FinishReasonTypeEnum.STOP, ) -> None: @@ -558,20 +684,40 @@ def finalize( # In case we failed to generate the assistant message and the generation prompt ids were already added to # input_ids, remove them from the end of input_ids - if self.input_ids[..., -self.generation_prompt_ids.shape[-1] :].eq(self.generation_prompt_ids).all(): - self.input_ids = self.input_ids[..., : -self.generation_prompt_ids.shape[-1]] - self.attention_mask = self.attention_mask[..., : -self.generation_prompt_ids.shape[-1]] - self.position_ids = self.position_ids[..., : -self.generation_prompt_ids.shape[-1]] - self.loss_mask = self.loss_mask[..., : -self.generation_prompt_ids.shape[-1]] - - self.response_ids = self.input_ids[..., self.prompt_ids.shape[-1] :] + if ( + self.input_ids[..., -self.generation_prompt_ids.shape[-1]:] + .eq(self.generation_prompt_ids) + .all() + ): + self.input_ids = self.input_ids[ + ..., : -self.generation_prompt_ids.shape[-1] + ] + self.attention_mask = self.attention_mask[ + ..., : -self.generation_prompt_ids.shape[-1] + ] + self.position_ids = self.position_ids[ + ..., : -self.generation_prompt_ids.shape[-1] + ] + self.loss_mask = self.loss_mask[ + ..., : -self.generation_prompt_ids.shape[-1] + ] + + self.response_ids = self.input_ids[..., self.prompt_ids.shape[-1]:] - if self.tokenization_sanity_check_mode != TokenizationSanityCheckModeEnum.DISABLE: - # When there is a diff, we log the diffs with diff_surrounding_chars context + if ( + self.tokenization_sanity_check_mode + != TokenizationSanityCheckModeEnum.DISABLE + ): + # When there is a diff, we log the diffs with + # diff_surrounding_chars context diff_surrounding_chars = 10 messages = [msg.model_dump() for msg in self.messages] - tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None + tools = ( + [tool.model_dump() for tool in self.tool_schemas] + if self.tool_schemas + else None + ) full_prompt_info = self._handle_apply_chat_template( processing_class, messages, @@ -609,20 +755,32 @@ def finalize( ) if diffs := self._get_prompt_diffs( - processing_class, full_prompt_ids, self.input_ids, diff_surrounding_chars=diff_surrounding_chars + processing_class, + full_prompt_ids, + self.input_ids, + diff_surrounding_chars=diff_surrounding_chars, ): log_warning = False - if self.tokenization_sanity_check_mode == TokenizationSanityCheckModeEnum.STRICT: + if ( + self.tokenization_sanity_check_mode + == TokenizationSanityCheckModeEnum.STRICT + ): log_warning = True - elif self.tokenization_sanity_check_mode == TokenizationSanityCheckModeEnum.IGNORE_STRIPPABLE: + elif ( + self.tokenization_sanity_check_mode + == TokenizationSanityCheckModeEnum.IGNORE_STRIPPABLE + ): non_strippable_diffs_exist = any( - d["full_prompt_chunk"].strip() or d["current_prompt_chunk"].strip() for d in diffs + d["full_prompt_chunk"].strip() + or d["current_prompt_chunk"].strip() + for d in diffs ) if non_strippable_diffs_exist: log_warning = True if log_warning: - mode_str = f" ({self.tokenization_sanity_check_mode.value})" + mode_str = f" ({ + self.tokenization_sanity_check_mode.value})" logger.warning( f"Inconsistent training and inference tokenization detected{mode_str}. This may lead to " f"unexpected behavior during training. Please review your chat template to determine if this " @@ -647,7 +805,8 @@ def finalize( elif finish_reason_type == FinishReasonTypeEnum.LENGTH: pass else: - raise ValueError(f"Unsupported finalize finish reason type: {finish_reason_type}") + raise ValueError( + f"Unsupported finalize finish reason type: {finish_reason_type}") self.truncate_output_ids(processing_class) assert ( @@ -655,21 +814,28 @@ def finalize( == self.attention_mask.shape[-1] == self.position_ids.shape[-1] == self.loss_mask.shape[-1] - ), f"""Request {self.request_id} has different length of {self.input_ids.shape[-1]=}, + ), f"""Request {self.request_id} has different length of {self.input_ids.shape[-1]=}, {self.attention_mask.shape[-1]=}, {self.position_ids.shape[-1]=}, {self.loss_mask.shape[-1]=}""" def truncate_output_ids( - self, processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + self, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), ) -> None: self.input_ids = self.input_ids[..., : self.max_model_len] self.attention_mask = self.attention_mask[..., : self.max_model_len] self.position_ids = self.position_ids[..., : self.max_model_len] self.loss_mask = self.loss_mask[..., : self.max_model_len] - self.response_ids = self.input_ids[..., self.prompt_ids.shape[-1] :][..., : self.max_response_len] - self.response_attention_mask = self.attention_mask[..., self.prompt_attention_mask.shape[-1] :][ - ..., : self.max_response_len - ] - self.response_position_ids = self.position_ids[..., self.prompt_position_ids.shape[-1] :][ + self.response_ids = self.input_ids[..., self.prompt_ids.shape[-1]:][ ..., : self.max_response_len ] - self.response_loss_mask = self.loss_mask[..., self.prompt_loss_mask.shape[-1] :][..., : self.max_response_len] + self.response_attention_mask = self.attention_mask[ + ..., self.prompt_attention_mask.shape[-1]: + ][..., : self.max_response_len] + self.response_position_ids = self.position_ids[ + ..., self.prompt_position_ids.shape[-1]: + ][..., : self.max_response_len] + self.response_loss_mask = self.loss_mask[ + ..., self.prompt_loss_mask.shape[-1]: + ][..., : self.max_response_len] diff --git a/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/__init__.py b/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/__init__.py index 43a1eeb..221a4f6 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/async_sglang_server.py b/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/async_sglang_server.py index df26765..a358891 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/async_sglang_server.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/async_sglang_server.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -28,10 +28,16 @@ @ray.remote(num_cpus=1) class AsyncSglangServer(AsyncServerBase): - def __init__(self, config: DictConfig, dp_size: int, dp_rank: int, wg_prefix: str): + def __init__( + self, + config: DictConfig, + dp_size: int, + dp_rank: int, + wg_prefix: str): super().__init__() self.config = config.actor_rollout_ref - self._tp_size = self.config.rollout.get("tensor_model_parallel_size", 1) + self._tp_size = self.config.rollout.get( + "tensor_model_parallel_size", 1) self._dp_size = dp_size self._dp_rank = dp_rank self.wg_prefix = wg_prefix @@ -44,18 +50,25 @@ async def init_engine(self): return all_actors = ray.util.list_named_actors(all_namespaces=True) matched_actors = [ - actor for actor in all_actors if actor.get("name", None).startswith(self.wg_prefix + "WorkerDict_") - ] + actor for actor in all_actors if actor.get( + "name", None).startswith( + self.wg_prefix + "WorkerDict_")] for matched_actor in matched_actors: fields = matched_actor["name"].split(":") - assert len(fields) == 2, f"invalid actor name: {matched_actor['name']}" - pg_index, local_rank = int(fields[0].split("_")[-1]), int(fields[1]) + assert len(fields) == 2, f"invalid actor name: { + matched_actor['name']}" + pg_index, local_rank = int( + fields[0].split("_")[-1]), int(fields[1]) - if (self._dp_size * pg_index + local_rank) // self._tp_size == self._dp_rank: + if ( + self._dp_size * pg_index + local_rank + ) // self._tp_size == self._dp_rank: worker = ray.get_actor(**matched_actor) self.workers.append(worker) - if (self._dp_size * pg_index + local_rank) / self._tp_size == self._dp_rank: + if ( + self._dp_size * pg_index + local_rank + ) / self._tp_size == self._dp_rank: self.master_worker = worker async def chat_completion(self, raw_request: Request): @@ -66,8 +79,14 @@ async def chat_completion(self, raw_request: Request): [outputs] = await asyncio.gather(output_future) return JSONResponse(outputs) - async def generate(self, prompt_ids: list[int], sampling_params: dict[str, Any], request_id: str) -> list[int]: - return await self.master_worker.generate.remote(prompt_ids, sampling_params, request_id) + async def generate(self, + prompt_ids: list[int], + sampling_params: dict[str, + Any], + request_id: str) -> list[int]: + return await self.master_worker.generate.remote( + prompt_ids, sampling_params, request_id + ) async def wake_up(self): if not self.config.rollout.free_cache_engine: diff --git a/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/sglang_rollout.py b/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/sglang_rollout.py index 3c66943..0dba8c4 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/sglang_rollout.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/sglang_rollout.py @@ -1,6 +1,6 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -54,10 +54,16 @@ from verl import DataProto from verl.interactions.base import BaseInteraction -from verl.interactions.utils.interaction_registry import initialize_interactions_from_config +from verl.interactions.utils.interaction_registry import ( + initialize_interactions_from_config, +) from verl.third_party.sglang import parallel_state as sglang_ps from verl.tools.base_tool import BaseTool -from verl.tools.schemas import OpenAIFunctionCallSchema, OpenAIFunctionParsedSchema, OpenAIFunctionToolCall +from verl.tools.schemas import ( + OpenAIFunctionCallSchema, + OpenAIFunctionParsedSchema, + OpenAIFunctionToolCall, +) from verl.tools.utils.tool_registry import initialize_tools_from_config from verl.utils.net_utils import is_ipv6 from verl.utils.profiler import GPUMemoryLogger @@ -105,7 +111,9 @@ def _set_envs_and_config(server_args: ServerArgs): # Fix triton bugs if server_args.tp_size * server_args.dp_size > 1: - # FIXME: remove this after https://github.com/triton-lang/triton/pull/4295 is used as a dependency. + # FIXME: remove this after + # https://github.com/triton-lang/triton/pull/4295 is used as a + # dependency. maybe_set_triton_cache_manager() # Check flashinfer version @@ -130,14 +138,17 @@ def _set_envs_and_config(server_args: ServerArgs): # because chatCompletion is an async method, it makes the whole ray actor be an async actor -# which can not call loop.run_until_complete. So we need to make the engine to be an async class +# which can not call loop.run_until_complete. So we need to make the +# engine to be an async class class AsyncEngine(sglang.srt.entrypoints.engine.Engine): def __init__(self, **kwargs): super().__init__(**kwargs) - # default to use dummy load format, which need to reload weights in first time + # default to use dummy load format, which need to reload weights in + # first time self._need_reload = True - async def release_memory_occupation(self, tags: Optional[list[str]] = None): + async def release_memory_occupation( + self, tags: Optional[list[str]] = None): """Release GPU occupation temporarily.""" if tags is None: obj = ReleaseMemoryOccupationReqInput() @@ -149,7 +160,8 @@ async def resume_memory_occupation(self, tags: Optional[list[str]] = None): """Resume GPU occupation.""" # because __init__ is a sync method, it can not call the async release_memory_occupation # have to move release_memory_occupation from __init__ to here - # For multi-stage awake, we run release weight and kv_cache when we resume weights for the first time. + # For multi-stage awake, we run release weight and kv_cache when we + # resume weights for the first time. if self._need_reload: await self.release_memory_occupation() self._need_reload = False @@ -170,7 +182,8 @@ async def update_weights_from_tensor( to avoid duplicated cache cleaning operation.""" obj = UpdateWeightsFromTensorReqInput( serialized_named_tensors=[ - MultiprocessingSerializer.serialize(named_tensors) for _ in range(self.server_args.tp_size) + MultiprocessingSerializer.serialize(named_tensors) + for _ in range(self.server_args.tp_size) ], load_format=load_format, flush_cache=flush_cache, @@ -188,7 +201,9 @@ def _pre_process_inputs( prompt_token_ids: torch.Tensor, ) -> torch.Tensor: # remove the left padding in the prompt token_id - non_pad_index = torch.nonzero(prompt_token_ids != pad_token_id, as_tuple=False)[0][0] + non_pad_index = torch.nonzero( + prompt_token_ids != pad_token_id, + as_tuple=False)[0][0] return prompt_token_ids[non_pad_index:] @@ -202,12 +217,17 @@ def _post_process_outputs(processing_class, output): # This is when processing_class is a tokenizer tokenizer = processing_class except AttributeError as e: - raise ValueError(f"Cannot get tokenizer from processing_class {processing_class}") from e + raise ValueError( + f"Cannot get tokenizer from processing_class {processing_class}") from e def _map_each_response(resp): output_token_logprobs = resp["meta_info"]["output_token_logprobs"] log_probs, output_token_ids = zip( - *[(log_prob, token_ids) for log_prob, token_ids, _ in output_token_logprobs], strict=True + *[ + (log_prob, token_ids) + for log_prob, token_ids, _ in output_token_logprobs + ], + strict=True, ) return torch.tensor(output_token_ids), torch.tensor(log_probs) @@ -217,10 +237,18 @@ def _map_each_response(resp): for output_token_ids, log_probs in out_map: batched_output_token_ids.append(output_token_ids) batched_logprobs.append(log_probs) - pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id - batched_output_token_ids = pad_sequence(batched_output_token_ids, batch_first=True, padding_value=pad_token_id) + pad_token_id = ( + tokenizer.pad_token_id + if tokenizer.pad_token_id is not None + else tokenizer.eos_token_id + ) + batched_output_token_ids = pad_sequence( + batched_output_token_ids, batch_first=True, padding_value=pad_token_id + ) if len(batched_logprobs) > 0: - batched_logprobs = pad_sequence(batched_logprobs, batch_first=True, padding_value=pad_token_id) + batched_logprobs = pad_sequence( + batched_logprobs, batch_first=True, padding_value=pad_token_id + ) return batched_output_token_ids, batched_logprobs @@ -238,14 +266,15 @@ def get_tool_call_parser_type( # This is when processing_class is a processor tokenizer_vocab = processing_class.tokenizer.get_vocab() except AttributeError as e: - raise ValueError(f"Cannot get vocab from processing_class {processing_class}") from e + raise ValueError( + f"Cannot get vocab from processing_class {processing_class}") from e if parser.bot_token.strip() in tokenizer_vocab and ( - parser.eot_token == "" or parser.eot_token.strip() in tokenizer_vocab - ): + parser.eot_token == "" or parser.eot_token.strip() in tokenizer_vocab): return parser_type else: - raise ValueError(f"No tool call parser found for processing_class {processing_class}") + raise ValueError( + f"No tool call parser found for processing_class {processing_class}") class SGLangRollout(BaseRollout): @@ -253,7 +282,9 @@ def __init__( self, actor_module: str, config: DictConfig, - processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin, + processing_class: ( + PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin + ), model_hf_config, port=None, trust_remote_code: bool = False, @@ -294,14 +325,18 @@ def __init__( self._sgl_tools, self._function_call_parser, ) = self._initialize_tools(config, processing_class) - self.interaction_map: dict[str, BaseInteraction] = self._initialize_interactions(config) + self.interaction_map: dict[str, BaseInteraction] = ( + self._initialize_interactions(config) + ) # If turn on `free_cache_engine`, SGLang engine's KV cache # will be freed after each `generate_sequences` call. logger.info( - f"tool_schemas: {self._tool_schemas}, tool_map: {self._tool_map}, tool_call_parser_type: " - f"{self._tool_call_parser_type}, sgl_tools: {self._sgl_tools}, function_call_parser: " - f"{self._function_call_parser}" - ) + f"tool_schemas: { + self._tool_schemas}, tool_map: { + self._tool_map}, tool_call_parser_type: " f"{ + self._tool_call_parser_type}, sgl_tools: { + self._sgl_tools}, function_call_parser: " f"{ + self._function_call_parser}") self._init_distributed_env(device_mesh_cpu=device_mesh, **kwargs) @@ -321,15 +356,18 @@ def __init__( # This is when processing_class is a processor self.pad_token_id = self.processing_class.tokenizer.pad_token_id except AttributeError as e: - raise ValueError(f"Cannot get pad_token_id from processing_class {self.processing_class}") from e + raise ValueError( + f"Cannot get pad_token_id from processing_class { + self.processing_class}") from e def _init_distributed_env(self, device_mesh_cpu, **kwargs): self._device_mesh_cpu = device_mesh_cpu os.environ.setdefault("SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK", "true") - self.tensor_parallel_size = self.config.get("tensor_model_parallel_size", 1) - assert self.tensor_parallel_size <= dist.get_world_size(), ( - "tensor parallel size should be less than or equal to the world size" - ) + self.tensor_parallel_size = self.config.get( + "tensor_model_parallel_size", 1) + assert ( + self.tensor_parallel_size <= dist.get_world_size() + ), "tensor parallel size should be less than or equal to the world size" self.train_tp = kwargs.get("train_tp", None) if self.train_tp is not None: # deployed with megatron @@ -352,45 +390,61 @@ def _init_distributed_env(self, device_mesh_cpu, **kwargs): mesh_dim_names=["dp", "tp", "pp"], ) - self._device_mesh_cpu = init_device_mesh("cpu", **device_mesh_kwargs) + self._device_mesh_cpu = init_device_mesh( + "cpu", **device_mesh_kwargs) self._rank = self._device_mesh_cpu.get_rank() self._tp_rank = self._device_mesh_cpu["tp"].get_local_rank() self._tp_size = self._device_mesh_cpu["tp"].size() if self._rank == 0: - logger.info(f"_init_distributed_env: :tp_world: {self._tp_size}, global_world: {world_size}") + logger.info( + f"_init_distributed_env: :tp_world: { + self._tp_size}, global_world: {world_size}") # get tp_rank of this process in this tp group visible_devices = [None] * self._device_mesh_cpu.size(1) torch.distributed.all_gather_object( - visible_devices, os.environ["CUDA_VISIBLE_DEVICES"], self._device_mesh_cpu.get_group("tp") + visible_devices, + os.environ["CUDA_VISIBLE_DEVICES"], + self._device_mesh_cpu.get_group("tp"), ) self.visible_devices_set = set(",".join(visible_devices).split(",")) - os.environ["CUDA_VISIBLE_DEVICES"] = ",".join(sorted(list(self.visible_devices_set))) + os.environ["CUDA_VISIBLE_DEVICES"] = ",".join( + sorted(list(self.visible_devices_set)) + ) def _verify_config(self, model_hf_config): if not self.config.get("max_model_len", None): - self.config.max_model_len = self.config.prompt_length + self.config.response_length + self.config.max_model_len = ( + self.config.prompt_length + self.config.response_length + ) assert ( - self.config.max_model_len >= self.config.prompt_length + self.config.response_length - ), f"""max_model_len should be greater than total sequence length (prompt_length + response_length): + self.config.max_model_len + >= self.config.prompt_length + self.config.response_length + ), f"""max_model_len should be greater than total sequence length (prompt_length + response_length): {self.config.max_model_len} >= {self.config.prompt_length} + {self.config.response_length}""" max_position_embeddings = None if hasattr(model_hf_config, "max_position_embeddings"): max_position_embeddings = model_hf_config.max_position_embeddings - elif hasattr(model_hf_config, "llm_config") and hasattr(model_hf_config.llm_config, "max_position_embeddings"): + elif hasattr(model_hf_config, "llm_config") and hasattr( + model_hf_config.llm_config, "max_position_embeddings" + ): max_position_embeddings = model_hf_config.llm_config.max_position_embeddings elif hasattr(model_hf_config, "text_config") and hasattr( model_hf_config.text_config, "max_position_embeddings" ): - max_position_embeddings = model_hf_config.text_config.max_position_embeddings + max_position_embeddings = ( + model_hf_config.text_config.max_position_embeddings + ) if max_position_embeddings is None: - raise ValueError("max_position_embeddings not found in model_hf_config") + raise ValueError( + "max_position_embeddings not found in model_hf_config") rope_scaling_config = getattr(model_hf_config, "rope_scaling", None) if not rope_scaling_config: - assert max_position_embeddings >= self.config.prompt_length + self.config.response_length, ( - "model context length should be greater than total sequence length" - ) + assert ( + max_position_embeddings + >= self.config.prompt_length + self.config.response_length + ), "model context length should be greater than total sequence length" else: # handle type where there's a length extend factor # see https://qwen.readthedocs.io/en/latest/deployment/vllm.html#extended-context-support @@ -425,11 +479,16 @@ def _init_inference_engine(self, trust_remote_code, actor_module, port): src=self._device_mesh_cpu["tp"].mesh[0].item(), force_cpu_device=False, ) - dist_init_addr = f"[{ip}]:{port}" if is_ipv6(ip) else f"{ip}:{port}" + dist_init_addr = f"[{ip}]:{port}" if is_ipv6( + ip) else f"{ip}:{port}" else: dist_init_addr = None - load_format = "dummy" if self.config.load_format.startswith("dummy") else self.config.load_format + load_format = ( + "dummy" + if self.config.load_format.startswith("dummy") + else self.config.load_format + ) tp_size_per_node = self._tp_size // nnodes node_rank = self._tp_rank // tp_size_per_node first_rank_in_node = self._tp_rank % tp_size_per_node == 0 @@ -516,11 +575,15 @@ def _initialize_tools(self, config, processing_class): tools_config_file = config.multi_turn.tool_config_path tool_list = initialize_tools_from_config(tools_config_file) - logger.info(f"Initialize tools from configuration.: tool_list: {tool_list}") - tool_schemas = [tool.get_openai_tool_schema().model_dump() for tool in tool_list] + logger.info( + f"Initialize tools from configuration.: tool_list: {tool_list}") + tool_schemas = [ + tool.get_openai_tool_schema().model_dump() for tool in tool_list + ] tool_map = {tool.name: tool for tool in tool_list} tool_call_parser_type = get_tool_call_parser_type(processing_class) - sgl_tools = [Tool.model_validate(tool_schema) for tool_schema in tool_schemas] + sgl_tools = [Tool.model_validate(tool_schema) + for tool_schema in tool_schemas] function_call_parser = FunctionCallParser( sgl_tools, tool_call_parser_type, @@ -544,9 +607,13 @@ def _initialize_interactions(self, config): return {} interaction_config_file = config.multi_turn.interaction_config_path - interaction_map = initialize_interactions_from_config(interaction_config_file) + interaction_map = initialize_interactions_from_config( + interaction_config_file) - logger.info(f"Initialize interactions from configuration: interaction_map: {list(interaction_map.keys())}") + logger.info( + f"Initialize interactions from configuration: interaction_map: { + list( + interaction_map.keys())}") return interaction_map @GPUMemoryLogger(role="sglang rollout", logger=logger) @@ -578,7 +645,9 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: @GPUMemoryLogger(role="sglang rollout", logger=logger) @torch.no_grad() - def _batch_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: + def _batch_level_generate_sequences( + self, prompts: DataProto, **kwargs + ) -> DataProto: """Generates single-turn sequences for a batch of prompts. For single-turn generation, all prompts are processed in one request. `_batch_level_generate_sequences` involves: @@ -635,7 +704,10 @@ def _batch_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataP non_tensor_batch = prompts.non_tensor_batch if "raw_prompt_ids" not in non_tensor_batch: non_tensor_batch["raw_prompt_ids"] = np.array( - [_pre_process_inputs(self.pad_token_id, idx[i]).tolist() for i in range(batch_size)], + [ + _pre_process_inputs(self.pad_token_id, idx[i]).tolist() + for i in range(batch_size) + ], dtype=object, ) @@ -651,13 +723,16 @@ def _batch_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataP "prompt_token_ids": raw_prompt_ids, "multi_modal_data": multi_modal_data, "image_data": ( - multi_modal_data.get("image", None) if isinstance(multi_modal_data, dict) else None + multi_modal_data.get("image", None) + if isinstance(multi_modal_data, dict) + else None ), } ) else: sglang_inputs = [ - {"prompt_token_ids": raw_prompt_ids} for raw_prompt_ids in non_tensor_batch.pop("raw_prompt_ids") + {"prompt_token_ids": raw_prompt_ids} + for raw_prompt_ids in non_tensor_batch.pop("raw_prompt_ids") ] # Ensure token IDs are lists or numpy arrays @@ -666,12 +741,16 @@ def _batch_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataP input_data["prompt_token_ids"] = input_data["prompt_token_ids"].tolist() elif not isinstance(input_data["prompt_token_ids"], list): raise TypeError( - f"prompt_token_ids must be a list or numpy array, got {type(input_data['prompt_token_ids'])}" - ) + f"prompt_token_ids must be a list or numpy array, got { + type( + input_data['prompt_token_ids'])}") # Extract token IDs and image data for SGLang Engine - idx_list = [input_data["prompt_token_ids"] for input_data in sglang_inputs] - image_list = [input_data.get("image_data", None) for input_data in sglang_inputs] + idx_list = [input_data["prompt_token_ids"] + for input_data in sglang_inputs] + image_list = [ + input_data.get("image_data", None) for input_data in sglang_inputs + ] do_sample = prompts.meta_info.get("do_sample", True) is_validate = prompts.meta_info.get("validate", False) @@ -722,7 +801,8 @@ def _batch_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataP else: output = None - # Most naive implementation, can extract tensor and send via gloo if too slow + # Most naive implementation, can extract tensor and send via gloo if + # too slow dist.barrier() [output] = broadcast_pyobj( data=[output], @@ -739,19 +819,24 @@ def _batch_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataP rollout_log_probs = out[1].to(idx.device) if response.shape[1] < self.config.response_length: - response = pad_sequence_to_length(response, self.config.response_length, self.pad_token_id) + response = pad_sequence_to_length( + response, self.config.response_length, self.pad_token_id + ) if self.config.calculate_log_probs: rollout_log_probs = pad_sequence_to_length( - rollout_log_probs, self.config.response_length, self.pad_token_id - ) + rollout_log_probs, self.config.response_length, self.pad_token_id) 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) + 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) if position_ids.dim() == 3: # qwen2vl mrope - delta_position_id = delta_position_id.view(batch_size, 1, -1).expand(batch_size, 3, -1) + delta_position_id = delta_position_id.view( + batch_size, 1, -1).expand(batch_size, 3, -1) # TODO(sgm): fix position_ids on right_pad # prompt: left pad + response: right pad @@ -760,11 +845,14 @@ def _batch_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataP response_position_ids = position_ids[..., -1:] + delta_position_id position_ids = torch.cat([position_ids, response_position_ids], dim=-1) response_attention_mask = get_response_mask( - response_id=response, eos_token=eos_token_id, dtype=attention_mask.dtype - ) - attention_mask = torch.cat((attention_mask, response_attention_mask), dim=-1) - - # all the tp ranks should contain the same data here. data in all ranks are valid + 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, @@ -846,25 +934,35 @@ async def _async_rollout_a_request( self._tool_map[tool_call.function.name].execute( _req.request_id, tool_call.function.arguments, - **_req.tools_kwargs[tool_call.function.name].get("execute_kwargs", {}), + **_req.tools_kwargs[tool_call.function.name].get( + "execute_kwargs", {} + ), ) for tool_call in parsed_tool_calls ] ) - _req.add_tool_response_messages(self.processing_class, [resp for resp, _, _ in tool_call_results]) - for tool_call, (resp, reward, metrics) in zip(parsed_tool_calls, tool_call_results, strict=True): + _req.add_tool_response_messages( + self.processing_class, + [resp for resp, _, _ in tool_call_results], + ) + for tool_call, (resp, reward, metrics) in zip( + parsed_tool_calls, tool_call_results, strict=True + ): _req.update_metrics(metrics, tool_call.function.name) if len(_req.input_ids) >= self.config.max_model_len: finish_reason_type = FinishReasonTypeEnum.STOP break _req.state = AsyncRolloutRequestStateEnum.RUNNING else: - raise ValueError(f"Unexpected tool calling last message state: {_req.messages[-1]}") + raise ValueError( + f"Unexpected tool calling last message state: {_req.messages[-1]}" + ) elif _req.state == AsyncRolloutRequestStateEnum.RUNNING: # Only continue the conversation if the prompt length is not greater than max_model_len - 1, # since SGLang raises an error when max_new_tokens + 1 is greater to max_model_len (the extra # token accounts for the EOS token). - if len(_req.get_generation_prompt_ids(self.processing_class)) + 1 >= self.config.max_model_len: + if (len(_req.get_generation_prompt_ids( + self.processing_class)) + 1 >= self.config.max_model_len): finish_reason_type = FinishReasonTypeEnum.LENGTH break @@ -881,22 +979,31 @@ async def _async_rollout_a_request( ) if video_data: logger.warning( - "video support is not implemented yet, current length of video data is %d", len(video_data) + "video support is not implemented yet, current length of video data is %d", + len(video_data), ) - output = await self._handle_engine_call(_req, request_sampling_params, image_data=image_data) + output = await self._handle_engine_call( + _req, request_sampling_params, image_data=image_data + ) content = output["text"] - finish_reason_type = FinishReasonTypeEnum.from_str(output["meta_info"]["finish_reason"]["type"]) + finish_reason_type = FinishReasonTypeEnum.from_str( + output["meta_info"]["finish_reason"]["type"] + ) current_turns += 1 if finish_reason_type == FinishReasonTypeEnum.LENGTH: _req.add_assistant_message(self.processing_class, content) break else: - if self._function_call_parser and self._function_call_parser.has_tool_call(content): + if ( + self._function_call_parser + and self._function_call_parser.has_tool_call(content) + ): finish_reason_type = FinishReasonTypeEnum.TOOL_CALL _req.state = AsyncRolloutRequestStateEnum.TOOL_CALLING try: - normed_content, tool_calls = self._function_call_parser.parse_non_stream(content) + normed_content, tool_calls = ( + self._function_call_parser.parse_non_stream(content)) except JSONDecodeError: normed_content = content tool_calls = [] @@ -905,13 +1012,12 @@ async def _async_rollout_a_request( tool_calls = [] parsed_tool_calls = [] for tool_call in tool_calls: - function, has_decode_error = OpenAIFunctionCallSchema.from_openai_function_parsed_schema( - OpenAIFunctionParsedSchema( - name=tool_call.name, - arguments=tool_call.parameters, - ) - ) - # Drop the tool call if its arguments has decode error + function, has_decode_error = ( + OpenAIFunctionCallSchema.from_openai_function_parsed_schema( + OpenAIFunctionParsedSchema( + name=tool_call.name, arguments=tool_call.parameters, ))) + # Drop the tool call if its arguments has decode + # error if has_decode_error: continue parsed_tool_calls.append( @@ -922,10 +1028,13 @@ async def _async_rollout_a_request( ) if len(parsed_tool_calls) > 0: _req.add_assistant_message( - self.processing_class, normed_content, tool_calls=parsed_tool_calls + self.processing_class, + normed_content, + tool_calls=parsed_tool_calls, ) else: - _req.add_assistant_message(self.processing_class, content) + _req.add_assistant_message( + self.processing_class, content) finish_reason_type = FinishReasonTypeEnum.STOP _req.state = AsyncRolloutRequestStateEnum.COMPLETED break @@ -938,14 +1047,16 @@ async def _async_rollout_a_request( _req.interaction_kwargs and self.interaction_map and user_turns < self.config.multi_turn.max_user_turns - and current_turns < self.config.multi_turn.max_assistant_turns + and current_turns + < self.config.multi_turn.max_assistant_turns ): _req.state = AsyncRolloutRequestStateEnum.INTERACTING else: break elif _req.state == AsyncRolloutRequestStateEnum.INTERACTING: user_turns += 1 - messages = [{"role": x.role, "content": x.content} for x in _req.messages] + messages = [{"role": x.role, "content": x.content} + for x in _req.messages] # Get interaction by name from interaction_kwargs interaction_name = _req.interaction_kwargs.get( @@ -958,8 +1069,10 @@ async def _async_rollout_a_request( ) interaction = self.interaction_map[interaction_name] - should_terminate_sequence, content, reward, metrics = await interaction.generate_response( - _req.request_id, messages, **_req.interaction_kwargs + should_terminate_sequence, content, reward, metrics = ( + await interaction.generate_response( + _req.request_id, messages, **_req.interaction_kwargs + ) ) user_turn_rewards.append(reward) if should_terminate_sequence: @@ -979,8 +1092,12 @@ async def _async_rollout_a_request( # Calculate the reward for each tool async def calc_reward_and_release_fn(name: str, tool: BaseTool): - reward = await tool.calc_reward(_req.request_id, **_req.tools_kwargs[name].get("calc_reward_kwargs", {})) - await tool.release(_req.request_id, **_req.tools_kwargs[name].get("release_kwargs", {})) + reward = await tool.calc_reward( + _req.request_id, **_req.tools_kwargs[name].get("calc_reward_kwargs", {}) + ) + await tool.release( + _req.request_id, **_req.tools_kwargs[name].get("release_kwargs", {}) + ) return name, reward tool_reward_tasks = [] @@ -989,21 +1106,34 @@ async def calc_reward_and_release_fn(name: str, tool: BaseTool): tool_reward_tasks.append(calc_reward_and_release_fn(name, tool)) tool_reward_scores = await asyncio.gather(*tool_reward_tasks) tool_reward_scores = dict(tool_reward_scores) - all_rewards = {**tool_reward_scores, **{"user_turn_rewards": user_turn_rewards}} + all_rewards = {**tool_reward_scores, + **{"user_turn_rewards": user_turn_rewards}} _req.finalize(self.processing_class, all_rewards, finish_reason_type) return _req async def _handle_engine_call( - self, _req: AsyncRolloutRequest, sampling_params: dict, image_data: Optional[list[Any]] = None + self, + _req: AsyncRolloutRequest, + sampling_params: dict, + image_data: Optional[list[Any]] = None, ) -> dict: - generation_prompt_ids = _req.get_generation_prompt_ids(self.processing_class) - return await self._handle_engine_generate(generation_prompt_ids, sampling_params, image_data) + generation_prompt_ids = _req.get_generation_prompt_ids( + self.processing_class) + return await self._handle_engine_generate( + generation_prompt_ids, sampling_params, image_data + ) async def _handle_engine_generate( - self, generation_prompt_ids: list[int], sampling_params: dict, image_data: Optional[list[Any]] = None + self, + generation_prompt_ids: list[int], + sampling_params: dict, + image_data: Optional[list[Any]] = None, ) -> dict: - max_new_tokens = min(self.config.response_length, self.config.max_model_len - len(generation_prompt_ids) - 1) + max_new_tokens = min( + self.config.response_length, + self.config.max_model_len - len(generation_prompt_ids) - 1, + ) kwargs = sampling_params.copy() kwargs["max_new_tokens"] = max_new_tokens kwargs["n"] = 1 # group size is supported in preprocess @@ -1015,18 +1145,25 @@ async def _handle_engine_generate( ) return output - async def _handle_pending_state(self, _req: AsyncRolloutRequest) -> AsyncRolloutRequest: + async def _handle_pending_state( + self, _req: AsyncRolloutRequest + ) -> AsyncRolloutRequest: if _req.tool_schemas is not None: tool_creation_coroutines = [] for tool_schema in _req.tool_schemas: tool = self._tool_map[tool_schema.function.name] - create_kwargs = _req.tools_kwargs[tool.name].get("create_kwargs", {}) - tool_creation_coroutines.append(tool.create(_req.request_id, **create_kwargs)) + create_kwargs = _req.tools_kwargs[tool.name].get( + "create_kwargs", {}) + tool_creation_coroutines.append( + tool.create(_req.request_id, **create_kwargs) + ) await asyncio.gather(*tool_creation_coroutines) if _req.interaction_kwargs and self.interaction_map: interaction_kwargs = _req.interaction_kwargs # Get interaction by name from interaction_kwargs - interaction_name = interaction_kwargs.get("name", "gsm8k") # Default to gsm8k for backward compatibility + interaction_name = interaction_kwargs.get( + "name", "gsm8k" + ) # Default to gsm8k for backward compatibility if interaction_name not in self.interaction_map: raise ValueError( f"Interaction '{interaction_name}' not found in interaction_map. Available interactions: " @@ -1038,7 +1175,10 @@ async def _handle_pending_state(self, _req: AsyncRolloutRequest) -> AsyncRollout @GPUMemoryLogger(role="sglang rollout", logger=logger) @torch.no_grad() - def generate_sequences_with_tools(self, prompts: DataProto, **kwargs) -> DataProto: + def generate_sequences_with_tools( + self, + prompts: DataProto, + **kwargs) -> DataProto: logger.warning( "`generate_sequences_with_tools` is deprecated, please use `generate_sequences(...)`", DeprecationWarning, @@ -1048,7 +1188,10 @@ def generate_sequences_with_tools(self, prompts: DataProto, **kwargs) -> DataPro @GPUMemoryLogger(role="sglang rollout", logger=logger) @torch.no_grad() - def _req_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: + def _req_level_generate_sequences( + self, + prompts: DataProto, + **kwargs) -> DataProto: """Generates multi-turn sequences for a batch of prompts. For multi-turn generation, each prompt is processed separately via `_req_level_generate_sequences` for better tool calling control. @@ -1066,10 +1209,17 @@ def _req_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataPro loop = asyncio.get_event_loop() output_req_list = loop.run_until_complete( asyncio.gather( - *[self._async_rollout_a_request(req, do_sample, is_validate, **kwargs) for req in req_list], + *[ + self._async_rollout_a_request( + req, do_sample, is_validate, **kwargs + ) + for req in req_list + ], ) ) - sorted_output_req_list = sorted(output_req_list, key=lambda x: (x.batch_data_id, x.rollout_offset)) + sorted_output_req_list = sorted( + output_req_list, key=lambda x: ( + x.batch_data_id, x.rollout_offset)) else: sorted_output_req_list = None @@ -1091,14 +1241,16 @@ def _req_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataPro multi_modal_inputs = [] for req in sorted_output_req_list: - assert req.state == AsyncRolloutRequestStateEnum.COMPLETED, f"Request {req.request_id} is not completed" + assert ( + req.state == AsyncRolloutRequestStateEnum.COMPLETED + ), f"Request {req.request_id} is not completed" assert ( req.input_ids.shape[-1] == req.attention_mask.shape[-1] == req.position_ids.shape[-1] == req.loss_mask.shape[-1] - ), f"""Request {req.request_id} has different length of - {req.input_ids.shape[-1]=}, {req.attention_mask.shape[-1]=}, + ), f"""Request {req.request_id} has different length of + {req.input_ids.shape[-1]=}, {req.attention_mask.shape[-1]=}, {req.position_ids.shape[-1]=}, {req.loss_mask.shape[-1]=}""" error_message_lines = [ f"""Request {req.request_id} has input_ids length {req.input_ids.shape[-1]} @@ -1116,15 +1268,24 @@ def _req_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataPro response_ids.append(req.response_ids.to(tgt_device).squeeze(0)) if req.response_ids.shape[-1] > self.config.response_length: logger.warning( - f"""{req.request_id=} has response_ids length {req.response_ids.shape[-1]} - greater than max_response_len {self.config.response_length},\n{req=}""" - ) - prompt_attention_mask.append(req.prompt_attention_mask.to(tgt_device).squeeze(0)) - response_attention_mask.append(req.response_attention_mask.to(tgt_device).squeeze(0)) - prompt_position_ids.append(req.prompt_position_ids.to(tgt_device).squeeze(0)) - response_position_ids.append(req.response_position_ids.to(tgt_device).squeeze(0)) - prompt_loss_mask.append(req.prompt_loss_mask.to(tgt_device).squeeze(0)) - response_loss_mask.append(req.response_loss_mask.to(tgt_device).squeeze(0)) + f"""{req.request_id=} has response_ids length {req.response_ids.shape[-1]} + greater than max_response_len {self.config.response_length},\n{req=}""") + prompt_attention_mask.append( + req.prompt_attention_mask.to(tgt_device).squeeze(0) + ) + response_attention_mask.append( + req.response_attention_mask.to(tgt_device).squeeze(0) + ) + prompt_position_ids.append( + req.prompt_position_ids.to(tgt_device).squeeze(0) + ) + response_position_ids.append( + req.response_position_ids.to(tgt_device).squeeze(0) + ) + prompt_loss_mask.append( + req.prompt_loss_mask.to(tgt_device).squeeze(0)) + response_loss_mask.append( + req.response_loss_mask.to(tgt_device).squeeze(0)) messages.append({"messages": req.messages}) reward_scores.append(req.reward_scores) multi_modal_inputs.append(req.multi_modal_inputs) @@ -1136,10 +1297,18 @@ def _req_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataPro padding_side="left", ) if prompt_ids.shape[-1] < self.config.prompt_length: - prompt_ids = pad_sequence_to_length(prompt_ids, self.config.prompt_length, self.pad_token_id, left_pad=True) - response_ids = pad_sequence(response_ids, batch_first=True, padding_value=self.pad_token_id) + prompt_ids = pad_sequence_to_length( + prompt_ids, + self.config.prompt_length, + self.pad_token_id, + left_pad=True) + response_ids = pad_sequence( + response_ids, batch_first=True, padding_value=self.pad_token_id + ) if response_ids.shape[-1] < self.config.response_length: - response_ids = pad_sequence_to_length(response_ids, self.config.response_length, self.pad_token_id) + response_ids = pad_sequence_to_length( + response_ids, self.config.response_length, self.pad_token_id + ) prompt_attention_mask = pad_sequence( prompt_attention_mask, batch_first=True, @@ -1148,54 +1317,86 @@ def _req_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataPro ) if prompt_attention_mask.shape[-1] < self.config.prompt_length: prompt_attention_mask = pad_sequence_to_length( - prompt_attention_mask, self.config.prompt_length, 0, left_pad=True - ) - response_attention_mask = pad_sequence(response_attention_mask, batch_first=True, padding_value=0) + prompt_attention_mask, self.config.prompt_length, 0, left_pad=True) + response_attention_mask = pad_sequence( + response_attention_mask, batch_first=True, padding_value=0 + ) if response_attention_mask.shape[-1] < self.config.response_length: - response_attention_mask = pad_sequence_to_length(response_attention_mask, self.config.response_length, 0) + response_attention_mask = pad_sequence_to_length( + response_attention_mask, self.config.response_length, 0 + ) # padding prompt_position_ids if prompt_position_ids[0].dim() == 2: # if prompt_position_ids is a 2D tensor # e.g. from qwen2vl, prompt_position_ids.shape = (3, seq_len) - transposed_prompt_position_ids = [p.transpose(0, 1) for p in prompt_position_ids] + transposed_prompt_position_ids = [ + p.transpose(0, 1) for p in prompt_position_ids + ] prompt_position_ids = pad_sequence( - transposed_prompt_position_ids, batch_first=True, padding_value=0, padding_side="left" + transposed_prompt_position_ids, + batch_first=True, + padding_value=0, + padding_side="left", ) prompt_position_ids = prompt_position_ids.transpose(1, 2) else: prompt_position_ids = pad_sequence( - prompt_position_ids, batch_first=True, padding_value=0, padding_side="left" + prompt_position_ids, + batch_first=True, + padding_value=0, + padding_side="left", ) if prompt_position_ids.shape[-1] < self.config.prompt_length: prompt_position_ids = pad_sequence_to_length( - prompt_position_ids, self.config.prompt_length, 0, left_pad=True - ) + prompt_position_ids, self.config.prompt_length, 0, left_pad=True) # padding response_position_ids if response_position_ids[0].dim() == 2: # if response_position_ids is a 2D tensor # e.g. from qwen2vl, response_position_ids.shape = (3, seq_len) - transposed_response_position_ids = [p.transpose(0, 1) for p in response_position_ids] + transposed_response_position_ids = [ + p.transpose(0, 1) for p in response_position_ids + ] response_position_ids = pad_sequence( - transposed_response_position_ids, batch_first=True, padding_value=0, padding_side="left" + transposed_response_position_ids, + batch_first=True, + padding_value=0, + padding_side="left", ) response_position_ids = response_position_ids.transpose(1, 2) else: - response_position_ids = pad_sequence(response_position_ids, batch_first=True, padding_value=0) + response_position_ids = pad_sequence( + response_position_ids, batch_first=True, padding_value=0 + ) if response_position_ids.shape[-1] < self.config.response_length: - response_position_ids = pad_sequence_to_length(response_position_ids, self.config.response_length, 0) + response_position_ids = pad_sequence_to_length( + response_position_ids, self.config.response_length, 0 + ) - prompt_loss_mask = pad_sequence(prompt_loss_mask, batch_first=True, padding_value=0, padding_side="left") + prompt_loss_mask = pad_sequence( + prompt_loss_mask, + batch_first=True, + padding_value=0, + padding_side="left") if prompt_loss_mask.shape[1] < self.config.prompt_length: - prompt_loss_mask = pad_sequence_to_length(prompt_loss_mask, self.config.prompt_length, 0, left_pad=True) - response_loss_mask = pad_sequence(response_loss_mask, batch_first=True, padding_value=0) + prompt_loss_mask = pad_sequence_to_length( + prompt_loss_mask, self.config.prompt_length, 0, left_pad=True + ) + response_loss_mask = pad_sequence( + response_loss_mask, batch_first=True, padding_value=0 + ) if response_loss_mask.shape[1] < self.config.response_length: - response_loss_mask = pad_sequence_to_length(response_loss_mask, self.config.response_length, 0) + response_loss_mask = pad_sequence_to_length( + response_loss_mask, self.config.response_length, 0 + ) input_ids = torch.cat((prompt_ids, response_ids), dim=-1) - attention_mask = torch.cat((prompt_attention_mask, response_attention_mask), dim=-1) - position_ids = torch.cat((prompt_position_ids, response_position_ids), dim=-1) + attention_mask = torch.cat( + (prompt_attention_mask, response_attention_mask), dim=-1 + ) + position_ids = torch.cat( + (prompt_position_ids, response_position_ids), dim=-1) # Construct the batch data batch = TensorDict( @@ -1220,14 +1421,18 @@ def _req_level_generate_sequences(self, prompts: DataProto, **kwargs) -> DataPro non_tensor_batch={ "messages": np.array(messages), "reward_scores": np.array(reward_scores), - "multi_modal_inputs": np.array(multi_modal_inputs, dtype=object), + "multi_modal_inputs": np.array( + multi_modal_inputs, + dtype=object), }, ) - def _preprocess_prompt_to_async_rollout_requests(self, prompts: DataProto, n: int = 1) -> list[AsyncRolloutRequest]: - assert "raw_prompt" in prompts.non_tensor_batch, ( - "need data.return_raw_chat=True, due to no official way do parse_messages" - ) + def _preprocess_prompt_to_async_rollout_requests( + self, prompts: DataProto, n: int = 1 + ) -> list[AsyncRolloutRequest]: + assert ( + "raw_prompt" in prompts.non_tensor_batch + ), "need data.return_raw_chat=True, due to no official way do parse_messages" logger.info( "n is deprecated for SGLang rollout since ray ppo trainer will repeat the prompts for rollout.n times" ) @@ -1237,21 +1442,34 @@ def _preprocess_prompt_to_async_rollout_requests(self, prompts: DataProto, n: in ) for data_idx, (raw_prompt, multi_modal_data) in enumerate( - zip(prompts.non_tensor_batch["raw_prompt"], multi_modal_data_list, strict=True) + zip( + prompts.non_tensor_batch["raw_prompt"], + multi_modal_data_list, + strict=True, + ) ): if self._tool_schemas: _tools_kwargs = prompts.non_tensor_batch["tools_kwargs"][data_idx] - _tool_schemas = [self._tool_map[k].get_openai_tool_schema() for k in _tools_kwargs.keys()] + _tool_schemas = [ + self._tool_map[k].get_openai_tool_schema() + for k in _tools_kwargs.keys() + ] _input_ids = None _attention_mask = None else: - _input_ids = _pre_process_inputs(self.pad_token_id, prompts.batch["input_ids"][data_idx]) - _attention_mask = _pre_process_inputs(0, prompts.batch["attention_mask"][data_idx]) + _input_ids = _pre_process_inputs( + self.pad_token_id, prompts.batch["input_ids"][data_idx] + ) + _attention_mask = _pre_process_inputs( + 0, prompts.batch["attention_mask"][data_idx] + ) _tools_kwargs = {} _tool_schemas = None if self.interaction_map: - _interaction_kwargs = prompts.non_tensor_batch["interaction_kwargs"][data_idx] + _interaction_kwargs = prompts.non_tensor_batch["interaction_kwargs"][ + data_idx + ] else: _interaction_kwargs = {} @@ -1274,15 +1492,18 @@ def _preprocess_prompt_to_async_rollout_requests(self, prompts: DataProto, n: in reward_scores={}, max_prompt_len=self.config.prompt_length, max_response_len=self.config.response_length, - max_model_len=min(self.config.max_model_len, self.config.prompt_length + self.config.response_length), + max_model_len=min( + self.config.max_model_len, + self.config.prompt_length + self.config.response_length, + ), use_inference_chat_template=self.config.multi_turn.use_inference_chat_template, tokenization_sanity_check_mode=self.config.multi_turn.tokenization_sanity_check_mode, processing_class=self.processing_class, ) - error_message = f"""Request {req.request_id} has mismatched lengths: - input_ids={req.input_ids.shape[-1]}, - attention_mask={req.attention_mask.shape[-1]}, - position_ids={req.position_ids.shape[-1]}, + error_message = f"""Request {req.request_id} has mismatched lengths: + input_ids={req.input_ids.shape[-1]}, + attention_mask={req.attention_mask.shape[-1]}, + position_ids={req.position_ids.shape[-1]}, loss_mask={req.loss_mask.shape[-1]}""" assert ( req.input_ids.shape[-1] @@ -1323,7 +1544,10 @@ async def chat_completion(self, json_request): reward_scores={}, max_prompt_len=self.config.prompt_length, max_response_len=self.config.response_length, - max_model_len=min(self.config.max_model_len, self.config.prompt_length + self.config.response_length), + max_model_len=min( + self.config.max_model_len, + self.config.prompt_length + self.config.response_length, + ), use_inference_chat_template=self.config.multi_turn.use_inference_chat_template, tokenization_sanity_check_mode=self.config.multi_turn.tokenization_sanity_check_mode, processing_class=self.processing_class, @@ -1332,9 +1556,13 @@ async def chat_completion(self, json_request): # json_request already contains sampling_params # Filter only valid SamplingParams arguments valid_sampling_params = {} - temp_sampling_params = SamplingParams() # Create temporary instance to check valid attributes + temp_sampling_params = ( + SamplingParams() + ) # Create temporary instance to check valid attributes for k, v in json_request.items(): - if k not in ["messages", "model", "tools"] and hasattr(temp_sampling_params, k): + if k not in ["messages", "model", "tools"] and hasattr( + temp_sampling_params, k + ): valid_sampling_params[k] = v output = await self._handle_engine_call(req, valid_sampling_params) # it can be Dict or AsyncIterator[Dict] @@ -1355,8 +1583,7 @@ async def chat_completion(self, json_request): "content": content["text"], }, "finish_reason": content["meta_info"]["finish_reason"]["type"], - } - ) + }) id = content["meta_info"]["id"] return { @@ -1369,9 +1596,11 @@ async def chat_completion(self, json_request): # this function is left for uniform train-inference resharding - async def generate( - self, prompt_ids: torch.Tensor, sampling_params: dict[str, Any], request_id: str - ) -> torch.Tensor: + async def generate(self, + prompt_ids: torch.Tensor, + sampling_params: dict[str, + Any], + request_id: str) -> torch.Tensor: request_sampling_params = self.sampling_params.copy() request_sampling_params.update(sampling_params) output = await self._handle_engine_generate(prompt_ids, request_sampling_params) diff --git a/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/utils.py b/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/utils.py index 776bd13..e4e1bd5 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/utils.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/sglang_rollout/utils.py @@ -1,5 +1,5 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -46,7 +46,9 @@ def broadcast_pyobj( serialized_data = pickle.dumps(data) size = len(serialized_data) - tensor_data = torch.ByteTensor(np.frombuffer(serialized_data, dtype=np.uint8)).to(device) + tensor_data = torch.ByteTensor( + np.frombuffer(serialized_data, dtype=np.uint8) + ).to(device) tensor_size = torch.tensor([size], dtype=torch.long, device=device) dist.broadcast(tensor_size, src=src, group=dist_group) diff --git a/Agent0/executor_train/verl/verl/workers/rollout/tokenizer.py b/Agent0/executor_train/verl/verl/workers/rollout/tokenizer.py index 1e1212e..a854d46 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/tokenizer.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/tokenizer.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -116,7 +116,9 @@ def decode( pass @abstractmethod - def convert_ids_to_tokens(self, ids: int | list[int], skip_special_tokens: bool = False) -> str | list[str]: + def convert_ids_to_tokens( + self, ids: int | list[int], skip_special_tokens: bool = False + ) -> str | list[str]: """ Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and added tokens. diff --git a/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/__init__.py b/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/__init__.py index dac55e0..d233c5b 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -29,8 +29,7 @@ def get_version(pkg): if vllm_package_version is None: raise PackageNotFoundError( "To use vllm rollout, please ensure the 'vllm' package is properly installed. See " - "https://verl.readthedocs.io/en/latest/start/install.html for more details" - ) + "https://verl.readthedocs.io/en/latest/start/install.html for more details") if "ROCM_PATH" in os.environ: import re @@ -39,6 +38,8 @@ def get_version(pkg): if match: vllm_package_version = match.group(1) else: - raise ValueError(f"Warning: Could not parse version format: {vllm_package_version}") + raise ValueError( + f"Warning: Could not parse version format: {vllm_package_version}" + ) vllm_mode = "spmd" diff --git a/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/vllm_async_server.py b/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/vllm_async_server.py index 988dac4..12dd1b7 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/vllm_async_server.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/vllm_async_server.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,7 +24,11 @@ from vllm import SamplingParams from vllm.engine.arg_utils import AsyncEngineArgs from vllm.entrypoints.logger import RequestLogger -from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ChatCompletionResponse, ErrorResponse +from vllm.entrypoints.openai.protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ErrorResponse, +) from vllm.entrypoints.openai.serving_chat import OpenAIServingChat from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels from vllm.inputs import TokensPrompt @@ -40,21 +44,30 @@ def _get_model_runner_workers(vllm_config, init_ray: bool = True): - assert vllm_config.instance_id is not None, "instance_id must be set for external ray actors." + assert ( + vllm_config.instance_id is not None + ), "instance_id must be set for external ray actors." fields = vllm_config.instance_id.split(":") assert len(fields) == 4, ( f"instance_id: {vllm_config.instance_id} must be in the format of " f":::." ) - namespace, wg_prefix, vllm_dp_size, vllm_dp_rank = fields[0], fields[1], int(fields[2]), int(fields[3]) + namespace, wg_prefix, vllm_dp_size, vllm_dp_rank = ( + fields[0], + fields[1], + int(fields[2]), + int(fields[3]), + ) # Make sure subprocess in same namespace as parent actor. # actor name format: {name_prefix}WorkerDict_{pg_idx}:{local_rank} if init_ray: ray.init(namespace=namespace) actor_names = [ - actor_name for actor_name in ray.util.list_named_actors() if actor_name.startswith(f"{wg_prefix}WorkerDict") + actor_name + for actor_name in ray.util.list_named_actors() + if actor_name.startswith(f"{wg_prefix}WorkerDict") ] vllm_tp_size = vllm_config.parallel_config.tensor_parallel_size @@ -71,9 +84,15 @@ def get_pg_index_and_local_rank(actor_name) -> tuple[int, int]: # sort actor names by pg_index and local_rank actor_names = sorted(actor_names, key=get_pg_index_and_local_rank) - actor_names = actor_names[vllm_dp_rank * vllm_tp_size : (vllm_dp_rank + 1) * vllm_tp_size] - workers: list[WorkerWrapperBase] = [ray.get_actor(actor_name) for actor_name in actor_names] - print(f"instance_id: {vllm_config.instance_id} initializes with external actors: {actor_names}") + actor_names = actor_names[ + vllm_dp_rank * vllm_tp_size: (vllm_dp_rank + 1) * vllm_tp_size + ] + workers: list[WorkerWrapperBase] = [ + ray.get_actor(actor_name) for actor_name in actor_names + ] + print( + f"instance_id: { + vllm_config.instance_id} initializes with external actors: {actor_names}") return workers @@ -84,7 +103,9 @@ class ExternalRayDistributedExecutor(Executor): uses_ray: bool = False def _init_executor(self) -> None: - self.workers = _get_model_runner_workers(vllm_config=self.vllm_config, init_ray=True) + self.workers = _get_model_runner_workers( + vllm_config=self.vllm_config, init_ray=True + ) kwargs = dict( vllm_config=self.vllm_config, @@ -96,7 +117,9 @@ def _init_executor(self) -> None: self.collective_rpc("init_worker", args=([kwargs],)) self.collective_rpc("init_device") self.collective_rpc("load_model") - print(f"instance_id: {self.vllm_config.instance_id} initializes finished.") + print( + f"instance_id: { + self.vllm_config.instance_id} initializes finished.") def collective_rpc( self, @@ -113,9 +136,8 @@ def collective_rpc( del method # ~3ms overhead per schedule step due to SchedulerOutput/ModelRunnerOutput serialization/deserialization. - outputs = ray.get( - [worker.execute_method.remote(sent_method, *args, **(kwargs or {})) for worker in self.workers] - ) + outputs = ray.get([worker.execute_method.remote( + sent_method, *args, **(kwargs or {})) for worker in self.workers]) return outputs def check_health(self): @@ -190,7 +212,12 @@ class AsyncvLLMServer(AsyncServerBase): For vLLM AsyncLLM design, see: https://github.com/vllm-project/vllm/pull/9826 """ - def __init__(self, config: DictConfig, vllm_dp_size: int, vllm_dp_rank: int, wg_prefix: str): + def __init__( + self, + config: DictConfig, + vllm_dp_size: int, + vllm_dp_rank: int, + wg_prefix: str): """ Args: config: DictConfig. @@ -217,7 +244,11 @@ async def init_engine(self): tensor_parallel_size = config.get("tensor_model_parallel_size", 1) max_num_batched_tokens = config.get("max_num_batched_tokens", 8192) - max_model_len = config.max_model_len if config.max_model_len else config.prompt_length + config.response_length + max_model_len = ( + config.max_model_len + if config.max_model_len + else config.prompt_length + config.response_length + ) self.max_model_len = int(max_model_len) # Override default generation config from hugging face model config, @@ -268,8 +299,12 @@ async def init_engine(self): # build serving chat model_config = self.engine.model_config - BASE_MODEL_PATHS = [BaseModelPath(name=model_name, model_path=model_path)] - models = OpenAIServingModels(self.engine, model_config, BASE_MODEL_PATHS) + BASE_MODEL_PATHS = [ + BaseModelPath( + name=model_name, + model_path=model_path)] + models = OpenAIServingModels( + self.engine, model_config, BASE_MODEL_PATHS) self.openai_serving_chat = OpenAIServingChat( self.engine, model_config, @@ -285,12 +320,22 @@ async def init_engine(self): def _create_engine_config(self, engine_args: AsyncEngineArgs): vllm_config = engine_args.create_engine_config() namespace = ray.get_runtime_context().namespace - vllm_config.instance_id = f"{namespace}:{self.wg_prefix}:{self.vllm_dp_size}:{self.vllm_dp_rank}" + vllm_config.instance_id = ( + f"{namespace}:{ + self.wg_prefix}:{ + self.vllm_dp_size}:{ + self.vllm_dp_rank}") # VERL_VLLM_ZMQ_ADDRESSES - if engine_args.distributed_executor_backend == ExternalZeroMQDistributedExecutor: - workers = _get_model_runner_workers(vllm_config=vllm_config, init_ray=False) - zmq_addresses = ray.get([worker.get_zeromq_address.remote() for worker in workers]) + if ( + engine_args.distributed_executor_backend + == ExternalZeroMQDistributedExecutor + ): + workers = _get_model_runner_workers( + vllm_config=vllm_config, init_ray=False) + zmq_addresses = ray.get( + [worker.get_zeromq_address.remote() for worker in workers] + ) print(f"VERL_VLLM_ZMQ_ADDRESSES: {zmq_addresses}") os.environ["VERL_VLLM_ZMQ_ADDRESSES"] = ",".join(zmq_addresses) @@ -303,21 +348,35 @@ async def chat_completion(self, raw_request: Request): """ request_json = await raw_request.json() request = ChatCompletionRequest(**request_json) - generator = await self.openai_serving_chat.create_chat_completion(request, raw_request) + generator = await self.openai_serving_chat.create_chat_completion( + request, raw_request + ) if isinstance(generator, ErrorResponse): - return JSONResponse(content=generator.model_dump(), status_code=generator.code) + return JSONResponse( + content=generator.model_dump(), status_code=generator.code + ) if request.stream: - return StreamingResponse(content=generator, media_type="text/event-stream") + return StreamingResponse( + content=generator, + media_type="text/event-stream") else: assert isinstance(generator, ChatCompletionResponse) return JSONResponse(content=generator.model_dump()) - async def generate(self, prompt_ids: list[int], sampling_params: dict[str, Any], request_id: str) -> list[int]: + async def generate(self, + prompt_ids: list[int], + sampling_params: dict[str, + Any], + request_id: str) -> list[int]: max_tokens = self.max_model_len - len(prompt_ids) - sampling_params = SamplingParams(max_tokens=max_tokens, **sampling_params) + sampling_params = SamplingParams( + max_tokens=max_tokens, **sampling_params) prompt = TokensPrompt(prompt_token_ids=prompt_ids) - generator = self.engine.generate(prompt=prompt, sampling_params=sampling_params, request_id=request_id) + generator = self.engine.generate( + prompt=prompt, + sampling_params=sampling_params, + request_id=request_id) # Get final response final_res: Optional[RequestOutput] = None diff --git a/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py b/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py index af637c1..cff8543 100644 --- a/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py +++ b/Agent0/executor_train/verl/verl/workers/rollout/vllm_rollout/vllm_rollout_spmd.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -64,18 +64,29 @@ # 3. simplify init logics -# NOTE(sgm): add for verl. We can optimize it by making the dataloader yield List[int] without padding. -def _pre_process_inputs(pad_token_id, prompt_token_ids: torch.Tensor) -> list[int]: +# NOTE(sgm): add for verl. We can optimize it by making the dataloader +# yield List[int] without padding. +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 - non_pad_index = torch.nonzero(prompt_token_ids != pad_token_id, as_tuple=False)[0][0] + 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 class vLLMRollout(BaseRollout): - def __init__(self, model_path: str, config: DictConfig, tokenizer, model_hf_config, **kwargs): + def __init__( + self, + model_path: str, + config: DictConfig, + tokenizer, + model_hf_config, + **kwargs): """A vLLM rollout. It requires the module is supported by the vllm. Args: @@ -89,10 +100,11 @@ def __init__(self, model_path: str, config: DictConfig, tokenizer, model_hf_conf self.config = config 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" - ) - max_num_batched_tokens = self.config.get("max_num_batched_tokens", 8192) + assert ( + tensor_parallel_size <= torch.distributed.get_world_size() + ), "tensor parallel size should be less than or equal to the world size" + max_num_batched_tokens = self.config.get( + "max_num_batched_tokens", 8192) if kwargs.get("train_tp") is not None: # deployed with megatron @@ -100,7 +112,9 @@ def __init__(self, model_path: str, config: DictConfig, tokenizer, model_hf_conf os.environ["CUDA_TIMER_STREAM_KAFKA_ENABLE"] = "0" os.environ["MEGATRON_IMPORT_TIMERS"] = "0" - vllm_ps.initialize_model_parallel(tensor_model_parallel_size=tensor_parallel_size) + vllm_ps.initialize_model_parallel( + tensor_model_parallel_size=tensor_parallel_size + ) rope_scaling_config = getattr(model_hf_config, "rope_scaling", None) if not rope_scaling_config: @@ -110,16 +124,21 @@ def __init__(self, model_path: str, config: DictConfig, tokenizer, model_hf_conf elif hasattr(model_hf_config, "llm_config") and hasattr( model_hf_config.llm_config, "max_position_embeddings" ): - max_position_embeddings = model_hf_config.llm_config.max_position_embeddings + max_position_embeddings = ( + model_hf_config.llm_config.max_position_embeddings + ) elif hasattr(model_hf_config, "text_config") and hasattr( model_hf_config.text_config, "max_position_embeddings" ): - max_position_embeddings = model_hf_config.text_config.max_position_embeddings + max_position_embeddings = ( + model_hf_config.text_config.max_position_embeddings + ) if max_position_embeddings is None: - raise ValueError("max_position_embeddings not found in model_hf_config") - assert max_position_embeddings >= config.prompt_length + config.response_length, ( - "model context length should be greater than total sequence length" - ) + raise ValueError( + "max_position_embeddings not found in model_hf_config") + assert ( + max_position_embeddings >= config.prompt_length + config.response_length + ), "model context length should be greater than total sequence length" else: # handle type where there's a length extend factor # see https://qwen.readthedocs.io/en/latest/deployment/vllm.html#extended-context-support @@ -135,16 +154,22 @@ def __init__(self, model_path: str, config: DictConfig, tokenizer, model_hf_conf + f"max_position_embeddings={model_hf_config.max_position_embeddings}" ) - max_model_len = int(config.max_model_len or config.prompt_length + config.response_length) + max_model_len = int( + config.max_model_len or config.prompt_length + + config.response_length) - if max_num_batched_tokens < max_model_len and self.config.enable_chunked_prefill: + if ( + max_num_batched_tokens < max_model_len + and self.config.enable_chunked_prefill + ): raise ValueError( "Enable chunked prefill, max_num_batched_tokens is smaller than max_model_len, \ please increase max_num_batched_tokens or disable chunked prefill" ) trust_remote_code = kwargs.get("trust_remote_code", False) - load_format = "dummy" if config.load_format.startswith("dummy") else config.load_format + load_format = ("dummy" if config.load_format.startswith( + "dummy") else config.load_format) lora_kwargs = kwargs.pop("lora_kwargs", {}) self.lora_kwargs = lora_kwargs @@ -158,9 +183,12 @@ def __init__(self, model_path: str, config: DictConfig, tokenizer, model_hf_conf # - `None` means not setting it, so we pop it, and leave it to vLLM default value # (which can vary across different vLLM versions); # - Otherwise it's the desired value we want to explicitly set. - engine_kwargs = {key: val for key, val in engine_kwargs.items() if val is not None} + engine_kwargs = { + key: val for key, val in engine_kwargs.items() if val is not None + } if config.get("limit_images", None): # support for multi-image data - engine_kwargs["limit_mm_per_prompt"] = {"image": config.get("limit_images")} + engine_kwargs["limit_mm_per_prompt"] = { + "image": config.get("limit_images")} self.inference_engine = LLM( model=model_path, @@ -258,7 +286,11 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: non_tensor_batch = prompts.non_tensor_batch if "raw_prompt_ids" not in non_tensor_batch: non_tensor_batch["raw_prompt_ids"] = np.array( - [_pre_process_inputs(self.pad_token_id, idx[i]) for i in range(batch_size)], dtype=object + [ + _pre_process_inputs(self.pad_token_id, idx[i]) + for i in range(batch_size) + ], + dtype=object, ) if batch_size != len(non_tensor_batch["raw_prompt_ids"]): @@ -267,12 +299,20 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: if "multi_modal_data" in non_tensor_batch: vllm_inputs = [] for raw_prompt_ids, multi_modal_data in zip( - non_tensor_batch.pop("raw_prompt_ids"), non_tensor_batch.pop("multi_modal_data"), strict=True + non_tensor_batch.pop("raw_prompt_ids"), + non_tensor_batch.pop("multi_modal_data"), + strict=True, ): - vllm_inputs.append({"prompt_token_ids": raw_prompt_ids, "multi_modal_data": multi_modal_data}) + vllm_inputs.append( + { + "prompt_token_ids": raw_prompt_ids, + "multi_modal_data": multi_modal_data, + } + ) else: vllm_inputs = [ - {"prompt_token_ids": raw_prompt_ids} for raw_prompt_ids in non_tensor_batch.pop("raw_prompt_ids") + {"prompt_token_ids": raw_prompt_ids} + for raw_prompt_ids in non_tensor_batch.pop("raw_prompt_ids") ] # ensure the type of `prompt_token_ids` passed to vllm is list[int] @@ -282,8 +322,9 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: input_data["prompt_token_ids"] = input_data["prompt_token_ids"].tolist() elif not isinstance(input_data["prompt_token_ids"], list): raise TypeError( - f"prompt_token_ids must be a list or numpy array, got {type(input_data['prompt_token_ids'])}" - ) + f"prompt_token_ids must be a list or numpy array, got { + type( + input_data['prompt_token_ids'])}") do_sample = prompts.meta_info.get("do_sample", True) is_validate = prompts.meta_info.get("validate", False) @@ -311,7 +352,11 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: if len(lora_int_ids) > 0: lora_int_id = lora_int_ids[0] lora_requests = [ - LoRARequest(lora_name=f"{lora_int_id}", lora_int_id=lora_int_id, lora_path="/simon-stub-path") + LoRARequest( + lora_name=f"{lora_int_id}", + lora_int_id=lora_int_id, + lora_path="/simon-stub-path", + ) ] * batch_size # users can customize different sampling_params at different run @@ -324,7 +369,8 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: ) # TODO(sgm): disable logprob when recompute_log_prob is enable - # if n = 1: (bs, response_length) ; if n > 1: (bs * n, response_length) + # if n = 1: (bs, response_length) ; if n > 1: (bs * n, + # response_length) response = [] rollout_log_probs = [] @@ -334,13 +380,17 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: response.append(response_ids) if self.config.calculate_log_probs: curr_log_prob = [] - for i, logprob in enumerate(output.outputs[sample_id].logprobs): - curr_log_prob.append(logprob[response_ids[i]].logprob) + for i, logprob in enumerate( + output.outputs[sample_id].logprobs): + curr_log_prob.append( + logprob[response_ids[i]].logprob) rollout_log_probs.append(curr_log_prob) - response = pad_2d_list_to_length(response, self.pad_token_id, max_length=self.config.response_length).to( - idx.device - ) + response = pad_2d_list_to_length( + response, + self.pad_token_id, + max_length=self.config.response_length).to( + idx.device) if self.config.calculate_log_probs: rollout_log_probs = pad_2d_list_to_length( rollout_log_probs, -1, max_length=self.config.response_length @@ -350,10 +400,14 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: 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).expand(batch_size, -1) + delta_position_id = torch.arange( + 1, response_length + 1, device=position_ids.device + ) + delta_position_id = delta_position_id.unsqueeze( + 0).expand(batch_size, -1) if position_ids.dim() == 3: # qwen2vl mrope - delta_position_id = delta_position_id.view(batch_size, 1, -1).expand(batch_size, 3, -1) + delta_position_id = delta_position_id.view( + batch_size, 1, -1).expand(batch_size, 3, -1) # TODO(sgm): fix position_ids on right_pad # prompt: left pad + response: right pad @@ -362,11 +416,14 @@ def generate_sequences(self, prompts: DataProto, **kwargs) -> DataProto: response_position_ids = position_ids[..., -1:] + delta_position_id position_ids = torch.cat([position_ids, response_position_ids], dim=-1) response_attention_mask = get_response_mask( - response_id=response, eos_token=eos_token_id, dtype=attention_mask.dtype - ) - attention_mask = torch.cat((attention_mask, response_attention_mask), dim=-1) - - # all the tp ranks should contain the same data here. data in all ranks are valid + 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, @@ -405,7 +462,13 @@ class vLLMAsyncRollout: which is engine in single worker process. """ - def __init__(self, model_path: str, config: DictConfig, tokenizer, model_hf_config, **kwargs): + def __init__( + self, + model_path: str, + config: DictConfig, + tokenizer, + model_hf_config, + **kwargs): self.tokenizer = tokenizer # Engine is deferred to be initialized in init_worker @@ -472,7 +535,9 @@ def load_model(self, *args, **kwargs): self.sharding_manager.inference_engine = self.inference_engine self.sharding_manager.model_runner = self.inference_engine.worker.model_runner - _monkey_patch_compute_logits(self.inference_engine.worker.model_runner.model, len(self.tokenizer)) + _monkey_patch_compute_logits( + self.inference_engine.worker.model_runner.model, len( + self.tokenizer)) def sleep(self, *args, **kwargs): """Offload model weights and discard kv cache.""" @@ -498,4 +563,5 @@ def execute_method(self, method: str | bytes, *args, **kwargs): elif method == "wake_up": return self.wake_up(*args, **kwargs) else: - return self.inference_engine.execute_method(method, *args, **kwargs) + return self.inference_engine.execute_method( + method, *args, **kwargs) diff --git a/Agent0/executor_train/verl/verl/workers/sharding_manager/__init__.py b/Agent0/executor_train/verl/verl/workers/sharding_manager/__init__.py index 1ce90c5..e40dc4f 100644 --- a/Agent0/executor_train/verl/verl/workers/sharding_manager/__init__.py +++ b/Agent0/executor_train/verl/verl/workers/sharding_manager/__init__.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/sharding_manager/base.py b/Agent0/executor_train/verl/verl/workers/sharding_manager/base.py index 59537be..17a0e4f 100644 --- a/Agent0/executor_train/verl/verl/workers/sharding_manager/base.py +++ b/Agent0/executor_train/verl/verl/workers/sharding_manager/base.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. diff --git a/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_sglang.py b/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_sglang.py index be74bbd..621e26b 100644 --- a/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_sglang.py +++ b/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_sglang.py @@ -1,6 +1,6 @@ -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -24,14 +24,24 @@ from sglang.srt.model_executor.model_runner import LocalSerializedTensor from sglang.srt.utils import MultiprocessingSerializer from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.fsdp.api import FullStateDictConfig, ShardedStateDictConfig, StateDictType -from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp.api import ( + FullStateDictConfig, + ShardedStateDictConfig, + StateDictType, +) +from torch.distributed.fsdp.fully_sharded_data_parallel import ( + FullyShardedDataParallel as FSDP, +) from torch.distributed.tensor import DTensor from verl import DataProto from verl.protocol import all_gather_data_proto from verl.utils.device import get_device_id, get_torch_device -from verl.utils.fsdp_utils import fsdp_version, load_fsdp_model_to_gpu, offload_fsdp_model_to_cpu +from verl.utils.fsdp_utils import ( + fsdp_version, + load_fsdp_model_to_gpu, + offload_fsdp_model_to_cpu, +) from verl.utils.model import convert_weight_keys from verl.utils.profiler import GPUMemoryLogger, log_gpu_memory_usage, simple_timer from verl.utils.torch_functional import check_device_is_available @@ -74,7 +84,9 @@ def __init__( self.full_params = full_params if full_params and fsdp_version(self.module) == 1: FSDP.set_state_dict_type( - self.module, state_dict_type=StateDictType.FULL_STATE_DICT, state_dict_config=FullStateDictConfig() + self.module, + state_dict_type=StateDictType.FULL_STATE_DICT, + state_dict_config=FullStateDictConfig(), ) elif fsdp_version(self.module) == 1: FSDP.set_state_dict_type( @@ -91,7 +103,9 @@ def __init__( # get a random rng states if self.device_mesh is not None: gen_dp_rank = self.device_mesh["dp"].get_local_rank() - get_torch_device().manual_seed(gen_dp_rank + 1000) # make sure all tp ranks have the same random states + get_torch_device().manual_seed( + gen_dp_rank + 1000 + ) # make sure all tp ranks have the same random states self.gen_random_states = get_torch_device().get_rng_state() get_torch_device().set_rng_state(self.torch_random_states) else: @@ -110,14 +124,18 @@ def __exit__(self, exc_type, exc_value, traceback): loop.run_until_complete(self.sleep()) async def update_weights(self, params): - # Most naive implementation, can optimize a lot if it is bottleneck from sglang Engine weight update + # Most naive implementation, can optimize a lot if it is bottleneck + # from sglang Engine weight update named_tensors = [(k, v) for k, v in params.items()] load_format = None for tensor_index, (name, tensor) in enumerate(named_tensors): - serialized_tensor = MultiprocessingSerializer.serialize(_preprocess_tensor_for_update_weights(tensor)) + serialized_tensor = MultiprocessingSerializer.serialize( + _preprocess_tensor_for_update_weights(tensor) + ) if self.device_mesh["infer_tp"].get_local_rank() == 0: - gathered_serialized_tensors = [None for _ in range(self.device_mesh["infer_tp"].mesh.size()[0])] + gathered_serialized_tensors = [None for _ in range( + self.device_mesh["infer_tp"].mesh.size()[0])] else: gathered_serialized_tensors = None dist.gather_object( @@ -140,43 +158,67 @@ async def update_weights(self, params): ) async def release_memory(self): - if self.device_mesh["infer_tp"].get_local_rank() == 0 and self.rollout_config.free_cache_engine: + if ( + self.device_mesh["infer_tp"].get_local_rank() == 0 + and self.rollout_config.free_cache_engine + ): await self.inference_engine.release_memory_occupation() @GPUMemoryLogger(role="FSDPSGLangShardingManager enter", logger=logger) async def wake_up(self): get_torch_device().empty_cache() - if self.device_mesh["infer_tp"].get_local_rank() == 0 and self.rollout_config.free_cache_engine: + if ( + self.device_mesh["infer_tp"].get_local_rank() == 0 + and self.rollout_config.free_cache_engine + ): if self.multi_stage_wake_up: await self.inference_engine.resume_memory_occupation(tags=["weights"]) - log_gpu_memory_usage("Before resume SGLang weights in sharding manager", logger=logger) + log_gpu_memory_usage( + "Before resume SGLang weights in sharding manager", + logger=logger) else: await self.inference_engine.resume_memory_occupation() - log_gpu_memory_usage("Before resume SGLang weights + kv_cache in sharding manager", logger=logger) + log_gpu_memory_usage( + "Before resume SGLang weights + kv_cache in sharding manager", + logger=logger, + ) - log_gpu_memory_usage("Before state_dict() in sharding manager memory", logger=logger) + log_gpu_memory_usage( + "Before state_dict() in sharding manager memory", logger=logger + ) if self.offload_param: load_fsdp_model_to_gpu(self.module) params = self.module.state_dict() - log_gpu_memory_usage("After state_dict() in sharding manager memory", logger=logger) + log_gpu_memory_usage( + "After state_dict() in sharding manager memory", logger=logger + ) device = get_device_id() # used when fsdp2 set cpu_offload_policy params = { - k: v.to(device, non_blocking=True) if fsdp_version(self.module) == 2 else v for k, v in params.items() - } + k: v.to( + device, + non_blocking=True) if fsdp_version( + self.module) == 2 else v for k, + v in params.items()} # convert weight keys to match the model config - params = convert_weight_keys(params, getattr(self.module, "_fsdp_wrapped_module", self.module)) + params = convert_weight_keys( + params, getattr(self.module, "_fsdp_wrapped_module", self.module) + ) # Copy, not share memory await self.update_weights(params) - log_gpu_memory_usage("After sync model weights in sharding manager", logger=logger) + log_gpu_memory_usage( + "After sync model weights in sharding manager", logger=logger + ) del params if self.offload_param: offload_fsdp_model_to_cpu(self.module) get_torch_device().empty_cache() - log_gpu_memory_usage("After del state_dict and empty_cache in sharding manager", logger=logger) + log_gpu_memory_usage( + "After del state_dict and empty_cache in sharding manager", + logger=logger) if ( self.multi_stage_wake_up @@ -184,9 +226,12 @@ async def wake_up(self): and self.device_mesh["infer_tp"].get_local_rank() == 0 ): await self.inference_engine.resume_memory_occupation(tags=["kv_cache"]) - log_gpu_memory_usage("After resume SGLang kv_cache in sharding manager", logger=logger) + log_gpu_memory_usage( + "After resume SGLang kv_cache in sharding manager", + logger=logger) - # important: need to manually set the random states of each tp to be identical. + # 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 = get_torch_device().get_rng_state() get_torch_device().set_rng_state(self.gen_random_states) @@ -194,9 +239,13 @@ async def wake_up(self): @GPUMemoryLogger(role="FSDPSGLangShardingManager exit", logger=logger) async def sleep(self): if self.rollout_config.free_cache_engine: - log_gpu_memory_usage("Before SGLang offload in sharding manager", logger=logger) + log_gpu_memory_usage( + "Before SGLang offload in sharding manager", logger=logger + ) await self.release_memory() - log_gpu_memory_usage("After SGLang offload in sharding manager", logger=logger) + log_gpu_memory_usage( + "After SGLang offload in sharding manager", logger=logger + ) self.module.train() diff --git a/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_ulysses.py b/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_ulysses.py index 39ccb77..1176129 100644 --- a/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_ulysses.py +++ b/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_ulysses.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,7 +19,10 @@ from verl import DataProto from verl.protocol import all_gather_data_proto -from verl.utils.ulysses import get_ulysses_sequence_parallel_group, set_ulysses_sequence_parallel_group +from verl.utils.ulysses import ( + get_ulysses_sequence_parallel_group, + set_ulysses_sequence_parallel_group, +) from .base import BaseShardingManager @@ -39,7 +42,8 @@ def __enter__(self): # 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()) + 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): diff --git a/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_vllm.py b/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_vllm.py index 1a9677d..8e9bf79 100644 --- a/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_vllm.py +++ b/Agent0/executor_train/verl/verl/workers/sharding_manager/fsdp_vllm.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -19,8 +19,14 @@ from collections import OrderedDict from torch.distributed.device_mesh import DeviceMesh -from torch.distributed.fsdp.api import FullStateDictConfig, ShardedStateDictConfig, StateDictType -from torch.distributed.fsdp.fully_sharded_data_parallel import FullyShardedDataParallel as FSDP +from torch.distributed.fsdp.api import ( + FullStateDictConfig, + ShardedStateDictConfig, + StateDictType, +) +from torch.distributed.fsdp.fully_sharded_data_parallel import ( + FullyShardedDataParallel as FSDP, +) try: # for torch 2.5+ @@ -41,10 +47,19 @@ load_fsdp_model_to_gpu, offload_fsdp_model_to_cpu, ) -from verl.utils.model import check_exclude_modules, check_target_modules, convert_weight_keys +from verl.utils.model import ( + check_exclude_modules, + check_target_modules, + convert_weight_keys, +) from verl.utils.profiler import GPUMemoryLogger, log_gpu_memory_usage, simple_timer from verl.utils.torch_functional import check_device_is_available -from verl.utils.vllm_utils import TensorLoRARequest, VLLMHijack, is_version_ge, patch_vllm_moe_model_weight_loader +from verl.utils.vllm_utils import ( + TensorLoRARequest, + VLLMHijack, + is_version_ge, + patch_vllm_moe_model_weight_loader, +) from .base import BaseShardingManager @@ -74,7 +89,8 @@ def __init__( layered_summon: bool = True, ): self.module = module - # For AsyncLLM, inference_engine and model_runner are defer initialized in vLLMAsyncRollout.load_model + # For AsyncLLM, inference_engine and model_runner are defer initialized + # in vLLMAsyncRollout.load_model self.inference_engine = inference_engine # self.model_runner = inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner if # inference_engine else None @@ -96,7 +112,9 @@ def __init__( self.full_params = full_params if full_params and fsdp_version(self.module) == 1: FSDP.set_state_dict_type( - self.module, state_dict_type=StateDictType.FULL_STATE_DICT, state_dict_config=FullStateDictConfig() + self.module, + state_dict_type=StateDictType.FULL_STATE_DICT, + state_dict_config=FullStateDictConfig(), ) elif fsdp_version(self.module) == 1: FSDP.set_state_dict_type( @@ -113,7 +131,9 @@ def __init__( # get a random rng states if self.device_mesh is not None: gen_dp_rank = self.device_mesh["dp"].get_local_rank() - get_torch_device().manual_seed(gen_dp_rank + 1000) # make sure all tp ranks have the same random states + get_torch_device().manual_seed( + gen_dp_rank + 1000 + ) # make sure all tp ranks have the same random states self.gen_random_states = get_torch_device().get_rng_state() get_torch_device().set_rng_state(self.torch_random_states) else: @@ -133,33 +153,46 @@ def __collect_lora_params() -> OrderedDict: from peft.utils.save_and_load import get_peft_model_state_dict lora_params = OrderedDict() - peft_model = getattr(self.module, "_fsdp_wrapped_module", self.module) + peft_model = getattr( + self.module, + "_fsdp_wrapped_module", + self.module) if fsdp_version(self.module) > 0: if self.layered_summon: if not self.base_sync_done: raise ValueError( "To use layered_summon, you must make sure base-model is preloaded in vllm, e.g. let " - "rollout.load_format=safetensors" - ) + "rollout.load_format=safetensors") lora_params = layered_summon_lora_params(self.module) else: with FSDP.summon_full_params(self.module, writeback=False): if self.base_sync_done: lora_params = get_peft_model_state_dict(peft_model) lora_params = { - name: param.full_tensor().detach().cpu() - if hasattr(param, "full_tensor") - else param.detach().cpu() + name: ( + param.full_tensor().detach().cpu() + if hasattr(param, "full_tensor") + else param.detach().cpu() + ) for name, param in lora_params.items() } else: model = peft_model.base_model.model - orig_dev = "cpu" if "cpu" in str(next(model.parameters()).device) else get_device_name() + orig_dev = ( + "cpu" + if "cpu" in str(next(model.parameters()).device) + else get_device_name() + ) model = model.to("cpu") for name, param in model.state_dict().items(): - if any(x in name for x in ["_flat_param", "lora_"]): + if any( + x in name for x in [ + "_flat_param", + "lora_"]): continue - name = name.replace("_fsdp_wrapped_module.", "").replace(".base_layer", "") + name = name.replace( + "_fsdp_wrapped_module.", "" + ).replace(".base_layer", "") lora_params[name] = ( param.full_tensor().detach().cpu() if hasattr(param, "full_tensor") @@ -172,12 +205,18 @@ def __collect_lora_params() -> OrderedDict: lora_params = get_peft_model_state_dict(peft_model) else: model = peft_model.base_model.model - orig_dev = "cpu" if "cpu" in str(next(model.parameters()).device) else get_device_name() + orig_dev = ( + "cpu" + if "cpu" in str(next(model.parameters()).device) + else get_device_name() + ) model = model.to("cpu") for name, param in model.state_dict().items(): if any(x in name for x in ["_flat_param", "lora_"]): continue - name = name.replace("_fsdp_wrapped_module.", "").replace(".base_layer", "") + name = name.replace( + "_fsdp_wrapped_module.", "").replace( + ".base_layer", "") lora_params[name] = param.detach().cpu() model = model.to(orig_dev) return lora_params @@ -188,34 +227,50 @@ def __collect_lora_params() -> OrderedDict: # to speed up memory allocations. # # pytorch: https://pytorch.org/docs/stable/notes/cuda.html#memory-management - # vllm: https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/device_allocator/cumem.py#L103 + # vllm: + # https://github.com/vllm-project/vllm/blob/v0.7.3/vllm/device_allocator/cumem.py#L103 self.timing = {} with simple_timer("reshard", self.timing): get_torch_device().empty_cache() - log_gpu_memory_usage("Before state_dict() in sharding manager memory", logger=logger) + log_gpu_memory_usage( + "Before state_dict() in sharding manager memory", logger=logger + ) if self.offload_param: load_fsdp_model_to_gpu(self.module) peft_config = None - peft_model = getattr(self.module, "_fsdp_wrapped_module", self.module) + peft_model = getattr( + self.module, + "_fsdp_wrapped_module", + self.module) if hasattr(peft_model, "peft_config"): peft_config = peft_model.peft_config.get("default", None) params = __collect_lora_params() else: params = self.module.state_dict() - params = convert_weight_keys(params, getattr(self.module, "_fsdp_wrapped_module", self.module)) - log_gpu_memory_usage("After state_dict() in sharding manager memory", logger=logger) + params = convert_weight_keys( + params, + getattr( + self.module, + "_fsdp_wrapped_module", + self.module)) + log_gpu_memory_usage( + "After state_dict() in sharding manager memory", logger=logger + ) if self.rollout_config.free_cache_engine: - if "tags" in inspect.signature(self.inference_engine.wake_up).parameters: + if ("tags" in inspect.signature( + self.inference_engine.wake_up).parameters): self.inference_engine.wake_up(tags=["weights"]) else: self.inference_engine.wake_up() # update model params self.update_params(params, peft_config=peft_config) - log_gpu_memory_usage("After sync model weights in sharding manager", logger=logger) + log_gpu_memory_usage( + "After sync model weights in sharding manager", logger=logger + ) del params if self.offload_param: offload_fsdp_model_to_cpu(self.module) @@ -223,13 +278,18 @@ def __collect_lora_params() -> OrderedDict: if ( self.rollout_config.free_cache_engine - and "tags" in inspect.signature(self.inference_engine.wake_up).parameters + and "tags" + in inspect.signature(self.inference_engine.wake_up).parameters ): self.inference_engine.wake_up(tags=["kv_cache"]) - log_gpu_memory_usage("After del state_dict and empty_cache in sharding manager", logger=logger) + log_gpu_memory_usage( + "After del state_dict and empty_cache in sharding manager", + logger=logger, + ) - # important: need to manually set the random states of each tp to be identical. + # 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 = get_torch_device().get_rng_state() get_torch_device().set_rng_state(self.gen_random_states) @@ -292,7 +352,9 @@ def update_params(self, updated_params, peft_config=None): lora_tensors=updated_params, ) self.inference_engine.llm_engine.add_lora(lora_reqest) - logger.info(f"vLLM load weights, loaded_params: {len(updated_params)}") + logger.info( + f"vLLM load weights, loaded_params: { + len(updated_params)}") return else: @@ -308,35 +370,54 @@ def replace_lora_wrapper(k): Returns: str: Transformed parameter key for base layer. """ - stacked_params = ["q_proj", "k_proj", "v_proj", "o_proj", "gate_proj", "up_proj", "down_proj"] + stacked_params = [ + "q_proj", + "k_proj", + "v_proj", + "o_proj", + "gate_proj", + "up_proj", + "down_proj", + ] if k.endswith(".weight"): module_k = k[: -len(".weight")] if check_exclude_modules(peft_config, module_k): return k - elif any([module_k.endswith(s) for s in stacked_params]) or check_target_modules( - peft_config, module_k - ): + elif any( + [module_k.endswith(s) for s in stacked_params] + ) or check_target_modules(peft_config, module_k): return f"{module_k}.base_layer.weight" if k.endswith(".bias"): module_k = k[: -len(".bias")] if check_exclude_modules(peft_config, module_k): return k - elif any([module_k.endswith(s) for s in stacked_params]) or check_target_modules( - peft_config, module_k - ): + elif any( + [module_k.endswith(s) for s in stacked_params] + ) or check_target_modules(peft_config, module_k): return f"{module_k}.base_layer.bias" return k - updated_params = {replace_lora_wrapper(k): v for k, v in updated_params.items()} + updated_params = { + replace_lora_wrapper(k): v for k, + v in updated_params.items()} patch_vllm_moe_model_weight_loader(model) device = get_device_id() # used when fsdp2 set cpu_offload_policy loaded_params = model.load_weights( ( - (name, param.to(device, non_blocking=True).full_tensor() if isinstance(param, DTensor) else param) + ( + name, + ( + param.to(device, non_blocking=True).full_tensor() + if isinstance(param, DTensor) + else param + ), + ) for name, param in updated_params.items() ) ) self.base_sync_done = True - logger.info(f"vLLM load weights, loaded_params: {len(loaded_params) if loaded_params else -1}") + logger.info( + f"vLLM load weights, loaded_params: { + len(loaded_params) if loaded_params else -1}") diff --git a/Agent0/executor_train/verl/verl/workers/sharding_manager/megatron_sglang.py b/Agent0/executor_train/verl/verl/workers/sharding_manager/megatron_sglang.py index 9bcc1f0..55b252d 100644 --- a/Agent0/executor_train/verl/verl/workers/sharding_manager/megatron_sglang.py +++ b/Agent0/executor_train/verl/verl/workers/sharding_manager/megatron_sglang.py @@ -1,6 +1,6 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates -# Copyright 2023-2024 SGLang Team -# Copyright 2025 ModelBest Inc. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates +# Copyright 2023-2026 SGLang Team +# Copyright 2025-2026 ModelBest Inc. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -47,7 +47,7 @@ """ 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 +- 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 @@ -111,7 +111,9 @@ def __init__( # get a random rng states if self.device_mesh is not None: gen_dp_rank = self.device_mesh["dp"].get_local_rank() - get_torch_device().manual_seed(gen_dp_rank + 1000) # make sure all tp ranks have the same random states + get_torch_device().manual_seed( + gen_dp_rank + 1000 + ) # make sure all tp ranks have the same random states self.gen_random_states = get_torch_device().get_rng_state() get_torch_device().set_rng_state(self.torch_random_states) else: @@ -130,15 +132,21 @@ def __exit__(self, exc_type, exc_value, traceback): loop.run_until_complete(self.sleep()) async def update_weights(self, params): - if self.device_mesh["tp"].get_local_rank() == 0 and self.rollout_config.free_cache_engine: + if ( + self.device_mesh["tp"].get_local_rank() == 0 + and self.rollout_config.free_cache_engine + ): await self.inference_engine.resume_memory_occupation() named_tensors = params load_format = None for tensor_index, (name, tensor) in enumerate(named_tensors): - serialized_tensor = MultiprocessingSerializer.serialize(tensor.detach()) + serialized_tensor = MultiprocessingSerializer.serialize( + tensor.detach()) if self.device_mesh["tp"].get_local_rank() == 0: - gathered_serialized_tensors = [None for _ in range(self.device_mesh["tp"].mesh.size()[0])] + gathered_serialized_tensors = [ + None for _ in range(self.device_mesh["tp"].mesh.size()[0]) + ] else: gathered_serialized_tensors = None dist.gather_object( @@ -163,7 +171,10 @@ async def update_weights(self, params): await self.inference_engine.flush_cache() async def release_memory(self): - if self.device_mesh["tp"].get_local_rank() == 0 and self.rollout_config.free_cache_engine: + if ( + self.device_mesh["tp"].get_local_rank() == 0 + and self.rollout_config.free_cache_engine + ): await self.inference_engine.release_memory_occupation() @GPUMemoryLogger(role="MegatronSGLangShardingManager enter", logger=logger) @@ -184,7 +195,8 @@ async def wake_up(self): if self.offload_param: offload_megatron_model_to_cpu(self.actor_module) get_torch_device().empty_cache() - # important: need to manually set the random states of each tp to be identical. + # 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 = get_torch_device().get_rng_state() get_torch_device().set_rng_state(self.gen_random_states) @@ -192,9 +204,13 @@ async def wake_up(self): @GPUMemoryLogger(role="MegatronSGLangShardingManager exit", logger=logger) async def sleep(self): if self.rollout_config.free_cache_engine: - log_gpu_memory_usage("Before SGLang offload in sharding manager", logger=logger) + log_gpu_memory_usage( + "Before SGLang offload in sharding manager", logger=logger + ) await self.release_memory() - log_gpu_memory_usage("After SGLang offload in sharding manager", logger=logger) + log_gpu_memory_usage( + "After SGLang offload in sharding manager", logger=logger + ) for model in self.actor_module: model.train() @@ -219,4 +235,6 @@ def postprocess_data(self, data: DataProto) -> DataProto: # DP_COMPUTE_PROTO: all training ranks are dp, the same as fsdp if self.infer_tp_size == 1: return data - return data.chunk(chunks=self.infer_tp_size)[self.device_mesh["tp"].get_local_rank()] + return data.chunk(chunks=self.infer_tp_size)[ + self.device_mesh["tp"].get_local_rank() + ] diff --git a/Agent0/executor_train/verl/verl/workers/sharding_manager/megatron_vllm.py b/Agent0/executor_train/verl/verl/workers/sharding_manager/megatron_vllm.py index b04352c..f004994 100644 --- a/Agent0/executor_train/verl/verl/workers/sharding_manager/megatron_vllm.py +++ b/Agent0/executor_train/verl/verl/workers/sharding_manager/megatron_vllm.py @@ -1,4 +1,4 @@ -# Copyright 2024 Bytedance Ltd. and/or its affiliates +# Copyright 2024-2026 Bytedance Ltd. and/or its affiliates # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -31,7 +31,11 @@ from verl.third_party.vllm import LLM from verl.third_party.vllm import parallel_state as vllm_ps from verl.utils.device import get_torch_device -from verl.utils.megatron_utils import load_megatron_model_to_gpu, offload_megatron_model_to_cpu, per_tensor_generator +from verl.utils.megatron_utils import ( + load_megatron_model_to_gpu, + offload_megatron_model_to_cpu, + per_tensor_generator, +) from verl.utils.profiler import GPUMemoryLogger, log_gpu_memory_usage from verl.utils.profiler.performance import simple_timer from verl.utils.torch_functional import check_device_is_available @@ -46,7 +50,7 @@ """ Megatron Hybrid Engine: - During training, only the current pp stage holds the parameters -- Before inference, broadcast the parameters of the current pp rank +- 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 @@ -97,7 +101,8 @@ def __init__( self.inference_engine = inference_engine self.offload_param = offload_param - # For AsyncLLM, inference_engine and model_runner are defer initialized in vLLMAsyncRollout.load_model + # For AsyncLLM, inference_engine and model_runner are defer initialized + # in vLLMAsyncRollout.load_model self.model_runner = ( self.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner if self.inference_engine @@ -133,7 +138,9 @@ def __init__( self.torch_random_states = get_torch_device().get_rng_state() if self.device_mesh is not None: gen_dp_rank = self.device_mesh["dp"].get_local_rank() - get_torch_device().manual_seed(gen_dp_rank + 1000) # make sure all tp ranks have the same random states + get_torch_device().manual_seed( + gen_dp_rank + 1000 + ) # make sure all tp ranks have the same random states self.gen_random_states = get_torch_device().get_rng_state() get_torch_device().set_rng_state(self.torch_random_states) else: @@ -145,17 +152,21 @@ def __enter__(self): with simple_timer("reshard", self.timing): get_torch_device().empty_cache() - log_gpu_memory_usage("Before state_dict() in sharding manager memory", logger=logger) + log_gpu_memory_usage( + "Before state_dict() in sharding manager memory", logger=logger + ) if self.offload_param: load_megatron_model_to_gpu(self.actor_module) if self.rollout_config.free_cache_engine: - if "tags" in inspect.signature(self.inference_engine.wake_up).parameters: + if ("tags" in inspect.signature( + self.inference_engine.wake_up).parameters): self.inference_engine.wake_up(tags=["weights"]) else: self.inference_engine.wake_up() if self.bridge is not None: - per_tensor_param = self.bridge.export_weights(self.actor_module) + per_tensor_param = self.bridge.export_weights( + self.actor_module) else: per_tensor_param = per_tensor_generator( self.actor_module, @@ -176,11 +187,13 @@ def __enter__(self): if ( self.rollout_config.free_cache_engine - and "tags" in inspect.signature(self.inference_engine.wake_up).parameters + and "tags" + in inspect.signature(self.inference_engine.wake_up).parameters ): self.inference_engine.wake_up(tags=["kv_cache"]) - # important: need to manually set the random states of each tp to be identical. + # 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 = get_torch_device().get_rng_state() get_torch_device().set_rng_state(self.gen_random_states) diff --git a/Agent0/executor_train/verl_tool/llm_agent/__init__.py b/Agent0/executor_train/verl_tool/llm_agent/__init__.py index 2766530..740673c 100644 --- a/Agent0/executor_train/verl_tool/llm_agent/__init__.py +++ b/Agent0/executor_train/verl_tool/llm_agent/__init__.py @@ -1,2 +1,2 @@ from .config import AgentActorConfig -from .manager import AgentActorManager \ No newline at end of file +from .manager import AgentActorManager diff --git a/Agent0/executor_train/verl_tool/llm_agent/config.py b/Agent0/executor_train/verl_tool/llm_agent/config.py index 30c481d..b5ca25f 100644 --- a/Agent0/executor_train/verl_tool/llm_agent/config.py +++ b/Agent0/executor_train/verl_tool/llm_agent/config.py @@ -1,35 +1,53 @@ from dataclasses import dataclass + @dataclass class AgentActorConfig: - enable_agent: bool=True - max_turns: int=0 - min_turns: int=0 - max_start_length: int=None - max_prompt_length: int=None - max_response_length: int=None - max_model_len: int=None # Maximum model length, used for async rollout to limit the input length. - max_obs_length: int=None - max_action_length: int=None + enable_agent: bool = True + max_turns: int = 0 + min_turns: int = 0 + max_start_length: int = None + max_prompt_length: int = None + max_response_length: int = None + max_model_len: int = ( + # Maximum model length, used for async rollout to limit the input + # length. + None + ) + max_obs_length: int = None + max_action_length: int = None tool_server_url: str = None - n: int=1 - truncate_obs_side: str='left' - truncate_response_side: str='left' - rolling_with_prompt: bool=False - call_tool_first: bool=False - action_stop_tokens: list=None - additional_eos_token_ids: list=None - mask_observations: bool=True - force_finish_for_last_turn: bool=False - enable_mtrl: bool=False - mtrl_role: str="user" - mtrl_sep: str=None # "\n<|im_start|>system\n{obs}<|im_end|>\n<|im_start|>assistant\n" - assistant_role: str="assistant" - turn_end_token: str="<|im_end|>" - rollout_mode: str="async" # "sync" or "async" - mask_overlong_loss: bool=False # whether to mask the overlong trajectory to not train on it - max_concurrent_trajectories: int=256 # Maximum number of concurrent trajectories for async rollout. If None, no limit is applied. - enable_tqdm: bool=True # Whether to enable tqdm for async rollout. - over_sampling: bool=False # Whether to over-sample the trajectories in async rollout. - tool_call_time_out: int=None # Timeout for tool calls in async rollout. - tool_call_max_retries: int=5 # Maximum number of retries for tool calls in async rollout. \ No newline at end of file + n: int = 1 + truncate_obs_side: str = "left" + truncate_response_side: str = "left" + rolling_with_prompt: bool = False + call_tool_first: bool = False + action_stop_tokens: list = None + additional_eos_token_ids: list = None + mask_observations: bool = True + force_finish_for_last_turn: bool = False + enable_mtrl: bool = False + mtrl_role: str = "user" + mtrl_sep: str = ( + # "\n<|im_start|>system\n{obs}<|im_end|>\n<|im_start|>assistant\n" + None + ) + assistant_role: str = "assistant" + turn_end_token: str = "<|im_end|>" + rollout_mode: str = "async" # "sync" or "async" + mask_overlong_loss: bool = ( + False # whether to mask the overlong trajectory to not train on it + ) + max_concurrent_trajectories: int = ( + # Maximum number of concurrent trajectories for async rollout. If None, + # no limit is applied. + 256 + ) + enable_tqdm: bool = True # Whether to enable tqdm for async rollout. + over_sampling: bool = ( + False # Whether to over-sample the trajectories in async rollout. + ) + tool_call_time_out: int = None # Timeout for tool calls in async rollout. + tool_call_max_retries: int = ( + 5 # Maximum number of retries for tool calls in async rollout. + ) diff --git a/Agent0/executor_train/verl_tool/llm_agent/manager.py b/Agent0/executor_train/verl_tool/llm_agent/manager.py index 8aff31a..b102cad 100644 --- a/Agent0/executor_train/verl_tool/llm_agent/manager.py +++ b/Agent0/executor_train/verl_tool/llm_agent/manager.py @@ -24,7 +24,13 @@ from .tensor_helper import TensorHelper, TensorConfig from PIL import Image from .utils import PerformanceTimer, nested_copy -from .vision_utils import encode_image, encode_image_url, encode_video_url, decode_image_url, decode_video_url +from .vision_utils import ( + encode_image, + encode_image_url, + encode_video_url, + decode_image_url, + decode_video_url, +) logger = logging.getLogger(__file__) @@ -32,9 +38,10 @@ # other C0 control characters except common whitespace). CONTROL_CHAR_RE = re.compile( # this matches U+0000 through U+001F, excluding tab(09), LF(0A), CR(0D) - r'[\x00-\x08\x0B\x0C\x0E-\x1F]' + r"[\x00-\x08\x0B\x0C\x0E-\x1F]" ) + def sanitize_request(obj: Any) -> Any: """ Recursively walk through obj and: @@ -46,13 +53,14 @@ def sanitize_request(obj: Any) -> Any: if isinstance(obj, np.ndarray): obj = obj.tolist() if isinstance(obj, dict): - return {sanitize_request(key): sanitize_request(val) for key, val in obj.items()} + return {sanitize_request(key): sanitize_request(val) + for key, val in obj.items()} elif isinstance(obj, (list, tuple)): return type(obj)(sanitize_request(item) for item in obj) elif isinstance(obj, str): # strip NUL (\x00) and other C0 control chars - return CONTROL_CHAR_RE.sub('', obj) - elif isinstance(obj,Image.Image): + return CONTROL_CHAR_RE.sub("", obj) + elif isinstance(obj, Image.Image): return encode_image(obj) else: return obj @@ -74,43 +82,74 @@ def __init__( self.config = config # self.logger = logger self.is_validation = is_validation - self.eos_token_id = self.generation_config.eos_token_id \ - if self.generation_config is not None else self.tokenizer.eos_token_id - self.tensor_fn = TensorHelper(TensorConfig( - pad_token_id=self.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, - max_response_length=config.max_response_length, - )) + self.eos_token_id = ( + self.generation_config.eos_token_id + if self.generation_config is not None + else self.tokenizer.eos_token_id + ) + self.tensor_fn = TensorHelper( + TensorConfig( + pad_token_id=self.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, + max_response_length=config.max_response_length, + ) + ) if self.config.action_stop_tokens is not None: if os.path.exists(self.config.action_stop_tokens): - with open(self.config.action_stop_tokens, 'r') as f: - self.action_stop_tokens = [x for x in f.read().split(',') if x] - logger.info(f"Using action stop tokens: {self.action_stop_tokens}") + with open(self.config.action_stop_tokens, "r") as f: + self.action_stop_tokens = [ + x for x in f.read().split(",") if x] + logger.info( + f"Using action stop tokens: { + self.action_stop_tokens}") else: - raise ValueError(f"action_stop_tokens file not found: {self.config.action_stop_tokens}") + raise ValueError( + f"action_stop_tokens file not found: { + self.config.action_stop_tokens}") else: self.action_stop_tokens = [] self.additional_eos_token_ids = self.config.additional_eos_token_ids if isinstance(self.additional_eos_token_ids, str): - self.additional_eos_token_ids = [int(x) for x in self.additional_eos_token_ids.split(',')] - elif isinstance(self.additional_eos_token_ids, list) or isinstance(self.additional_eos_token_ids, omegaconf.listconfig.ListConfig): - self.additional_eos_token_ids = [int(x) for x in self.additional_eos_token_ids] + self.additional_eos_token_ids = [ + int(x) for x in self.additional_eos_token_ids.split(",") + ] + elif isinstance(self.additional_eos_token_ids, list) or isinstance( + self.additional_eos_token_ids, omegaconf.listconfig.ListConfig + ): + self.additional_eos_token_ids = [ + int(x) for x in self.additional_eos_token_ids + ] elif self.additional_eos_token_ids is None: self.additional_eos_token_ids = [] if self.config.mtrl_sep is None: messages = [{"role": "system", "content": "{obs}"}] - self.config.mtrl_sep = "\n" + self.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) - self.config.mtrl_sep = self.config.mtrl_sep.replace("system", self.config.mtrl_role) - self.max_action_length = self.config.max_action_length if self.config.max_action_length is not None else 0 - self.max_model_len = int(config.max_model_len or config.max_prompt_length + config.max_response_length) + self.config.mtrl_sep = "\n" + self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + self.config.mtrl_sep = self.config.mtrl_sep.replace( + "system", self.config.mtrl_role + ) + self.max_action_length = ( + self.config.max_action_length + if self.config.max_action_length is not None + else 0 + ) + self.max_model_len = int( + config.max_model_len + or config.max_prompt_length + config.max_response_length + ) self.tokenizer_lock = asyncio.Lock() # for multimodal processing if self.processor: self.mm_prefix, self.mm_postfix = self.processor.apply_chat_template( [{"role": "system", "content": [{"type": "text", "text": "|||"}]}], - tokenize=False, add_generation_prompt=False).split("|||") # this is used to create the correct multi-modal prompt + tokenize=False, + add_generation_prompt=False, + ).split( + "|||" + ) # this is used to create the correct multi-modal prompt else: self.mm_prefix = "" self.mm_postfix = "" @@ -120,40 +159,48 @@ def __init__( logger.setLevel(logging.WARNING) @classmethod - def from_rollout_config(cls, actor_rollout_wg, rollout_config, rollout_mode="async"): + def from_rollout_config( + cls, actor_rollout_wg, rollout_config, rollout_mode="async" + ): agent_config = AgentActorConfig() - for key in getattr(rollout_config, 'agent', {}).keys(): + for key in getattr(rollout_config, "agent", {}).keys(): if key in agent_config.__dict__.keys(): setattr(agent_config, key, rollout_config.agent[key]) - setattr(agent_config, 'n', rollout_config.rollout.n) - setattr(agent_config, 'max_model_len', rollout_config.rollout.max_model_len) + setattr(agent_config, "n", rollout_config.rollout.n) + setattr( + agent_config, + "max_model_len", + rollout_config.rollout.max_model_len) model_path = rollout_config.model.path agent_config.rollout_mode = rollout_mode print(f"AgentAsyncActorRolloutRefWorker: {agent_config}") agent_actor_manager = cls(model_path, actor_rollout_wg, agent_config) return agent_actor_manager - + 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'] - + return_tensors="pt", + padding="longest")["input_ids"] + def repeat_inputs_by_n(self, inputs: DataProto, n=None, force=False): """ this version verl do not repeat the input by n times, so we manually repeat the input by n times """ if inputs.meta_info.get("is_repeated_by_n", False) and not force: - # if the inputs are already repeated by n times, we do not need to repeat again + # if the inputs are already repeated by n times, we do not need to + # repeat again return inputs - # we manually repeat the input by n times if needed since every trajectory is independent + # we manually repeat the input by n times if needed since every + # trajectory is independent do_sample = inputs.meta_info.get("do_sample", True) - assert 'traj_ids' in inputs.non_tensor_batch, "traj_ids should be claimed univerally in the ray trainer" - ori_len = len(inputs.batch['input_ids']) + assert ( + "traj_ids" in inputs.non_tensor_batch + ), "traj_ids should be claimed univerally in the ray trainer" + ori_len = len(inputs.batch["input_ids"]) if not do_sample: n = 1 else: @@ -162,22 +209,29 @@ def repeat_inputs_by_n(self, inputs: DataProto, n=None, force=False): n = self.config.val_kwargs.n else: n = self.config.n - + inputs = inputs.repeat(n, interleave=True) # add "_{i}" for each trajectory to the traj_ids for i in range(ori_len): for j in range(n): - inputs.non_tensor_batch['traj_ids'][i*n+j] += f"_{j}" + inputs.non_tensor_batch["traj_ids"][i * n + j] += f"_{j}" # deepcopy to avoid reference bug for key in inputs.non_tensor_batch.keys(): - if key == 'traj_ids': + if key == "traj_ids": continue # # check if it's the same reference as the inputs.non_tensor_batch[key][i] - inputs.non_tensor_batch[key][i*n+j] = nested_copy(inputs.non_tensor_batch[key][i*n]) - inputs.meta_info['is_repeated_by_n'] = True + inputs.non_tensor_batch[key][i * n + j] = nested_copy( + inputs.non_tensor_batch[key][i * n] + ) + inputs.meta_info["is_repeated_by_n"] = True return inputs - async def _postprocess_responses(self, responses: Union[torch.Tensor, List[str]], action_step: int, rollout_messages: list) -> torch.Tensor: + async def _postprocess_responses( + self, + responses: Union[torch.Tensor, List[str]], + action_step: int, + rollout_messages: list, + ) -> torch.Tensor: """Process responses to stop at python operation or answer operation. Args: responses (Union[torch.Tensor, List[str]]): Responses from the model, either as a tensor or a list of strings. of length sum(active_mask), which <= batch_size @@ -190,13 +244,13 @@ async def _postprocess_responses(self, responses: Union[torch.Tensor, List[str]] do_actions (List[bool]): List indicating whether to perform actions based on the responses. rollings (DataProto): Updated rolling state with new responses. """ - effective_lens = self.tensor_fn.create_attention_mask(responses).sum(dim=1) + effective_lens = self.tensor_fn.create_attention_mask( + responses).sum(dim=1) do_actions = [] async with self.tokenizer_lock: if isinstance(responses, torch.Tensor): responses_str = self.tokenizer.batch_decode( - responses, - skip_special_tokens=True + responses, skip_special_tokens=True ) else: responses_str = responses @@ -206,48 +260,73 @@ async def _postprocess_responses(self, responses: Union[torch.Tensor, List[str]] rollout_messages[i].update_rollout_messages( { "role": self.config.assistant_role, - "content": responses_str[i] + "content": responses_str[i], } ) - + for i in range(len(responses_str)): # check if the response contains action stop tokens has_action = False for j in range(len(self.action_stop_tokens)): if self.action_stop_tokens[j] in responses_str[i]: - responses_str[i] = responses_str[i].split(self.action_stop_tokens[j])[0] + self.action_stop_tokens[j] + responses_str[i] = ( + responses_str[i].split( + self.action_stop_tokens[j])[0] + + self.action_stop_tokens[j]) has_action = True break - + # judge whether do action or not if action_step >= self.config.min_turns: # do action if there are action stop tokens in the response - do_action = has_action or (self.config.enable_mtrl and not self.action_stop_tokens) + do_action = has_action or ( + self.config.enable_mtrl and not self.action_stop_tokens + ) else: - # always do action, decided by the server about whether an action stops + # always do action, decided by the server about whether an + # action stops do_action = True if self.action_stop_tokens and not has_action: - # force add a action stop token for those responses that do not have action stop tokens - turn_end_token_idx = responses_str[i].rfind(self.config.turn_end_token) + # force add a action stop token for those responses + # that do not have action stop tokens + turn_end_token_idx = responses_str[i].rfind( + self.config.turn_end_token + ) if turn_end_token_idx != -1: - responses_str[i] = responses_str[i][:turn_end_token_idx] + self.action_stop_tokens[0] + responses_str[i] = ( + responses_str[i][:turn_end_token_idx] + + self.action_stop_tokens[0] + ) else: - responses_str[i] = responses_str[i] + self.action_stop_tokens[0] - - # now if do action, responses_str[i] should end with a action stop token, if not do action, we use the original response + responses_str[i] = ( + responses_str[i] + self.action_stop_tokens[0] + ) + + # now if do action, responses_str[i] should end with a action + # stop token, if not do action, we use the original response if do_action: if self.config.enable_mtrl: # add turn end token responses_str[i] += self.config.turn_end_token else: # preserve eos token - responses_str[i] = self.tokenizer.decode(responses[i][:effective_lens[i]], skip_special_tokens=False) - do_actions.append(do_action) + responses_str[i] = self.tokenizer.decode( + responses[i][: effective_lens[i]], skip_special_tokens=False + ) + do_actions.append(do_action) responses = self._batch_tokenize(responses_str).to(torch.int64) return responses, responses_str, do_actions, rollout_messages - async def _process_next_obs(self, next_obs: List[str], dones: List[bool], valid_action: List[bool], finishs: List[bool], tool_interact_info: List[dict], rollings: DataProto) -> Tuple[torch.Tensor, List[dict]]: + async def _process_next_obs( + self, + next_obs: List[str], + dones: List[bool], + valid_action: List[bool], + finishs: List[bool], + tool_interact_info: List[dict], + rollings: DataProto, + ) -> Tuple[torch.Tensor, List[dict]]: """Process next observations from environment. Args: next_obs (List[str]): List of next observations, only the text part. @@ -260,70 +339,98 @@ async def _process_next_obs(self, next_obs: List[str], dones: List[bool], valid_ next_obs_ids (torch.Tensor): Tokenized next observations. rollings (DataProto): Updated rolling state with new observations. """ - has_multi_modal_data = "multi_modal_data" in rollings.non_tensor_batch and rollings.non_tensor_batch['multi_modal_data'] is not None + has_multi_modal_data = ( + "multi_modal_data" in rollings.non_tensor_batch + and rollings.non_tensor_batch["multi_modal_data"] is not None + ) mm_data_list = None async with self.tokenizer_lock: mtrl_sep = self.config.mtrl_sep - next_obs = [obs if not done else "" for obs, done in zip(next_obs, dones)] - if self.config.truncate_obs_side == 'left': + next_obs = [ + obs if not done else "" for obs, + done in zip( + next_obs, + dones)] + if self.config.truncate_obs_side == "left": next_obs_ids = self.tokenizer( next_obs, - padding='longest', - return_tensors='pt', + padding="longest", + return_tensors="pt", add_special_tokens=False, # Prevents adding special tokens - padding_side='left', - )['input_ids'].to(torch.int64) + padding_side="left", + )["input_ids"].to(torch.int64) if next_obs_ids.shape[1] > self.config.max_obs_length: - logger.warning(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:] - elif self.config.truncate_obs_side == 'right': + logger.warning( + 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:] + elif self.config.truncate_obs_side == "right": next_obs_ids = self.tokenizer( next_obs, - padding='longest', - return_tensors='pt', + padding="longest", + return_tensors="pt", add_special_tokens=False, # Prevents adding special tokens - padding_side='right', - )['input_ids'].to(torch.int64) + padding_side="right", + )["input_ids"].to(torch.int64) if next_obs_ids.shape[1] > self.config.max_obs_length: - logger.warning(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] + logger.warning( + 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] else: - raise ValueError(f"Invalid truncate_obs_side: {self.config.truncate_obs_side}") + raise ValueError( + f"Invalid truncate_obs_side: { + self.config.truncate_obs_side}") next_obs = self.tokenizer.batch_decode( - next_obs_ids, - skip_special_tokens=True + next_obs_ids, skip_special_tokens=True ) if not has_multi_modal_data: - + if self.config.enable_mtrl: processed_next_obs = [] for i in range(len(next_obs)): if finishs[i] or dones[i]: # do action is false - assert next_obs[i] == "", f"next_obs should be empty when finishs is True, but got {next_obs[i]}" + assert ( + next_obs[i] == ""), f"next_obs should be empty when finishs is True, but got { + next_obs[i]}" processed_next_obs.append("") elif valid_action[i]: - processed_next_obs.append(mtrl_sep.format(obs=next_obs[i])) + processed_next_obs.append( + mtrl_sep.format(obs=next_obs[i])) else: - processed_next_obs.append(mtrl_sep.format(obs="Your action is not valid, please check the format and try again." + next_obs[i])) + processed_next_obs.append( + mtrl_sep.format( + obs="Your action is not valid, please check the format and try again." + + next_obs[i])) next_obs = processed_next_obs next_obs_ids = self.tokenizer( next_obs, - padding='longest', - return_tensors='pt', + padding="longest", + return_tensors="pt", add_special_tokens=False, # Prevents adding special tokens - )['input_ids'].to(torch.int64) + )["input_ids"].to(torch.int64) # update rollout messages with next_obs if "rollout_messages" in rollings.non_tensor_batch: for i in range(len(next_obs)): if next_obs[i]: - rollings.non_tensor_batch['rollout_messages'][i].update_rollout_messages( + rollings.non_tensor_batch["rollout_messages"][ + i + ].update_rollout_messages( { - "role": self.config.mtrl_role if self.config.enable_mtrl else self.config.assistant_role, - "content": next_obs[i] + "role": ( + self.config.mtrl_role + if self.config.enable_mtrl + else self.config.assistant_role + ), + "content": next_obs[i], } ) else: @@ -331,39 +438,61 @@ async def _process_next_obs(self, next_obs: List[str], dones: List[bool], valid_ raw_prompts = [] import traceback - + for k, tool_interact_info_k in enumerate(tool_interact_info): try: multi_modal_data = {} - next_obs_image = tool_interact_info_k.get('image', []) + next_obs_image = tool_interact_info_k.get("image", []) if not isinstance(next_obs_image, list): next_obs_image = [next_obs_image] - next_obs_image = [decode_image_url(img) for img in next_obs_image] + next_obs_image = [ + decode_image_url(img) for img in next_obs_image + ] multi_modal_data["image"] = next_obs_image - - next_obs_video = tool_interact_info_k.get('video', []) + + next_obs_video = tool_interact_info_k.get("video", []) if not isinstance(next_obs_video, list): next_obs_video = [next_obs_video] - next_obs_video = [decode_video_url(video) for video in next_obs_video] - multi_modal_data["video"] = [video.numpy() for video in next_obs_video] - - # add additional and