diff --git a/Dockerfile b/Dockerfile index e01ac04566..b753242e06 100644 --- a/Dockerfile +++ b/Dockerfile @@ -39,7 +39,6 @@ RUN cd /vllm && \ python3 use_existing_torch.py && \ pip3 install -r requirements-build.txt && \ MAX_JOBS=64 pip3 install -e . --no-build-isolation -RUN yes | pip3 uninstall uvloop RUN pip3 install opencv-python-headless==4.5.4.58 RUN apt-get update && apt-get install -y python3.10-venv diff --git a/evaluation/eval_and_aggregate.py b/evaluation/eval_and_aggregate.py index fe57f1faf7..be5d042735 100644 --- a/evaluation/eval_and_aggregate.py +++ b/evaluation/eval_and_aggregate.py @@ -5,6 +5,7 @@ from glob import glob import numpy as np +import wandb from rm_maj_eval import group_pred from tqdm import tqdm from transformers import AutoTokenizer @@ -29,6 +30,7 @@ def parse_args(): parser.add_argument("--overwrite", action="store_true") parser.add_argument("--evaluate_train", action="store_true") parser.add_argument("--max_gen_tokens", default=32768, type=int) + args = parser.parse_args() if args.output_path is None: args.output_path = args.model_path @@ -145,7 +147,7 @@ def process_single_data_name(args, data_name, base_dir, tokenizer): if __name__ == "__main__": args = parse_args() - + print(f"Evaluation output to {args.output_path}") assert args.num_sample_nodes * args.samples_per_node >= args.n_sampling eval_dir = ( @@ -155,6 +157,7 @@ def process_single_data_name(args, data_name, base_dir, tokenizer): ) base_dir = os.path.join(args.output_path, eval_dir) + os.makedirs(base_dir, exist_ok=True) tokenizer = AutoTokenizer.from_pretrained(args.model_path) result_path = os.path.join(base_dir, f"aggregate_parallel_{args.prompt_type}.json") @@ -228,10 +231,10 @@ def process_single_data_name(args, data_name, base_dir, tokenizer): from prettytable import PrettyTable table = PrettyTable() - filed_names = ["dataset"] + list(all_results[args.data_names[0]].keys()) - table.field_names = filed_names + field_names = ["dataset"] + list(all_results[args.data_names[0]].keys()) + table.field_names = field_names for k, v in all_results.items(): - table.add_row([k, *[round(v[x], 1) for x in filed_names[1:]]]) + table.add_row([k, *[round(v[x], 1) for x in field_names[1:]]]) print(table) except: diff --git a/evaluation/math_eval.py b/evaluation/math_eval.py index d4511ad04b..0b31a08f57 100644 --- a/evaluation/math_eval.py +++ b/evaluation/math_eval.py @@ -137,7 +137,7 @@ def generate_in_parallel(requests, model_args, sampling_params, data_parallel_si def run_inference_one_model( model_args: dict, sampling_params, requests, cuda_visisble_devices ): - os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG" + os.environ["VLLM_LOGGING_LEVEL"] = "INFO" os.environ["CUDA_VISIBLE_DEVICES"] = ",".join( [str(x) for x in cuda_visisble_devices] ) diff --git a/evaluation/requirements.txt b/evaluation/requirements.txt index 46cd1bf5e7..49a45d510f 100644 --- a/evaluation/requirements.txt +++ b/evaluation/requirements.txt @@ -3,7 +3,7 @@ vllm tqdm datasets torch -transformers +transformers==4.47.0 python_dateutil flash_attn @@ -12,4 +12,6 @@ sympy==1.12 antlr4-python3-runtime==4.11.1 # ! The version needs to be compatible with sympy. word2number Pebble -timeout-decorator \ No newline at end of file +prettytable +timeout-decorator +wandb diff --git a/evaluation/sh/install_deps_and_eval.sh b/evaluation/sh/install_deps_and_eval.sh new file mode 100644 index 0000000000..40e1b9faad --- /dev/null +++ b/evaluation/sh/install_deps_and_eval.sh @@ -0,0 +1,7 @@ +#!/bin/bash +# Users should run this script under AReaL directory + +/usr/bin/python3 -m pip install -e evaluation/latex2sympy/ +/usr/bin/python3 -m pip install -r evaluation/requirements.txt + +cd evaluation && /usr/bin/python3 eval_and_aggregate.py --model_path $1 --output_path $2 --data_names $3 --max_gen_tokens $4 --prompt_type $5 diff --git a/examples/README.md b/examples/README.md index 2f4c63fe7c..d86db6c1cf 100644 --- a/examples/README.md +++ b/examples/README.md @@ -13,14 +13,14 @@ Check if your hardware meets these minimum requirements: | Memory | 1 TB |1 TB per node|1 TB per node| 1 TB per node |1 TB per node| | Network | NVSwitch |NVSwitch + RoCE 3.2 Tbps|NVSwitch + RoCE 3.2 Tbps| NVSwitch + RoCE 3.2 Tbps |NVSwitch + RoCE 3.2 Tbps| | Storage | 1TB |Shared storage (NAS) 10TB|Shared storage (NAS) 10TB| Shared storage (NAS) 10TB |Shared storage (NAS) 10TB| -| **Total Time (Hours)** | **520** | **150** | **50** | **410** | **130** | +| **Total Time (Hours)** | **520** | **150** | **50** | **680** | **200** | Notes: - GPUs need to have 80GB memory. Other GPU models with similar specs are acceptable. - Single-node training can use local storage, but multi-node training requires shared storage. - Total Training Time = Number of Epochs × Number of Steps per Epoch × Training Time per Step - Number of Epochs defaults to 10. - - Number of Steps per Epoch depends on the dataset size. In this tutorial, the dataset approximately requires 40 steps per epoch for the 1.5B model and 20 steps for the 7B model. The batch sizes for 1.5B and 7B models are different. + - Number of steps per epoch depends on the dataset size and batch size. For example, with a dataset of 40,315 samples and a batch size of 1024, each epoch requires training for 40,315 / 1024 = 39.37 steps. Ultimately, an epoch will involve a minimum of 39 steps and a maximum of 40 steps of training. - Training Time per Step depends on the number of GPUs used. ## Software Requirements @@ -358,7 +358,7 @@ The last entry is used to explain the meaning of key fields: - `importance_weight`: The average importance sampling ratio across all tokens in the PPO loss. This value is typically close to 1. - `actor_clip_ratio`: The ratio of tokens clipped in the PPO loss to the total number of tokens. This is usually less than 0.1. - `actor_loss`: The PPO loss. **It does not show a clear upward or downward trend during training** and should not be used as a reference for model performance. - - `avg_seq_len`: The average length of all sampled answers in this step. In a full multi-stage training process, this value will first decrease and then increase. + - `avg_seq_len`: The average length of all sequences (i.e., prompts with sampled answers) in this step. In a full multi-stage training process, this value will first decrease and then increase. - `no_eos_ratio`: The ratio of sampled answers truncated due to exceeding the maximum generation length. An increase in this value indicates that the average length of answers is increasing. # Evaluation diff --git a/examples/README_zh.md b/examples/README_zh.md index 7a7c6c866e..3b792547f5 100644 --- a/examples/README_zh.md +++ b/examples/README_zh.md @@ -14,7 +14,7 @@ | 内存 | 1 TB |每节点 1 TB|每节点 1 TB |每节点 1 TB |每节点 1 TB | | 通信 | NVSwitch |NVSwitch+RoCE 带宽 3.2 Tbps|NVSwitch+RoCE 带宽 3.2 Tbps|NVSwitch+RoCE 带宽 3.2 Tbps|NVSwitch+RoCE 带宽 3.2 Tbps| | 存储 | 1TB |共享存储（NAS）10TB |共享存储（NAS）10TB |共享存储（NAS）10TB |共享存储（NAS）10TB | -|总训练时间（小时）|520|150|50|410|130| +|总训练时间（小时）|520|150|50|680|200| 关于硬件要求的说明： @@ -27,7 +27,7 @@ - 总训练时间 = Epoch 数量 * 每个 Epoch 的 Step 数量 * 单步训练时间 - Epoch 数量默认为 10 - - 每个 Epoch 的 Step 数量与数据集大小有关。利用我们的数据集， 1.5B 模型每个epoch需要训练 39 步，7B 模型每个epoch需要训练 19 步。不同大小模型使用的batch size也不同 + - 每个 Epoch 的 Step 数量与数据集大小和 Batch Size 有关。比如数据集为 40315 条，Batch Size 为 1024 时，每个 Epoch 需要训练 40315 / 1024 = 39.37 步，最终一个 Epoch 最少训练 39 步，最多训练 40 步。 - 单步训练时间与 GPU 卡数有关 ## 软件要求 @@ -380,7 +380,7 @@ python3 -m realhf.apps.quickstart ppo-math --show-args + `importance_weight`: PPO loss中重要性采样比率在所有token上的平均值，通常接近1。 + `actor_clip_ratio`: PPO loss中被clip掉的token占所有token的比率，通常小于0.1。 + `actor_loss`: PPO loss，**不会随着训练过程有明显的上升或下降趋势**，不应作为模型表现的参考。 -+ `avg_seq_len`: 这一步中采样的所有答案的平均长度。在完整的多阶段训练中，这个值会先下降再上升。 ++ `avg_seq_len`: 这一步中采样的所有序列（即提示词和答案相加）的平均长度。在完整的多阶段训练中，这个值会先下降再上升。 + `no_eos_ratio`: 这一步中采样的所有答案因为超出最大生成长度被截断的比率。这个值上升也代表了答案的平均长度在上升。 # 评估 diff --git a/examples/train_code_batch_1.5B_n16.sh b/examples/train_code_batch_1.5B_n16.sh new file mode 100644 index 0000000000..e888056ec7 --- /dev/null +++ b/examples/train_code_batch_1.5B_n16.sh @@ -0,0 +1,52 @@ +#!/bin/bash + +SCRIPT_DIR=$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd) + +EXP_NAME=ppo-zero-distill-1.5B-n16-jun1 +MODEL_NAME="DeepSeek-R1-Distill-Qwen-1.5B" +DATASET_NAME="codeparrot-apps-test.jsonl" +NODES=16 +ALLOCATION_MODE="vllm.d64p1m1+d32p2m1" + +LOG_DIR="/storage/ray/train_batch_logs/${EXP_NAME}/$(date +'%Y%m%d-%H%M%S')" +mkdir -p ${LOG_DIR} +echo "Log Dir: ${LOG_DIR}" + +MAX_WORKERS=$(expr 16 / ${NODES}) + +FIFO_NAME=$(mktemp -u) +mkfifo "$FIFO_NAME" +exec 3<>"$FIFO_NAME" +rm -f "$FIFO_NAME" + +for ((i=0; i&3 +done + + +ALL_PARAMS=( + "${EXP_NAME} ${MODEL_NAME} ${DATASET_NAME} 1024 8 ${NODES} ${ALLOCATION_MODE} 16384 128 1 0.001" + #"${EXP_NAME} ${MODEL_NAME} ${DATASET_NAME} 1024 8 ${NODES} ${ALLOCATION_MODE} 16384 128 1 0.001" + #"${EXP_NAME} ${MODEL_NAME} ${DATASET_NAME} 1024 8 ${NODES} ${ALLOCATION_MODE} 16384 128 1 0.001" +) + +echo "Task Count: ${#ALL_PARAMS[@]}" + +for ((i=0; i<${#ALL_PARAMS[@]}; i++)); do + read -u3 + + { + echo "$(date +"%Y-%m-%d %H:%M.%S") Task $i started: ${ALL_PARAMS[$i]}" + bash -c "bash ${SCRIPT_DIR}/train_code_small_on_ray.sh ${ALL_PARAMS[$i]} &> ${LOG_DIR}/${i}.log" + echo "$(date +"%Y-%m-%d %H:%M.%S") Task $i completed with exit code: $?, ${ALL_PARAMS[$i]}" + sleep 120 + echo >&3 + } & + + sleep 120 +done + +wait + +exec 3>&- +echo "All tasks completed" diff --git a/examples/train_code_small_on_ray.sh b/examples/train_code_small_on_ray.sh new file mode 100644 index 0000000000..8f40376952 --- /dev/null +++ b/examples/train_code_small_on_ray.sh @@ -0,0 +1,123 @@ +#!/bin/sh +MODEL_FAMILY=qwen2 + +EXP_NAME="$1" +MODEL_NAME="$2" +DATASET_NAME="$3" +TRAIN_BATCH_SIZE="$4" +GROUP_SIZE="$5" +NODES="$6" +ALLOCATION_MODE="$7" +MAX_NEW_TOKENS=$8 +MAX_NUM_SEQS=$9 +PPO_MBS=${10} +KL_CTL=${11} + +MAX_TOKEN_PER_MB=$(expr 2048 + ${MAX_NEW_TOKENS} + 1024) +MAX_SEQ_LEN_TO_CAPTURE=$(expr 2048 + ${MAX_NEW_TOKENS}) + +BASE_MODEL_PATH="/storage/models/${MODEL_NAME}" + +# original data +DATA_PATH="/storage/datasets/${DATASET_NAME}" +REAL_CODE_METADATA_PATH="/storage/datasets/codeparrot-apps-test.jsonl" + +# Option 1: The experiment runs locally with subprocesses. +# MODE=local +# Option 2: The experiment runs in a Ray cluster +# MODE=ray +# Option 3: The experiment runs in a SLURM + pyxis cluster +# Using the slurm mode requires a cluster spec file +# and setting CLUSTER_SPEC_PATH to the path of it. +MODE=ray + +# `experiment_name` and `trial_name` can be arbitrary. +# Logs and saved checkpoints will be indexed by them. +#EXP_NAME=ppo-zero--${MODEL_NAME}--${DATASET_NAME} +#EXP_NAME=ppo-zero-distill-1.5B-default +TRIAL_NAME="${TRAIN_BATCH_SIZE}x${GROUP_SIZE}-n${NODES}" + +# We use the "heuristic" allocation mode here to automatically determine the parallelism strategy +# for each model function call, i.e., actor generation, critic inference, actor train, etc. +# The number of GPUs is `n_nodes` * `n_gpus_per_node` (not set explictly here, defaults to 8). +# ReaL will make full use of these available GPUs to design allocations. +# This does not ensure the optimal throughput, but it is a good starting point. + +# The `heuristic` allocation mode is not ensured to run with every model configurations. +# For example, if the vocabulary size is an odd number, the model parallelism may not work. +# In these cases, you can use the `ppo_manual.sh` to specify the parallelism strategy manually. + +# The `ppo` subcommand specifies that this is a PPO experiment. +# The `save_freq_steps` is set to `null` to disable saving checkpoints. +# Enable it if you want to save checkpoints. +# The `ppo` option is used to control the generation and PPO algorithm hyperparameters. +# Note that the performance of PPO is sensitive to the the pre-trained model and hyperparameters. +# It's the user's responsibility to tune them appropriately. +unset CLUSTER_SPEC_PATH +CLUSTER_SPEC_PATH=/storage/ray/cluster_config_on_ray.json \ +REAL_CODE_METADATA_PATH=${REAL_CODE_METADATA_PATH} \ +FUNCTIONCALL_SERVICE_DOMAIN="" \ +REAL_GPU_MEMORY_KILL_THRESHOLD=1 \ +python3 -m realhf.apps.quickstart ppo-code \ + mode=$MODE \ + experiment_name=$EXP_NAME \ + trial_name=$TRIAL_NAME \ + wandb.mode=disabled \ + exp_ctrl.total_train_epochs=1 \ + exp_ctrl.save_freq_epochs=1 \ + exp_ctrl.ckpt_freq_secs=600 \ + group_size=${GROUP_SIZE} \ + group_adv_norm=False \ + use_dense_reward=False \ + reward_delta=True \ + rw_type=sparse \ + check_xml_format=False \ + actor.type._class=$MODEL_FAMILY \ + actor.path=$BASE_MODEL_PATH \ + actor.vllm.hybrid_train=False \ + actor.vllm.enforce_eager=False \ + actor.vllm.max_seq_len_to_capture=${MAX_SEQ_LEN_TO_CAPTURE} \ + actor.vllm.max_num_seqs=${MAX_NUM_SEQS} \ + actor.vllm.gpu_memory_utilization=1 \ + actor.vllm.swap_space=64 \ + critic.type._class=$MODEL_FAMILY \ + critic.type.is_critic=True \ + critic.init_critic_from_actor=True \ + critic.path=$BASE_MODEL_PATH\ + ref.type._class=$MODEL_FAMILY \ + ref.path=$BASE_MODEL_PATH \ + rew.type._class=$MODEL_FAMILY \ + rew.type.is_critic=True \ + rew.init_critic_from_actor=True \ + rew.path=$BASE_MODEL_PATH \ + dataset.path=$DATA_PATH \ + dataset.max_prompt_len=2048 \ + dataset.train_bs_n_seqs=${TRAIN_BATCH_SIZE} \ + ppo.gen.max_new_tokens=${MAX_NEW_TOKENS} \ + ppo.gen.min_new_tokens=0 \ + ppo.disable_value=True \ + ppo.gen.top_p=1 ppo.gen.top_k=1000000 \ + ppo.ppo_n_minibatches=${PPO_MBS} \ + ppo.gen.temperature=0.6 \ + ppo.kl_ctl=${KL_CTL} \ + ppo.value_eps_clip=0.2 \ + ppo.reward_output_scaling=5 \ + ppo.reward_output_bias=0.0 \ + ppo.adv_norm=True ppo.value_norm=True \ + mask_too_long=False \ + ppo.discount=1.0 \ + actor.optimizer.lr=1e-6 \ + critic.optimizer.lr=5e-6 \ + actor.optimizer.lr_scheduler_type=constant \ + actor_gen.mb_spec.max_tokens_per_mb=${MAX_TOKEN_PER_MB} \ + ref_inf.mb_spec.max_tokens_per_mb=${MAX_TOKEN_PER_MB} \ + rew_inf.mb_spec.max_tokens_per_mb=${MAX_TOKEN_PER_MB} \ + critic_inf.mb_spec.max_tokens_per_mb=${MAX_TOKEN_PER_MB} \ + actor_train.mb_spec.max_tokens_per_mb=${MAX_TOKEN_PER_MB} \ + critic_train.mb_spec.max_tokens_per_mb=${MAX_TOKEN_PER_MB} \ + cache_clear_freq=1 \ + n_nodes=${NODES} \ + allocation_mode="'${ALLOCATION_MODE}'" n_gpus_per_node=8 \ + recover_mode=auto \ + recover_retries=10 \ + torch_cache_mysophobia=True diff --git a/examples/train_small_on_ray.sh b/examples/train_small_on_ray.sh index c82af0dc9c..20747649d6 100644 --- a/examples/train_small_on_ray.sh +++ b/examples/train_small_on_ray.sh @@ -77,7 +77,7 @@ python3 -m realhf.apps.quickstart ppo-math \ actor.vllm.enforce_eager=False \ actor.vllm.max_seq_len_to_capture=${MAX_SEQ_LEN_TO_CAPTURE} \ actor.vllm.max_num_seqs=${MAX_NUM_SEQS} \ - actor.vllm.gpu_memory_utilization=1 \ + actor.vllm.gpu_memory_utilization=0.85 \ actor.vllm.swap_space=64 \ critic.type._class=$MODEL_FAMILY \ critic.type.is_critic=True \ diff --git a/functioncall/__init__.py b/functioncall/__init__.py new file mode 100644 index 0000000000..5a8a3c25f9 --- /dev/null +++ b/functioncall/__init__.py @@ -0,0 +1 @@ +# Copyright 2025 Ant Group Inc. diff --git a/functioncall/base/__init__.py b/functioncall/base/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/functioncall/base/call.py b/functioncall/base/call.py new file mode 100644 index 0000000000..77c2b76015 --- /dev/null +++ b/functioncall/base/call.py @@ -0,0 +1,145 @@ +import asyncio +import logging +import os +import random +import time +import traceback +from statistics import median +from typing import Any, Dict + +import aiohttp + +from functioncall.base import logging + +logger = logging.getLogger("Functioncall") + + +FUNCTIONCALL_SERVICE_DOMAIN = os.getenv( + "FUNCTIONCALL_SERVICE_DOMAIN", + "", +) + + +def calculate_percentile(elapsed_times, percentile): + sorted_times = sorted(elapsed_times) + index = int(len(sorted_times) * (percentile / 100)) + return sorted_times[min(index, len(sorted_times) - 1)] + + +async def async_invoke_function( + session: aiohttp.ClientSession, + function_name: str, + timeout: aiohttp.ClientTimeout, + payload: Dict[str, Any] = None, + max_retries: int = 3, + initial_retry_interval: float = 0.1, + max_retry_interval: float = 10.0, +): + if payload is None: + payload = {} + url = f"{FUNCTIONCALL_SERVICE_DOMAIN}/hapis/faas.hcs.io/v1/functions/{function_name}/invoke" + params = {"invocationType": "RequestResponse"} + + retries = 0 + while retries < max_retries: + try: + async with session.post( + url, + params=params, + json=payload, + timeout=timeout, + ) as response: + if response.status != 200: + text = await response.text() + raise Exception( + f"HTTP Error {response.status}: {text} : {response.headers}" + ) + + try: + result = await response.json() + return result, response.headers + except aiohttp.ContentTypeError as e: + raise Exception("Invalid JSON response") from e + + except asyncio.TimeoutError as e: + logger.warning( + f"Request timeout after {timeout}s, URL: {url}, Headers: {session.headers}, payload: {payload}" + ) + break + + except Exception as e: + logger.error(f"Async invocation failed on attempt {retries + 1}:{str(e)}, URL: {url}, Headers: {session.headers}") + + retries += 1 + if retries > max_retries: + return None, None + + # 指数退避 + 随机抖动 + sleep_time = min( + initial_retry_interval * (2**retries) + random.uniform(0, 0.1), + max_retry_interval, + ) + await asyncio.sleep(sleep_time) + + +async def batch_function_call_async( + payload_list, function_name, timeout, concurrency=1000 +): + connector = aiohttp.TCPConnector(limit=0) + async with aiohttp.ClientSession(connector=connector) as session: + semaphore = asyncio.Semaphore(concurrency) + + async def limited_task(payload): + if not payload: + return None + async with semaphore: + st = time.monotonic() + result = await async_invoke_function( + session, function_name, timeout, payload + ) + return result, time.monotonic() - st + + tasks = [limited_task(payload) for payload in payload_list] + results = await asyncio.gather(*tasks, return_exceptions=True) + + data_list = [] + elapsed_times = [] + max_elapsed = -1 + max_elapsed_header = None + for (data, header), elapsed in results: + if elapsed > max_elapsed: + max_elapsed = elapsed + max_elapsed_header = header + data_list.append(data) + elapsed_times.append(elapsed) + # logger.debug(f"functioncall took {elapsed:.4f} seconds, header: {header}.)") + + p50 = median(elapsed_times) + p90 = calculate_percentile(elapsed_times, 90) + p99 = calculate_percentile(elapsed_times, 99) + logger.info( + f"Longest functioncall took {max_elapsed:.4f} seconds, header: {max_elapsed_header}, p50: {p50}, p90: {p90}, p99: {p99}" + ) + + return data_list + + +def get_function_name(runtime_type): + if runtime_type == "python_code": + return "realhf_code_verify" + elif runtime_type == "python_math": + return "realhf_math_verify" + return "empty_code" + + +def batch_function_call(payload_list, runtime_type, timeout=30): + start_time = time.time() + function_name = get_function_name(runtime_type) + result = asyncio.run( + batch_function_call_async(payload_list, function_name, timeout) + ) + execution_time = time.time() - start_time + logger.debug( + f"Batch function call done, runtime type: {runtime_type}, batch size: {len(payload_list)}, cost: {execution_time * 1000:.0f} ms" + ) + return result diff --git a/functioncall/base/logging.py b/functioncall/base/logging.py new file mode 100644 index 0000000000..4051606894 --- /dev/null +++ b/functioncall/base/logging.py @@ -0,0 +1,30 @@ +import logging +from typing import Optional + +import colorlog + +LOG_FORMAT = "%(asctime)s.%(msecs)03d %(name)s %(levelname)s: %(message)s" +DATE_FORMAT = "%Y%m%d-%H:%M:%S" +LOGLEVEL = logging.DEBUG + +formatter = colorlog.ColoredFormatter( + fmt="%(log_color)s" + LOG_FORMAT, + datefmt=DATE_FORMAT, + log_colors={ + "DEBUG": "blue", + "INFO": "light_purple", + "WARNING": "yellow", + "ERROR": "red", + "CRITICAL": "bold_white,bg_red", + }, +) + +handler = logging.StreamHandler() +handler.setLevel(LOGLEVEL) +handler.setFormatter(formatter) + +logging.basicConfig(level=LOGLEVEL, handlers=[handler]) + + +def getLogger(name: Optional[str] = None): + return logging.getLogger(name) diff --git a/functioncall/code/__init__.py b/functioncall/code/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/functioncall/code/function/handler.py b/functioncall/code/function/handler.py new file mode 100644 index 0000000000..cb8a8132fa --- /dev/null +++ b/functioncall/code/function/handler.py @@ -0,0 +1,29 @@ +import json + +from testing_util import run_test + + +def handle(event, context): + problem = event.get("problem", "") + code = event.get("code", "") + debug = event.get("debug", False) + timeout = event.get("timeout", 6) + + if isinstance(problem, str): + problem = json.loads(problem) + + # print(f"problem:{problem}, code:{code}, debug: {debug}\n") + result, metadata = run_test(problem, test=code, debug=debug, timeout=timeout) + return {"result": result, "metadata": metadata} + + +# handle( +# { +# "problem": { +# "input_output": '{"inputs": ["1\\n", "2\\n"], "outputs": ["1\\n", "2\\n"]}' +# }, +# "code": "s = input()\nprint(s)\n", +# "debug": True, +# }, +# None, +# ) diff --git a/functioncall/code/function/testing_util.py b/functioncall/code/function/testing_util.py new file mode 100644 index 0000000000..2d4cff73d7 --- /dev/null +++ b/functioncall/code/function/testing_util.py @@ -0,0 +1,771 @@ +import ast +import faulthandler +import json +import platform + +# to run the solution files we're using a timing based approach +import signal +import sys +import types + +# used for debugging to time steps +from datetime import datetime +from enum import Enum + +# for capturing the stdout +from io import StringIO + +# used for testing the code that reads from input +from unittest.mock import mock_open, patch + +import numpy as np + + +# Note: adapted from https://github.com/LiveCodeBench/LiveCodeBench +class RuntimeModule: + @staticmethod + def from_string(name, _, code): + module = types.ModuleType(name) + exec(code, module.__dict__) + return module + + +def truncatefn(s, length=300): + assert isinstance(s, str) + if len(s) <= length: + return s + + return s[: length // 2] + "...(truncated) ..." + s[-length // 2 :] + + +class CODE_TYPE(Enum): + call_based = 0 + standard_input = 1 + + +# stuff for setting up signal timer +class TimeoutException(Exception): + pass + + +def timeout_handler(signum, frame): + print("alarm went off") + # return + raise TimeoutException + + +signal.signal(signal.SIGALRM, timeout_handler) +# timeout = 6 # seconds + + +# used to capture stdout as a list +# from https://stackoverflow.com/a/16571630/6416660 +# alternative use redirect_stdout() from contextlib +class Capturing(list): + def __enter__(self): + self._stdout = sys.stdout + sys.stdout = self._stringio = StringIO() + # Make closing the StringIO a no-op + self._stringio.close = lambda x: 1 + return self + + def __exit__(self, *args): + self.append(self._stringio.getvalue()) + del self._stringio # free up some memory + sys.stdout = self._stdout + + +def only_int_check(val): + return isinstance(val, int) + + +def string_int_check(val): + return isinstance(val, str) and val.isdigit() + + +def combined_int_check(val): + return only_int_check(val) or string_int_check(val) + + +def run_test(sample, test=None, debug=False, timeout=6): + """ + if test(generated_code) is not None it'll try to run the code. + otherwise it'll just return an input and output pair. + """ + # Disable functionalities that can make destructive changes to the test. + reliability_guard() + + if debug: + print(f"start = {datetime.now().time()}") + + try: + in_outs = json.loads(sample["input_output"]) + except ValueError: + in_outs = None + + which_type = CODE_TYPE.call_based + if in_outs: + if in_outs.get("fn_name") is None: + which_type = CODE_TYPE.standard_input # Standard input + method_name = None + else: + which_type = CODE_TYPE.call_based # Call-based + method_name = in_outs["fn_name"] + + if debug: + print(f"loaded input_output = {datetime.now().time()}") + + if test is None: + assert False, "should not happen: test code is none" + return in_outs, {"error": "no test code provided"} + elif test is not None: + results = [] + sol = "from string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\nfrom copy import *\nfrom math import *\nfrom random import *\nfrom statistics import *\nfrom itertools import *\nfrom functools import *\nfrom operator import *\nfrom io import *\nfrom sys import *\nfrom json import *\nfrom builtins import *\nfrom typing import *\nimport string\nimport re\nimport datetime\nimport collections\nimport heapq\nimport bisect\nimport copy\nimport math\nimport random\nimport statistics\nimport itertools\nimport functools\nimport operator\nimport io\nimport sys\nimport json\nsys.setrecursionlimit(6*10**5)\n" + if debug: + print(f"loading test code = {datetime.now().time()}") + + if which_type == CODE_TYPE.call_based: + + sol += test + if debug: + print(f"sol = {sol}") + signal.alarm(timeout) + try: + tmp_sol = RuntimeModule.from_string("tmp_sol", "", sol) + if "class Solution" not in test: + tmp = tmp_sol + else: + tmp = tmp_sol.Solution() + signal.alarm(0) + except Exception as e: + signal.alarm(0) + if debug: + print(f"type 0 compilation error = {e}") + results.append(-2) + return results, { + "error": repr(e), + "error_code": -1, + "error_message": "Compilation Error", + } + signal.alarm(0) + + elif which_type == CODE_TYPE.standard_input: + # sol + # if code has if __name__ == "__main__": then remove it + try: + astree = ast.parse(test) + 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) + ) + except: + pass + + tmp_test = test.split("\n") + + new_test = [] + for x in tmp_test: + if (not x.startswith("from ")) and (not x.startswith("import ")): + new_test.append("\t" + x + "\n") + else: + new_test.append(x + "\n") + tmp_test = new_test + + new_test = "" + started = False + for i in tmp_test: + if i.startswith("\t") and not started: + new_test += "stdin = sys.stdin\nstdout = sys.stdout\n" + new_test += "def code():\n" + new_test += i + started = True + elif started and ((i.startswith("from ")) or (i.startswith("import "))): + new_test += "\t" + i + else: + new_test += i + tmp_test = new_test + + sol += tmp_test + if debug: + print(f"sol = {sol}") + method_name = "code" + signal.alarm(timeout) + try: + tmp_sol = RuntimeModule.from_string("tmp_sol", "", sol) + tmp = tmp_sol + signal.alarm(0) + except Exception as e: + signal.alarm(0) + if debug: + print(f"type 1 compilation error = {e}") + results.append(-2) + return results, { + "error": repr(e), + "error_code": -1, + "error_message": "Compilation Error", + } + signal.alarm(0) + if debug: + print(f"get method = {datetime.now().time()}") + + try: + method = getattr(tmp, method_name) # get_attr second arg must be str + except: + signal.alarm(0) + e = sys.exc_info() + print(f"unable to get function error = {e}") + results.append(-2) + return results, { + "error": repr(e), + "error_code": -1, + "error_message": "Unable to extract code", + } + + for index, inputs in enumerate(in_outs["inputs"]): + raw_inputs = inputs + 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]) + + 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") + ] + ) + 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) + try: + if isinstance(inputs[0], dict): + inputs = [{int(k): v for k, v in inputs[0].items()}] + except: + True + try: + if isinstance(in_outs["outputs"][index], dict): + in_outs["outputs"][index] = [ + {int(k): v for k, v in in_outs["outputs"][index].items()} + ] + except: + True + 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()} + ] + except: + True + + if debug: + print( + f"time: {datetime.now().time()} testing index = {index} inputs = {inputs}, {type(inputs)}. type = {which_type}" + ) + if which_type == CODE_TYPE.call_based: # Call-based + signal.alarm(timeout) + faulthandler.enable() + try: + output = method(*inputs) + raw_true_output = output + + raw_true_output_copy = json.dumps(output) + 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] + ) + + # 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] + ) + except: + True + results.append(tmp_result) + if tmp_result != True: + return results, { + "output": raw_true_output_copy, + "expected": raw_outputs, + "inputs": raw_inputs, + "error_code": -2, + "error_message": "Wrong Answer", + } + # reset the alarm + signal.alarm(0) + except Exception as e: + signal.alarm(0) + faulthandler.disable() + if debug: + print( + f"Standard input runtime error or time limit exceeded error = {e}" + ) + results.append(-1) + if "timeoutexception" in repr(e).lower(): + return results, { + "error": repr(e), + "error_code": -3, + "error_message": "Time Limit Exceeded", + "inputs": raw_inputs, + "expected": raw_outputs, + } + else: + return results, { + "error": repr(e), + "error_code": -4, + "error_message": "Runtime Error", + "inputs": raw_inputs, + "expected": raw_outputs, + } + faulthandler.disable() + signal.alarm(0) + if debug: + print( + f"outputs = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + elif which_type == CODE_TYPE.standard_input: # Standard input + faulthandler.enable() + passed = False + + 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]) + + signal.alarm(timeout) + with Capturing() as output: + try: + call_method(method, inputs) + # reset the alarm + signal.alarm(0) + passed = True + except Exception as e: + # runtime error or took too long + signal.alarm(0) + print( + f"Call-based runtime error or time limit exceeded error = {repr(e)}{e}" + ) + results.append(-1) + if "timeoutexception" in repr(e).lower(): + return results, { + "error": repr(e), + "error_code": -3, + "error_message": "Time Limit Exceeded", + "inputs": raw_inputs, + "expected": raw_outputs, + } + else: + return results, { + "error": repr(e), + "error_code": -4, + "error_message": "Runtime Error", + "inputs": raw_inputs, + "expected": raw_outputs, + } + signal.alarm(0) + raw_true_output = output[0] + raw_true_output_copy = truncatefn(raw_true_output, 200) + output = raw_true_output.splitlines() + if not passed: + if debug: + nl = "\n" + if not isinstance(inputs, list): + print( + f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + else: + print( + f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + continue + + if passed and debug: + print( + f"==> output = {output}, test outputs = {in_outs['outputs'][index]}" + ) + + if custom_compare_(output, in_outs["outputs"][index]): + tmp_result = True + results.append(tmp_result) + continue + + # ground truth sequences are expressed as lists not tuples + if isinstance(output, tuple): + output = list(output) + + tmp_result = False + 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]) + if isinstance(output[0], str): + 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}") + pass + + if tmp_result == True: + results.append(tmp_result) + continue + + # try one more time without \n + if isinstance(in_outs["outputs"][index], list): + 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 + ] + 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]) + ) + + 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]) + except Exception as e: + if debug: + print(f"Failed check2 exception = {e}") + pass + + if tmp_result == True: + results.append(tmp_result) + continue + + # try by converting the output into a split up list too + if isinstance(output, list): + output = list(filter(len, output)) + + if debug: + nl = "\n" + if not isinstance(inputs, list): + print( + f"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]} {tmp_result=}" + ) + else: + print( + f"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]} {tmp_result=}" + ) + + if tmp_result == True: + results.append(tmp_result) + continue + + if debug: + print(f"{tmp_result=} @a") + + 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]) + except Exception as e: + if debug: + print(f"Failed check3 exception = {e}") + pass + + if debug: + print(f"{tmp_result=} @b") + + try: + all_ints = all( + combined_int_check(e1) and combined_int_check(e2) + for e1, e2 in zip(output, in_outs["outputs"][index]) + ) + 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]) + ] + ) + output_float = [float(e) for e in output] + 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) + ) + except Exception as e: + pass + + if debug: + print(f"{tmp_result=} @c") + + try: + 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]) + ) + 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]] + tmp_result = tmp_result or ( + (len(output_float) == len(gt_float)) + and np.allclose(output_float, gt_float) + ) + except Exception as e: + pass + + if tmp_result == True: + results.append(tmp_result) + continue + + if debug: + print(f"{tmp_result=} @d") + # try by converting the stuff into split up list + if isinstance(in_outs["outputs"][index], list): + 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()) + + if debug: + print(f"{tmp_result=} @e") + + try: + tmp_result = output == in_outs["outputs"][index] + except Exception as e: + if debug: + print(f"Failed check4 exception = {e}") + continue + + if tmp_result == True: + results.append(tmp_result) + continue + + if debug: + print(f"{tmp_result=} @f") + + # try by converting the output into a split up list too + if isinstance(output, list): + for tmp_index, i in enumerate(output): + output[tmp_index] = i.split() + output = list(filter(len, output)) + for tmp_index, i in enumerate(output): + output[tmp_index] = set(i) + else: + output = output.split() + output = list(filter(len, output)) + output = set(output) + + if debug: + print(f"{tmp_result=} @g") + # try: + # tmp_result = set(frozenset(s) for s in output) == set( + # frozenset(s) for s in in_outs["outputs"][index] + # ) + # except Exception as e: + # if debug: + # print(f"Failed check5 exception = {e}") + + # if they are all numbers, round so that similar numbers are treated as identical + # try: + # all_ints = all( + # combined_int_check(e1) and combined_int_check(e2) + # for e1, e2 in zip(output, in_outs["outputs"][index]) + # ) + # tmp_result = tmp_result or ( + # set(frozenset(round(float(t), 3) for t in s) for s in output) + # == set( + # frozenset(round(float(t), 3) for t in s) + # for s in in_outs["outputs"][index] + # ) + # ) + # except Exception as e: + # if debug: + # print(f"Failed check6 exception = {e}") + + if debug: + print(f"{tmp_result=} @h") + + if tmp_result == True and debug: + print("PASSED") + + results.append(tmp_result) + if tmp_result != True: + return results, { + "output": raw_true_output_copy, + "expected": raw_outputs, + "inputs": raw_inputs, + "error_code": -2, + "error_message": "Wrong Answer", + } + + if debug: + nl = "\n" + if not isinstance(inputs, list): + print( + f"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl,' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + else: + print( + f"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + + print(f"results = {results}") + + return results, {} + + +def custom_compare_(output, ground_truth): + + if isinstance(output, list): + output_1 = "\n".join(output) + if stripped_string_compare(output_1, ground_truth): + return True + + if isinstance(output, list): + output_2 = [o.lstrip().rstrip() for o in output] + output_2 = "\n".join(output_2) + if stripped_string_compare(output_2, ground_truth): + return True + + return False + + +def stripped_string_compare(s1, s2): + s1 = s1.lstrip().rstrip() + s2 = s2.lstrip().rstrip() + return s1 == s2 + + +def call_method(method, inputs): + + if isinstance(inputs, list): + inputs = "\n".join(inputs) + + inputs_line_iterator = iter(inputs.split("\n")) + + # sys.setrecursionlimit(10000) + + # @patch('builtins.input', side_effect=inputs.split("\n")) + @patch("builtins.open", mock_open(read_data=inputs)) + @patch("sys.stdin", StringIO(inputs)) + @patch("sys.stdin.readline", lambda *args: next(inputs_line_iterator)) + @patch("sys.stdin.readlines", lambda *args: inputs.split("\n")) + @patch("sys.stdin.read", lambda *args: inputs) + # @patch('sys.stdout.write', print) + def _inner_call_method(_method): + try: + return _method() + except SystemExit as e: + pass + finally: + pass + + return _inner_call_method(method) + + +def reliability_guard(maximum_memory_bytes=None): + """ + This disables various destructive functions and prevents the generated code + from interfering with the test (e.g. fork bomb, killing other processes, + removing filesystem files, etc.) + WARNING + This function is NOT a security sandbox. Untrusted code, including, model- + generated code, should not be blindly executed outside of one. See the + Codex paper for more information about OpenAI's code sandbox, and proceed + with caution. + """ + + 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) + ) + if not platform.uname().system == "Darwin": + resource.setrlimit( + resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes) + ) + + faulthandler.disable() + + import builtins + + builtins.exit = None + builtins.quit = None + + import os + + if os.putenv: + os.environ["OMP_NUM_THREADS"] = "20" + + os.kill = None + os.system = None + os.putenv = None + os.remove = None + os.removedirs = None + os.rmdir = None + os.fchdir = None + os.setuid = None + os.fork = None + os.forkpty = None + os.killpg = None + os.rename = None + os.renames = None + os.truncate = None + os.replace = None + os.unlink = None + os.fchmod = None + os.fchown = None + os.chmod = None + os.chown = None + os.chroot = None + os.fchdir = None + os.lchflags = None + os.lchmod = None + os.lchown = None + os.getcwd = None + os.chdir = None + + import shutil + + shutil.rmtree = None + shutil.move = None + shutil.chown = None + + import subprocess + + subprocess.Popen = None # type: ignore + + __builtins__["help"] = None + + import sys + + sys.modules["ipdb"] = None + sys.modules["joblib"] = None + sys.modules["resource"] = None + sys.modules["psutil"] = None + sys.modules["tkinter"] = None diff --git a/functioncall/code/local_verify.py b/functioncall/code/local_verify.py new file mode 100644 index 0000000000..50ccd91055 --- /dev/null +++ b/functioncall/code/local_verify.py @@ -0,0 +1,194 @@ +import json +import multiprocessing +import os +import sys +import time +import traceback +from io import StringIO + +from function.testing_util import run_test + +from functioncall.base import logging + +logger = logging.getLogger("Functioncall") + + +def try_load_json(data): + """ + Attempts to load the given data as JSON. If successful, returns the parsed JSON. + Otherwise, returns None and logs an error. + """ + try: + loaded_data = json.loads(data) + return loaded_data + except json.JSONDecodeError as e: + logger.info(f"Failed to load JSON: {e}") + return data + + +def capture_stdout(code): + original_stdout = sys.stdout + fake_stdout = StringIO() + + try: + sys.stdout = fake_stdout # 重定向输出 + exec(code, {"__builtins__": __builtins__}) # 在隔离环境中执行 + except Exception as e: + return f"error: {str(e)}, traceback: {traceback.format_exc()}" + finally: + sys.stdout = original_stdout # 恢复原stdout + return fake_stdout.getvalue() + + +def _temp_run(problem, generation, debug, result): + start_time = time.time() + + try: + if debug: + logger.debug(f"Running test for problem: {problem}") + result.append(run_test(sample=problem, test=generation, debug=debug)) + if debug: + logger.debug(f"Test completed with result: {result}") + except Exception as e: + if debug: + logger.error(f"Error in _temp_run: {e}, problem:{problem}", exe_info=True) + + execution_time = time.time() - start_time + logger.info( + f'[_temp_run] query_id: {problem["problem_id"]}, start_time: {str(start_time)}, Time elapsed: {execution_time * 1000:.0f} ms' + ) + + +def check_correctness(problem, generation, timeout, debug=False): + """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`""" + if debug: + result = capture_stdout( + "from function.testing_util import run_test\n" + + "run_test(sample=problem, test=generation, debug=debug)" + ) + return result[0], result[1] + + start_time = time.time() + manager = multiprocessing.Manager() + result = manager.list() + p = multiprocessing.Process( + target=_temp_run, args=(problem, generation, debug, result) + ) + p.start() + p.join(timeout=timeout + 1) + if p.is_alive(): + if debug: + logger.debug(f"Process is still alive. Killing the process.") + p.kill() + if not result: + # Remark: ideally we would consider that all tests failed but we can't access number of tests here easily + # so we use 21=the average number of tests for a smaple in the test split instead + avg_number_tests = 21 + result = [[-1] * avg_number_tests] + if debug: + logger.debug(f"Global timeout occurred, returning default result.") + if debug: + logger.debug(f"Final result: {result}") + + execution_time = time.time() - start_time + logger.info( + f'[check_correctness] query_id: {problem["problem_id"]}, start_time: {str(start_time)}, Time elapsed: {execution_time * 1000:.0f} ms' + ) + return result[0] + + +def load_problems(path): + problem_map = {} + for idx, line in enumerate(open(path, "rb")): + if line is not None: + line = line.strip().decode("utf-8") + row = json.loads(line) + + query_id = str(row["id"]) + problem_map[query_id] = { + "problem_id": query_id, + # "question": row["question"], + "solutions": row["solutions"], + "input_output": row["input_output"], + # "difficulty": level, + # "url": row["url"], + # "starter_code": row["starter_code"], + } + + return problem_map + + +global_problems = load_problems( + os.getenv( + "REAL_CODE_METADATA_PATH", + "/storage/datasets/codeparrot-apps-test.jsonl", + ) +) + + +def code_verify(generateds, query_ids, debug=False): + assert len(generateds) == len(query_ids), ( + len(generateds), + len(query_ids), + ) + + result = [] + global global_problems + + for idx, query_id in enumerate(query_ids): + if query_id not in global_problems: + continue + + problem = global_problems[query_id] + test_code = generateds[idx] + logger.debug(f"run_batch_code, query_id: {query_id}") + + try: + curr_res, metadata = check_correctness( + problem=problem, generation=test_code, timeout=6000, debug=debug + ) + + if any(x != True for x in curr_res): + logger.debug( + f'id:{problem["problem_id"]}, Results were not all True: {metadata}' + ) + result.append(f"{query_id} failed") + else: + # print(f"id:{problem["problem_id"]}, result : {curr_res}") + result.append(f"{query_id} success") + + except Exception as e: + logger.error(f"test framework exception = {repr(e)}{e}\n", exe_info=True) + result.append(f"{query_id} failed") + break + finally: + # print(f"id:{problem["problem_id"]}, result : {curr_res}") + # assert isinstance(curr_res, list) + pass + + return result + + +if __name__ == "__main__": + + def create_test_params(index_list): + global global_problems + codes, query_ids = [], [] + + for index in index_list: + if str(index) not in global_problems: + continue + + problem = global_problems[str(index)] + if "solutions" not in problem or not problem["solutions"]: + continue + + codes.append(try_load_json(problem["solutions"])[0]) + query_ids.append(str(index)) + return codes, query_ids + + codes, query_ids = create_test_params(list(range(805, 806))) + result = code_verify(codes, query_ids, True) + print(result) diff --git a/functioncall/code/verify.py b/functioncall/code/verify.py new file mode 100644 index 0000000000..0730be3022 --- /dev/null +++ b/functioncall/code/verify.py @@ -0,0 +1,154 @@ +import json +import os +from collections import defaultdict + +from functioncall.base import logging +from functioncall.base.call import batch_function_call + +logger = logging.getLogger("Functioncall") + + +def try_load_json(data): + """ + Attempts to load the given data as JSON. If successful, returns the parsed JSON. + Otherwise, returns None and logs an error. + """ + try: + loaded_data = json.loads(data) + return loaded_data + except json.JSONDecodeError as e: + # print(f"Failed to load JSON: {e}") + return data + + +def load_problems_with_testcase_batch(path, debug=False, test_case_batch_size=None): + problem_map = defaultdict(list) + for idx, line in enumerate(open(path, "rb")): + if line is None: + continue + + # parse one problem + row = json.loads(line.strip().decode("utf-8")) + query_id = str(row["id"]) + input_output = json.loads(row["input_output"]) if "input_output" in row else {} + inputs = input_output.get("inputs", []) + outputs = input_output.get("outputs", []) + assert len(inputs) == len( + outputs + ), f"Inputs({len(inputs)}) and outputs({len(outputs)}) mismatch for {query_id}" + + other_io_fields = { + k: v for k, v in input_output.items() if k not in ["inputs", "outputs"] + } + # create batches for testcases + if not test_case_batch_size or test_case_batch_size <= 0: + test_case_batch_size = len(inputs) + + for batch_idx in range(0, len(inputs), test_case_batch_size): + batch_io = { + **other_io_fields, + "inputs": inputs[batch_idx : batch_idx + test_case_batch_size], + "outputs": outputs[batch_idx : batch_idx + test_case_batch_size], + } + + sub_problem = { + "problem_id": query_id, + "input_output": json.dumps(batch_io), + "batche_index": batch_idx, + } + if debug: + sub_problem["solutions"] = row.get("solutions", []) + problem_map[query_id].append(sub_problem) + + return problem_map + + +global_problems = None + + +def code_verify(generateds, query_ids, debug=False, timeout=20, timeout_for_testcase=6): + assert len(generateds) == len(query_ids), ( + len(generateds), + len(query_ids), + ) + payload_list = [] + + global global_problems + if global_problems is None: + global_problems = load_problems_with_testcase_batch( + os.getenv( + "REAL_CODE_METADATA_PATH", + "/storage/datasets/codeparrot-apps-test.jsonl", + ), + debug=True, + test_case_batch_size=20, + ) + for idx, query_id in enumerate(query_ids): + if query_id not in global_problems: + payload_list.append(None) + logger.warning( + f"Invalid query id : {query_id}, type: {type(query_id)}, should be in problem dataset!" + ) + continue + + problems = global_problems[query_id] + for problem in problems: + payload_list.append( + { + "problem": problem, + "code": generateds[idx], + "debug": debug, + "timeout": timeout_for_testcase, + "query_index": idx, + } + ) + + logger.debug( + f"code_verify, payload_list size: {len(payload_list)}, query size: {len(query_ids)}, query_id_0: {query_ids[0]}" + ) + rsp_list = batch_function_call(payload_list, "python_code", timeout=timeout) + + results = [1] * len(query_ids) + for idx, rsp in enumerate(rsp_list): + query_index = payload_list[idx]["query_index"] + query_id = query_ids[query_index] + + value = 0 + if rsp and "result" in rsp and not any(x != True for x in rsp["result"]): + value = 1 + + else: + logger.debug( + f"Functioncall code verify failed, query index: {query_index}, query id: {query_id}, results: {rsp}" + ) + + # logger.debug(f"query index: {idx}, value: {value}, results[query_index]: {results[query_index]}") + + results[query_index] = results[query_index] and value + + return results + + +if __name__ == "__main__": + + def create_test_params(count=10): + global global_problems + codes, query_ids = [], [] + idx = 0 + for query_id, problems in global_problems.items(): + if idx >= count: + break + + problem = problems[0] + if "solutions" not in problem or not problem["solutions"]: + continue + + codes.append(try_load_json(problem["solutions"])[0]) + query_ids.append(query_id) + idx += 1 + + return codes, query_ids + + codes, query_ids = create_test_params(1000) + result = code_verify(codes, query_ids, True, 100) + print(result) diff --git a/functioncall/math/__init__.py b/functioncall/math/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/functioncall/math/function/grader.py b/functioncall/math/function/grader.py new file mode 100644 index 0000000000..19b83cf520 --- /dev/null +++ b/functioncall/math/function/grader.py @@ -0,0 +1,396 @@ +""" +This logic is largely copied from the Hendrycks' MATH release (math_equivalence), and borrowed from: +- https://github.com/microsoft/ProphetNet/tree/master/CRITIC +- https://github.com/openai/prm800k +- https://github.com/microsoft/ToRA/blob/main/src/eval/grader.py +- https://github.com/deepseek-ai/DeepSeek-Math/blob/main/evaluation/eval/eval_utils.py +""" + +import multiprocessing +import re +from collections import defaultdict +from math import isclose +from typing import Union + +import regex +from latex2sympy2 import latex2sympy +from sympy import N, simplify +from sympy.parsing.latex import parse_latex +from sympy.parsing.sympy_parser import parse_expr + +# from .parser import choice_answer_clean, strip_string +# from parser import choice_answer_clean + + +def choice_answer_clean(pred: str): + pred = pred.strip("\n").rstrip(".").rstrip("/").strip(" ").lstrip(":") + # Clean the answer based on the dataset + tmp = re.findall(r"\b(A|B|C|D|E)\b", pred.upper()) + if tmp: + pred = tmp + else: + pred = [pred.strip().strip(".")] + pred = pred[-1] + # Remove the period at the end, again! + pred = pred.rstrip(".").rstrip("/") + return pred + + +def parse_digits(num): + num = regex.sub(",", "", str(num)) + try: + return float(num) + except: + if num.endswith("%"): + num = num[:-1] + if num.endswith("\\"): + num = num[:-1] + try: + return float(num) / 100 + except: + pass + return None + + +def is_digit(num): + # paired with parse_digits + return parse_digits(num) is not None + + +def str_to_pmatrix(input_str): + input_str = input_str.strip() + matrix_str = re.findall(r"\{.*,.*\}", input_str) + pmatrix_list = [] + + for m in matrix_str: + m = m.strip("{}") + pmatrix = r"\begin{pmatrix}" + m.replace(",", "\\") + r"\end{pmatrix}" + pmatrix_list.append(pmatrix) + + return ", ".join(pmatrix_list) + + +def math_equal( + prediction: Union[bool, float, str], + reference: Union[float, str], + include_percentage: bool = True, + is_close: bool = True, + timeout: bool = False, +) -> bool: + """ + Exact match of math if and only if: + 1. numerical equal: both can convert to float and are equal + 2. symbolic equal: both can convert to sympy expression and are equal + """ + # print("Judge:", prediction, reference) + if prediction is None or reference is None: + return False + if str(prediction.strip().lower()) == str(reference.strip().lower()): + return True + if ( + reference in ["A", "B", "C", "D", "E"] + and choice_answer_clean(prediction) == reference + ): + return True + + try: # 1. numerical equal + if is_digit(prediction) and is_digit(reference): + prediction = parse_digits(prediction) + reference = parse_digits(reference) + # number questions + if include_percentage: + gt_result = [reference / 100, reference, reference * 100] + else: + gt_result = [reference] + for item in gt_result: + try: + if is_close: + if numeric_equal(prediction, item): + return True + else: + if item == prediction: + return True + except Exception: + continue + return False + except: + pass + + if not prediction and prediction not in [0, False]: + return False + + # 2. symbolic equal + reference = str(reference).strip() + prediction = str(prediction).strip() + + ## pmatrix (amps) + if "pmatrix" in prediction and not "pmatrix" in reference: + reference = str_to_pmatrix(reference) + + ## deal with [], (), {} + 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("[") + ): + pred_str = pred_str.strip("[]()") + ref_str = ref_str.strip("[]()") + for s in ["{", "}", "(", ")"]: + ref_str = ref_str.replace(s, "") + pred_str = pred_str.replace(s, "") + if pred_str.lower() == ref_str.lower(): + return True + + ## [a, b] vs. [c, d], return a==c and b==d + if ( + regex.match(r"(\(|\[).+(\)|\])", prediction) is not None + and regex.match(r"(\(|\[).+(\)|\])", reference) is not None + ): + pred_parts = prediction[1:-1].split(",") + ref_parts = reference[1:-1].split(",") + if len(pred_parts) == len(ref_parts): + if all( + [ + math_equal( + pred_parts[i], ref_parts[i], include_percentage, is_close + ) + for i in range(len(pred_parts)) + ] + ): + return True + if ( + ( + prediction.startswith("\\begin{pmatrix}") + or prediction.startswith("\\begin{bmatrix}") + ) + and ( + prediction.endswith("\\end{pmatrix}") + or prediction.endswith("\\end{bmatrix}") + ) + and ( + reference.startswith("\\begin{pmatrix}") + or reference.startswith("\\begin{bmatrix}") + ) + and ( + reference.endswith("\\end{pmatrix}") or reference.endswith("\\end{bmatrix}") + ) + ): + pred_lines = [ + line.strip() + for line in prediction[ + len("\\begin{pmatrix}") : -len("\\end{pmatrix}") + ].split("\\\\") + if line.strip() + ] + ref_lines = [ + line.strip() + for line in reference[ + len("\\begin{pmatrix}") : -len("\\end{pmatrix}") + ].split("\\\\") + if line.strip() + ] + matched = True + if len(pred_lines) == len(ref_lines): + for pred_line, ref_line in zip(pred_lines, ref_lines): + pred_parts = pred_line.split("&") + ref_parts = ref_line.split("&") + if len(pred_parts) == len(ref_parts): + if not all( + [ + math_equal( + pred_parts[i], + ref_parts[i], + include_percentage, + is_close, + ) + for i in range(len(pred_parts)) + ] + ): + matched = False + break + else: + matched = False + if not matched: + break + else: + matched = False + if matched: + return True + + if prediction.count("=") == 1 and reference.count("=") == 1: + pred = prediction.split("=") + pred = f"{pred[0].strip()} - ({pred[1].strip()})" + ref = reference.split("=") + ref = f"{ref[0].strip()} - ({ref[1].strip()})" + if symbolic_equal(pred, ref) or symbolic_equal(f"-({pred})", ref): + return True + elif ( + prediction.count("=") == 1 + and len(prediction.split("=")[0].strip()) <= 2 + and "=" not in reference + ): + if math_equal( + prediction.split("=")[1], reference, include_percentage, is_close + ): + return True + elif ( + reference.count("=") == 1 + and len(reference.split("=")[0].strip()) <= 2 + and "=" not in prediction + ): + if math_equal( + prediction, reference.split("=")[1], include_percentage, is_close + ): + return True + + # symbolic equal with sympy + if timeout: + if call_with_timeout(symbolic_equal_process, prediction, reference): + return True + else: + if symbolic_equal(prediction, reference): + return True + + return False + + +def math_equal_process(param): + return math_equal(param[-2], param[-1]) + + +def numeric_equal(prediction: float, reference: float): + # Note that relative tolerance has significant impact + # on the result of the synthesized GSM-Hard dataset + # if reference.is_integer(): + # return isclose(reference, round(prediction), abs_tol=1e-4) + # else: + # prediction = round(prediction, len(str(reference).split(".")[-1])) + return isclose(reference, prediction, rel_tol=1e-4) + + +def symbolic_equal(a, b): + def _parse(s): + for f in [parse_latex, parse_expr, latex2sympy]: + try: + return f(s.replace("\\\\", "\\")) + except: + try: + return f(s) + except: + pass + return s + + a = _parse(a) + b = _parse(b) + + # direct equal + try: + if str(a) == str(b) or a == b: + return True + except: + pass + + # simplify equal + try: + if a.equals(b) or simplify(a - b) == 0: + return True + except: + pass + + # equation equal + try: + if (abs(a.lhs - a.rhs)).equals(abs(b.lhs - b.rhs)): + return True + except: + pass + + try: + if numeric_equal(float(N(a)), float(N(b))): + return True + except: + pass + + # matrix + try: + # if a and b are matrix + if a.shape == b.shape: + _a = a.applyfunc(lambda x: round(x, 3)) + _b = b.applyfunc(lambda x: round(x, 3)) + if _a.equals(_b): + return True + except: + pass + + return False + + +def symbolic_equal_process(a, b, output_queue): + result = symbolic_equal(a, b) + output_queue.put(result) + + +def call_with_timeout(func, *args, timeout=3, **kwargs): + output_queue = multiprocessing.Queue() + process_args = args + (output_queue,) + process = multiprocessing.Process(target=func, args=process_args, kwargs=kwargs) + process.start() + process.join(timeout) + + if process.is_alive(): + process.terminate() + process.join() + return False + + return output_queue.get() + + +def _test_math_equal(): + # print(math_equal("0.0833333333333333", "\\frac{1}{12}")) + # print(math_equal("(1,4.5)", "(1,\\frac{9}{2})")) + # print(math_equal("\\frac{x}{7}+\\frac{2}{7}", "\\frac{x+2}{7}", timeout=True)) + # print(math_equal("\\sec^2(y)", "\\tan^2(y)+1", timeout=True)) + # print(math_equal("\\begin{pmatrix}-\\frac{7}{4}&-2\\\\4&\\frac{1}{4}\\end{pmatrix}", "(\\begin{pmatrix}-\\frac{7}{4}&-2\\\\4&\\frac{1}{4}\\\\\\end{pmatrix})", timeout=True)) + + # pred = '\\begin{pmatrix}\\frac{1}{3x^{2/3}}&0&0\\\\0&1&0\\\\-\\sin(x)&0&0\\end{pmatrix}' + # gt = '(\\begin{pmatrix}\\frac{1}{3\\sqrt[3]{x}^2}&0&0\\\\0&1&0\\\\-\\sin(x)&0&0\\\\\\end{pmatrix})' + + # pred= '-\\frac{8x^2}{9(x^2-2)^{5/3}}+\\frac{2}{3(x^2-2)^{2/3}}' + # gt= '-\\frac{2(x^2+6)}{9(x^2-2)\\sqrt[3]{x^2-2}^2}' + + # pred = '-34x-45y+20z-100=0' + # gt = '34x+45y-20z+100=0' + + # pred = '\\frac{100}{3}' + # gt = '33.3' + + # pred = '\\begin{pmatrix}0.290243531202435\\\\0.196008371385084\\\\-0.186381278538813\\end{pmatrix}' + # gt = '(\\begin{pmatrix}0.29\\\\0.196\\\\-0.186\\\\\\end{pmatrix})' + + # pred = '\\frac{\\sqrt{\\sqrt{11}+\\sqrt{194}}}{2\\sqrt{33}+15}' + # gt = '\\frac{\\sqrt{\\sqrt{11}+\\sqrt{194}}}{15+2\\sqrt{33}}' + + # pred = '(+5)(b+2)' + # gt = '(a+5)(b+2)' + + # pred = '\\frac{1+\\sqrt{5}}{2}' + # gt = '2' + + # pred = '\\frac{34}{16}+\\frac{\\sqrt{1358}}{16}', gt = '4' + # pred = '1', gt = '1\\\\sqrt{19}' + + # pred = "(0.6,2.6667]" + # gt = "(\\frac{3}{5},\\frac{8}{3}]" + + gt = "x+0.5+0.5" + pred = "x+1" + + print(math_equal(pred, gt, timeout=True)) + + +if __name__ == "__main__": + _test_math_equal() diff --git a/functioncall/math/function/handler.py b/functioncall/math/function/handler.py new file mode 100644 index 0000000000..dca346ed08 --- /dev/null +++ b/functioncall/math/function/handler.py @@ -0,0 +1,39 @@ +import json +from parser import extract_answer + +from grader import math_equal + + +def process_results(answer, solution): + extracted_answer = extract_answer(answer, "math", use_last_number=False) + extracted_solution = extract_answer(solution, "math", use_last_number=True) + + print( + f"extracted_answer: {extracted_answer}, extracted_solution: {extracted_solution}, equal: {math_equal(extracted_answer, extracted_solution)}" + ) + if extracted_answer is None or extracted_answer.strip() in ["None", "none", ""]: + retval = 0 + elif math_equal(extracted_answer, extracted_solution, timeout=True): + retval = 1 + else: + retval = 0 + + return retval, (extracted_answer, extracted_solution) + + +def handle(event, context): + answers = event.get("answers", "") + solutions = event.get("solutions", "") + + #print(f"math payload:{event}\n") + # answers and solutions are json lists, and call process_results then collect result into a list + if isinstance(answers, str): + answers = json.loads(answers) + if isinstance(solutions, str): + solutions = json.loads(solutions) + + results = [] + for answer, solution in zip(answers, solutions): + results.append(process_results(answer, solution)) + + return results diff --git a/functioncall/math/function/parser.py b/functioncall/math/function/parser.py new file mode 100644 index 0000000000..6798304cf5 --- /dev/null +++ b/functioncall/math/function/parser.py @@ -0,0 +1,765 @@ +import random +import re +from typing import Any, Dict, Iterable, List, TypeVar, Union + +import regex +import sympy +from latex2sympy2 import latex2sympy +from word2number import w2n + +# from utils import * + + +def _fix_fracs(string): + substrs = string.split("\\frac") + new_str = substrs[0] + if len(substrs) > 1: + substrs = substrs[1:] + for substr in substrs: + new_str += "\\frac" + if len(substr) > 0 and substr[0] == "{": + new_str += substr + else: + try: + assert len(substr) >= 2 + except: + return string + a = substr[0] + b = substr[1] + if b != "{": + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}{" + b + "}" + post_substr + else: + new_str += "{" + a + "}{" + b + "}" + else: + if len(substr) > 2: + post_substr = substr[2:] + new_str += "{" + a + "}" + b + post_substr + else: + new_str += "{" + a + "}" + b + string = new_str + return string + + +def _fix_a_slash_b(string): + if len(string.split("/")) != 2: + return string + a = string.split("/")[0] + b = string.split("/")[1] + try: + if "sqrt" not in a: + a = int(a) + if "sqrt" not in b: + b = int(b) + assert string == "{}/{}".format(a, b) + new_string = "\\frac{" + str(a) + "}{" + str(b) + "}" + return new_string + except: + return string + + +def _fix_sqrt(string): + _string = re.sub(r"\\sqrt(\w+)", r"\\sqrt{\1}", string) + return _string + + +def convert_word_number(text: str) -> str: + try: + text = str(w2n.word_to_num(text)) + except: + pass + return text + + +# units mainly from MathQA +unit_texts = [ + "east", + "degree", + "mph", + "kmph", + "ft", + "m sqaure", + " m east", + "sq m", + "deg", + "mile", + "q .", + "monkey", + "prime", + "ratio", + "profit of rs", + "rd", + "o", + "gm", + "p . m", + "lb", + "tile", + "per", + "dm", + "lt", + "gain", + "ab", + "way", + "west", + "a .", + "b .", + "c .", + "d .", + "e .", + "f .", + "g .", + "h .", + "t", + "a", + "h", + "no change", + "men", + "soldier", + "pie", + "bc", + "excess", + "st", + "inches", + "noon", + "percent", + "by", + "gal", + "kmh", + "c", + "acre", + "rise", + "a . m", + "th", + "π r 2", + "sq", + "mark", + "l", + "toy", + "coin", + "sq . m", + "gallon", + "° f", + "profit", + "minw", + "yr", + "women", + "feet", + "am", + "pm", + "hr", + "cu cm", + "square", + "v â € ™", + "are", + "rupee", + "rounds", + "cubic", + "cc", + "mtr", + "s", + "ohm", + "number", + "kmph", + "day", + "hour", + "minute", + "min", + "second", + "man", + "woman", + "sec", + "cube", + "mt", + "sq inch", + "mp", + "∏ cm ³", + "hectare", + "more", + "sec", + "unit", + "cu . m", + "cm 2", + "rs .", + "rs", + "kg", + "g", + "month", + "km", + "m", + "cm", + "mm", + "apple", + "liter", + "loss", + "yard", + "pure", + "year", + "increase", + "decrease", + "d", + "less", + "Surface", + "litre", + "pi sq m", + "s .", + "metre", + "meter", + "inch", +] + +unit_texts.extend([t + "s" for t in unit_texts]) + + +def strip_string(string, skip_unit=False): + string = str(string).strip() + # linebreaks + string = string.replace("\n", "") + + # right "." + string = string.rstrip(".") + + # remove inverse spaces + # replace \\ with \ + string = string.replace("\\!", "") + # string = string.replace("\\ ", "") + # string = string.replace("\\\\", "\\") + + # matrix + string = re.sub(r"\\begin\{array\}\{.*?\}", r"\\begin{pmatrix}", string) + string = re.sub(r"\\end\{array\}", r"\\end{pmatrix}", string) + string = string.replace("bmatrix", "pmatrix") + + # 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") + ) + + # remove \left and \right + string = string.replace("\\left", "") + string = string.replace("\\right", "") + string = string.replace("\\{", "{") + string = string.replace("\\}", "}") + + # Remove unit: miles, dollars if after is not none + _string = re.sub(r"\\text{.*?}$", "", string).strip() + if _string != "" and _string != string: + # print("Warning: unit not removed: '{}' -> '{}'".format(string, _string)) + string = _string + + if not skip_unit: + # Remove unit: texts + 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) + if _string != "": + string = _string + + # Remove circ (degrees) + string = string.replace("^{\\circ}", "") + string = string.replace("^\\circ", "") + + # remove dollar signs + string = string.replace("\\$", "") + string = string.replace("$", "") + string = string.replace("\\(", "").replace("\\)", "") + + # convert word number to digit + string = convert_word_number(string) + + # replace "\\text{...}" to "..." + string = re.sub(r"\\text\{(.*?)\}", r"\1", string) + for key in ["x=", "y=", "z=", "x\\in", "y\\in", "z\\in", "x\\to", "y\\to", "z\\to"]: + string = string.replace(key, "") + string = string.replace("\\emptyset", r"{}") + string = string.replace("(-\\infty,\\infty)", "\\mathbb{R}") + + # 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.") + string = string.replace("{.", "{0.") + + # cdot + # string = string.replace("\\cdot", "") + if ( + string.startswith("{") + and string.endswith("}") + and string.isalnum() + or string.startswith("(") + and string.endswith(")") + and string.isalnum() + or string.startswith("[") + and string.endswith("]") + and string.isalnum() + ): + string = string[1:-1] + + # inf + string = string.replace("infinity", "\\infty") + if "\\infty" not in string: + string = string.replace("inf", "\\infty") + string = string.replace("+\\inity", "\\infty") + + # and + string = string.replace("and", "") + string = string.replace("\\mathbf", "") + + # use regex to remove \mbox{...} + string = re.sub(r"\\mbox{.*?}", "", string) + + # quote + string.replace("'", "") + string.replace('"', "") + + # i, j + if "j" in string and "i" not in string: + string = string.replace("j", "i") + + # replace a.000b where b is not number or b is end, with ab, use regex + string = re.sub(r"(\d+)\.0*([^\d])", r"\1\2", string) + string = re.sub(r"(\d+)\.0*$", r"\1", string) + + # if empty, return empty string + if len(string) == 0: + return string + if string[0] == ".": + string = "0" + string + + # to consider: get rid of e.g. "k = " or "q = " at beginning + if len(string.split("=")) == 2: + if len(string.split("=")[0]) <= 2: + string = string.split("=")[1] + + string = _fix_sqrt(string) + string = string.replace(" ", "") + + # \frac1b or \frac12 --> \frac{1}{b} and \frac{1}{2}, etc. Even works with \frac1{72} (but not \frac{72}1). Also does a/b --> \\frac{a}{b} + string = _fix_fracs(string) + + # 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 + + +def extract_multi_choice_answer(pred_str): + # TODO: SFT models + if "Problem:" in pred_str: + pred_str = pred_str.split("Problem:", 1)[0] + pred_str = pred_str.replace("choice is", "answer is") + patt = regex.search(r"answer is \(?(?P[abcde])\)?", pred_str.lower()) + if patt is not None: + return patt.group("ans").upper() + return "placeholder" + + +direct_answer_trigger_for_fewshot = ("choice is", "answer is") + + +def choice_answer_clean(pred: str): + pred = pred.strip("\n") + + # Determine if this is ICL, if so, use \n\n to split the first chunk. + ICL = False + for trigger in direct_answer_trigger_for_fewshot: + if pred.count(trigger) > 1: + ICL = True + if ICL: + pred = pred.split("\n\n")[0] + + # Split the trigger to find the answer. + preds = re.split("|".join(direct_answer_trigger_for_fewshot), pred) + if len(preds) > 1: + answer_flag = True + pred = preds[-1] + else: + answer_flag = False + + pred = pred.strip("\n").rstrip(".").rstrip("/").strip(" ").lstrip(":") + + # Clean the answer based on the dataset + tmp = re.findall(r"\b(A|B|C|D|E)\b", pred.upper()) + if tmp: + pred = tmp + else: + pred = [pred.strip().strip(".")] + + if len(pred) == 0: + pred = "" + else: + if answer_flag: + # choose the first element in list ... + pred = pred[0] + else: + # choose the last e + pred = pred[-1] + + # Remove the period at the end, again! + pred = pred.rstrip(".").rstrip("/") + + return pred + + +def find_box(pred_str: str): + ans = pred_str.split("boxed")[-1] + if not ans: + return "" + if ans[0] == "{": + stack = 1 + a = "" + for c in ans[1:]: + if c == "{": + stack += 1 + a += c + elif c == "}": + stack -= 1 + if stack == 0: + break + a += c + else: + a += c + else: + a = ans.split("$")[0].strip() + return a + + +def clean_units(pred_str: str): + """Clean the units in the number.""" + + def convert_pi_to_number(code_string): + code_string = code_string.replace("\\pi", "π") + # Replace \pi or π not preceded by a digit or } with 3.14 + code_string = re.sub(r"(? "3*3.14" + code_string = re.sub(r"(\d)(\\?π)", r"\1*3.14", code_string) + # Handle cases where π is within braces or followed by a multiplication symbol + # This replaces "{π}" with "3.14" directly and "3*π" with "3*3.14" + code_string = re.sub(r"\{(\\?π)\}", "3.14", code_string) + code_string = re.sub(r"\*(\\?π)", "*3.14", code_string) + return code_string + + pred_str = convert_pi_to_number(pred_str) + pred_str = pred_str.replace("%", "/100") + pred_str = pred_str.replace("$", "") + pred_str = pred_str.replace("¥", "") + pred_str = pred_str.replace("°C", "") + pred_str = pred_str.replace(" C", "") + pred_str = pred_str.replace("°", "") + return pred_str + + +def extract_theoremqa_answer(pred: str, answer_flag: bool = True): + if any([option in pred.lower() for option in ["yes", "true"]]): + pred = "True" + elif any([option in pred.lower() for option in ["no", "false"]]): + pred = "False" + elif any( + [ + option in pred.lower() + for option in ["(a)", "(b)", "(c)", "(d)", "(e)", "(f)"] + ] + ): + pass + else: + # Some of the models somehow get used to boxed output from pre-training + if "boxed" in pred: + pred = find_box(pred) + + if answer_flag: + # Extract the numbers out of the string + pred = pred.split("=")[-1].strip() + pred = clean_units(pred) + try: + tmp = str(latex2sympy(pred)) + pred = str(eval(tmp)) + except Exception: + if re.match(r"-?[\d\.]+\s\D+$", pred): + pred = pred.split(" ")[0] + elif re.match(r"-?[\d\.]+\s[^\s]+$", pred): + pred = pred.split(" ")[0] + else: + # desparate search over the last number + preds = re.findall(r"-?\d*\.?\d+", pred) + if len(preds) >= 1: + pred = preds[-1] + else: + pred = "" + + return pred + + +def extract_answer(pred_str, data_name, use_last_number=True): + pred_str = pred_str.replace("\u043a\u0438", "") + if data_name in ["mmlu_stem", "sat_math", "aqua", "gaokao2023"]: + # TODO check multiple choice + return choice_answer_clean(pred_str) + + if "final answer is $" in pred_str and "$. I hope" in pred_str: + # minerva_math + tmp = pred_str.split("final answer is $", 1)[1] + pred = tmp.split("$. I hope", 1)[0].strip() + elif "boxed" in pred_str: + ans = pred_str.split("boxed")[-1] + if len(ans) == 0: + return "" + elif ans[0] == "{": + stack = 1 + a = "" + for c in ans[1:]: + if c == "{": + stack += 1 + a += c + elif c == "}": + stack -= 1 + if stack == 0: + break + a += c + else: + a += c + else: + a = ans.split("$")[0].strip() + pred = a + elif "he answer is" in pred_str: + pred = pred_str.split("he answer is")[-1].strip() + elif "final answer is" in pred_str: + pred = pred_str.split("final answer is")[-1].strip() + elif "答案是" in pred_str: + # Handle Chinese few-shot multiple choice problem answer extraction + pred = pred_str.split("答案是")[1].strip().split("\n\n")[0].strip() + else: # use the last number + if use_last_number: + pattern = "-?\d*\.?\d+" + pred = re.findall(pattern, pred_str.replace(",", "")) + if len(pred) >= 1: + pred = pred[-1] + else: + pred = "" + else: + pred = "" + + # choice answer + if data_name in ["sat_math", "aqua"] or "mmlu" in data_name: + tmp = re.findall(r"\b(A|B|C|D|E)\b", pred.upper()) + if tmp: + pred = tmp[-1] + else: + pred = pred.strip().strip(".") + + # multiple line + # pred = pred.split("\n")[0] + pred = re.sub(r"\n\s*", "", pred) + if pred != "" and pred[0] == ":": + pred = pred[1:] + if pred != "" and pred[-1] == ".": + pred = pred[:-1] + if pred != "" and pred[-1] == "/": + pred = pred[:-1] + pred = strip_string(pred, skip_unit=data_name in ["carp_en", "minerva_math"]) + return pred + + +STRIP_EXCEPTIONS = ["carp_en", "minerva_math"] + + +def parse_ground_truth(example: Dict[str, Any], data_name): + if "gt_cot" in example and "gt" in example: + if data_name in ["math"]: + gt_ans = extract_answer(example["gt_cot"], data_name) + elif data_name in STRIP_EXCEPTIONS: + gt_ans = example["gt"] + else: + gt_ans = strip_string(example["gt"]) + return example["gt_cot"], gt_ans + + # parse ground truth + if data_name in ["math", "minerva_math", "math_500"]: + gt_cot = example["solution"] + gt_ans = extract_answer(gt_cot, data_name) + elif data_name == "gsm8k": + gt_cot, gt_ans = example["answer"].split("####") + elif data_name == "svamp": + gt_cot, gt_ans = example["Equation"], example["Answer"] + elif data_name == "asdiv": + gt_cot = example["formula"] + gt_ans = re.sub(r"\(.*?\)", "", example["answer"]) + elif data_name == "mawps": + gt_cot, gt_ans = None, example["target"] + elif data_name == "tabmwp": + gt_cot = example["solution"] + gt_ans = example["answer"] + if example["ans_type"] in ["integer_number", "decimal_number"]: + if "/" in gt_ans: + gt_ans = int(gt_ans.split("/")[0]) / int(gt_ans.split("/")[1]) + elif "," in gt_ans: + gt_ans = float(gt_ans.replace(",", "")) + elif "%" in gt_ans: + gt_ans = float(gt_ans.split("%")[0]) / 100 + else: + gt_ans = float(gt_ans) + elif data_name == "carp_en": + gt_cot, gt_ans = example["steps"], example["answer"] + elif data_name == "mmlu_stem": + abcd = "ABCD" + gt_cot, gt_ans = None, abcd[example["answer"]] + elif data_name == "sat_math": + gt_cot, gt_ans = None, example["Answer"] + elif data_name == "aqua": + gt_cot, gt_ans = None, example["correct"] + elif data_name in ["gaokao2023en", "college_math", "gaokao_math_cloze"]: + gt_cot, gt_ans = None, example["answer"].replace("$", "").strip() + elif data_name == "gaokao_math_qa": + gt_cot, gt_ans = None, example["label"] + elif data_name in ["gaokao2024_mix", "cn_middle_school"]: + if len(example["choice_answer"]) > 0: + gt_cot, gt_ans = None, example["choice_answer"] + else: + gt_cot, gt_ans = None, example["answer"] + elif data_name == "olympiadbench": + gt_cot, gt_ans = None, example["final_answer"][0].strip("$") + elif data_name in [ + "aime24", + "amc23", + "cmath", + "gaokao2024_I", + "gaokao2024_II", + "imo2024", + ]: + gt_cot, gt_ans = None, example["answer"] + elif data_name.startswith("train"): # train_amc_aime + gt_cot, gt_ans = None, example["final_answer"] + else: + raise NotImplementedError(f"`{data_name}`") + # post process + gt_cot = str(gt_cot).strip() + if data_name not in STRIP_EXCEPTIONS: + gt_ans = strip_string(gt_ans, skip_unit=data_name == "carp_en") + else: + gt_ans = ( + gt_ans.replace("\\neq", "\\ne") + .replace("\\leq", "\\le") + .replace("\\geq", "\\ge") + ) + return gt_cot, gt_ans + + +def parse_question(example, data_name): + question = "" + if data_name == "asdiv": + question = f"{example['body'].strip()} {example['question'].strip()}" + elif data_name == "svamp": + body = example["Body"].strip() + if not body.endswith("."): + body = body + "." + question = f'{body} {example["Question"].strip()}' + elif data_name == "tabmwp": + title_str = ( + f'regarding "{example["table_title"]}" ' if example["table_title"] else "" + ) + question = f"Read the following table {title_str}and answer a question:\n" + question += f'{example["table"]}\n{example["question"]}' + if example["choices"]: + question += ( + f' Please select from the following options: {example["choices"]}' + ) + elif data_name == "carp_en": + question = example["content"] + elif data_name == "mmlu_stem": + options = example["choices"] + assert len(options) == 4 + for i, (label, option) in enumerate(zip("ABCD", options)): + options[i] = f"({label}) {str(option).strip()}" + options = " ".join(options) + # question = f"{example['question'].strip()}\nWhat of the following is the right choice? Explain your answer.\n{options}" + question = f"{example['question'].strip()}\nAnswer Choices: {options}" + elif data_name == "sat_math": + options = example["options"].strip() + assert "A" == options[0] + options = "(" + options + for ch in "BCD": + if f" {ch}) " in options: + options = regex.sub(f" {ch}\) ", f" ({ch}) ", options) + # question = f"{example['question'].strip()}\nWhat of the following is the right choice? Explain your answer.\n{options.strip()}" + question = f"{example['question'].strip()}\nAnswer Choices: {options}" + elif "aqua" in data_name: + options = example["options"] + choice = "(" + "(".join(options) + choice = choice.replace("(", " (").replace(")", ") ").strip() + choice = "\nAnswer Choices: " + choice + question = example["question"].strip() + choice + elif data_name == "gaokao_math_qa": + options_dict = example["options"] + options = [] + for key in options_dict: + options.append(f"({key}) {options_dict[key]}") + options = " ".join(options) + question = f"{example['question'].strip()}\n选项: {options}" + elif data_name.startswith("train"): # train_amc_aime: + question = example["prompt"] + else: + for key in ["question", "problem", "Question", "input"]: + if key in example: + question = example[key] + break + # assert question != "" + # Yes or No question + _, gt_ans = parse_ground_truth(example, data_name) + if isinstance(gt_ans, str): + gt_lower = gt_ans.lower() + if gt_lower in ["true", "false"]: + question += " (True or False)" + if gt_lower in ["yes", "no"]: + question += " (Yes or No)" + return question.strip() + + +# def run_execute(executor, result, prompt_type, data_name, execute=False): +# if not result or result == "error": +# return None, None +# report = None + +# if "program_only" in prompt_type: +# prediction = extract_program_output(result) +# elif prompt_type in ["pot", "pal"] and execute: +# code = extract_program(result) +# prediction, report = executor.apply(code) +# else: +# prediction = extract_answer(result, data_name) + +# # prediction = strip_string(prediction, skip_unit=data_name == "carp_en") +# prediction = strip_string(prediction, skip_unit=data_name in STRIP_EXCEPTIONS) +# return prediction, report + + +def _test_extract_answer(): + text = """ +This is still not equal to $0$, so we must have made another mistake. + +When we subtracted $7$ from $\frac{386}{64}$, we should have subtracted $7 \cdot 64$ from $386$, not the other way around. Let's correct that: + +\[\frac{386}{64} - 7 = \frac{386}{64} - \frac{7 \cdot 64}{1 \cdot 64} = \frac{386 - 448}{64} = \frac{-62}{64}.\] + +This is still not equal to $0$, so we must have made another mistake. + +When we subtracted $7$ from $\frac{386}{64}$, we should have subtracted $7 \cdot 64$ from $386$, not the other way around. Let's correct that: + +\[\frac{386}{64} 025 +""" + print(extract_answer(text, "math-oai", use_last_number=True)) + print(choice_answer_clean("\mathrm{(D)\}1,008,016")) + # should output a dict + + print(extract_answer("The product of the inner terms is \\($ 15x $\\).", "math")) + + +if __name__ == "__main__": + _test_extract_answer() diff --git a/functioncall/math/verify.py b/functioncall/math/verify.py new file mode 100644 index 0000000000..7ee14edb55 --- /dev/null +++ b/functioncall/math/verify.py @@ -0,0 +1,112 @@ +import json +import os +import time +from typing import List + +from functioncall.base import logging +from functioncall.base.call import batch_function_call + +logger = logging.getLogger("Functioncall") + + +def loadJson(dataDir): + with open(dataDir, "r") as f: + if dataDir.endswith(".jsonl"): + samples = [json.loads(line) for line in f.readlines()] + else: + samples = json.load(f) + + return samples + + +id2info = None + + +def math_verify(generateds: List, query_ids: List, batch_size=10, timeout=5) -> List: + global id2info + if id2info is None: + id2info = loadJson( + os.getenv( + "REAL_MATH_MEATADATA_PATH", + "/storage/datasets/id2info.json", + ) + ) + assert len(generateds) == len(query_ids), ( + len(generateds), + len(query_ids), + ) + + st = time.monotonic() + query_indices = [] + parameters = [] + # Collect all (generated, solution) pairs with their original indices + for idx, (query_id, generated) in enumerate(zip(query_ids, generateds)): + base_query_id = query_id.split("@idx:")[0] + info = id2info[base_query_id] + for cur_solution in info["solutions"]: + parameters.append((generated, cur_solution, idx)) + query_indices.append(idx) + + # Process in batches + start_time = time.time() + batch_args_list = [] + for i in range(0, len(parameters), batch_size): + answers, solutions, indices = zip(*parameters[i : i + batch_size]) + batch_args = { + "answers": list(answers), + "solutions": list(solutions), + "query_ids": [query_ids[i] for i in indices], + } + + #print(batch_args) + batch_args_list.append(batch_args) + + results_batch = batch_function_call(batch_args_list, "python_math", timeout) + + labels = [0] * len(query_ids) + # Map results back to original indices + index = 0 + for batch_idx, results in enumerate(results_batch): + query_index = query_indices[index] + if not isinstance(results, list) or len(results) == 0: + index += len(batch_args_list[batch_idx]["answers"]) + logger.warning( + f"Invalid functioncall math results: {results}, batch index:{batch_idx}, query index: {query_index}, params: {batch_args_list[batch_idx]['answers']}." + ) + continue + + for result in results: + if ( + isinstance(result, list) + and len(result) > 0 + and (isinstance(result[0], int) and result[0] in [0, 1]) + ): + labels[query_index] = result[0] or labels[query_index] + else: + logger.warning( + f"Invalid functioncall math result: {result}, index:{index}, qeury_id: {query_ids[query_index]}." + ) + + index += 1 + + logger.info( + f"verify math with query size={len(query_ids)}, takes {time.time() - start_time:.4f} seconds, result: {labels}" + ) + return labels + + +if __name__ == "__main__": + sample = { + "prompt": "", + "query_id": "fe11b471-1aa9-4867-958f-a0a811c85f92", + "answer": "\\boxed{-\\frac{2}{3}}", + } + start_time = time.time() + batch_size = 10 + result = math_verify( + [sample["answer"]] * batch_size, [sample["query_id"] for _ in range(batch_size)] + ) + + hint = f"batch_size: {batch_size}, total time : {(time.time() - start_time) * 1000:.0f} ms" + print(result) + print(hint) diff --git a/realhf/api/core/data_api.py b/realhf/api/core/data_api.py index 3cb7c2c89a..393bd7f7d8 100644 --- a/realhf/api/core/data_api.py +++ b/realhf/api/core/data_api.py @@ -396,6 +396,7 @@ def split_with_lengths( group_indices = datapack.ffd_allocate( lens, mb_spec.max_tokens_per_mb, min_groups=mb_spec.n_mbs ) + group_indices = sorted([sorted(g) for g in group_indices]) forward_indices = datapack.flat2d(group_indices) sample = SequenceSample.reorder(self, forward_indices) @@ -828,3 +829,25 @@ def make_dataset( logger.info(f"Dataset creation/loading time: {time.perf_counter() - tik:.3f}s") return dataset + + +def gather_stat(src: List[Dict]) -> Dict: + cnt, stats = {}, {} + for reply in src: + # FIXME: understand why the reply can be None + if not reply: + continue + for k, v in reply.items(): + cnt[k] = cnt.get(k, 0) + 1 + stats[k] = stats.get(k, 0) + v + res = {k: v / cnt for k, v, cnt in zip(stats.keys(), stats.values(), cnt.values())} + for k, c in cnt.items(): + if c != len(src): + logger.warning(f"Gathered `{k}` is not present in every returned stats.") + for k, v in res.items(): + if any(abs(v - x.get(k, None)) > 1e-4 for x in src): + logger.warning( + f"Gathered `{k}` is not all-reduced " + f"before returning: ({[x.get(k, None) for x in src]}, {v})." + ) + return res diff --git a/realhf/api/core/system_api.py b/realhf/api/core/system_api.py index 29dbc6b5d3..20943eba05 100644 --- a/realhf/api/core/system_api.py +++ b/realhf/api/core/system_api.py @@ -2,16 +2,9 @@ # Copyright 2024 Wei Fu & Zhiyu Mei # Licensed under the Apache License, Version 2.0 (the "License"). -import collections -import copy import dataclasses -import enum -import getpass -import itertools -import math import os -import sys -from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union +from typing import Dict, List, Optional, Tuple, Union import realhf.api.core.dfg as dfg from realhf.api.core.config import ( @@ -234,6 +227,36 @@ class ExperimentScheduling: controller_image: str = _LLM_CPU_IMAGE +@dataclasses.dataclass +class AutomaticEvaluator: + """Configuration for automatic evaluation. + :param data_names: Dataset for evaluation seperated by comma. Currently support datasets stored under ./evaluation/data, + including "aime24", "amc23" and "math_500". For example, if "aime24" and "amc23" are required for evaluation, + this field should be set to "aime24,amc23". + :type data_names: str + :param max_gen_tokens: Maximum number of tokens to be generated in evaluation. + :type max_gen_tokens: int + :param max_concurrent_jobs: Maximum number of concurrent evaluation jobs to submit. If number of existing jobs is equal to + `max_concurrent_jobs` and a new checkpoint is saved, the evaluation job will wait until former jobs complete. + :type max_concurrent_jobs: int + :param eval_job_image: Container image used to launch evaluation job. If set to None, evaluation jobs will use + GPU image for training. + :type eval_job_image: Optional[str] + :param initial_checkpoint_path: Initial checkpoint path to evaluate. If specified, this initial checkpoint will be evaluated, + results will be stored as global_step = 0. + :type initial_checkpoint_path: Optional[str] + :param prompt_type: Prompt format used in evaluation. + :type prompt_type: str + """ + + data_names: str = "aime24" + max_gen_tokens: int = 32768 + max_concurrent_jobs: int = 3 + eval_job_image: Optional[str] = None + initial_checkpoint_path: Optional[str] = None + prompt_type: str = "deepscaler" + + @dataclasses.dataclass class WandBConfig: mode: str = "disabled" @@ -247,15 +270,26 @@ class WandBConfig: config: Optional[Dict] = None +@dataclasses.dataclass +class TensorBoardConfig: + path: Optional[str] = None + + @dataclasses.dataclass class ExperimentConfig: exp_ctrl: ExperimentSaveEvalControl wandb: WandBConfig + tensorboard: TensorBoardConfig # dataflow model_rpcs: List[dfg.MFCDef] model_worker: List[ModelWorker] = dataclasses.field(default_factory=list) # master_worker will be set automatically master_worker: Optional[List[MasterWorker]] = None + # automatic evaluation + auto_eval: bool = False + evaluator: AutomaticEvaluator = dataclasses.field( + default_factory=AutomaticEvaluator + ) def __post_init__(self): self.master_worker = [ @@ -429,13 +463,12 @@ def _resolve_data_transfer_pairs( data_transfer_pairs.append((mn1, mn2)) src_rpcs = [rpc for rpc in self.model_rpcs if rpc.is_src] data_src_rpc = src_rpcs[0] - for r in src_rpcs[1:]: + for r in src_rpcs: if ( data_src_rpc.model_name, r.model_name, ) not in data_transfer_pairs: data_transfer_pairs.append((data_src_rpc.model_name, r.model_name)) - data_transfer_pairs += [(mn, mn) for mn in model_names] return data_transfer_pairs def _resolve_param_realloc_pairs( diff --git a/realhf/apps/main.py b/realhf/apps/main.py index 9a767b518a..a3ff135451 100644 --- a/realhf/apps/main.py +++ b/realhf/apps/main.py @@ -18,6 +18,7 @@ import realhf.scheduler.client as sched_client import realhf.system as system from realhf.scheduler.client import JobException, JobState +from realhf.scheduler.evaluator import AutomaticEvaluator logger = logging.getLogger("main", "system") @@ -85,10 +86,17 @@ def main_start(args, recover_count: int = 0): # Run initial_setup to go through all sanity checks. try: exp_cfg = experiment.initial_setup() + assert isinstance(exp_cfg, config_package.ExperimentConfig) exp_cfg.lazy_init() except Exception as e: raise RuntimeError("Experiment initial setup failed.") from e + evaluator = ( + AutomaticEvaluator(exp_cfg.evaluator, exp_cfg.wandb) + if exp_cfg.auto_eval + else None + ) + if args.mode == "local": assert ( args.recover_mode == "disabled" @@ -153,6 +161,8 @@ def main_start(args, recover_count: int = 0): REAL_RECOVER_RUN="1" if is_recover_run else "0", REAL_SAVE_RECOVER_STATES="1" if save_recover_states else "0", REAL_MATH_METADATA_PATH=os.getenv("REAL_MATH_METADATA_PATH", ""), + REAL_CODE_METADATA_PATH=os.getenv("REAL_CODE_METADATA_PATH", ""), + FUNCTIONCALL_SERVICE_DOMAIN=os.getenv("FUNCTIONCALL_SERVICE_DOMAIN", ""), ) for k, v in BASE_ENVIRONS.items(): os.environ[k] = v @@ -167,6 +177,7 @@ def main_start(args, recover_count: int = 0): expr_name=expr_name, trial_name=trial_name, schedule_strategy=args.schedule_strategy, + evaluator=evaluator, ) setup = experiment.scheduling_setup() diff --git a/realhf/base/asyncio_utils.py b/realhf/base/asyncio_utils.py deleted file mode 100644 index f4bfb9ca96..0000000000 --- a/realhf/base/asyncio_utils.py +++ /dev/null @@ -1,80 +0,0 @@ -# Copyright 2025 Ant Group Inc. -# Copyright 2024 Wei Fu & Zhiyu Mei -# Licensed under the Apache License, Version 2.0 (the "License"). - -import asyncio -import dataclasses -import sys -import threading -from asyncio.base_events import _run_until_complete_cb - - -@dataclasses.dataclass -class AsyncRunUntilCompleteContext: - loop: asyncio.BaseEventLoop - future: asyncio.Future - new_task: bool - - -def setup_run_until_complete( - loop: asyncio.BaseEventLoop, - future: asyncio.Future, -) -> AsyncRunUntilCompleteContext: - loop._check_closed() - loop._check_running() - - new_task = not asyncio.futures.isfuture(future) - future = asyncio.tasks.ensure_future(future, loop=loop) - if new_task: - # An exception is raised if the future didn't complete, so there - # is no need to log the "destroy pending task" message - future._log_destroy_pending = False - - future.add_done_callback(_run_until_complete_cb) - - # set up run forever - loop._set_coroutine_origin_tracking(loop._debug) - - loop._old_agen_hooks = sys.get_asyncgen_hooks() - loop._thread_id = threading.get_ident() - sys.set_asyncgen_hooks( - firstiter=loop._asyncgen_firstiter_hook, - finalizer=loop._asyncgen_finalizer_hook, - ) - asyncio.events._set_running_loop(loop) - return AsyncRunUntilCompleteContext(loop=loop, future=future, new_task=new_task) - - -def teardown_run_util_complete(ctx: AsyncRunUntilCompleteContext): - ctx.loop._stopping = False - ctx.loop._thread_id = None - asyncio.events._set_running_loop(None) - ctx.loop._set_coroutine_origin_tracking(False) - # Restore any pre-existing async generator hooks. - if ctx.loop._old_agen_hooks is not None: - sys.set_asyncgen_hooks(*ctx.loop._old_agen_hooks) - ctx.loop._old_agen_hooks = None - - ctx.future.remove_done_callback(_run_until_complete_cb) - - if not ctx.future.done(): - raise RuntimeError("Event loop stopped before Future completed.") - - -def raise_asyncio_exception( - ctx: AsyncRunUntilCompleteContext, raise_error: bool = True -): - if ctx.new_task and ctx.future.done() and not ctx.future.cancelled(): - # The coroutine raised a BaseException. Consume the exception - # to not log a warning, the caller doesn't have access to the - # local task. - ctx.future.exception() - - try: - teardown_run_util_complete(ctx) - except RuntimeError as e: - if raise_error: - raise e - - if raise_error: - raise diff --git a/realhf/base/constants.py b/realhf/base/constants.py index 4ead5a6623..d8aba6b16c 100644 --- a/realhf/base/constants.py +++ b/realhf/base/constants.py @@ -105,7 +105,6 @@ def get_tensor(self, tensor_shape, dtype, name, force_zero: bool = False): # "TORCH_SHOW_CPP_STACKTRACES": "1", # "RAY_DEDUP_LOGS": "0", # disable ray log deduplication "CUDA_DEVICE_MAX_CONNECTIONS": "1", - "PYTHONUSERBASE": "/nonsense", # a random PYTHONUSERBASE to avoid local user site-packages interference "OMP_NUM_THREADS": str(min(os.cpu_count(), 32)), # torch.distributed.all_reduce does not free the input tensor until # the synchronization point. This causes the memory usage to grow diff --git a/realhf/base/timeutil.py b/realhf/base/timeutil.py index ca58cd7090..85f2e6c4ea 100644 --- a/realhf/base/timeutil.py +++ b/realhf/base/timeutil.py @@ -5,9 +5,9 @@ import dataclasses import math import threading -import time from abc import ABC -from typing import Callable, List +from datetime import datetime +from typing import List INFINITE_DURATION = 60 * 60 * 24 * 365 * 1000 @@ -33,9 +33,9 @@ def __init__( """ self.frequency_seconds = frequency_seconds self.frequency_steps = frequency_steps - self.__start_time = time.monotonic() + self.__start_time = datetime.now() self.__steps = 0 - self.__last_time = time.monotonic() + self.__last_time = datetime.now() self.__last_steps = 0 self.__interval_seconds = self.__interval_steps = None self.__initial_value = initial_value @@ -67,7 +67,8 @@ def load_state_dict(self, state_dict): @property def total_seconds(self): - return time.monotonic() - self.__start_time + now = datetime.now() + return (now - self.__start_time).total_seconds() @property def total_steps(self): @@ -90,7 +91,7 @@ def check(self, steps=1): flag: True if condition is met, False other wise """ with self.__lock: - now = time.monotonic() + now = datetime.now() self.__steps += steps if self.__initial_value: @@ -99,7 +100,7 @@ def check(self, steps=1): self.__initial_value = False return True - self.__interval_seconds = now - self.__last_time + self.__interval_seconds = (now - self.__last_time).total_seconds() self.__interval_steps = self.__steps - self.__last_steps if self.frequency_steps is None and self.frequency_seconds is None: return False @@ -119,7 +120,7 @@ def check(self, steps=1): return True def reset_time(self): - self.__last_time = time.monotonic() + self.__last_time = datetime.now() @dataclasses.dataclass diff --git a/realhf/experiments/common/common.py b/realhf/experiments/common/common.py index 2347e09728..6ef970535b 100644 --- a/realhf/experiments/common/common.py +++ b/realhf/experiments/common/common.py @@ -2,20 +2,14 @@ # Copyright 2024 Wei Fu & Zhiyu Mei # Licensed under the Apache License, Version 2.0 (the "License"). -import contextlib import dataclasses -import functools import itertools -import os -import pprint -import re from collections import defaultdict from typing import * import numpy as np import transformers from omegaconf import MISSING, OmegaConf -from transformers.utils import is_accelerate_available import realhf.base.logging as logging from realhf.api.core.config import ( @@ -26,9 +20,10 @@ ModelShardID, StandaloneModelShardAbstraction, ) -from realhf.api.core.dfg import MFCDef, ModelInterfaceType, build_graph +from realhf.api.core.dfg import MFCDef, ModelInterfaceType from realhf.api.core.model_api import HF_MODEL_FAMILY_REGISTRY from realhf.api.core.system_api import ( + AutomaticEvaluator, Experiment, ExperimentConfig, ExperimentSaveEvalControl, @@ -36,6 +31,7 @@ ModelWorker, Scheduling, TasksGroup, + TensorBoardConfig, WandBConfig, ) from realhf.api.quickstart.device_mesh import ( @@ -60,6 +56,7 @@ ) from realhf.experiments.common.utils import ( AllocationMode, + asdict, get_topo, make_inf_backend_config, make_train_backend_config, @@ -137,6 +134,9 @@ class CommonExperimentConfig(Experiment): :param wandb: The WandB initialization config. See https://docs.wandb.ai/ref/python/init/ for details. :type wandb: WandbConfig + :param tensorboard: The tensorboard initialization config. + Only the field of `path` is needed to specify the directory of saving the tensorboard events. + :type tensorboard: TensorBoardConfig :param image_name: The name of the Docker image used by the controller. This parameter is only used in SLURM mode. :type image_name: str or None @@ -183,6 +183,16 @@ class CommonExperimentConfig(Experiment): torch.cuda.empty_cache() before each RPC in model worker If enabled, there will be a ~0.1s overhead per RPC. :type torch_cache_mysophobia: bool + :param auto_eval: Whether to automatic evaluation in training. When enabled, an evaluation + job is submitted whenever a checkpoint is saved, and the result will be logged on disk and + on wandb if wandb is active. + :type auto_eval: bool + :param auto_eval_config: Configuration for automatic evaluation. + :type auto_eval_config: AutomaticEvaluator + :param cpus_per_master_worker: The number of CPUs for each master worker. + :param mem_per_master_worker: The size of memory for each master worker, measured in MB. + :param cpus_per_model_worker: The number of CPUs for each model worker. + :param mem_per_model_worker: The size of memory for each model worker, measured in MB. """ experiment_name: str = MISSING @@ -194,6 +204,9 @@ class CommonExperimentConfig(Experiment): partition: str = "dev" schedule_strategy: str = "empty_first" wandb: WandBConfig = dataclasses.field(default_factory=WandBConfig) + tensorboard: TensorBoardConfig = dataclasses.field( + default_factory=TensorBoardConfig + ) image_name: Optional[str] = None recover_mode: str = "disabled" recover_retries: int = 1 @@ -210,6 +223,16 @@ class CommonExperimentConfig(Experiment): default_factory=ExperimentSaveEvalControl ) torch_cache_mysophobia: bool = True + # Options for automatic evaluation + auto_eval: bool = False + auto_eval_config: AutomaticEvaluator = dataclasses.field( + default_factory=AutomaticEvaluator + ) + # Options for worker resources + cpus_per_master_worker: int = 4 + mem_per_master_worker: int = 20000 + cpus_per_model_worker: int = 4 + mem_per_model_worker: int = 90000 @property def models(self) -> Dict[str, ModelTrainEvalConfig]: @@ -324,18 +347,18 @@ def scheduling_setup(self) -> ExperimentScheduling: master_worker=TasksGroup( count=1, scheduling=Scheduling.master_worker_default( - cpu=4, - mem=20000, + cpu=self.cpus_per_master_worker, + mem=self.mem_per_master_worker, nodelist=self.nodelist, ), ), model_worker=TasksGroup( count=self.n_nodes * self.n_gpus_per_node, scheduling=Scheduling.model_worker_default( - cpu=4, + cpu=self.cpus_per_model_worker, gpu=1, gpu_type=cluster_spec.gpu_type, - mem=90000, + mem=self.mem_per_model_worker, nodelist=self.nodelist, ), ), @@ -541,9 +564,7 @@ def _get_model_worker_configs( and self.gen_device_mesh.mapping[i, j] ): gen_rpc_alloc = next( - alloc - for alloc in rpc_allocs - if alloc.rpc.interface_type == ModelInterfaceType.GENERATE + alloc for alloc in rpc_allocs if alloc.rpc.is_generate() ) model_name = gen_rpc_alloc.rpc.model_name topo = get_topo( @@ -600,6 +621,10 @@ def _get_model_worker_configs( model_rpc_allocs, ) in model_name_to_rpc_allocs.items(): rpcs = [rpc_alloc.rpc for rpc_alloc in model_rpc_allocs] + if self._allocation_mode.is_decoupled() and all( + rpc.is_generate() for rpc in rpcs + ): + continue rpc_alloc = model_rpc_allocs[0] model_cfg = self.models[model_name.role] model = get_real_model_config( @@ -657,9 +682,7 @@ def _get_model_worker_configs( rpc.is_generate() for rpc in rpcs ): assert len(rpcs) == 1 and rpcs[0].is_generate(), rpcs - vllm_dict_args: Dict[str, Any] = OmegaConf.to_container( - model_cfg.vllm, resolve=True - ) + vllm_dict_args: Dict[str, Any] = asdict(model_cfg.vllm) backend = ModelBackendAbstraction( "vllm", args=dict( @@ -669,6 +692,17 @@ def _get_model_worker_configs( ) else: backend = make_inf_backend_config(model_cfg, rpc_alloc.parallel) + if any(rpc.is_generate() for rpc in rpcs) and backend.type_ not in [ + "vllm", + "sglang", + ]: + print(rpcs, model_name, backend.type_) + raise ValueError( + "vLLM or SGLang is not enabled for generation. " + "This behavior has been deprecated. " + "Please set model.vllm.hybrid_train=True " + "or model.sglang.hybrid_train=True." + ) check_valid_vllm(model_name.role, model_cfg.vllm, rpc_allocs) if mapping[i, j]: @@ -706,8 +740,11 @@ def initial_setup(self) -> ExperimentConfig: return ExperimentConfig( exp_ctrl=self.exp_ctrl, wandb=self.wandb, + tensorboard=self.tensorboard, model_rpcs=[rpc_alloc.rpc for rpc_alloc in rpc_allocs], model_worker=model_worker, + auto_eval=self.auto_eval, + evaluator=self.auto_eval_config, ) def __check_legal_allocation_options(self): diff --git a/realhf/experiments/common/ppo_code_exp.py b/realhf/experiments/common/ppo_code_exp.py new file mode 100644 index 0000000000..01c9c2422d --- /dev/null +++ b/realhf/experiments/common/ppo_code_exp.py @@ -0,0 +1,564 @@ +# Copyright 2025 Ant Group Inc. +# Copyright 2024 Wei Fu & Zhiyu Mei +# Licensed under the Apache License, Version 2.0 (the "License"). +import copy +import dataclasses +import math +import os +import pprint +from typing import * + +import numpy as np +from omegaconf import DictConfig, OmegaConf + +import realhf.base.logging as logging +from realhf.api.core.config import ( + DatasetAbstraction, + ModelInterfaceAbstraction, + ModelInterfaceType, +) +from realhf.api.core.dfg import MFCDef, ParamReallocHook +from realhf.api.core.model_api import GenerationHyperparameters +from realhf.api.core.system_api import ExperimentConfig +from realhf.api.quickstart.dataset import PromptOnlyDatasetConfig +from realhf.api.quickstart.device_mesh import DeviceMesh, MFCConfig, RPCAllocation +from realhf.api.quickstart.entrypoint import register_quickstart_exp +from realhf.api.quickstart.model import ModelTrainEvalConfig, ParallelismConfig +from realhf.experiments.common.common import CommonExperimentConfig +from realhf.experiments.common.utils import resolve_replica_ids, resolve_rpc_hooks + +logger = logging.getLogger("PPO Code exp", "colored") + + +@dataclasses.dataclass +class PPOHyperparameters: + """Configuration of PPO hyperparameters. + + :param gen: Generation hyperparameters. + :type gen: GenerationHyperparameters + :param ppo_n_minibatches: Number of minibatches in each PPO update. + :type ppo_n_minibatches: int + :param kl_ctl: Coefficient of KL divergence rewards. + :type kl_ctl: float + :param discount: Discount factor. + :type discount: float + :param gae_lambda: Lambda factor in GAE. + :type gae_lambda: float + :param eps_clip: PPO actor probability ratio clipping factor. + :type eps_clip: float + :param value_eps_clip: PPO value clipping factor. + :type value_eps_clip: float + :param max_reward_clip: Maximum reward value. + :type max_reward_clip: float + :param reward_output_scaling: Scaling factor of the reward model output. + :type reward_output_scaling: float + :param reward_output_bias: Bias of the reward model output. + The number outputed by the reward model will be + CLIP((x - bias) * scaling, -max_reward_clip, max_reward_clip). + :type reward_output_bias: float + :param early_stop_imp_ratio: PPO update will be early stopped if importance ratio + exceeds this maximum value. + :type early_stop_imp_ratio: float + :param use_adaptive_kl_ctl: Whether to use adaptive KL divergence coefficient. + :type use_adaptive_kl_ctl: bool + :param adv_norm: Whether to use advantage normalization. + :type adv_norm: bool + :param value_norm: Whether to denormalize valued and normalize return predictions. + :type value_norm: bool + :param value_norm_type: Type of value normalization. + Either exponential moving average ("exp") or moving average ("ma"). + :type value_norm_type: str + :param value_norm_beta: Exponential decay factor + in exponential moving average. + :type value_norm_beta: float + :param value_norm_eps: Epsilon factor in the + denominator of exponential moving average. + :type value_norm_eps: float + :param disable_value: A shortcut option to disable the critic model. + :type disable_value: bool + """ + + gen: GenerationHyperparameters = dataclasses.field( + default_factory=GenerationHyperparameters + ) + ppo_n_minibatches: int = 4 + kl_ctl: float = 0.1 + discount: float = 1.0 + gae_lambda: float = 1.0 + eps_clip: float = 0.2 + value_eps_clip: float = 0.2 + disable_value: bool = False + max_reward_clip: float = 20.0 + reward_output_scaling: float = 1.0 + reward_output_bias: float = 0.0 + early_stop_imp_ratio: float = 5.0 + use_adaptive_kl_ctl: bool = False + adv_norm: bool = True + value_norm: bool = True + value_norm_type: str = dataclasses.field( + metadata={"choices": ["exp", "ma"]}, default="exp" + ) + value_norm_beta: float = 0.99995 + value_norm_eps: float = 1e-5 + + +@dataclasses.dataclass +class PPOCODEConfig(CommonExperimentConfig): + """PPO experiment configuration. + + It is a subclass of :class:`CommonExperimentConfig`, + so all CLI options in the base class are available. + + We don't implement runtime evaluation for PPO. + + We identify that the RLHF process is composed of four + distinct models with independent parameters and six + *model function calls* upon these models. + + The four models are\: + + - Actor\: The primary LLM that generates text. + - Critic\: The value function that estimates the value of a state. + - Ref\: The reference LLM that provides KL regularization. + - Rew\: The reward model that provides reward signals. + + The four model function calls and their dependencies are\: + + - Rollout\: Generate text from the actor model. + - InfReward\: Infer rewards from the reward model given generated text. + - InfRef\: Infer log probabilities from the reference model given generated text. + - InfValues\: Infer values from the critic model given generated text. + - TrainActor\: Train the actor model given generated text, rewards, values, and reference log probabilities. + - TrainCritic\: Train the critic model given generated text, rewards, values, and reference log probabilities. + + This class resolves these dependencies under the hood. + What the users should specify are the runtime configurations + of models and allocations of *each model function call*. + + :param actor: Runtime configuration of the primary LLM. + :type actor: ModelTrainEvalConfig + :param critic: Runtime configuration of the critic model of PPO. + :type critic: ModelTrainEvalConfig + :param ref: Runtime configuration of the reference LLM. + :type ref: ModelTrainEvalConfig + :param rew: Runtime configuration of the reward LLM. + :type rew: ModelTrainEvalConfig + :param actor_train: :class:`MFCConfig` for TrainActor. + :type actor_train: MFCConfig + :param critic_train: :class:`MFCConfig` for TrainCritic. + :type critic_train: MFCConfig + :param actor_gen: :class:`MFCConfig` for Rollout. + :type actor_gen: MFCConfig + :param critic_inf: :class:`MFCConfig` for InfValues. + :type critic_inf: MFCConfig + :param rew_inf: :class:`MFCConfig` for InfReward. + :type rew_inf: MFCConfig + :param ref_inf: :class:`MFCConfig` for InfRef. + :type ref_inf: MFCConfig + :param dataset: Dataset configuration. + :type dataset: PromptOnlyDatasetConfig + :param ppo: Configuration for the PPO algorithm. + :type ppo: PPOHyperparameters + :param group_size: The number of answers remained for each prompt. + :type group_size: int + :param generation_size: The number of answers sampled for each prompt. + Among them, only `group_size` samples are remained according to + the reward score, aka best-of-n sampling. If None, use `group_size`. + :type generation_size: Optional[int] + :param mask_no_eos_with_zero: Whether to mask out the reward if an + answer is truncated due to exceeding the length limit. + :type mask_no_eos_with_zero: bool + :param mask_too_long: Whether to mask out the PPO loss if an + answer is truncated due to exceeding the length limit. + :type mask_too_long: bool + :param check_verifier_status: If True, raise an error + when the reward is all-zero. This usually indicates a bug + of the verifier. + :type check_verifier_status: bool + :param group_adv_norm: Whther to use grouped advantage + normaliztion in GRPO. + :type group_adv_norm: bool + """ + + actor: ModelTrainEvalConfig = dataclasses.field( + default_factory=ModelTrainEvalConfig + ) + critic: ModelTrainEvalConfig = dataclasses.field( + default_factory=ModelTrainEvalConfig + ) + ref: ModelTrainEvalConfig = dataclasses.field(default_factory=ModelTrainEvalConfig) + rew: ModelTrainEvalConfig = dataclasses.field(default_factory=ModelTrainEvalConfig) + + # for manual allocation only + actor_train: MFCConfig = dataclasses.field(default_factory=MFCConfig) + critic_train: MFCConfig = dataclasses.field(default_factory=MFCConfig) + actor_gen: MFCConfig = dataclasses.field(default_factory=MFCConfig) + critic_inf: MFCConfig = dataclasses.field(default_factory=MFCConfig) + rew_inf: MFCConfig = dataclasses.field(default_factory=MFCConfig) + ref_inf: MFCConfig = dataclasses.field(default_factory=MFCConfig) + + dataset: PromptOnlyDatasetConfig = dataclasses.field( + default_factory=PromptOnlyDatasetConfig + ) + + ppo: PPOHyperparameters = dataclasses.field(default_factory=PPOHyperparameters) + + group_size: int = 1 + generation_size: Optional[int] = None + mask_no_eos_with_zero: bool = False + rm_output_scaling: Optional[float] = None + ref_ema_eta: Optional[float] = None + group_adv_norm: bool = False + mask_too_long: bool = False + rw_type: Optional[str] = "sparse" + task: str = "code" + check_xml_format: bool = False + use_dense_reward: bool = False + reward_delta: bool = True + + check_verifier_status: bool = False + + dataset_filter_threshold: float = 100.0 + dataset_max_filter_percentage: float = 0.0 + + def __post_init__(self): + + self.ppo_kwargs = dict( + n_minibatches=self.ppo.ppo_n_minibatches, + kl_ctl=self.ppo.kl_ctl, + discount=self.ppo.discount, + gae_lambda=self.ppo.gae_lambda, + eps_clip=self.ppo.eps_clip, + value_eps_clip=self.ppo.value_eps_clip, + max_reward_clip=self.ppo.max_reward_clip, + adaptive_kl_ctl=self.ppo.use_adaptive_kl_ctl, + value_norm=self.ppo.value_norm, + value_norm_type=self.ppo.value_norm_type, + value_norm_beta=self.ppo.value_norm_beta, + value_norm_eps=self.ppo.value_norm_eps, + disable_value=self.ppo.disable_value, + mask_no_eos_with_zero=self.mask_no_eos_with_zero, + ) + + if self.rm_output_scaling is None: + self.rm_output_scaling = self.ppo.reward_output_scaling + + @property + def models(self) -> Dict[str, ModelTrainEvalConfig]: + # role to config + if self.ppo.disable_value: + return { + "actor": self.actor, + # "critic": self.critic, + "ref": self.ref, + "reward": self.rew, + } + else: + return { + "actor": self.actor, + "critic": self.critic, + "ref": self.ref, + "reward": self.rew, + } + + @property + def rpcs(self): + if ( + self.dataset.max_prompt_len + self.ppo.gen.max_new_tokens + > self.actor.vllm.max_seq_len_to_capture + ): + raise RuntimeError( + f"vllm max seq len to capture {self.actor.vllm.max_seq_len_to_capture} is " + f"smaller than the prompt length + generation length {self.dataset.max_prompt_len + self.ppo.gen.max_new_tokens}" + ) + if not os.path.exists(os.getenv("REAL_CODE_METADATA_PATH")): + raise RuntimeError( + "Dataset json path REAL_CODE_METADATA_PATH does not exist." + ) + + domain = os.getenv("FUNCTIONCALL_SERVICE_DOMAIN", "") + if not (domain.startswith("http://") and ":" in domain): + raise RuntimeError( + "function call address FUNCTIONCALL_SERVICE_DOMAIN is invalid." + ) + + + + # interfaces + actor_interface = ModelInterfaceAbstraction( + "ppo_actor", + args={ + **copy.deepcopy(self.ppo_kwargs), + # NOTE: to_container converts the object to a dict + # It is used for unifying the profiling API, which requires to + # pass external interface configurations in the launch command. + # Customized dataclass objects will not work in that case. + "generation_config": ( + OmegaConf.to_container(self.ppo.gen, resolve=True) + if isinstance(self.ppo.gen, (OmegaConf, DictConfig)) + else dataclasses.asdict(self.ppo.gen) + ), + "early_stop_imp_ratio": self.ppo.early_stop_imp_ratio, + "adv_norm": self.ppo.adv_norm, + "group_size": self.group_size, + "generation_size": self.generation_size, + "group_adv_norm": self.group_adv_norm, + "mask_too_long": self.mask_too_long, + "use_dense_reward": self.use_dense_reward, + "reward_delta": self.reward_delta, + }, + ) + ref_interface = copy.deepcopy(actor_interface) + ref_interface.args["enable_save"] = False + + critic_interface = ModelInterfaceAbstraction( + "ppo_critic", + args={ + **copy.deepcopy(self.ppo_kwargs), + "group_size": self.group_size, + "mask_too_long": self.mask_too_long, + "use_dense_reward": self.use_dense_reward, + "reward_delta": self.reward_delta, + }, + ) + critic_interface.args.pop("eps_clip") + rw_interface = ModelInterfaceAbstraction( + "rw_code", + args=dict( + rw_type=self.rw_type, + task=self.task, + check_xml_format=self.check_xml_format, + tokenizer_path=self.actor.path, + enable_save=False, + output_scaling=self.ppo.reward_output_scaling, + rm_output_scaling=self.rm_output_scaling, + output_bias=self.ppo.reward_output_bias, + group_size=self.group_size, + check_verifier_status=self.check_verifier_status, + max_sync_length=self.ppo.gen.max_new_tokens + + self.dataset.max_prompt_len + + 128, + ), + ) + rollout = MFCDef( + name="actor_gen", + model_name="actor", + mb_spec=self.actor_gen.mb_spec, + interface_type=ModelInterfaceType.GENERATE, + model_type=self.actor.type, + model_path=self.actor.path, + interface_impl=actor_interface, + input_keys=["packed_prompts"], + output_keys=[ + "seq_no_eos_mask", + "packed_input_ids", + "packed_logprobs", + "prompt_mask", + ], + n_seqs=self.dataset.train_bs_n_seqs, + ) + + inf_reward = MFCDef( + name="rew_inf", + model_name="reward", + mb_spec=self.rew_inf.mb_spec, + interface_type=ModelInterfaceType.INFERENCE, + interface_impl=rw_interface, + model_type=self.rew.type, + model_path=self.rew.path, + min_n_seqs_per_pass=1 / self.group_size, + input_keys=["packed_input_ids", "packed_prompts"], + output_keys=["rewards", "dense_rewards"], + n_seqs=self.dataset.train_bs_n_seqs, + ) + + inf_ref_inputs = ["packed_input_ids"] + inf_ref_logits = MFCDef( + name="ref_inf", + model_name="ref", + mb_spec=self.ref_inf.mb_spec, + interface_type=ModelInterfaceType.INFERENCE, + model_type=self.ref.type, + model_path=self.ref.path, + interface_impl=ref_interface, + min_n_seqs_per_pass=1 / self.group_size, + input_keys=inf_ref_inputs, + output_keys=["packed_ref_logprobs"], + n_seqs=self.dataset.train_bs_n_seqs, + ) + + inf_values = MFCDef( + name="critic_inf", + model_name="critic", + mb_spec=self.critic_inf.mb_spec, + interface_type=ModelInterfaceType.INFERENCE, + interface_impl=critic_interface, + model_type=self.critic.type, + model_path=self.critic.path, + min_n_seqs_per_pass=1 / self.group_size, + input_keys=["packed_input_ids", "seq_no_eos_mask"], + output_keys=["values"], + n_seqs=self.dataset.train_bs_n_seqs, + ) + + train_actor_inputs = [ + "packed_input_ids", + "packed_logprobs", + "packed_ref_logprobs", + "rewards", + "dense_rewards", + "values", + "prompt_mask", + "seq_no_eos_mask", + ] + if self.ppo.disable_value: + train_actor_inputs.remove("values") + train_actor = MFCDef( + name="actor_train", + model_name="actor", + mb_spec=self.actor_train.mb_spec, + interface_type=ModelInterfaceType.TRAIN_STEP, + model_type=self.actor.type, + model_path=self.actor.path, + interface_impl=actor_interface, + input_keys=train_actor_inputs, + log_return_value=True, + min_n_seqs_per_pass=self.ppo.ppo_n_minibatches / self.group_size, + n_seqs=self.dataset.train_bs_n_seqs, + ) + + train_critic = MFCDef( + name="critic_train", + model_name="critic", + mb_spec=self.critic_train.mb_spec, + interface_type=ModelInterfaceType.TRAIN_STEP, + interface_impl=critic_interface, + model_type=self.critic.type, + model_path=self.critic.path, + input_keys=[ + "packed_input_ids", + "packed_logprobs", + "packed_ref_logprobs", + "rewards", + "dense_rewards", + "values", + "prompt_mask", + "seq_no_eos_mask", + ], + log_return_value=True, + min_n_seqs_per_pass=self.ppo.ppo_n_minibatches / self.group_size, + n_seqs=self.dataset.train_bs_n_seqs, + ) + if self.ppo.disable_value: + return { + "actor_gen": rollout, + "actor_train": train_actor, + # "critic_inf": inf_values, + # "critic_train": train_critic, + "ref_inf": inf_ref_logits, + "rew_inf": inf_reward, + } + else: + return { + "actor_gen": rollout, + "actor_train": train_actor, + "critic_inf": inf_values, + "critic_train": train_critic, + "ref_inf": inf_ref_logits, + "rew_inf": inf_reward, + } + + @property + def allocations(self): + if self.ppo.disable_value: + return { + "actor_gen": self.actor_gen, + "actor_train": self.actor_train, + # "critic_inf": self.critic_inf, + # "critic_train": self.critic_train, + "ref_inf": self.ref_inf, + "rew_inf": self.rew_inf, + } + else: + return { + "actor_gen": self.actor_gen, + "actor_train": self.actor_train, + "critic_inf": self.critic_inf, + "critic_train": self.critic_train, + "ref_inf": self.ref_inf, + "rew_inf": self.rew_inf, + } + + @property + def datasets(self): + return [ + DatasetAbstraction( + "code_prompt", + args=dict( + dataset_path=self.dataset.path, + max_length=self.dataset.max_prompt_len, + fill_to_max_length=self.dataset.fill_to_max_length, + filter_threshold=self.dataset_filter_threshold, + max_filter_percentage=self.dataset_max_filter_percentage, + ), + ) + ] + + @property + def tokenizer_name_or_path(self) -> str: + return self.actor.path + + @property + def search_kwargs(self): + return { + "num_gen_tokens": self.ppo.gen.max_new_tokens, + "n_ppo_minibatches": self.ppo.ppo_n_minibatches, + "seq_len": self.dataset.max_prompt_len, + } + + @property + def max_prompt_len(self): + return self.dataset.max_prompt_len + + def initial_setup(self) -> ExperimentConfig: + rpc_allocs = self._get_rpc_allocations() + + resolve_replica_ids(rpc_allocs, self.models) + resolve_rpc_hooks( + rpc_allocs, self.models + ) # inplace modify MFCDefs in rpc allocations + + pprint.pprint(rpc_allocs) + + ######### update ref model using ema, ref_ema_eta = 0 means fixed ref model ######### + def _find_rpc(name): + return next(alloc.rpc for alloc in rpc_allocs if alloc.rpc.name == name) + + # Remove the offload hook of ref_inf, because + # we need to receive parameters from peer GPUs and update it immediately. + if self.ref_ema_eta is not None: + + ref_inf = _find_rpc("ref_inf") + ref_inf._post_hooks = [] + + # Add an unidirectional parameter reallocation hook. + actor_train = _find_rpc("actor_train") + actor_train.add_post_hook( + ParamReallocHook( + target=ref_inf.model_name, + eta=self.ref_ema_eta, + ) + ) + ######### The main difference from normal PPO ######### + + model_worker = self._get_model_worker_configs(rpc_allocs) + + return ExperimentConfig( + exp_ctrl=self.exp_ctrl, + wandb=self.wandb, + model_rpcs=[rpc_alloc.rpc for rpc_alloc in rpc_allocs], + model_worker=model_worker, + ) + + +register_quickstart_exp("ppo-code", PPOCODEConfig) diff --git a/realhf/experiments/common/ppo_math_exp.py b/realhf/experiments/common/ppo_math_exp.py index 3f297e9f65..d70fd4d1a7 100644 --- a/realhf/experiments/common/ppo_math_exp.py +++ b/realhf/experiments/common/ppo_math_exp.py @@ -3,12 +3,10 @@ # Licensed under the Apache License, Version 2.0 (the "License"). import copy import dataclasses -import math import os import pprint from typing import * -import numpy as np from omegaconf import DictConfig, OmegaConf import realhf.base.logging as logging @@ -21,11 +19,15 @@ from realhf.api.core.model_api import GenerationHyperparameters from realhf.api.core.system_api import ExperimentConfig from realhf.api.quickstart.dataset import PromptOnlyDatasetConfig -from realhf.api.quickstart.device_mesh import DeviceMesh, MFCConfig, RPCAllocation +from realhf.api.quickstart.device_mesh import MFCConfig from realhf.api.quickstart.entrypoint import register_quickstart_exp -from realhf.api.quickstart.model import ModelTrainEvalConfig, ParallelismConfig +from realhf.api.quickstart.model import ModelTrainEvalConfig from realhf.experiments.common.common import CommonExperimentConfig -from realhf.experiments.common.utils import resolve_replica_ids, resolve_rpc_hooks +from realhf.experiments.common.utils import ( + asdict, + resolve_replica_ids, + resolve_rpc_hooks, +) logger = logging.getLogger("PPO Math exp", "colored") @@ -275,6 +277,13 @@ def rpcs(self): raise RuntimeError( "Dataset json path REAL_MATH_METADATA_PATH does not exist." ) + + domain = os.getenv("FUNCTIONCALL_SERVICE_DOMAIN", "") + if domain and (not (domain.startswith("http://") and ":" in domain)): + raise RuntimeError( + "function call address FUNCTIONCALL_SERVICE_DOMAIN is invalid." + ) + # interfaces actor_interface = ModelInterfaceAbstraction( @@ -285,11 +294,7 @@ def rpcs(self): # It is used for unifying the profiling API, which requires to # pass external interface configurations in the launch command. # Customized dataclass objects will not work in that case. - "generation_config": ( - OmegaConf.to_container(self.ppo.gen, resolve=True) - if isinstance(self.ppo.gen, (OmegaConf, DictConfig)) - else dataclasses.asdict(self.ppo.gen) - ), + "generation_config": asdict(self.ppo.gen), "early_stop_imp_ratio": self.ppo.early_stop_imp_ratio, "adv_norm": self.ppo.adv_norm, "group_size": self.group_size, @@ -544,12 +549,16 @@ def _find_rpc(name): ######### The main difference from normal PPO ######### model_worker = self._get_model_worker_configs(rpc_allocs) + self.auto_eval_config.initial_checkpoint_path = self.actor.path return ExperimentConfig( exp_ctrl=self.exp_ctrl, wandb=self.wandb, + tensorboard=self.tensorboard, model_rpcs=[rpc_alloc.rpc for rpc_alloc in rpc_allocs], model_worker=model_worker, + auto_eval=self.auto_eval, + evaluator=self.auto_eval_config, ) diff --git a/realhf/experiments/common/utils.py b/realhf/experiments/common/utils.py index 8de0dca0ee..7bc94dcb6d 100644 --- a/realhf/experiments/common/utils.py +++ b/realhf/experiments/common/utils.py @@ -10,7 +10,7 @@ from typing import * import numpy as np -from omegaconf import OmegaConf +from omegaconf import DictConfig, OmegaConf from realhf.api.core.config import ( ModelBackendAbstraction, @@ -229,6 +229,8 @@ def resolve_rpc_hooks( class AllocationType(enum.Enum): DECOUPLED = 1 GLOBAL_HYBRID = 2 + MANUAL = 3 + SEARCH = 4 @dataclasses.dataclass @@ -244,6 +246,10 @@ def is_global_hybrid(self): @classmethod def from_str(cls, allocation_mode: str): + if allocation_mode == "manual": + return cls(AllocationType.MANUAL, None) + if allocation_mode == "search": + return cls(AllocationType.SEARCH, None) alloc_3d = AllocationMode.extract_3d_alloc(allocation_mode) alloc_hybrid = AllocationMode.extract_key_value_alloc(allocation_mode) alloc_decoupled = AllocationMode.extract_decoupled_alloc(allocation_mode) @@ -317,3 +323,9 @@ def parse_key_value_pairs(s: str): return allocs[k] = v["*"] return allocs + + +def asdict(cfg): + if isinstance(cfg, (OmegaConf, DictConfig)): + return OmegaConf.to_container(cfg, resolve=True) + return dataclasses.asdict(cfg) diff --git a/realhf/impl/dataset/prompt_dataset.py b/realhf/impl/dataset/prompt_dataset.py index 18dd67a6f1..16004b30e1 100644 --- a/realhf/impl/dataset/prompt_dataset.py +++ b/realhf/impl/dataset/prompt_dataset.py @@ -82,7 +82,6 @@ def __getitem__(self, idx): ids=[self.ids[idx]], seqlens=[self.prompt_lengths[idx]], data=dict(packed_prompts=torch.tensor(self.prompts[idx], dtype=torch.long)), - metadata=dict(random_id=[uuid.uuid4()]), ) @@ -168,6 +167,144 @@ def __len__(self): def __getitem__(self, idx): # print(self.base_scores) idx = self.active_indices[idx] + if hasattr(self, "base_scores"): + return data_api.SequenceSample.from_default( + ids=[self.ids[idx]], + seqlens=[self.prompt_lengths[idx]], + data=dict( + packed_prompts=torch.tensor(self.prompts[idx], dtype=torch.long), + base_scores=torch.tensor( + [self.base_scores[idx]], dtype=torch.float32 + ), + ), + ) + else: + return data_api.SequenceSample.from_default( + ids=[self.ids[idx]], + seqlens=[self.prompt_lengths[idx]], + data=dict( + packed_prompts=torch.tensor(self.prompts[idx], dtype=torch.long) + ), + ) + + def filter(self, eval_scores: Dict[Hashable, float]): + # Get all data indices that have a higher score than the threshold. + idx2scores_to_remove = {} + for pop_idx, idx in enumerate(self.active_indices): + data_id = self.ids[idx] + if data_id not in eval_scores: + continue + if eval_scores[data_id] > self.filter_threshold: + idx2scores_to_remove[pop_idx] = eval_scores[data_id] + + # Control the number of samples to be removed according to max_filter_percentage. + n = int(len(self.active_indices) * self.max_filter_percentage) + indices_to_remove = sorted( + idx2scores_to_remove.keys(), + key=lambda x: idx2scores_to_remove[x], + reverse=True, + )[:n] + + for pop_idx in sorted(indices_to_remove, reverse=True): + self.active_indices.pop(pop_idx) + logger.info( + f"Math prompt dataset DP rank {self.util.dp_rank} filtered" + f" {len(indices_to_remove)} samples, {len(self.active_indices)} samples remain. " + f"Original dataset size: {len(self.prompts)}. " + f"Filter threshold: {self.filter_threshold}. " + f"Max filter percentage: {self.max_filter_percentage}. " + f"Current number of eval scores: {len(eval_scores)}." + ) + + +class CODEPromptDataset(torch.utils.data.Dataset): + + def __init__( + self, + util: data_api.DatasetUtility, + max_length: Optional[int] = None, + dataset_path: Optional[str] = None, + dataset_builder: Optional[Callable[[], List[Dict]]] = None, + fill_to_max_length: bool = False, + filter_threshold: float = 1e4, + max_filter_percentage: float = 0.0, + ): + """A dataset with prompts. Usually used for PPO. + + Args: + util (api.data.DatasetUtility): . + max_length (Optional[int], optional): The maximum length of each sequence in the batch. + dataset_path (Optional[str], optional): Path to the dataset json/jsonl file. + The json/jsonl file should be a list of dictionary. Each element in the list should have + a key "prompt". Defaults to None. + dataset_builder (Optional[Callable[[], List[Dict]]], optional): Alternative to dataset_path. + A callable that returns a list of dictionary. Defaults to None. + """ + self._util = util + self.max_length = max_length + + data = data_api.load_shuffle_split_dataset(util, dataset_path, dataset_builder) + + prompts_str = [x["question"] for x in data] + self.ids = [x["id"] for x in data] + if "scores" in data[0]: + self.base_scores = [np.mean(x["scores"]) for x in data] + util.tokenizer.padding_side = "left" + prompt_encodings = util.tokenizer( + prompts_str, + truncation=True, + # max_length=max_length, + padding=False, + return_length=True, + return_attention_mask=False, + ) + + if fill_to_max_length: + for i in range(len(prompt_encodings["length"])): + x = prompt_encodings["input_ids"][i] + if max_length > len(x): + # Fill with the final non-pad token to the left. + prompt_encodings["input_ids"][i] = [x[-1]] * ( + max_length - len(x) + ) + x + prompt_encodings["length"][i] = max_length + + logger.info(f"{len(data)} samples, checking lengths (max_length={max_length})") + indices = [ + i for i, x in enumerate(prompt_encodings["length"]) if x <= max_length + ] + logger.info(f"{len(indices)} samples remain") + + self.prompt_lengths = [int(prompt_encodings["length"][idx]) for idx in indices] + self.prompts = [prompt_encodings["input_ids"][idx] for idx in indices] + self.ids = [ + str(self.ids[idx]) + f"@idx:{idx}-{util.dp_rank}" for idx in indices + ] + if "scores" in data[0]: + self.base_scores = [self.base_scores[idx] for idx in indices] + + assert all(len(x) == l for x, l in zip(self.prompts, self.prompt_lengths)) + + logger.info(f"Number of prompts in the dataset: {len(self.prompts)}") + + self.active_indices = list(range(len(self.prompts))) + self.filter_threshold = filter_threshold + self.max_filter_percentage = max_filter_percentage + + @property + def util(self): + return self._util + + def __len__(self): + return len(self.active_indices) + + def __getitem__(self, idx): + # print(self.base_scores) + print( + f"CODEPromptDataset idx{idx}, use idx: {self.active_indices[idx]}, size: {len(self.ids)}" + ) + idx = self.active_indices[idx] + if hasattr(self, "base_scores"): return data_api.SequenceSample.from_default( ids=[self.ids[idx]], @@ -211,7 +348,7 @@ def filter(self, eval_scores: Dict[Hashable, float]): for pop_idx in sorted(indices_to_remove, reverse=True): self.active_indices.pop(pop_idx) logger.info( - f"Math prompt dataset DP rank {self.util.dp_rank} filtered" + f"Code prompt dataset DP rank {self.util.dp_rank} filtered" f" {len(indices_to_remove)} samples, {len(self.active_indices)} samples remain. " f"Original dataset size: {len(self.prompts)}. " f"Filter threshold: {self.filter_threshold}. " @@ -223,6 +360,7 @@ def filter(self, eval_scores: Dict[Hashable, float]): if not __name__ == "__main__": data_api.register_dataset("prompt", PromptDataset) data_api.register_dataset("math_prompt", MATHPromptDataset) + data_api.register_dataset("code_prompt", CODEPromptDataset) else: from transformers import AutoTokenizer @@ -238,3 +376,21 @@ def filter(self, eval_scores: Dict[Hashable, float]): max_length=512, dataset_path="/storage/openpsi/data/math/Qwen_RL_training_xss/0101/train_rl@0101_with_qwqsft-7b_score.jsonl", ) + + dataset = CODEPromptDataset( + data_api.DatasetUtility( + seed=0, + dp_rank=0, + world_size=1, + tokenizer=AutoTokenizer.from_pretrained( + "/storage/openpsi/models/Qwen__Qwen2-1.5B-Instruct/" + ), + ), + max_length=512, + dataset_path="/storage/datasets/codeparrot-apps-test.jsonl", + ) + data_generator = enumerate(dataset) + dataset_batch_counter, cur_sample = next(data_generator) + print( + f"size: {len(dataset)}, index: {dataset_batch_counter}, cur_sample: {cur_sample}" + ) diff --git a/realhf/impl/model/backend/mock_train.py b/realhf/impl/model/backend/mock_train.py index 8e6ff36ed7..541cae574b 100644 --- a/realhf/impl/model/backend/mock_train.py +++ b/realhf/impl/model/backend/mock_train.py @@ -32,7 +32,7 @@ class MockPipeTrainInstrSet(PipeTrainInstrSet): Used for testing only. """ - optimizer: torch.optim.Optimizer + optim: torch.optim.Optimizer def _exec_backward_pass( self, @@ -78,14 +78,19 @@ def _exec_optimizer_step( micro_batch_id: int, step_id: int, ): - self.optimizer.step() + self.optim.step() + + +class AdamWithLossScale(torch.optim.Adam): + def get_loss_scale(self) -> torch.Tensor: + return torch.tensor([1.0], device=constants.current_device()) class MockTrainEngine(model_api.PipelinableEngine): - def __init__(self, module: ReaLModel, optimizer: torch.optim.Optimizer): + def __init__(self, module: ReaLModel, optimizer: AdamWithLossScale): self.module = module - self.optimizer = optimizer + self.optim = optimizer self.inf_engine = PipelinableInferenceEngine(module) if constants.pipe_parallel_world_size() > 1: @@ -111,10 +116,10 @@ def train_batch( token_normalize_scope: str, version_steps: int, ): - self.optimizer.zero_grad() + self.optim.zero_grad() if constants.pipe_parallel_world_size() > 1: # Fusing the minibatched forward-backward in a pipeline training schedule. - instr_set = MockPipeTrainInstrSet(self.optimizer) + instr_set = MockPipeTrainInstrSet(self, self.optim) # NOTE: When training with pipeline parallel, num micro batches should be # larger than 2 x num_pipeline_stages to avoid idle time. return self.pipe_runner.train_batch( @@ -222,7 +227,7 @@ def _initialize( raise ValueError("MegatronTrainBackend only supports ReaLModel.") if self.optimizer_name == "adam": - optimizer = torch.optim.Adam(module.parameters(), **self.optimizer_config) + optimizer = AdamWithLossScale(module.parameters(), **self.optimizer_config) else: raise NotImplementedError( f"Optimizer {self.optimizer_name} not implemented for testing." diff --git a/realhf/impl/model/comm/data_transfer.py b/realhf/impl/model/comm/data_transfer.py deleted file mode 100644 index 00d735e819..0000000000 --- a/realhf/impl/model/comm/data_transfer.py +++ /dev/null @@ -1,344 +0,0 @@ -# Copyright 2025 Ant Group Inc. -# Copyright 2024 Wei Fu & Zhiyu Mei -# Licensed under the Apache License, Version 2.0 (the "License"). - -import dataclasses -import itertools -from typing import * - -import torch -import torch.distributed as dist - -from realhf.api.core.config import ModelName, ModelShardID -from realhf.api.core.data_api import SequenceSample -from realhf.base import constants, topology -from realhf.impl.model.comm.global_comm import filter_match_mwids -from realhf.impl.model.comm.param_realloc import pipeline_repartition_strategy - -BCAST_CHUNK_SIZE_BYTES = int(4e6) - - -@dataclasses.dataclass(unsafe_hash=True) -class DataTransferPair: - src: ModelName - src_dp_rank: int - dst: ModelName - dst_dp_rank: int - - -@dataclasses.dataclass -class DataTransferInfo: - # Groups for data transfer among model workers. - data_transfer_groups: Dict[DataTransferPair, dist.ProcessGroup] - data_transfer_src_ranks: Dict[DataTransferPair, int] - data_transfer_dst_ranks: Dict[DataTransferPair, List[int]] - - -def setup_data_transfer( - model_topos: Optional[Dict[str, topology.PipeModelDataParallelTopology]] = None, - msid2mwid: Optional[Dict[ModelShardID, int]] = None, - data_transfer_pairs: Optional[List[Tuple[ModelName, ModelName]]] = None, -) -> DataTransferInfo: - - # Stores the ranks given a (model_name, dp_rank) pair. - # These workers correspond to a complete set of model parameters sharded by TP+PP. - mw_dp_ranks: Dict[Tuple[ModelName, int], List[int]] = {} - - # Stores the dp_head (i.e., mp_rank=0, pp_rank=-1) ranks given a model_name. - mw_dp_head_ranks: Dict[ModelName, List[int]] = {} - - if model_topos is not None: - assert msid2mwid is not None - for model_name, topo in model_topos.items(): - mw_dp_head_ranks[model_name] = filter_match_mwids( - model_name, - topo, - msid2mwid, - pipe=topo.get_dim("pipe") - 1, - model=0, - ) - dp_size = topo.get_dim("data") - for dp_i in range(dp_size): - mw_dp_ranks[model_name, dp_i] = filter_match_mwids( - model_name, - topo, - msid2mwid, - data=dp_i, - ) - - data_transfer_groups, data_transfer_src_ranks = {}, {} - data_transfer_dst_ranks = {} - if data_transfer_pairs is not None: - for src, dst in data_transfer_pairs: - src_topo = model_topos[src] - dst_topo = model_topos[dst] - - # Construct all src-dst pairs, from any src dp rank to any dst dp rank. - # Note that a dp rank corresponds to multiple parameter shards (TP+PP), - # so each pair is a group-to-group communication. - # Since the models in the source group have duplicate data (TP+PP), - # we just use its "head" as the broadcast source, - # and broadcast to all the ranks in the destination group. - for src_dp, dst_dp in itertools.product( - range(src_topo.get_dim("data")), range(dst_topo.get_dim("data")) - ): - key = DataTransferPair( - src=src, src_dp_rank=src_dp, dst=dst, dst_dp_rank=dst_dp - ) - src_mw_rank = mw_dp_head_ranks[src][src_dp] - dst_mw_ranks = mw_dp_ranks[dst, dst_dp] - data_transfer_dst_ranks[key] = dst_mw_ranks - # The src and dst groups can be disjoint or overlapped. - # If they are disjoint, we need to include the src_mw_rank in the group. - # Otherwise, we only need to include the dst_mw_ranks. - if src_mw_rank not in dst_mw_ranks: - _ranks = [src_mw_rank] + dst_mw_ranks - else: - _ranks = dst_mw_ranks - data_transfer_groups[key] = topology.new_or_get_group( - _ranks, backend="nccl" if constants.use_cuda() else "gloo" - ) - data_transfer_src_ranks[key] = src_mw_rank - - return DataTransferInfo( - data_transfer_groups=data_transfer_groups, - data_transfer_src_ranks=data_transfer_src_ranks, - data_transfer_dst_ranks=data_transfer_dst_ranks, - ) - - -@dataclasses.dataclass -class DataTransferSenderStep: - rank: int - dst_ranks: List[int] - group: dist.ProcessGroup - key: str - ids: List[int] - - -@dataclasses.dataclass -class DataTransferReceiverStep: - rank: int - src: int - dst_ranks: List[int] - group: dist.ProcessGroup - key: str - ids: List[int] - - -def chunked_bcast(buf, src_rank, group, chunk_size_bytes=BCAST_CHUNK_SIZE_BYTES): - shape = buf.shape - chunk_size = chunk_size_bytes // buf.dtype.itemsize - buf = buf.flatten() - for i in range(0, buf.numel(), chunk_size): - s = slice(i, i + chunk_size) - dist.broadcast(buf[s], src=src_rank, group=group) - return buf.view(*shape) - - -def derive_data_transfer_plan( - keys: List[str], - global_ids: List[int], - consumer_name: ModelName, - consumer_mapping: Dict[int, List[int]], - producer_names: Dict[str, ModelName], - producer_mappings: Dict[Tuple[ModelName, str], Dict[int, List[int]]], - data_transfer_info: DataTransferInfo, -) -> List[DataTransferReceiverStep | DataTransferSenderStep]: - comm_plan = [] - - for k in keys: - producer_name = producer_names[k] - producer_mapping = producer_mappings[(producer_name, k)] - - # partition mapping starts from zero, which is different from buffer indices - repart_strat = pipeline_repartition_strategy(producer_mapping, consumer_mapping) - - for (dp_i, dp_j), comm_slots in repart_strat.items(): - if len(comm_slots) == 0: - continue - - group_key = DataTransferPair( - src=producer_name, - src_dp_rank=dp_i, - dst=consumer_name, - dst_dp_rank=dp_j, - ) - bcast_src = data_transfer_info.data_transfer_src_ranks[group_key] - group = data_transfer_info.data_transfer_groups[group_key] - dst_ranks = data_transfer_info.data_transfer_dst_ranks[group_key] - - ids = [global_ids[_i] for _i in comm_slots] - - for dst_rank in dst_ranks: - comm_plan.append( - DataTransferReceiverStep( - rank=dst_rank, - src=bcast_src, - dst_ranks=dst_ranks, - group=group, - key=k, - ids=ids, - ) - ) - - comm_plan.append( - DataTransferSenderStep( - rank=bcast_src, - dst_ranks=dst_ranks, - group=group, - key=k, - ids=ids, - ) - ) - - return comm_plan - - -def run_data_transfer( - comm_plan: List[DataTransferReceiverStep | DataTransferSenderStep], - meta_samples: Dict[int, SequenceSample], - storage: Dict[int, SequenceSample], - sent_worker_idx_table: Dict[int, Dict[str, Set[int]]], - received_worker_idx_table: Dict[int, Dict[str, Set[int]]], -) -> Tuple[Set[int], Set[str]]: - device = "cuda" if constants.use_cuda() else "cpu" - for step in comm_plan: - - if isinstance(step, DataTransferReceiverStep) and step.rank == dist.get_rank(): - ids = step.ids - if step.src == dist.get_rank(): - # The receiver is also a sender. - # We can directly use the data without comm. - for _id in step.ids: - if storage[_id].data[step.key] is not None: - storage[_id].data[step.key] = ( - storage[_id].data[step.key].to(device) - ) - else: - # If we have to receive remote data, we first check whether - # the data has been sent here in previous function calls. - # If so, just fetch it from the cache. - cached = all( - [ - set(step.dst_ranks).issubset( - set(received_worker_idx_table[_id][step.key]) - ) - for _id in ids - ] - ) - - metadata_cached = all( - [ - set(step.dst_ranks).issubset( - set(received_worker_idx_table[_id]["__metadata__"]) - ) - for _id in ids - ] - ) - if cached: - pass - else: - dtype = meta_samples[ids[0]].dtypes[step.key] - total_len = sum( - sum(meta_samples[_id].seqlens[step.key][0]) for _id in ids - ) - trailing_shape = meta_samples[ids[0]].trailing_shapes[step.key] - - # Receive data if it is not None. - if trailing_shape is not None: - buf = torch.zeros( - (total_len, *trailing_shape), - dtype=dtype, - device=constants.current_device(), - ) - chunked_bcast(buf, step.src, step.group) - else: - buf = None - - # Receive metadata if not cached. - if not metadata_cached: - metadatas = [{} for _ in step.ids] - dist.broadcast_object_list( - metadatas, src=step.src, group=step.group - ) - - # Mark that the data has been received. - for _id in ids: - received_worker_idx_table[_id][step.key].union(step.dst_ranks) - received_worker_idx_table[_id]["__metadata__"].union( - step.dst_ranks - ) - - # Split the received data and put it into the storage. - offset = 0 - for _id, metadata in zip(ids, metadatas): - seqlens = meta_samples[_id].seqlens[step.key] - assert len(seqlens) == 1 - seqlen = sum(seqlens[0]) - if buf is not None: - vs = buf[offset : offset + seqlen] - else: - vs = None - offset = offset + seqlen - with SequenceSample.disable_validation(): - s = SequenceSample( - keys=[step.key], - dtypes={step.key: vs.dtype if vs is not None else None}, - trailing_shapes={ - step.key: vs.shape[1:] if vs is not None else None - }, - ids=[_id], - seqlens={step.key: seqlens}, - data={step.key: vs}, - metadata=metadata, - ) - if _id in storage: - storage[_id].update_(s) - else: - storage[_id] = s - - if isinstance(step, DataTransferSenderStep) and step.rank == dist.get_rank(): - # Similar to the receiver, we first check whether the data has been sent to all destinations. - cached = all( - [ - set(step.dst_ranks).issubset( - set(sent_worker_idx_table[_id][step.key]) - ) - for _id in step.ids - ] - ) - metadata_cached = all( - [ - set(step.dst_ranks).issubset( - set(sent_worker_idx_table[_id]["__metadata__"]) - ) - for _id in step.ids - ] - ) - if cached: - pass - else: - # If not cached, we fetch the data from the storage and send it to all destinations. - for _id in step.ids: - if storage[_id].data[step.key] is not None: - storage[_id].data[step.key] = ( - storage[_id].data[step.key].to(device) - ) - if all([storage[_id].data[step.key] is not None for _id in step.ids]): - vs = torch.cat( - [storage[_id].data[step.key] for _id in step.ids], - dim=0, - ) - chunked_bcast(vs, step.rank, step.group) - - if not metadata_cached: - dist.broadcast_object_list( - [storage[_id].metadata for _id in step.ids], - src=step.rank, - group=step.group, - ) - - for _id in step.ids: - sent_worker_idx_table[_id][step.key].union(step.dst_ranks) - sent_worker_idx_table[_id]["__metadata__"].union(step.dst_ranks) diff --git a/realhf/impl/model/interface/code_rw_interface.py b/realhf/impl/model/interface/code_rw_interface.py new file mode 100644 index 0000000000..2f83acae97 --- /dev/null +++ b/realhf/impl/model/interface/code_rw_interface.py @@ -0,0 +1,429 @@ +# Copyright 2025 Ant Group Inc. +import collections +import dataclasses +import html +import json +import os +import re +import xml.etree.ElementTree as ET +from ast import parse +from concurrent.futures import ThreadPoolExecutor, as_completed + +import colorama +import numpy as np +import torch +import torch.distributed as dist + +import realhf.api.core.model_api as model_api +import realhf.base.logging as logging +from functioncall.code.verify import code_verify +from realhf.api.core.data_api import SequenceSample, load_hf_tokenizer +from realhf.base import constants + +logger = logging.getLogger("Packed Reward Modeling Interface", "benchmark") + + +class CodeVerifierException(Exception): + pass + + +def extract_python_code(text, min_length=20, strict_syntax=True): + code_pattern = r"(?i)```(?:python|py)?\s*\n?(.*?)\n?```" + code_blocks = re.findall(code_pattern, text, re.DOTALL) + valid_blocks = [] + for block in code_blocks: + clean_block = block.strip() + if len(clean_block) < min_length: + continue + + # verify code syntax + if strict_syntax: + try: + parse(clean_block, mode="exec") + except (SyntaxError, IndentationError): + continue + + valid_blocks.append(clean_block) + + if not valid_blocks: + return None + # return the last code block + return valid_blocks[-1] + + +def check_with_elementtree(text): + def escape_between_tags(text, tags=["think", "answer"]): + """转义标签之间的内容，但保留标签本身.""" + # 构建标签模式 + tag_pattern = "|".join(tags) + parts = [] + current_pos = 0 + + # 匹配开始和结束标签 + pattern = f"]*>" + + for match in re.finditer(pattern, text): + # 添加标签之前的内容（需要转义） + if current_pos < match.start(): + parts.append(html.escape(text[current_pos : match.start()])) + + # 添加标签本身（不转义） + parts.append(match.group()) + current_pos = match.end() + + # 添加最后剩余的内容 + if current_pos < len(text): + parts.append(html.escape(text[current_pos:])) + + return "".join(parts) + + text = escape_between_tags(text) + if not text.strip().startswith(""): + text = "" + text + try: + xml_text = f"{text}" + x = ET.fromstring(xml_text) + if x.text is not None and x.text.strip() != "": + return False, f"Error: extra content before . {x.text}" + if len(x) != 2: + return False, f"Error: there are {len(x)} tags." + if x[0].tag is None or x[0].tag != "think": + return False, f"Error: tag is missing. {x[0].tag}" + if x[0].tail is not None and x[0].tail.strip() != "": + return ( + False, + f"Error: extra content between and . {x[0].tail}", + ) + if x[1].tag is None or x[1].tag != "answer": + return False, f"Error: tag is missing. {x[1].tag}" + if x[1].tail is not None and x[1].tail.strip() != "": + return False, f"Error: extra content after , {x[1].tail}" + + return True, x[1].text if x[1].text is not None else "" + except ET.ParseError as e: + return False, f"Error: XML格式错误, {str(e)}" + + +def retokenize( + task, + tokenizer, + packed_input_ids, + input_cu_seqlens, + prompts, + prompt_cu_seqlens, + query_ids, + check_xml_format=False, + do_eval=False, +): + input_ids = [ + packed_input_ids[start:end] + for start, end in zip(input_cu_seqlens[:-1], input_cu_seqlens[1:]) + ] + prompt_ids = [ + prompts[start:end] + for start, end in zip(prompt_cu_seqlens[:-1], prompt_cu_seqlens[1:]) + ] + seq_strs = tokenizer.batch_decode( + input_ids, clean_up_tokenization_spaces=False, skip_special_tokens=True + ) + prompt_strs = tokenizer.batch_decode( + prompt_ids, clean_up_tokenization_spaces=False, skip_special_tokens=True + ) + # query_id_strs = query_ids + query_id_strs = [query_id.split("@")[0] for query_id in query_ids] + + format_rewards = [] + + queryid_to_results = collections.defaultdict(list) + # 8 processes on each node, with 10 subprocesses each + if do_eval == True: + _answers = [ + seq_str.split(prompt_str)[1] + for seq_str, prompt_str in zip(seq_strs, prompt_strs) + ] + + codes = [extract_python_code(_answer) for _answer in _answers] + logger.info( + f"code_rw_interface, size: {len(query_id_strs)}, valid code size: {len(codes)}, query_id_0: {query_id_strs[0]}" + ) + format_rewards = code_verify(codes, query_id_strs) + + if check_xml_format: + with ThreadPoolExecutor(max_workers=22) as executor: + futures = [ + executor.submit(check_with_elementtree, answer_str) + for answer_str in _answers + ] + # xml_rewards = [] + for idx, future in enumerate(futures): + xml_reward, _ = future.result() + # xml_rewards.append(xml_reward) + if xml_reward == 1 and format_rewards[idx] == 0: + format_rewards[idx] = -0.8 + elif xml_reward == 0 and format_rewards[idx] == 0: + format_rewards[idx] = -1 + + for query_id_str, format_reward in zip(query_id_strs, format_rewards): + if query_id_str not in queryid_to_results: + queryid_to_results[query_id_str] = [] + queryid_to_results[query_id_str].append(format_reward) + else: + for query_id_str in query_id_strs: + if query_id_str not in queryid_to_results: + queryid_to_results[query_id_str] = [] + queryid_to_results[query_id_str].append(0) + format_rewards.append(0) + + return format_rewards, prompt_strs, prompt_ids, seq_strs, queryid_to_results + + +@dataclasses.dataclass +class PackedCodeRewardInterface(model_api.ModelInterface): + + enable_save: bool = False + tokenizer_path: str = "/storage/models/Qwen__Qwen2.5-1.5B" + output_scaling: float = 1.0 + rm_output_scaling: float = 1.0 + rm_output_bias: float = 0.0 + output_bias: float = 0.0 + loss_fun = torch.nn.CrossEntropyLoss(reduction="none") + max_sync_length: int = 2048 + rw_type: str = "sparse" + task: str = "code" # math or countdown or code + check_xml_format: bool = False + post_process: str = "sigmoid" + group_size: int = 1 + check_verifier_status: bool = False + + _call_count: int = 0 + + def __post_init__(self): + self.tokenizer = load_hf_tokenizer(self.tokenizer_path) + logger.info(f"rm_output_scaling: {self.rm_output_scaling}") + logger.info(f"rm_output_bias: {self.rm_output_bias}") + logger.info(f"output_scaling: {self.output_scaling}") + logger.info(f"output_bias: {self.output_bias}") + logger.info(f"max_sync_length: {self.max_sync_length}") + logger.info(f"rw_type: {self.rw_type}") + logger.info(f"post_process: {self.post_process}") + + while True: + gen_file_path = os.path.join( + constants.LOG_ROOT, + constants.experiment_name(), + constants.trial_name(), + "generated", + f"v{self._call_count}r{dist.get_rank()}.txt", + ) + if os.path.exists(gen_file_path): + self._call_count += 1 + else: + break + logger.info(f"call_count: {self._call_count}") + + def inference( + self, + model: model_api.Model, + data_: SequenceSample, + mb_spec, + ) -> SequenceSample: + + packed_input_ids: torch.Tensor = data_.data["packed_input_ids"].squeeze() + + input_seqlens = torch.tensor(data_.seqlens["packed_input_ids"]).view(-1) + input_cu_seqlens = torch.nn.functional.pad( + input_seqlens.cumsum(0), (1, 0) + ).int() + + packed_prompts = data_.data["packed_prompts"] + prompts = [] + prompt_seqlens = [] + offset = 0 + for x in data_.seqlens["packed_prompts"]: + prompts += [packed_prompts[offset : offset + x[0]]] * self.group_size + offset += x[0] + prompt_seqlens.extend(x * self.group_size) + + assert offset == sum(x[0] for x in data_.seqlens["packed_prompts"]) + # non_packed_prompts = copy.deepcopy(prompts) + prompts = torch.cat(prompts) + prompt_seqlens = torch.tensor(prompt_seqlens).view(-1) + prompt_cu_seqlens = torch.nn.functional.pad( + prompt_seqlens.cumsum(0), (1, 0) + ).int() + + query_ids = [data_id for data_id in data_.ids for _ in range(self.group_size)] + + format_rewards, prompt_strs, prompt_ids, seq_strs, queryid_to_results = ( + retokenize( + self.task, + self.tokenizer, + packed_input_ids, + input_cu_seqlens, + prompts, + prompt_cu_seqlens, + query_ids, + check_xml_format=self.check_xml_format, + do_eval=constants.is_last_pipe_stage(), + ) + ) + + assert self.rw_type == "sparse" + dense_scores = torch.zeros_like(packed_input_ids).float() + scores = torch.FloatTensor(format_rewards).to(packed_input_ids.device) + scores[scores == 0] = -1 + + if len(scores) == 0: + return None + + assert dense_scores.shape == packed_input_ids.shape + scores = ( + scores.to(packed_input_ids.device) - self.output_bias + ) * self.output_scaling + + logger.info(f"Code reward logging info @v{model.version.global_step}") + logger.info( + f"before: Format success rate: {torch.FloatTensor(format_rewards).mean().item()}" + ) + logger.info( + f"number of samples: {len(scores)}, {scores.shape}, group_size: {self.group_size}" + ) + + gen_file_path = os.path.join( + constants.LOG_ROOT, + constants.experiment_name(), + constants.trial_name(), + "generated", + f"v{self._call_count}r{dist.get_rank()}.txt", + ) + os.makedirs(os.path.dirname(gen_file_path), exist_ok=True) + logger.info(f"Generated samples and rewards will be dumped to: {gen_file_path}") + with open(gen_file_path, "w") as _f: + for idx, (score, prompt_str, seq_str) in enumerate( + zip(scores, prompt_strs, seq_strs) + ): + info = "\n".join( + [ + f"idx: {idx} / {len(scores)}", + f"reward is {score.item()}, prompt is {colorama.Fore.YELLOW + colorama.Style.DIM}{prompt_str}{colorama.Style.RESET_ALL}", + f"sequence is: {colorama.Fore.YELLOW + colorama.Style.DIM}{seq_str.split(prompt_str)[1]}{colorama.Style.RESET_ALL}.", + ] + ) + _f.write(info + "\n") + + gen_file_path = os.path.join( + constants.LOG_ROOT, + constants.experiment_name(), + constants.trial_name(), + "generated_jsonl", + f"v{self._call_count}r{dist.get_rank()}.jsonl", + ) + os.makedirs(os.path.dirname(gen_file_path), exist_ok=True) + logger.info(f"Generated samples and rewards will be dumped to: {gen_file_path}") + with open(gen_file_path, "w") as _f: + for idx, (score, prompt_str, seq_str) in enumerate( + zip(scores, prompt_strs, seq_strs) + ): + _f.write( + json.dumps( + { + "prompt": prompt_str, + "generated": seq_str.split(prompt_str)[1], + "reward": score.item(), + }, + ensure_ascii=False, + ) + + "\n" + ) + + logger.info( + f"Format success rate: {torch.FloatTensor(format_rewards).mean().item()}" + ) + + pass_at_k = np.mean( + [sum([xx == 1 for xx in x]) > 0 for x in queryid_to_results.values()] + ) + avg_num_samples = np.mean([len(x) for x in queryid_to_results.values()]) + logger.info(f"pass@k: {pass_at_k}, num_samples: {avg_num_samples}") + + logger.info(f"number of samples: {len(scores)}, {scores.shape}") + + logger.info(f"reward: {sum(scores) / len(scores)}") + + train_pass_monitor_file_path = os.path.join( + constants.LOG_ROOT, + constants.experiment_name(), + constants.trial_name(), + "training_monitor", + f"v{self._call_count}r{dist.get_rank()}.jsonl", + ) + os.makedirs(os.path.dirname(train_pass_monitor_file_path), exist_ok=True) + logger.info( + f"pass monitor result will be dumped to: {train_pass_monitor_file_path}" + ) + with open(train_pass_monitor_file_path, "w") as monitor_file: + for key, value in queryid_to_results.items(): + pass1 = sum(value) / len(value) + pass8 = int(sum(value) > 0) + monitor_file.write( + json.dumps( + {"query_id": key, "pass1": pass1, "pass8": pass8}, + ensure_ascii=False, + ) + + "\n" + ) + + self._call_count += 1 + + if scores.dtype != torch.float32: + scores = scores.to(torch.float32) + if dense_scores.dtype != torch.float32: + dense_scores = dense_scores.to(torch.float32) + + res = SequenceSample( + keys=["rewards", "dense_rewards"], + trailing_shapes=dict(rewards=(), dense_rewards=()), + dtypes=dict(rewards=torch.float32, dense_rewards=torch.float32), + ids=data_.ids, + seqlens=dict( + rewards=[ + torch.tensor([1 for _ in range(len(x))], dtype=torch.int32) + for x in data_.seqlens["packed_input_ids"] + ], + dense_rewards=data_.seqlens["packed_input_ids"], + ), + data=dict(rewards=scores, dense_rewards=dense_scores), + ) + + # record rewards for each piece of data + avg_scores = [] + offset = 0 + for i in range(data_.bs): + score_lis = scores[ + offset : offset + len(data_.seqlens["packed_input_ids"][i]) + ] + avg_scores.append(score_lis.mean().item()) + offset += len(data_.seqlens["packed_input_ids"][i]) + assert offset == sum(len(x) for x in data_.seqlens["packed_input_ids"]) + + res.metadata["scores"] = avg_scores + + if self.check_verifier_status: + avg_score = torch.tensor( + np.mean(avg_scores), device=constants.current_device() + ) + dist.all_reduce( + avg_score, op=dist.ReduceOp.SUM, group=constants.parallelism_group() + ) + avg_score /= constants.parallelism_group_size() + avg_score = avg_score.item() + minimal_score = (-1 - self.output_bias) * self.rm_output_scaling + + if avg_score <= minimal_score or np.isclose(avg_score, minimal_score): + raise CodeVerifierException( + "All rewards are at minimal value. Probably there are something wrong with the verifier!" + ) + return res + + +model_api.register_interface("rw_code", PackedCodeRewardInterface) diff --git a/realhf/impl/model/interface/math_parser.py b/realhf/impl/model/interface/math_parser.py index d9af254963..eef806a8cf 100644 --- a/realhf/impl/model/interface/math_parser.py +++ b/realhf/impl/model/interface/math_parser.py @@ -4,11 +4,11 @@ import os import signal import subprocess +import time +import traceback import uuid from typing import * -import requests - from realhf.base import logging from realhf.base.constants import parallelism_rank @@ -112,10 +112,9 @@ def parse_line(prompt_str, generated, query_id): def parse_lines_in_parallel( - prompt_strs: List, generateds: List, query_ids: List, - max_workers: int, + max_workers=22, check_xml_format=False, ) -> List: global id2info @@ -124,12 +123,11 @@ def parse_lines_in_parallel( id2info = loadJson(os.environ["REAL_MATH_METADATA_PATH"]) except KeyError as e: raise KeyError("The json file REAL_MATH_METADATA_PATH is not set") from e - assert len(prompt_strs) == len(generateds) == len(query_ids), ( - len(prompt_strs), + assert len(generateds) == len(query_ids), ( len(generateds), len(query_ids), ) - bs = len(prompt_strs) + bs = len(query_ids) mbs = (bs + max_workers - 1) // max_workers tmp_ids = [] @@ -183,7 +181,7 @@ def parse_lines_in_parallel( except ProcessLookupError: pass - labels = [0 for _ in prompt_strs] + labels = [0 for _ in query_ids] for i, (tmp_id, query_indices) in enumerate(zip(tmp_ids, all_query_indices)): try: with open(f"/tmp/{tmp_id}-output.jsonl", "r") as f: @@ -214,7 +212,6 @@ def parse_lines_in_parallel( print( parse_lines_in_parallel( - [sample["prompt"] for _ in range(100)], # [sample["answer_in_box"][0] for _ in range(50)] + [sample["answer_in_box"][1] for _ in range(50)], [sample["answer"]] * 100, [sample["query_id"] for _ in range(100)], diff --git a/realhf/impl/model/interface/math_rw_interface.py b/realhf/impl/model/interface/math_rw_interface.py index d20d14539f..d8425a0975 100644 --- a/realhf/impl/model/interface/math_rw_interface.py +++ b/realhf/impl/model/interface/math_rw_interface.py @@ -23,6 +23,7 @@ import realhf.api.core.model_api as model_api import realhf.base.logging as logging +from functioncall.math.verify import math_verify from realhf.api.core.data_api import SequenceSample, load_hf_tokenizer from realhf.base import constants from realhf.base.constants import data_parallel_group, data_parallel_world_size @@ -32,6 +33,9 @@ logger = logging.getLogger("Packed Reward Modeling Interface", "benchmark") +ENABLE_FUNCTION_CALL = True if os.getenv("FUNCTIONCALL_SERVICE_DOMAIN", "") else False +math_verify_call = math_verify if ENABLE_FUNCTION_CALL else parse_lines_in_parallel + class MathVerifierException(Exception): pass @@ -127,13 +131,7 @@ def retokenize( for seq_str, prompt_str in zip(seq_strs, prompt_strs) ] - format_rewards = parse_lines_in_parallel( - prompt_strs, - _answers, - query_id_strs, - max_workers=22, - check_xml_format=False, - ) + format_rewards = math_verify_call(_answers, query_id_strs) if check_xml_format: with ThreadPoolExecutor(max_workers=22) as executor: futures = [ @@ -181,8 +179,6 @@ class PackedMathRewardInterface(model_api.ModelInterface): group_size: int = 1 check_verifier_status: bool = False - _call_count: int = 0 - def __post_init__(self): self.tokenizer = load_hf_tokenizer(self.tokenizer_path) logger.info(f"rm_output_scaling: {self.rm_output_scaling}") @@ -193,20 +189,6 @@ def __post_init__(self): logger.info(f"rw_type: {self.rw_type}") logger.info(f"post_process: {self.post_process}") - while True: - gen_file_path = os.path.join( - constants.LOG_ROOT, - constants.experiment_name(), - constants.trial_name(), - "generated", - f"v{self._call_count}r{dist.get_rank()}.txt", - ) - if os.path.exists(gen_file_path): - self._call_count += 1 - else: - break - logger.info(f"call_count: {self._call_count}") - def inference( self, model: model_api.Model, @@ -272,92 +254,98 @@ def inference( f"before: Format success rate: {torch.FloatTensor(format_rewards).mean().item()}" ) logger.info(f"number of samples: {len(scores)}, {scores.shape}") + if constants.is_last_pipe_stage(): - gen_file_path = os.path.join( - constants.LOG_ROOT, - constants.experiment_name(), - constants.trial_name(), - "generated", - f"v{self._call_count}r{dist.get_rank()}.txt", - ) - os.makedirs(os.path.dirname(gen_file_path), exist_ok=True) - logger.info(f"Generated samples and rewards will be dumped to: {gen_file_path}") - with open(gen_file_path, "w") as _f: - for idx, (score, prompt_str, seq_str) in enumerate( - zip(scores, prompt_strs, seq_strs) - ): - info = "\n".join( - [ - f"idx: {idx} / {len(scores)}", - f"reward is {score.item()}, prompt is {colorama.Fore.YELLOW + colorama.Style.DIM}{prompt_str}{colorama.Style.RESET_ALL}", - f"sequence is: {colorama.Fore.YELLOW + colorama.Style.DIM}{seq_str.split(prompt_str)[1]}{colorama.Style.RESET_ALL}.", - ] - ) - _f.write(info + "\n") - - gen_file_path = os.path.join( - constants.LOG_ROOT, - constants.experiment_name(), - constants.trial_name(), - "generated_jsonl", - f"v{self._call_count}r{dist.get_rank()}.jsonl", - ) - os.makedirs(os.path.dirname(gen_file_path), exist_ok=True) - logger.info(f"Generated samples and rewards will be dumped to: {gen_file_path}") - with open(gen_file_path, "w") as _f: - for idx, (score, prompt_str, seq_str) in enumerate( - zip(scores, prompt_strs, seq_strs) - ): - _f.write( - json.dumps( - { - "prompt": prompt_str, - "generated": seq_str.split(prompt_str)[1], - "reward": score.item(), - }, - ensure_ascii=False, + gen_file_path = os.path.join( + constants.LOG_ROOT, + constants.experiment_name(), + constants.trial_name(), + "generated", + f"v{model.version.global_step}r{dist.get_rank()}.txt", + ) + + os.makedirs(os.path.dirname(gen_file_path), exist_ok=True) + logger.info( + f"Generated samples and rewards will be dumped to: {gen_file_path}" + ) + with open(gen_file_path, "w") as _f: + for idx, (score, prompt_str, seq_str) in enumerate( + zip(scores, prompt_strs, seq_strs) + ): + info = "\n".join( + [ + f"idx: {idx} / {len(scores)}", + f"reward is {score.item()}, prompt is {colorama.Fore.YELLOW + colorama.Style.DIM}{prompt_str}{colorama.Style.RESET_ALL}", + f"sequence is: {colorama.Fore.YELLOW + colorama.Style.DIM}{seq_str.split(prompt_str)[1]}{colorama.Style.RESET_ALL}.", + ] ) - + "\n" - ) + _f.write(info + "\n") - logger.info( - f"Format success rate: {torch.FloatTensor(format_rewards).mean().item()}" - ) + gen_file_path = os.path.join( + constants.LOG_ROOT, + constants.experiment_name(), + constants.trial_name(), + "generated_jsonl", + f"v{model.version.global_step}r{dist.get_rank()}.jsonl", + ) + os.makedirs(os.path.dirname(gen_file_path), exist_ok=True) + logger.info( + f"Generated samples and rewards will be dumped to: {gen_file_path}" + ) + with open(gen_file_path, "w") as _f: + for idx, (score, prompt_str, seq_str) in enumerate( + zip(scores, prompt_strs, seq_strs) + ): + _f.write( + json.dumps( + { + "prompt": prompt_str, + "generated": seq_str.split(prompt_str)[1], + "reward": score.item(), + }, + ensure_ascii=False, + ) + + "\n" + ) - pass_at_k = np.mean( - [sum([xx == 1 for xx in x]) > 0 for x in queryid_to_results.values()] - ) - avg_num_samples = np.mean([len(x) for x in queryid_to_results.values()]) - logger.info(f"pass@k: {pass_at_k}, num_samples: {avg_num_samples}") + logger.info( + f"Format success rate: {torch.FloatTensor(format_rewards).mean().item()}" + ) - logger.info(f"number of samples: {len(scores)}, {scores.shape}") + pass_at_k = np.mean( + [sum([xx == 1 for xx in x]) > 0 for x in queryid_to_results.values()] + ) + avg_num_samples = np.mean([len(x) for x in queryid_to_results.values()]) + logger.info(f"pass@k: {pass_at_k}, num_samples: {avg_num_samples}") - logger.info(f"reward: {sum(scores) / len(scores)}") + logger.info(f"number of samples: {len(scores)}, {scores.shape}") - train_pass_monitor_file_path = os.path.join( - constants.LOG_ROOT, - constants.experiment_name(), - constants.trial_name(), - "training_monitor", - f"v{self._call_count}r{dist.get_rank()}.jsonl", - ) - os.makedirs(os.path.dirname(train_pass_monitor_file_path), exist_ok=True) - logger.info( - f"pass monitor result will be dumped to: {train_pass_monitor_file_path}" - ) - with open(train_pass_monitor_file_path, "w") as monitor_file: - for key, value in queryid_to_results.items(): - pass1 = sum(value) / len(value) - pass8 = int(sum(value) > 0) - monitor_file.write( - json.dumps( - {"query_id": key, "pass1": pass1, "pass8": pass8}, - ensure_ascii=False, + logger.info(f"reward: {sum(scores) / len(scores)}") + + train_pass_monitor_file_path = os.path.join( + constants.LOG_ROOT, + constants.experiment_name(), + constants.trial_name(), + "training_monitor", + f"v{model.version.global_step}r{dist.get_rank()}.jsonl", + ) + os.makedirs(os.path.dirname(train_pass_monitor_file_path), exist_ok=True) + logger.info( + f"pass monitor result will be dumped to: {train_pass_monitor_file_path}" + ) + with open(train_pass_monitor_file_path, "w") as monitor_file: + for key, value in queryid_to_results.items(): + pass1 = sum(value) / len(value) + pass8 = int(sum(value) > 0) + monitor_file.write( + json.dumps( + {"query_id": key, "pass1": pass1, "pass8": pass8}, + ensure_ascii=False, + ) + + "\n" ) - + "\n" - ) - self._call_count += 1 + model.inc_version() if scores.dtype != torch.float32: scores = scores.to(torch.float32) @@ -407,6 +395,9 @@ def inference( raise MathVerifierException( "All rewards are at minimal value. Probably there are something wrong with the verifier!" ) + + if not constants.is_last_pipe_stage(): + return None return res diff --git a/realhf/impl/model/interface/ppo_interface.py b/realhf/impl/model/interface/ppo_interface.py index 760f2a0ae3..4a1456c18f 100644 --- a/realhf/impl/model/interface/ppo_interface.py +++ b/realhf/impl/model/interface/ppo_interface.py @@ -37,7 +37,7 @@ def get_score(prompt_ids, generated, query_ids, tokenizer): generated, clean_up_tokenization_spaces=False, skip_special_tokens=True ) query_id_strs = [query_id.split("@")[0] for query_id in query_ids] - return parse_lines_in_parallel(prompt_strs, seq_strs, query_id_strs, max_workers=22) + return parse_lines_in_parallel(seq_strs, query_id_strs) def topk(scores, gen_lengths, k) -> list: diff --git a/realhf/scheduler/client.py b/realhf/scheduler/client.py index cec32ff091..ea22cbcba9 100644 --- a/realhf/scheduler/client.py +++ b/realhf/scheduler/client.py @@ -123,8 +123,6 @@ def wait(self, timeout=None, **kwargs): def get_python3_path(): if cluster_spec.cluster_type == "ray": return subprocess.check_output(["which", "python3"]).decode("utf-8").strip() - if cluster_spec.name == "na132": - return "/usr/bin/python3" return "python3" @@ -162,7 +160,8 @@ def make(mode, expr_name, trial_name, **kwargs) -> SchedulerClient: from realhf.scheduler.slurm.client import SlurmSchedulerClient schedule_strategy = kwargs.get("schedule_strategy", "empty_first") - return SlurmSchedulerClient(expr_name, trial_name, schedule_strategy) + evaluator = kwargs.get("evaluator", None) + return SlurmSchedulerClient(expr_name, trial_name, schedule_strategy, evaluator) elif mode == "local": from realhf.scheduler.local.client import LocalSchedulerClient diff --git a/realhf/scheduler/evaluator.py b/realhf/scheduler/evaluator.py new file mode 100644 index 0000000000..997a0fe878 --- /dev/null +++ b/realhf/scheduler/evaluator.py @@ -0,0 +1,246 @@ +import json +import os +import re +import subprocess +import time +from typing import Dict + +import wandb + +import realhf.api.core.system_api as config_pkg +from realhf.base import cluster, constants, logging + +logger = logging.getLogger("AutomaticEvaluator", "colored") + + +class AutomaticEvaluator: + + def __init__( + self, + config: config_pkg.AutomaticEvaluator, + wandb_config: config_pkg.WandBConfig, + ): + self.__running_processes: Dict[int, subprocess.Popen] = {} + self.__start_time = {} + self.__done_steps = [] + self.__wandb_log_steps = [] + self.__pending_ckpts = {} + self.__image = config.eval_job_image or cluster.spec.gpu_image + self.__data_names = config.data_names + self.__max_gen_tokens = config.max_gen_tokens + self.__max_concurrent_jobs = config.max_concurrent_jobs + self.__prompt_type = config.prompt_type + self.__wandb_config = wandb_config + self.__config = config + self.__wandb_inited = False + + # Check evaluated checkpoints by logs + former_output_dir = os.path.join( + constants.LOG_ROOT, + constants.experiment_name(), + constants.trial_name(), + "eval_output", + ) + if os.path.exists(former_output_dir): + for log_dir in os.listdir(former_output_dir): + match = re.match(r"globalstep(\d+)", log_dir) + if not match: + continue + global_step = int(match.groups()[0]) + self.__done_steps.append(global_step) + + logger.info( + f"Initializing AutomaticEvaluator: \n" + f"eval_job_image: {config.eval_job_image}\n" + f"data_names: {config.data_names}\n" + f"max_gen_tokens: {config.max_gen_tokens}\n" + f"max_concurrent_jobs: {config.max_concurrent_jobs}\n" + f"Existing eval outputs for global steps: {self.__done_steps}" + ) + if self.__config.initial_checkpoint_path and 0 not in self.__done_steps: + self.__pending_ckpts[0] = self.__config.initial_checkpoint_path + + if not cluster.spec.cluster_type == "slurm": + raise NotImplementedError( + "Currently only support automatic evaluation for slurm" + ) + + def __lazy_wandb_init(self): + # Initializing wandb for evaluator + wandb.login() + wandb.init( + mode=self.__wandb_config.mode, + entity=self.__wandb_config.entity, + project=self.__wandb_config.project or constants.experiment_name(), + name=self.__wandb_config.name or f"{constants.trial_name()}_eval", + job_type=self.__wandb_config.job_type, + group=self.__wandb_config.group + or f"{constants.experiment_name()}_{constants.trial_name()}", + notes=self.__wandb_config.notes, + tags=self.__wandb_config.tags, + config=self.__wandb_config.config, + dir=os.path.join( + constants.LOG_ROOT, constants.experiment_name(), constants.trial_name() + ), + force=True, + id=f"{constants.experiment_name()}_{constants.trial_name()}_eval", + resume="allow", + settings=wandb.Settings(start_method="fork"), + ) + + def __check_new_ckpts(self): + save_path = os.path.join( + constants.MODEL_SAVE_ROOT, + constants.experiment_name(), + constants.trial_name(), + "actor", + ) + if not os.path.exists(save_path): + return + for ckpt_dir in os.listdir(save_path): + match = re.match(r"epoch(\d+)epochstep(\d+)globalstep(\d+)", ckpt_dir) + if not match: + continue + _, _, global_step = map(int, match.groups()) + if not global_step in ( + list(self.__running_processes.keys()) + + list(self.__pending_ckpts.keys()) + + self.__done_steps + ): + abs_ckpt_dir = os.path.join(save_path, ckpt_dir) + logger.info( + f"Found new checkpoint (globalstep{global_step}) at {abs_ckpt_dir}" + ) + self.__pending_ckpts[global_step] = os.path.join( + save_path, abs_ckpt_dir + ) + + def __check_and_maybe_submit_jobs(self): + for global_step, process in self.__running_processes.items(): + result = process.poll() + if not result is None: + self.__done_steps.append(global_step) + start_time = self.__start_time[global_step] + logger.info( + f"Evaluation of checkpoint (globalstep{global_step}) is done, returncode={process.returncode}, " + f"time passed {time.perf_counter() - start_time:.3f} s." + ) + + for done in self.__done_steps: + if done in self.__running_processes: + self.__running_processes.pop(done) + self.__start_time.pop(done) + + submitted = [] + # Jobs should be submitted by the order of global steps + ordered_steps = sorted(self.__pending_ckpts.keys()) + for global_step in ordered_steps: + ckpt_path = self.__pending_ckpts[global_step] + if len(self.__running_processes) >= self.__max_concurrent_jobs: + return + self.__submit_one(global_step, ckpt_path) + submitted.append(global_step) + + for global_step in submitted: + self.__pending_ckpts.pop(global_step) + + def __submit_one(self, global_step, ckpt_path): + output_path = self.eval_output_path(global_step) + os.makedirs(output_path, exist_ok=True) + log_file = open(os.path.join(output_path, "output.log"), "w") + if cluster.spec.cluster_type == "slurm": + cmd = self.slurm_eval_cmd(global_step, ckpt_path) + else: + raise NotImplementedError( + "AutomaticEvaluator does only support slurm cluster." + ) + + logger.info( + f"Submitting evaluation job of checkpoint at {ckpt_path} (globalstep{global_step}), " + f"command: {cmd}" + ) + self.__running_processes[global_step] = subprocess.Popen( + cmd, + stdout=log_file, + stderr=log_file, + shell=True, + ) + self.__start_time[global_step] = time.perf_counter() + + def __maybe_parse_and_log_to_wandb(self): + # Note that after recover, all previous done steps will be + # logged again in case some data points are missing. + # If the data point is already logged, wandb will raise + # a warning. + to_log = list( + filter( + lambda x: x not in self.__wandb_log_steps, + ( + self.__done_steps + + list(self.__running_processes.keys()) + + list(self.__pending_ckpts.keys()) + ), + ) + ) + while to_log: + # The wandb should always log the minimal global step + # whose eval job has been submitted but not logged. + # If this minimal step is not logged, other steps should wait. + global_step = min(to_log) + result_path = os.path.join( + self.eval_output_path(global_step), + f"math_eval_{self.__max_gen_tokens}", + f"aggregate_parallel_{self.__prompt_type}.json", + ) + if not os.path.exists(result_path): + break + + if not self.__wandb_inited: + self.__lazy_wandb_init() + self.__wandb_inited = True + + try: + with open(result_path, "r") as fp: + data = json.load(fp) + except json.JSONDecodeError: + logger.warning(f"JSON decoding for eval result in {result_path} failed") + continue + except FileNotFoundError: + logger.warning( + f"{result_path} not found, but the eval job is done. " + "Maybe the eval job abnormally exited and did not output the result." + ) + continue + wandb_data = {} + for data_name, d in data.items(): + for k, v in d.items(): + wandb_data[f"{data_name}_{k}"] = v + wandb.log(wandb_data, step=global_step) + logger.info(f"Logging eval result {wandb_data} to step {global_step}") + self.__wandb_log_steps.append(global_step) + to_log.remove(global_step) + + def step(self): + self.__check_new_ckpts() + self.__check_and_maybe_submit_jobs() + self.__maybe_parse_and_log_to_wandb() + + def slurm_eval_cmd(self, global_step, ckpt_path): + slurm_job_name = f"{constants.experiment_name()}_{constants.trial_name()}:eval_globalstep{global_step}" + cmd = ( + f"srun --mpi=pmi2 -J {slurm_job_name} --ntasks=1 --cpus-per-task=128 --gres=gpu:8 --mem-per-cpu=12G " + f"singularity exec --nv --pid --writable-tmpfs --bind /storage:/storage " + f"{self.__image} " + f"bash ./evaluation/sh/install_deps_and_eval.sh {ckpt_path} {self.eval_output_path(global_step)} " + f"{self.__data_names} {self.__max_gen_tokens} {self.__prompt_type}" + ) + return cmd + + def eval_output_path(self, global_step): + return os.path.join( + constants.LOG_ROOT, + constants.experiment_name(), + constants.trial_name(), + "eval_output", + f"globalstep{global_step}", + ) diff --git a/realhf/scheduler/slurm/client.py b/realhf/scheduler/slurm/client.py index 5cb1ae3dca..43234e81dd 100644 --- a/realhf/scheduler/slurm/client.py +++ b/realhf/scheduler/slurm/client.py @@ -14,6 +14,7 @@ from realhf.base.cluster import spec as cluster_spec from realhf.base.constants import SLURM_LOCK_FILE_NAME as LOCK_FILE_NAME from realhf.scheduler.client import JobException, JobInfo, JobState, SchedulerClient +from realhf.scheduler.evaluator import AutomaticEvaluator from realhf.scheduler.slurm.utils import ( SlurmLaunchInfo, SlurmResource, @@ -25,12 +26,19 @@ SCHEDULING_RETRY_INTERVAL_SECONDS = 30 SCHEDULING_TIMEOUT_MAX_SECONDS = 3600 * 24 +SCHEDULER_WAIT_CHECK_TIME_INTERVAL = 5 class SlurmSchedulerClient(SchedulerClient): """Uses Slurm (https://slurm.schedmd.com/overview.html).""" - def __init__(self, expr_name, trial_name, schedule_strategy): + def __init__( + self, + expr_name: str, + trial_name: str, + schedule_strategy: str, + evaluator: Optional[AutomaticEvaluator], + ): super().__init__(expr_name, trial_name) self.__schedule_strategy = schedule_strategy @@ -40,6 +48,7 @@ def __init__(self, expr_name, trial_name, schedule_strategy): self.__submission_counter = defaultdict(int) self.__wprocs_counter = defaultdict(int) + self.__evaluator = evaluator def submit(self, worker_type, cmd, **kwargs): self.submit_array(worker_type, cmd, count=1, **kwargs) @@ -118,7 +127,7 @@ def __resolve_multiprog_file(self, launch_info: SlurmLaunchInfo): ) wrap_cmd = "singularity exec " if cluster_spec.name == "na132": - wrap_cmd += "--pid " + wrap_cmd += "--pid --no-home --writable-tmpfs " if cluster_spec.gpu_type == "tesla": wrap_cmd += "--nv " if len(launch_info.env_vars) > 0: @@ -245,6 +254,8 @@ def wait( if len(left) < num_jobs_left: num_jobs_left = len(left) logger.info(f"Waiting for {num_jobs_left} jobs.") + if self.__evaluator is not None: + self.__evaluator.step() if deadline is not None and time.time() > deadline: raise TimeoutError( f"Timeout waiting for {self.run_name}: {', '.join(sorted(left))}" @@ -276,7 +287,7 @@ def wait( left.remove(job_slurm_name) if update: self.__committed_jobs.pop(job_slurm_name) - time.sleep(2) + time.sleep(SCHEDULER_WAIT_CHECK_TIME_INTERVAL) def __update_all(self): states = [] diff --git a/realhf/system/buffer.py b/realhf/system/buffer.py index 6f375fa8ee..896baea957 100644 --- a/realhf/system/buffer.py +++ b/realhf/system/buffer.py @@ -17,32 +17,6 @@ logger = logging.getLogger("buffer") -def _extract_intervals(arr): - if len(arr) == 0: - return [] - - # Initialize the list to hold the intervals - intervals = [] - - # Start of the first interval - start = arr[0] - - for i in range(1, len(arr)): - # Check if the current element is not contiguous with the previous one - if arr[i] != arr[i - 1] + 1: - # End of the current interval - end = arr[i - 1] - # Add the interval as a tuple - intervals.append((start, end + 1)) - # Start a new interval - start = arr[i] - - # Add the last interval - intervals.append((start, arr[-1] + 1)) - - return intervals - - class BufferFull(Exception): pass @@ -147,6 +121,7 @@ def __init__( rpcs: List[dfg.MFCDef], max_size: int, ): + self.rpcs = rpcs self._lock = asyncio.Condition(asyncio.Lock()) # Buffer indicators, should be locked by self._lock. @@ -269,6 +244,62 @@ def put_batch_synced(self, samples: List[SequenceSample]): ) return indices + async def put_batch(self, samples: List[SequenceSample]): + n = len(samples) + + if n == 0: + return np.array([], dtype=np.int64) + + async with self._lock: + self._assert_valid_indicator() + + indices = np.where(self._is_empty)[0][:n] + + if len(indices) < n: + raise BufferFull( + "You are probably using a large dataset. " + "The default buffer size 1M is not large enough. " + "Please set a larger buffer size by setting " + "the environment variable, e.g., REAL_MASTER_BUFFER_SIZE=3000000." + ) + self._is_empty[indices] = False + self._is_being_put[indices] = True + + self.__buffer.put_batch(indices, samples) + + # Set a slight difference in birth time to let the order + # be deterministic. + self._birth_time[indices] = time.monotonic_ns() + np.arange( + len(indices), dtype=np.int64 + ) + + async with self._lock: + self.__buffer._update_has_keys(indices) + + has_keys = self.__buffer._get_has_keys(indices) # [bs, #keys] + rpc_key_mask = self._rpc_key_mask # [#keys, #rpcs] + self._ready_for_rpcs[indices] = ( + has_keys[:, :, None] >= rpc_key_mask[None, :, :] + ).all(axis=1) + + self._is_being_put[indices] = False + self._is_idle[indices] = True + + self._buf_size += len(samples) + if self._buf_size >= 0.95 * self.__max_size: + logger.warning( + f"Buffer is 95% full. The current buffer size is {self._buf_size} " + f"while the maximum size is {self.__max_size}. " + f"If your dataset has more than 1M sequences, consider enlarge " + f"the default batch size in the master worker." + ) + + can_do_rpcs = {rpc.name: self._can_do_rpc(rpc) for rpc in self.rpcs} + logger.info(f"After putting batch, can do RPCs? {can_do_rpcs}.") + + self._lock.notify(len(self._rpc_names)) + return indices + async def amend_batch(self, indices: List[int], samples: List[SequenceSample]): async with self._lock: await self._lock.wait_for( @@ -298,6 +329,17 @@ async def amend_batch(self, indices: List[int], samples: List[SequenceSample]): if self._is_idle[indices].any(): self._lock.notify(len(self._rpc_names)) + def _can_do_rpc(self, rpc: dfg.MFCDef) -> bool: + rpc_idx = self._rpc_names.index(rpc.name) + ready_indices = np.nonzero( + (self._is_idle | self._is_being_read) + & self._ready_for_rpcs[:, rpc_idx] + & ~self._completed_rpc[:, rpc_idx] + )[0] + if len(ready_indices) < rpc.n_seqs: + return False + return True + async def get_batch_for_rpc( self, rpc: dfg.MFCDef ) -> Tuple[List[int], SequenceSample]: @@ -306,20 +348,10 @@ async def get_batch_for_rpc( ) rpc_idx = self._rpc_names.index(rpc.name) - def _can_do_rpc() -> bool: - ready_indices = np.nonzero( - (self._is_idle | self._is_being_read) - & self._ready_for_rpcs[:, rpc_idx] - & ~self._completed_rpc[:, rpc_idx] - )[0] - if len(ready_indices) < rpc.n_seqs: - return False - return True - async with self._lock: # await self._lock.wait_for(_can_do_rpc) - while not _can_do_rpc(): + while not self._can_do_rpc(rpc): await self._lock.wait() logger.info(f"Input keys ({rpc.input_keys}) for MFC {rpc.name} are ready!") diff --git a/realhf/system/controller.py b/realhf/system/controller.py index c84db5ae2e..859fbd1de6 100644 --- a/realhf/system/controller.py +++ b/realhf/system/controller.py @@ -450,7 +450,7 @@ class RayController: It uses the basic Controller to configure workers. Besides, it launchs all remote workers using Ray, instead of submitting them to - the scheduelr. + the scheduler. """ def __init__(self, experiment_name, trial_name): @@ -619,6 +619,8 @@ def start(self, experiment: system_api.Experiment, ignore_worker_error=False): REAL_RECOVER_RUN=os.environ.get("REAL_RECOVER_RUN", ""), REAL_SAVE_RECOVER_STATES=os.environ.get("REAL_SAVE_RECOVER_STATES", ""), REAL_MATH_METADATA_PATH=os.environ.get("REAL_MATH_METADATA_PATH", ""), + REAL_CODE_METADATA_PATH=os.getenv("REAL_CODE_METADATA_PATH", ""), + FUNCTIONCALL_SERVICE_DOMAIN=os.getenv("FUNCTIONCALL_SERVICE_DOMAIN", ""), ) runtime_env = { "env_vars": env_vars, diff --git a/realhf/system/data_manager.py b/realhf/system/data_manager.py new file mode 100644 index 0000000000..365baebbe9 --- /dev/null +++ b/realhf/system/data_manager.py @@ -0,0 +1,395 @@ +# Copyright 2025 Ant Group Inc. + +import bisect +import dataclasses +import itertools +from collections import defaultdict +from typing import * + +import numpy as np +import torch +import torch.distributed as dist + +from realhf import SequenceSample +from realhf.api.core.config import ModelName, ModelShardID +from realhf.base import constants +from realhf.base.topology import PipeModelDataParallelTopology, new_or_get_group +from realhf.impl.model.comm.global_comm import filter_match_mwids +from realhf.system.redistributor import RedistribStep + +BCAST_GROUPS = {} +GATHER_GROUPS = {} +SCATTER_GROUPS = {} + + +class DataManager: + + def __init__( + self, + model_topos: Dict[ModelName, PipeModelDataParallelTopology], + msid2mwid: Optional[Dict[ModelShardID, int]] = None, + data_transfer_pairs: Optional[List[Tuple[ModelName, ModelName]]] = None, + ): + self.model_topos = model_topos + self.msid2mwid = msid2mwid + self.data_transfer_pairs = data_transfer_pairs + + self.storage: Dict[Hashable, SequenceSample] = {} + + def setup_process_groups(self): + if self.msid2mwid is None or self.data_transfer_pairs is None: + return + + model_topos = self.model_topos + msid2mwid = self.msid2mwid + data_transfer_pairs = self.data_transfer_pairs + + # Stores the ranks given a (model_name, dp_rank) pair. + # These workers correspond to a complete set of model parameters sharded by TP+PP. + mw_dp_ranks: Dict[Tuple[ModelName, int], List[int]] = {} + + mw_ranks: Dict[ModelName, List[int]] = {} + + # Stores the dp_head (i.e., mp_rank=0, pp_rank=-1) ranks given a model_name. + mw_dp_head_ranks: Dict[ModelName, List[int]] = defaultdict(list) + + assert msid2mwid is not None + for model_name, topo in model_topos.items(): + mw_ranks[model_name] = filter_match_mwids( + model_name, + topo, + msid2mwid, + ) + mw_dp_head_ranks[model_name] = filter_match_mwids( + model_name, + topo, + msid2mwid, + pipe=topo.get_dim("pipe") - 1, + model=0, + ) + dp_size = topo.get_dim("data") + for dp_i in range(dp_size): + mw_dp_ranks[model_name, dp_i] = filter_match_mwids( + model_name, + topo, + msid2mwid, + data=dp_i, + ) + + for src, dst in data_transfer_pairs: + src_topo = model_topos[src] + dst_topo = model_topos[dst] + + ranks = tuple(sorted(mw_dp_head_ranks[src])) + GATHER_GROUPS[ranks] = new_or_get_group( + list(ranks), backend="nccl" if constants.use_cuda() else "gloo" + ) + + scatter_ranks = tuple(sorted(set([ranks[0]] + mw_ranks[dst]))) + SCATTER_GROUPS[scatter_ranks] = new_or_get_group( + list(scatter_ranks), + backend="nccl" if constants.use_cuda() else "gloo", + ) + + # Construct all src-dst pairs, from any src dp rank to any dst dp rank. + # Note that a dp rank corresponds to multiple parameter shards (TP+PP), + # so each pair is a group-to-group communication. + # Since the models in the source group have duplicate data (TP+PP), + # we just use its "head" as the broadcast source, + # and broadcast to all the ranks in the destination group. + for src_dp, dst_dp in itertools.product( + range(src_topo.get_dim("data")), range(dst_topo.get_dim("data")) + ): + src_mw_rank = mw_dp_head_ranks[src][src_dp] + dst_mw_ranks = mw_dp_ranks[dst, dst_dp] + # The src and dst groups can be disjoint or overlapped. + # If they are disjoint, we need to include the src_mw_rank in the group. + # Otherwise, we only need to include the dst_mw_ranks. + if src_mw_rank not in dst_mw_ranks: + _ranks = [src_mw_rank] + dst_mw_ranks + else: + _ranks = dst_mw_ranks + key = tuple(sorted(_ranks)) + BCAST_GROUPS[key] = new_or_get_group( + _ranks, backend="nccl" if constants.use_cuda() else "gloo" + ) + + def storage_size(self): + return len(self.storage) + + def store(self, x: SequenceSample): + assert len(x.ids) == 1 + assert x.ids[0] not in self.storage + self.storage[x.ids[0]] = x + + def update(self, x: SequenceSample): + self.storage[x.ids[0]].update_(x) + + def get(self, data_id: Hashable): + return self.storage[data_id] + + def has_data(self, data_id: Hashable): + return data_id in self.storage + + def remove(self, ids: List[Hashable]): + for data_id in ids: + if data_id in self.storage: + del self.storage[data_id] + + def clear_data(self): + self.storage.clear() + + def _bcast_recv( + self, + step: RedistribStep, + data_infos: Dict[Hashable, SequenceSample], + ): + assert len(step.keys) == 1 + ids = step.ids + key = step.keys[0] + dtype = data_infos[ids[0]].dtypes[key] + total_len = sum(sum(data_infos[_id].seqlens[key][0]) for _id in ids) + trailing_shape = data_infos[ids[0]].trailing_shapes[key] + + buf = torch.zeros( + (total_len, *trailing_shape), + dtype=dtype, + device=constants.current_device(), + ) + + if len(step.dsts) == 1: + dist.recv(buf, src=step.root) + else: + global BCAST_GROUPS + group = BCAST_GROUPS[tuple(sorted([step.root] + list(step.dsts)))] + dist.broadcast(buf, src=step.root, group=group) + + # Split the received data and put it into the storage. + offset = 0 + for _id in ids: + seqlens = data_infos[_id].seqlens[key] + assert len(seqlens) == 1 + seqlen = sum(seqlens[0]) + if buf is not None: + vs = buf[offset : offset + seqlen] + else: + vs = None + offset = offset + seqlen + with SequenceSample.disable_validation(): + s = SequenceSample( + keys=[key], + dtypes={key: vs.dtype if vs is not None else None}, + trailing_shapes={key: vs.shape[1:] if vs is not None else None}, + ids=[_id], + seqlens={key: seqlens}, + data={key: vs}, + metadata={}, + ) + if _id in self.storage: + self.storage[_id].update_(s) + else: + self.storage[_id] = s + + def _bcast_send(self, step: RedistribStep): + ids = step.ids + for _id in ids: + self.storage[_id].to_device(constants.current_device()) + assert len(step.keys) == 1 + key = step.keys[0] + vs = torch.cat( + [self.storage[_id].data[key] for _id in ids], + dim=0, + ) + if len(step.dsts) == 1: + dist.send(vs, dst=step.dsts[0]) + else: + global BCAST_GROUPS + group = BCAST_GROUPS[tuple(sorted([step.root] + list(step.dsts)))] + dist.broadcast(vs, src=step.root, group=group) + + def _run_bcast( + self, step: RedistribStep, data_infos: Dict[Hashable, SequenceSample] + ): + if dist.get_rank() in step.dsts: + self._bcast_recv(step, data_infos=data_infos) + + if dist.get_rank() == step.root: + self._bcast_send(step) + + def _pad_data(self, x: torch.Tensor, maxlen: int): + assert x.dtype == torch.float32 + assert len(x.shape) == 1 + if maxlen > x.numel(): + return torch.nn.functional.pad(x, (0, maxlen - x.numel()), value=0.0) + return x + + def _run_gather( + self, step: RedistribStep, data_infos: Dict[Hashable, SequenceSample] + ): + if dist.get_rank() not in step.srcs: + return + + maxlen = 0 + for ids in step.ids: + infos = [data_infos[i] for i in ids] + maxlen = max( + maxlen, + sum( + [ + sum([sum(info.seqlens[key][0]) for info in infos]) + for key in step.keys + ] + ), + ) + + if dist.get_rank() == step.root: + gather_list = [ + torch.empty( + maxlen, device=constants.current_device(), dtype=torch.float32 + ) + for _ in range(len(step.srcs)) + ] + else: + gather_list = None + + local_gather_idx = step.srcs.index(dist.get_rank()) + ids = step.ids[local_gather_idx] + for i in ids: + self.storage[i].to_device(constants.current_device()) + samples = [self.storage[i] for i in ids] + data = torch.cat( + [ + sample.data[key].float().flatten() + for sample in samples + for key in step.keys + ] + ) + data = self._pad_data(data, maxlen) + + dist.gather( + data, + gather_list, + dst=step.root, + group=GATHER_GROUPS[tuple(sorted(step.srcs))], + ) + + if dist.get_rank() != step.root: + return + + for ids, buf in zip(step.ids, gather_list): + offset = 0 + for i in ids: + for key in step.keys: + seqlen = data_infos[i].seqlens[key][0] + dtype = data_infos[i].dtypes[key] + trailing_shape = data_infos[i].trailing_shapes[key] + size = int(np.prod(trailing_shape) * sum(seqlen)) + data = buf[offset : offset + size].to(dtype) + offset += size + # with SequenceSample.disable_validation(): + s = SequenceSample( + keys=[key], + dtypes={key: dtype}, + trailing_shapes={key: trailing_shape}, + ids=[i], + seqlens={key: [seqlen]}, + data={key: data}, + metadata={}, + ) + if i in self.storage: + self.storage[i].update_(s) + else: + self.storage[i] = s + + def _run_scatter( + self, step: RedistribStep, data_infos: Dict[Hashable, SequenceSample] + ): + if dist.get_rank() != step.root and dist.get_rank() not in step.dsts: + return + + maxlen = 0 + for ids in step.ids: + infos = [data_infos[i] for i in ids] + maxlen = max( + maxlen, + sum( + [ + sum([sum(info.seqlens[key][0]) for info in infos]) + for key in step.keys + ] + ), + ) + + buf = torch.empty( + maxlen, device=constants.current_device(), dtype=torch.float32 + ) + + if dist.get_rank() == step.root: + scatter_list = [] + for ids in step.ids: + for i in ids: + self.storage[i].to_device(constants.current_device()) + samples = [self.storage[i] for i in ids] + data = torch.cat( + [ + sample.data[key].float().flatten() + for sample in samples + for key in step.keys + ] + ) + scatter_list.append(data) + maxlen = max([x.shape[0] for x in scatter_list]) + scatter_list = [self._pad_data(x, maxlen) for x in scatter_list] + if step.root not in step.dsts: + + idx = bisect.bisect(step.dsts, step.root) + scatter_list.insert(idx, buf) + else: + scatter_list = None + + key = tuple(sorted(set([step.root] + step.dsts))) + dist.scatter(buf, scatter_list, src=step.root, group=SCATTER_GROUPS[key]) + + if dist.get_rank() not in step.dsts: + return + + local_dst_idx = step.dsts.index(dist.get_rank()) + ids = step.ids[local_dst_idx] + offset = 0 + for i in ids: + for key in step.keys: + seqlen = data_infos[i].seqlens[key][0] + dtype = data_infos[i].dtypes[key] + trailing_shape = data_infos[i].trailing_shapes[key] + size = int(np.prod(trailing_shape) * sum(seqlen)) + data = buf[offset : offset + size].to(dtype) + offset += size + # with SequenceSample.disable_validation(): + s = SequenceSample( + keys=[key], + dtypes={key: dtype}, + trailing_shapes={key: trailing_shape}, + ids=[i], + seqlens={key: [seqlen]}, + data={key: data}, + metadata={}, + ) + if i in self.storage: + self.storage[i].update_(s) + else: + self.storage[i] = s + + def redistribute( + self, + data_info: SequenceSample, + plan: List[RedistribStep], + ): + data_infos = {x.ids[0]: x for x in data_info.unpack()} + + for step in plan: + if step.comm_type == "bcast": + self._run_bcast(step, data_infos) + elif step.comm_type == "gather": + self._run_gather(step, data_infos) + elif step.comm_type == "scatter": + self._run_scatter(step, data_infos) diff --git a/realhf/system/function_executor.py b/realhf/system/function_executor.py new file mode 100644 index 0000000000..54dc3bab26 --- /dev/null +++ b/realhf/system/function_executor.py @@ -0,0 +1,187 @@ +# Copyright 2025 Ant Group Inc. +import asyncio +import random +from typing import * + +import networkx as nx +from tensorboardX import SummaryWriter + +from realhf.api.core.config import ModelName, ModelShardID +from realhf.api.core.data_api import DataBatchMeta, SequenceSample +from realhf.api.core.dfg import MFCDef +from realhf.api.core.model_api import ReaLModelConfig +from realhf.base import logging +from realhf.base.topology import PipeModelDataParallelTopology +from realhf.system.buffer import AsyncIOSequenceBuffer +from realhf.system.model_function_call import ModelFunctionCall, RPCCorountineControl +from realhf.system.redistributor import GlobalStorageTracker, RedistribPlanner +from realhf.system.request_reply_stream import NameResolvingRequestClient + +logger = logging.getLogger(__name__, "system") +blogger = logging.getLogger("benchmark") + + +class FunctionExecutor: + def __init__( + self, + rpcs: List[MFCDef], + msid2mwid: Dict[ModelShardID, int], + stream: NameResolvingRequestClient, + buffer: AsyncIOSequenceBuffer, + model_topos: Dict[str, PipeModelDataParallelTopology], + model_configs: Dict[str, None | ReaLModelConfig], + ctrl: RPCCorountineControl, + summary_writer: SummaryWriter | None, + ): + + self.func_calls: Dict[str, ModelFunctionCall] = {} + self.ctrl = ctrl + + self.n_model_workers = len(set(msid2mwid.values())) + self.msid2mwid = msid2mwid + + self.storage_tracker = GlobalStorageTracker(self.n_model_workers) + self.redistrib_planner = RedistribPlanner(self.storage_tracker) + + self.rpcs = rpcs + self.src_rpc = list(filter(lambda rpc: rpc.is_src, rpcs))[0] + self.src_dp_size = model_topos[self.src_rpc.model_name].get_dim("data") + + # Create model function calls. + for rpc in self.rpcs: + func_call = ModelFunctionCall( + rpc=rpc, + src_rpc=self.src_rpc, + stream=stream, + msid2mwid=msid2mwid, + model_topos=model_topos, + model_configs=model_configs, + ctrl=ctrl, + buffer=buffer, + redistrib_planner=self.redistrib_planner, + summary_writer=summary_writer, + ) + self.func_calls[rpc.name] = func_call + + self.stream = stream + self.buffer = buffer + + # Sort all MFCs in the topological order and + # calculate the width of each level. + # These numbers will determine when to flush MFC requests. + self.topo_widths = [] + for generation in nx.topological_generations(rpcs[0]._G): + self.topo_widths.append(len(generation)) + + def get_leaf_tasks(self) -> List[str]: + dst_rpcs = list(filter(lambda rpc: rpc.is_dst, self.rpcs)) + return [rpc.name for rpc in dst_rpcs] + + async def flush_calls(self): + for level, w in enumerate(self.topo_widths): + for _ in range(w): + await self.ctrl.topo_level_count.get() + logger.info(f"DFG level {level}. Flushing {w} function calls.") + self.stream.request( + handlers=list(range(self.n_model_workers)), handle_type="flush" + ) + + async def finish_traverse(self): + for _ in range(len(self.get_leaf_tasks())): + await self.ctrl.train_count.get() + + async def load_data(self): + src_rpc = self.src_rpc + src_rpc_model_name = src_rpc.model_name + buffer = self.buffer + ctrl = self.ctrl + + dp_idx = -1 + received_ids = set() + + while self.buffer.size < max(rpc.n_seqs for rpc in self.rpcs): + + dp_idx += 1 + dp_idx %= self.src_dp_size + + resps = await self.stream.call_async( + handlers=[f"__data{dp_idx}__"], + handle_type="fetch", + datas=[None], + verbose=False, + ) + x: DataBatchMeta | None = resps[0] + + if x is None: + continue + if x.meta_sample is None: + continue + + # RPCs corountines will use this information to + # determine the src and dst of data transfer. + for xx in x.meta_sample.unpack(): + if xx.ids[0] in received_ids: + raise ValueError(f"Duplicate data id {xx.ids[0]}.") + received_ids.add(xx.ids[0]) + all_data = x.meta_sample.unpack() + + filtered_data = [] + for xx in x.meta_sample.unpack(): + if xx.ids[0] in ctrl.hash_vals_to_ignore_in_recover: + ctrl.hash_vals_to_ignore_in_recover.remove(xx.ids[0]) + ctrl.ids_to_clear.add(xx.ids[0]) + else: + filtered_data.append(xx) + all_data = filtered_data + + # We load data in a round-robin manner across different DP ranks, + # so we also need to shuffle the data to fuse different dataset splits. + random.shuffle(all_data) + + # Update resource tracker for planning data redistribution. + gpu_id = self.stream.route_to(f"__data{dp_idx}__") + for k in all_data[0].keys: + self.storage_tracker.add_data( + gpu_id, + [x.ids[0] for x in all_data], + k, + is_owner=True, + ) + + # Store into buffer! + buffer_indices = await buffer.put_batch(all_data) + assert len(buffer_indices) == len(all_data) + + blogger.info( + f"Master worker loaded {len(all_data)} pieces of data. " + f"Remaining number of data to ignore: {len(self.ctrl.hash_vals_to_ignore_in_recover)}. " + f"Current buffer size: {buffer.size}/{buffer.max_size}. " + ) + + def execute_step(self): + logger.info("Waiting for the finish of the execution graph.") + loop = asyncio.get_event_loop() + + tasks = [loop.create_task(fc.run()) for fc in self.func_calls.values()] + [ + loop.create_task(self.flush_calls()), + loop.create_task(self.load_data()), + loop.create_task(self.finish_traverse()), + ] + + loop.run_until_complete(asyncio.gather(*tasks)) + + logger.info("Execution finished!") + + self.clear_gpu_cache() + + def clear_gpu_cache(self): + self.stream.request( + handlers=list(range(self.n_model_workers)), + handle_type="clear_data_cache", + datas=[self.ctrl.ids_to_clear for _ in list(range(self.n_model_workers))], + no_syn=True, + ) + # Clear resource tracker as well. + self.storage_tracker.clear_data(self.ctrl.ids_to_clear) + + self.ctrl.ids_to_clear.clear() diff --git a/realhf/system/master_worker.py b/realhf/system/master_worker.py index 445ff3effa..29a88b2560 100644 --- a/realhf/system/master_worker.py +++ b/realhf/system/master_worker.py @@ -3,30 +3,25 @@ # Licensed under the Apache License, Version 2.0 (the "License"). import asyncio -import collections import copy -import dataclasses import gc -import getpass -import itertools import os -import pprint -import random -import re import time -import uuid -from collections import defaultdict -from typing import Dict, List, Optional, Set, Tuple, Union +from typing import Dict import colorama import networkx as nx import numpy as np -import torch -import torch.distributed import wandb +from tensorboardX import SummaryWriter + +try: + import uvloop + + uvloop.install() +except (ModuleNotFoundError, ImportError): + pass -import realhf.api.core.config as config_api -import realhf.api.core.data_api as data_api import realhf.api.core.dfg as dfg import realhf.api.core.model_api as model_api import realhf.api.core.system_api as config_pkg @@ -35,512 +30,16 @@ import realhf.system.worker_base as worker_base from realhf.api.core.config import ModelName from realhf.api.core.model_api import ReaLModelConfig -from realhf.base import ( - constants, - datapack, - logging, - name_resolve, - names, - seeding, - timeutil, - topology, -) -from realhf.base.asyncio_utils import ( - raise_asyncio_exception, - setup_run_until_complete, - teardown_run_util_complete, -) +from realhf.base import constants, logging, name_resolve, names, timeutil, topology from realhf.system.buffer import AsyncIOSequenceBuffer -from realhf.system.flops_counter import FlopsCounter +from realhf.system.function_executor import FunctionExecutor +from realhf.system.model_function_call import RPCCorountineControl logger = logging.getLogger("master worker", "system") blogger = logging.getLogger("benchmark") -def _attach_param_realloc_hooks( - payload: request_reply_stream.Payload, - msid2mwid: Dict[config_pkg.ModelShardID, int], - from_model_name: ModelName, - to_model_name: ModelName, - from_topo: topology.PipeModelDataParallelTopology, - to_topo: topology.PipeModelDataParallelTopology, - to_model_config: ReaLModelConfig, - pre: bool, -) -> request_reply_stream.Payload: - - model_name = from_model_name - target = to_model_name - # Prioritize handlers of `from_model`, then handlers of `to_model`. - # As a result, if both `from_model` and `to_model` reside in a model worker, - # the handler in the received request will be `from_model`. Layers will also built in `from_model`. - # After that, we assign layers of the `from_model` to `to_model`. - handlers = [ - config_pkg.ModelShardID.from_parallelism_rank(model_name, from_topo, j) - for j in range(from_topo.world_size()) - ] - all_handler_mwids = set([msid2mwid[h] for h in handlers]) - dst_handlers = [ - config_pkg.ModelShardID.from_parallelism_rank(target, to_topo, j) - for j in range(to_topo.world_size()) - ] - for h in dst_handlers: - if msid2mwid[h] not in all_handler_mwids: - handlers.append(h) - all_handler_mwids.add(msid2mwid[h]) - - ps_data = { - "from_model_name": model_name, - "to_model_name": target, - "from_topo": from_topo, - "to_topo": to_topo, - "to_model_config": to_model_config, - "eta": 1.0, - } - if pre: - payload.pre_hooks.append("param_realloc") - payload.pre_hook_data.append(ps_data) - else: - payload.post_hooks.append("param_realloc") - payload.post_hook_data.append(ps_data) - return payload - - -@dataclasses.dataclass -class RPCCorountineControl: - ## Shared resources ## - stop: asyncio.Event - # for counting the number of finished training steps - # one training step corresponds to traversal of the whole DFG - train_count: asyncio.Queue - # For flushing requests - topo_level_count: asyncio.Queue - - ## Per-coroutine resources ## - # Used for counting the number of concurrent calls. - rpc_concurrency: Dict[str, asyncio.Semaphore] - rpc_traversal: Dict[str, int] - # for synchronizing req ids between req and reply coroutines - request_queues: Dict[str, asyncio.Queue] - - # for training data management and data cleaning after each step - ids_to_clear: Set[int] = dataclasses.field(default_factory=set) - flops_counter: FlopsCounter = dataclasses.field(default_factory=FlopsCounter) - - should_save: bool = False - should_eval: bool = False - should_ckpt: bool = False - step_info: recover.StepInfo = dataclasses.field(default_factory=recover.StepInfo) - - # recover information - used_hash_vals_this_epoch: List[int] = dataclasses.field(default_factory=list) - hash_vals_to_ignore_in_recover: List[int] = dataclasses.field(default_factory=list) - - -def _attach_payloads_with_hooks( - rpc: dfg.MFCDef, - payloads: Dict[config_api.ModelShardID, request_reply_stream.Payload], - mwids: List[int], - msid2mwid: Dict[config_pkg.ModelShardID, int], - model_configs: Dict[str, None | ReaLModelConfig], - model_topos: Dict[str, topology.PipeModelDataParallelTopology], - main_handlers: List[config_pkg.ModelShardID], - hook_type: str, -) -> Tuple[Dict[config_api.ModelShardID, request_reply_stream.Payload], List[int]]: - assert hook_type in ["pre", "post"], hook_type - - main_mwids = set([msid2mwid[h] for h in main_handlers]) - for hook in getattr(rpc, f"_{hook_type}_hooks"): - if isinstance(hook, dfg.ParamReallocHook): - assert (hook.source is None) != (hook.target is None), hook - if hook.source is None: - src_topo = model_topos[rpc.model_name] - dst_topo = model_topos[hook.target] - dst_config = model_configs[hook.target] - src_model_name, dst_model_name = rpc.model_name, hook.target - other_model_name = hook.target - other_topo = dst_topo - else: - src_topo = model_topos[hook.source] - dst_topo = model_topos[rpc.model_name] - dst_config = model_configs[rpc.model_name] - src_model_name, dst_model_name = hook.source, rpc.model_name - other_model_name = hook.source - other_topo = src_topo - - ps_data = { - "from_model_name": src_model_name, - "to_model_name": dst_model_name, - "from_topo": src_topo, - "to_topo": dst_topo, - "to_model_config": dst_config, - "eta": hook.eta, - } - for h in main_handlers: - getattr(payloads[h], f"{hook_type}_hooks").append("param_realloc") - getattr(payloads[h], f"{hook_type}_hook_data").append(ps_data) - other_handlers = [ - config_api.ModelShardID.from_parallelism_rank( - other_model_name, other_topo, j - ) - for j in range(other_topo.world_size()) - ] - for h in other_handlers: - if msid2mwid[h] not in mwids: - payloads[h] = request_reply_stream.Payload( - handler=h, - handle_name="empty", - ) - setattr(payloads[h], f"{hook_type}_hooks", ["param_realloc"]) - setattr(payloads[h], f"{hook_type}_hook_data", [ps_data]) - mwids.append(msid2mwid[h]) - elif msid2mwid[h] not in main_mwids: - hh = next(hh for hh in payloads if msid2mwid[hh] == msid2mwid[h]) - getattr(payloads[hh], f"{hook_type}_hooks").append("param_realloc") - getattr(payloads[hh], f"{hook_type}_hook_data").append(ps_data) - - elif isinstance(hook, dfg.OffloadHook): - for h in main_handlers: - getattr(payloads[h], f"{hook_type}_hooks").append("offload") - getattr(payloads[h], f"{hook_type}_hook_data").append( - dict(model_name=h.model_name) - ) - else: - raise NotImplementedError(f"Unknown hook type: {hook}") - return payloads, mwids - - -def _request_model_function_call( - rpc: dfg.MFCDef, - stream: request_reply_stream.NameResolvingRequestClient, - msid2mwid: Dict[config_pkg.ModelShardID, int], - model_topos: Dict[str, topology.PipeModelDataParallelTopology], - model_configs: Dict[str, None | ReaLModelConfig], - producer_names: Dict[str, str], - producer_name2producer_handlers: Dict[str, List[config_pkg.ModelShardID]], - producer_mappings: Dict[str, Dict[str, List[int]]], - target_mapping: Dict[str, List[int]], - meta_sample: data_api.SequenceSample, - handlers: List[config_pkg.ModelShardID], - ctrl: RPCCorountineControl, - model_save_root: str, -) -> List[uuid.UUID]: - - dt_data = { - "keys": rpc.input_keys, - "target": rpc.model_name, - "producer_names": producer_names, - "producer_mappings": producer_mappings, - "target_mapping": target_mapping, - "handle_name": rpc.interface_type.value, - "rpc_name": rpc.name, - "meta_sample": meta_sample, - } - - payloads = { - handler: request_reply_stream.Payload( - handler=handler, - handle_name=rpc.interface_type.value, - pre_hooks=["data_transfer"], - pre_hook_data=[dt_data], - data=rpc.name, - ) - for handler in handlers - } - if ctrl.should_eval: - for p in payloads.values(): - p.post_hooks.append("evaluate") - p.post_hook_data.append(dict(model_name=rpc.model_name)) - if ( - ctrl.should_save or ctrl.should_ckpt - ) and rpc.interface_type == dfg.ModelInterfaceType.TRAIN_STEP: - for p in payloads.values(): - p.post_hooks.append("save") - save_dir = os.path.join( - model_save_root, - rpc.model_name.role, - f"epoch{ctrl.step_info.epoch + 1}" - f"epochstep{ctrl.step_info.epoch_step + 1}" - f"globalstep{ctrl.step_info.global_step + 1}", - ) - p.post_hook_data.append( - dict( - model_name=rpc.model_name, - save_dir=save_dir, - recover_only=not ctrl.should_save, - ) - ) - mwids = [msid2mwid[h] for h in handlers] - assert len(mwids) == len(set(mwids)) - - for producer_name in producer_names.values(): - for h in producer_name2producer_handlers[producer_name]: - if msid2mwid[h] not in mwids: - payloads[h] = request_reply_stream.Payload( - handler=h, - handle_name="empty", - pre_hooks=["data_transfer"], - pre_hook_data=[dt_data], - ) - mwids.append(msid2mwid[h]) - - payloads, mwids = _attach_payloads_with_hooks( - rpc, - payloads, - mwids, - msid2mwid=msid2mwid, - model_configs=model_configs, - model_topos=model_topos, - main_handlers=handlers, - hook_type="pre", - ) - payloads, mwids = _attach_payloads_with_hooks( - rpc, - payloads, - mwids, - msid2mwid=msid2mwid, - model_configs=model_configs, - model_topos=model_topos, - main_handlers=handlers, - hook_type="post", - ) - main_payloads = [p for h, p in payloads.items() if h in handlers] - other_payloads = [p for h, p in payloads.items() if h not in handlers] - all_req_ids = stream.request( - payloads=main_payloads + other_payloads, - ) - return all_req_ids[: len(main_payloads)], all_req_ids[len(main_payloads) :] - - -async def model_rpc_request_func( - rpc: dfg.MFCDef, - msid2mwid: Dict[config_pkg.ModelShardID, int], - src_rpc_model_name: ModelName, - stream: request_reply_stream.NameResolvingRequestClient, - buffer: AsyncIOSequenceBuffer, - data_owner: Dict[Tuple[int, str], Tuple[ModelName, int]], - model_topos: Dict[str, topology.PipeModelDataParallelTopology], - model_configs: Dict[str, None | ReaLModelConfig], - model_save_root: str, - ctrl: RPCCorountineControl, -): - """The corountine for sending requests to model workers.""" - topo = model_topos[rpc.model_name] - logger.info( - f"Requesting Model RPC, interface_type=#{rpc.interface_type}# " - f"(dp,mp,pp) = *({topo.get_dim('data')},{topo.get_dim('model')},{topo.get_dim('pipe')})*" - ) - - topo = model_topos[rpc.model_name] - handlers = [ - config_pkg.ModelShardID.from_parallelism_rank(rpc.model_name, topo, j) - for j in range(topo.world_size()) - ] - - producer_names = {} # data key -> model name - for k in rpc.input_keys: - if k in rpc.data_producers: - producer_names[k] = rpc.data_producers[k] - else: - producer_names[k] = src_rpc_model_name - keys_to_send = defaultdict(list) # model name -> List[keys] to send - for k in producer_names: - keys_to_send[producer_names[k]].append(k) - - # convert producer model name to ModelShardID - producer_name2producer_handlers = {} - for producer_name in keys_to_send: - producer_name2producer_handlers[producer_name] = [ - config_pkg.ModelShardID.from_parallelism_rank( - producer_name, model_topos[producer_name], j - ) - for j in range(model_topos[producer_name].world_size()) - ] - - request_queue = ctrl.request_queues[rpc.name] - rpc_concurrency = ctrl.rpc_concurrency[rpc.name] - - this_rpc_consumed_seqs = 0 - while not ctrl.stop.is_set(): - - await rpc_concurrency.acquire() - - # Ensure that parent RPCs will not be over-consumed. - while any( - this_rpc_consumed_seqs >= (ctrl.rpc_traversal[c.name] + 1) * c.n_seqs - for c in rpc.all_successors() - ): - await asyncio.sleep(0.1) - - buf_indices, sample = await buffer.get_batch_for_rpc(rpc) - - ctrl.flops_counter.add_rpc(rpc, sample, model_configs[rpc.model_name]) - - this_rpc_consumed_seqs += sample.bs - - # logger.info(f"Model rpc {rpc.name} requesting.") - - # Dispatch data to different data parallel ranks. - dp_size = topo.get_dim("data") - if rpc.is_generate(): - # The workload of generation is decided by batch size, instead of the generated length. - samples, forward_indices, _ = sample.split_with_lengths( - mb_spec=data_api.MicroBatchSpec(n_mbs=dp_size), - lens=[1 for _ in range(sample.bs)], - ) - else: - samples, forward_indices, _ = sample.split( - data_api.MicroBatchSpec(n_mbs=dp_size) - ) - blogger.info( - f"DP split (DP size {dp_size}) for RPC {rpc.name}: " - f"#seqs: {[s.bs for s in samples]}, " - f"#tokens: {[sum([sum(lens) for lens in s.seqlens[s._get_split_key()]]) for s in samples]}" - ) - sample = data_api.SequenceSample.gather(samples) - buf_indices = [buf_indices[i] for i in forward_indices] - - partitions = data_api.SequenceSplitSpec( - sizes=[s.bs for s in samples] - ).partitions - target_mapping = {i: list(range(v[0], v[1])) for i, v in enumerate(partitions)} - - # Set data owner of produced data by this RPC, such that downstream RPCs can know - # where to fetch these data. - for dp_idx, (st, ed) in enumerate(partitions): - for i in range(st, ed): - for k in rpc.output_keys: - data_owner[sample.ids[i], k] = (rpc.model_name, dp_idx) - - # Get the data owner of this RPC's input data. - # We use it to determine the source of data transfer. - producer_mappings = {} - for k in rpc.input_keys: - names, dp_indices = [], [] - for sample_id in sample.ids: - owner_name, dp_idx = data_owner[(sample_id, k)] - names.append(owner_name) - dp_indices.append(dp_idx) - assert len(set(names)) == 1 - producer_mapping = defaultdict(list) - for i, dp_idx in enumerate(dp_indices): - producer_mapping[dp_idx].append(i) - producer_mapping = {k: sorted(v) for k, v in producer_mapping.items()} - producer_mappings[names[0], k] = producer_mapping - - # send partitioned data to model workers - req_ids, other_req_ids = _request_model_function_call( - rpc=rpc, - stream=stream, - msid2mwid=msid2mwid, - model_topos=model_topos, - model_configs=model_configs, - producer_names=producer_names, - producer_name2producer_handlers=producer_name2producer_handlers, - producer_mappings=producer_mappings, - target_mapping=target_mapping, - meta_sample=sample, - handlers=handlers, - ctrl=ctrl, - model_save_root=model_save_root, - ) - await request_queue.put( - (buf_indices, sample.ids, req_ids, other_req_ids, time.perf_counter()) - ) - await ctrl.topo_level_count.put(1) - logger.info(f"Model rpc {rpc.name} requested.") - - -async def model_rpc_reply_func( - rpc: dfg.MFCDef, - stream: request_reply_stream.NameResolvingRequestClient, - buffer: AsyncIOSequenceBuffer, - model_topos: Dict[str, topology.PipeModelDataParallelTopology], - ctrl: RPCCorountineControl, -): - topo = model_topos[rpc.model_name] - dp_size = topo.get_dim("data") - dp_head_indices = [ - topo.get_rank(data=i, pipe=topo.get_dim("pipe") - 1, model=0) - for i in range(dp_size) - ] - - request_queue = ctrl.request_queues[rpc.name] - rpc_concurrency = ctrl.rpc_concurrency[rpc.name] - - while not ctrl.stop.is_set(): - # Wait for master worker's request. - buf_indices, ids, req_ids, other_req_ids, tik = await request_queue.get() - - # Then, wait for all main requests to finish. - responses = await stream.gather_async(request_ids=req_ids) - # logger.info(f"rpc {rpc.name} received responses {req_ids}") - - # Filter out responses other than DP heads. - # Other repsonses are duplicated or None. - responses: List[request_reply_stream.Payload] = [ - responses[i] for i in dp_head_indices - ] - - # If the returned data is a SequenceSample, it is the data returned by - # model function calls. The data shoulbe be amended into buffer. - # Otherwise, it's the train statistics and should be reduced and logged. - if isinstance(responses[-1], data_api.SequenceSample): - res = data_api.SequenceSample.gather(responses) - else: - res = _gather_stat(responses) - - if rpc.log_return_value: - logger.info(f"RPC name {rpc.name} returns {res}") - - if isinstance(res, Dict): - wandb.log(res, step=ctrl.step_info.global_step) - - logger.info( - f"Model rpc {rpc.name} finished. Run time {time.perf_counter() - tik:.4f}s." - ) - - # Release the semaphore to let the request corountine continue running. - rpc_concurrency.release() - ctrl.rpc_traversal[rpc.name] += 1 - - # If this RPC is the final node in the dataflow graph, - # update the train counter. - # Otherwise, amend data in the buffer. - if rpc.is_dst: - ctrl.ids_to_clear = ctrl.ids_to_clear.union(ids) - await ctrl.train_count.put(1) - else: - logger.info(f"Amending RPC {rpc.name} output keys: {res.keys}") - await buffer.amend_batch(buf_indices, res.unpack()) - - # Wait for all side-effect requests to finish. - # Side-effect or empty requests are required for data transfer - # and parameter synchronization. - # Wait them after the main request to log the oorrect MFC time. - await stream.gather_async(other_req_ids) - - -def _gather_stat(src: List[Dict]) -> Dict: - cnt, stats = {}, {} - for reply in src: - for k, v in reply.items(): - cnt[k] = cnt.get(k, 0) + 1 - stats[k] = stats.get(k, 0) + v - res = {k: v / cnt for k, v, cnt in zip(stats.keys(), stats.values(), cnt.values())} - for k, c in cnt.items(): - if c != len(src): - logger.warning(f"Gathered `{k}` is not present in every returned stats.") - for k, v in res.items(): - if any(abs(v - x.get(k, None)) > 1e-4 for x in src): - logger.warning( - f"Gathered `{k}` is not all-reduced " - f"before returning: ({[x.get(k, None) for x in src]}, {v})." - ) - return res - - class MasterWorker(worker_base.Worker): - os.makedirs(constants.MODEL_SAVE_ROOT, exist_ok=True) global_exp_tik = time.perf_counter() def _configure(self, config: config_pkg.MasterWorker): @@ -563,14 +62,8 @@ def _configure(self, config: config_pkg.MasterWorker): self.__topo_widths.append(len(generation)) logger.info("Topological widths: " + str(self.__topo_widths)) - self.__mwid2msids = defaultdict(list) - for msid, mwid in self.config.msid2mwid.items(): - self.__mwid2msids[mwid].append(msid) - self.__rpc_srcs = list(filter(lambda rpc: rpc.is_src, self.__model_rpcs)) self.__rpc_dsts = list(filter(lambda rpc: rpc.is_dst, self.__model_rpcs)) - self.__n_rpc_srcs = len(self.__rpc_srcs) - self.__n_rpc_dsts = len(self.__rpc_dsts) # Save and eval control. self.__total_train_epochs = config.exp_ctrl.total_train_epochs @@ -618,14 +111,8 @@ def _configure(self, config: config_pkg.MasterWorker): # Create corountine control objects for running the dataflow graph. self.__rpc_ctrl = RPCCorountineControl( - stop=asyncio.Event(), train_count=asyncio.Queue(maxsize=len(self.__rpc_dsts)), topo_level_count=asyncio.Queue(maxsize=sum(self.__topo_widths)), - rpc_traversal={rpc.name: 0 for rpc in self.__model_rpcs}, - request_queues={rpc.name: asyncio.Queue(1) for rpc in self.__model_rpcs}, - rpc_concurrency={ - rpc.name: asyncio.Semaphore(1) for rpc in self.__model_rpcs - }, # NOTE: We should accumulate the used data hashes in the same epoch # to prevent loading data used before. used_hash_vals_this_epoch=( @@ -660,158 +147,10 @@ def _configure(self, config: config_pkg.MasterWorker): return config.worker_info - def __lazy_init(self): - # Set up streams. - handler_routing = copy.deepcopy(self.config.msid2mwid) - src_rpc = self.__rpc_srcs[0] - src_rpc_topo = self.config.model_topos[src_rpc.model_name] - src_rpc_dp_size = src_rpc_topo.get_dim("data") - src_rpc_pp_size = src_rpc_topo.get_dim("pipe") - for i in range(src_rpc_dp_size): - rank = src_rpc_topo.get_rank(data=i, pipe=src_rpc_pp_size - 1, model=0) - handler_routing[f"__data{i}__"] = self.config.msid2mwid[ - config_pkg.ModelShardID.from_parallelism_rank( - model_name=src_rpc.model_name, - topo=src_rpc_topo, - parallelism_rank=rank, - ) - ] - self.__stream = request_reply_stream.make_master_stream( - self.config.worker_info, - n_subscribers=self.config.n_model_workers, - handler_routing=handler_routing, - ) - self.__stream: request_reply_stream.NameResolvingRequestClient - - self.__src_rpc = src_rpc = [ - rpc for rpc in self.config.model_rpcs if rpc.is_src - ][0] - src_rpc_model_name = src_rpc.model_name - self.__src_rpc_dp_size = src_rpc_dp_size = self.config.model_topos[ - src_rpc.model_name - ].get_dim("data") - - # Request training specification from data workers. - all_data = sum( - self.__stream.call( - handlers=[f"__data{i}__" for i in range(src_rpc_dp_size)], - datas=[None for i in range(src_rpc_dp_size)], - handle_type="spec", - ), - [], - ) - - # NOTE: For dynamic datasets, we still count epoch according to the initial number of data, - # such that the learning rate decay is not affected. - seqlens = [max(sum(v[0]) for v in x.seqlens.values()) for x in all_data] - self._dataset_size = len(all_data) - self._steps_per_epoch = self._dataset_size // src_rpc.n_seqs - self._avg_tokens_per_batch = sum(seqlens) / self._steps_per_epoch - self._dataset_ids = [copy.deepcopy(x.ids[0]) for x in all_data] - - event_loop = asyncio.new_event_loop() - asyncio.set_event_loop(event_loop) - - # Build some data required for subsequent model function calls. - self.__all_model_handlers: List[config_pkg.ModelShardID] = [] - self.__all_mw_handlers: List[config_pkg.ModelShardID] = [] - _covered_mws = set() - self.__dp0_model_handlers: List[config_pkg.ModelShardID] = [] - self.__trainable_model_handlers: List[config_pkg.ModelShardID] = [] - for model_name, topo in self.config.model_topos.items(): - for j in range(topo.world_size()): - h = config_pkg.ModelShardID.from_parallelism_rank(model_name, topo, j) - _mw_id = self.config.msid2mwid[h] - if _mw_id not in _covered_mws: - _covered_mws.add(_mw_id) - self.__all_mw_handlers.append(h) - num_dp = topo.get_dim("data") - self.__all_model_handlers += [ - config_pkg.ModelShardID.from_parallelism_rank(model_name, topo, j) - for j in range(topo.world_size()) - ] - - if any( - rpc.model_name == model_name - and rpc.interface_type == dfg.ModelInterfaceType.TRAIN_STEP - for rpc in self.__model_rpcs - ): - self.__trainable_model_handlers += [ - config_pkg.ModelShardID.from_parallelism_rank(model_name, topo, j) - for j in range(topo.world_size()) - ] - self.__dp0_model_handlers += [ - config_pkg.ModelShardID.from_parallelism_rank(model_name, topo, j) - for j in topo.filter_match(data=0) - ] - - # Request model configs from model workers. - # Return None if the model is not a ReaLModel. - self.__model_configs: Dict[ModelName, None | ReaLModelConfig] = {} - for model_name in self.config.model_topos: - h = config_pkg.ModelShardID.from_parallelism_rank(model_name, topo, 0) - self.__model_configs[model_name] = self.__stream.call( - handlers=[h], - datas=[None], - handle_type="model_config", - )[0] - + def initialize_models(self): # Initialize model backends. - # For models with the same role, they share the same model parameters. - # Therefore, we must call reallocate parameters from A to B - # before we send requests to initialize B. - _param_senders = [v[0] for v in self.config.sync_param_pairs] - _param_receivers = [v[1] for v in self.config.sync_param_pairs] - - # The parameters are by default held by the trainable model. - # If all replicas are not trainable, the parameters are held in replica 0. - _model_is_trainable = collections.defaultdict(list) - for rpc in self.__model_rpcs: - _model_is_trainable[rpc.model_name].append( - rpc.interface_type == dfg.ModelInterfaceType.TRAIN_STEP - ) - - _model_is_trainable = { - model_name: any(values) - for model_name, values in _model_is_trainable.items() - } - - _roles = set([rpc.model_name.role for rpc in self.__model_rpcs]) - _role_cnt = { - role: len( - set( - [ - rpc.model_name - for rpc in self.__model_rpcs - if rpc.model_name.role == role - ] - ) - ) - for role in _roles - } - _reordered_model_names = [] - for role in sorted(_roles): - if _role_cnt[role] == 1: - _reordered_model_names.append(ModelName(role, 0)) - continue - _indices = list(range(_role_cnt[role])) - _trainable_this_role = [ - _model_is_trainable[ModelName(role, i)] for i in range(_role_cnt[role]) - ] - if any(_trainable_this_role): - assert ( - sum(_trainable_this_role) == 1 - ), "only one train for each model is allowed" - _trainable_idx = _trainable_this_role.index(True) - _reordered_model_names.append(ModelName(role, _trainable_idx)) - _indices.remove(_trainable_idx) - for i in _indices: - _reordered_model_names.append(ModelName(role, i)) - - # Send initialization requests. - self.logger.info( - f"Initialize model backends with order: {_reordered_model_names}." - ) + model_names = list(self.__model_topos.keys()) + self.logger.info(f"Initialize model backends with order: {model_names}.") train_rpcs = list( filter( lambda rpc: rpc.interface_type == dfg.ModelInterfaceType.TRAIN_STEP, @@ -832,7 +171,7 @@ def __lazy_init(self): train_batch_size=self.__src_rpc.n_seqs, ) _initialized_roles = [] - for model_name in _reordered_model_names: + for model_name in model_names: topo = self.config.model_topos[model_name] # Build FinetuneSpec, which is required to initialize backends. _handlers = [ @@ -849,55 +188,6 @@ def __lazy_init(self): for _h in _handlers ] - # Reallocate parameters if necessary. - if model_name.role in _initialized_roles and model_name in _param_receivers: - _param_realloc_src = _param_senders[_param_receivers.index(model_name)] - - # Update handler and payloads to send empty requests - # with param realloc hooks to source workers. - src_topo = self.config.model_topos[_param_realloc_src] - other_handlers = [ - config_pkg.ModelShardID.from_parallelism_rank( - _param_realloc_src, src_topo, j - ) - for j in range(src_topo.world_size()) - ] - main_mw_ids = [self.config.msid2mwid[_h] for _h in _handlers] - _other_hanlders = [] - for other_h in other_handlers: - if self.config.msid2mwid[other_h] not in main_mw_ids: - _other_hanlders.append(other_h) - init_payloads += [ - request_reply_stream.Payload( - handler=_h, - handle_name="empty", - ) - for _h in _other_hanlders - ] - _handlers = _handlers + _other_hanlders - - for i, p in enumerate(init_payloads): - p = _attach_param_realloc_hooks( - payload=p, - msid2mwid=self.config.msid2mwid, - from_model_name=_param_realloc_src, - to_model_name=model_name, - from_topo=self.config.model_topos[_param_realloc_src], - to_topo=self.config.model_topos[model_name], - to_model_config=self.__model_configs[model_name], - pre=True, - ) - init_payloads[i] = _attach_param_realloc_hooks( - payload=p, - msid2mwid=self.config.msid2mwid, - from_model_name=model_name, - to_model_name=_param_realloc_src, - to_topo=self.config.model_topos[_param_realloc_src], - from_topo=self.config.model_topos[model_name], - to_model_config=self.__model_configs[_param_realloc_src], - pre=False, - ) - # Send initialization requests then immediately flush them. self.__stream.request( payloads=init_payloads, @@ -913,48 +203,75 @@ def __lazy_init(self): self._ft_spec = ft_spec logger.info("Initializations of models and backends complete.") - self.__seqbuffer = AsyncIOSequenceBuffer( - self.__model_rpcs, - max_size=int(os.getenv("REAL_MASTER_BUFFER_SIZE", str(int(1e7)))), + def get_dataset_model_info(self): + src_rpc = self.__rpc_srcs[0] + src_rpc_topo = self.config.model_topos[src_rpc.model_name] + src_rpc_dp_size = src_rpc_topo.get_dim("data") + + # Request training specification from data workers. + all_data = sum( + self.__stream.call( + handlers=[f"__data{i}__" for i in range(src_rpc_dp_size)], + datas=[None for i in range(src_rpc_dp_size)], + handle_type="spec", + ), + [], ) - self.__data_owner = {} + # NOTE: For dynamic datasets, we still count epoch according to the initial number of data, + # such that the learning rate decay is not affected. + seqlens = [max(sum(v[0]) for v in x.seqlens.values()) for x in all_data] + self._dataset_size = len(all_data) + self._steps_per_epoch = self._dataset_size // src_rpc.n_seqs + self._avg_tokens_per_batch = sum(seqlens) / self._steps_per_epoch + self._dataset_ids = [copy.deepcopy(x.ids[0]) for x in all_data] - logger.info(f"Creating asyncio coroutines...") + # Request model configs from model workers. + # Return None if the model is not a ReaLModel. + self.__model_configs: Dict[ModelName, None | ReaLModelConfig] = {} + for model_name, topo in self.config.model_topos.items(): + h = config_pkg.ModelShardID.from_parallelism_rank(model_name, topo, 0) + self.__model_configs[model_name] = self.__stream.call( + handlers=[h], + datas=[None], + handle_type="model_config", + )[0] - # Create coroutines for model RPCs. - coroutine_tasks = [] - for rpc in self.__model_rpcs: - request_task = event_loop.create_task( - model_rpc_request_func( - rpc=rpc, - msid2mwid=self.config.msid2mwid, - src_rpc_model_name=src_rpc_model_name, - data_owner=self.__data_owner, - stream=self.__stream, - buffer=self.__seqbuffer, - model_topos=self.__model_topos, - model_configs=self.__model_configs, - ctrl=self.__rpc_ctrl, - model_save_root=self.MODEL_SAVE_ROOT, - ) - ) - reply_task = event_loop.create_task( - model_rpc_reply_func( - rpc=rpc, - stream=self.__stream, - buffer=self.__seqbuffer, - model_topos=self.__model_topos, - ctrl=self.__rpc_ctrl, + def __lazy_init(self): + # Set up streams. + handler_routing = copy.deepcopy(self.config.msid2mwid) + src_rpc = self.__rpc_srcs[0] + src_rpc_topo = self.config.model_topos[src_rpc.model_name] + src_rpc_dp_size = src_rpc_topo.get_dim("data") + src_rpc_pp_size = src_rpc_topo.get_dim("pipe") + for i in range(src_rpc_dp_size): + rank = src_rpc_topo.get_rank(data=i, pipe=src_rpc_pp_size - 1, model=0) + handler_routing[f"__data{i}__"] = self.config.msid2mwid[ + config_pkg.ModelShardID.from_parallelism_rank( + model_name=src_rpc.model_name, + topo=src_rpc_topo, + parallelism_rank=rank, ) - ) - coroutine_tasks += [request_task, reply_task] + ] + handler_routing.update({i: i for i in range(self.config.n_model_workers)}) + self.__stream = request_reply_stream.make_master_stream( + self.config.worker_info, + n_subscribers=self.config.n_model_workers, + handler_routing=handler_routing, + ) + self.__stream: request_reply_stream.NameResolvingRequestClient - # Set up a run context of EventLoop.run_util_complete, baiscally copy-paste from cpython. - # With this context, we can call the non-block EventLoop._run_once (similar to worker._poll). - self.__asyncio_tasks: List[asyncio.Task] = coroutine_tasks - self.__asyncio_ctx = setup_run_until_complete( - event_loop, asyncio.gather(*coroutine_tasks) + self.__src_rpc = src_rpc = [ + rpc for rpc in self.config.model_rpcs if rpc.is_src + ][0] + + self.get_dataset_model_info() + + self.initialize_models() + + self.__seqbuffer = AsyncIOSequenceBuffer( + self.__model_rpcs, + max_size=int(os.getenv("REAL_MASTER_BUFFER_SIZE", str(int(1e7)))), ) # wandb init, connect to remote wandb host @@ -976,7 +293,23 @@ def __lazy_init(self): resume="allow", settings=wandb.Settings(start_method="fork"), ) + # tensorboard logging + self.__summary_writer = None + if self.tensorboard_config.path is not None: + self.__summary_writer = SummaryWriter(log_dir=self.tensorboard_config.path) + # Create coroutines for model RPCs. + logger.info(f"Creating asyncio coroutines...") + self.func_executor = FunctionExecutor( + rpcs=self.__model_rpcs, + msid2mwid=self.config.msid2mwid, + stream=self.__stream, + buffer=self.__seqbuffer, + model_topos=self.__model_topos, + model_configs=self.__model_configs, + ctrl=self.__rpc_ctrl, + summary_writer=self.__summary_writer, + ) logger.info(f"Coroutines created. The master worker is ready to run.") self.__initialized = True @@ -990,30 +323,13 @@ def __lazy_init(self): def _poll(self): is_new_epoch = False - first_poll = not self.__initialized if not self.__initialized: self.__lazy_init() - self._maybe_request_load_data(first_poll=first_poll) # Main execution steps. The graph runs under-the-hood in RPC & stream threads. # Wait for the finish of the traversal of the execution graph. execution_start = time.perf_counter() - logger.info("Master worker is waiting for the finish of the execution graph.") - if self.__rpc_ctrl.ids_to_clear: - # Send clear cache requests to model workers. - # Clearing the data used in the last step. - self._clear_gpu_cache() - - for _ in range(10): - self._maybe_request_load_data(first_poll=first_poll) - if self.__seqbuffer.size < self.__src_rpc.n_seqs: - raise RuntimeError( - f"Buffer size {self.__seqbuffer.size} smaller than " - f"required batch size {self.__src_rpc.n_seqs} after loading data. " - "This should not happen, but we raise an error to stop the experiment. " - "Is your dataset size larger than the configured batch size?" - ) is_new_epoch = self._ft_spec.is_new_epoch(self.__rpc_ctrl.step_info) is_epoch_last_step = self._ft_spec.is_epoch_last_step(self.__rpc_ctrl.step_info) @@ -1029,49 +345,10 @@ def _poll(self): epochs=int(is_epoch_last_step), steps=1 ) - _rpc_dst_cnt = 0 - _topo_level_cnt, _topo_level_idx = 0, 0 - while _rpc_dst_cnt < self.__n_rpc_dsts: - try: - self.__rpc_ctrl.train_count.get_nowait() - _rpc_dst_cnt += 1 - continue - except asyncio.QueueEmpty: - pass - try: - self.__rpc_ctrl.topo_level_count.get_nowait() - _topo_level_cnt += 1 - if _topo_level_cnt >= self.__topo_widths[_topo_level_idx]: - logger.info( - f"Flushing the current level of the DFG with {self.__topo_widths[_topo_level_idx]} vertices." - ) - self.__stream.request( - handlers=self.__all_mw_handlers, - handle_type="flush", - ) - _topo_level_idx += 1 - _topo_level_cnt = 0 - continue - except asyncio.QueueEmpty: - pass - try: - # Similar to worker._poll. Run multiple times until a train step is finished. - self.__asyncio_ctx.loop._run_once() - # NOTE: The following line will propagate errors in corountines back to the main thread. - # It raises asyncio.exceptions.InvalidStateError if the result is not ready. - # (In our use cases, the result will never be ready because corountines run while-loops.) - # We just ignore this error and continue running. - self.__asyncio_ctx.future.result() - except asyncio.exceptions.InvalidStateError: - # Catch the exception when future.result() is not ready. - pass - except KeyboardInterrupt as e: - raise_asyncio_exception(self.__asyncio_ctx, raise_error=False) - raise e - except: - raise_asyncio_exception(self.__asyncio_ctx) - logger.info("Execution finished!") + # Traverse over the dataflow graph for once. + self.func_executor.execute_step() + # Post-process. if self.__rpc_ctrl.should_save or self.__rpc_ctrl.should_ckpt: self.__last_step_info = copy.deepcopy(self.__rpc_ctrl.step_info) @@ -1120,107 +397,6 @@ def _poll(self): return worker_base.PollResult(sample_count=1, batch_count=1) - def _maybe_request_load_data(self, first_poll: bool): - should_load_data = self.__seqbuffer.size < self.__src_rpc.n_seqs - if not should_load_data: - return - blogger.info( - f"Current buffer size {self.__seqbuffer.size}/{self.__seqbuffer.max_size}. " - f"The batch size of the source MFC is {self.__src_rpc.n_seqs}." - ) - - src_rpc_dp_size = self.__src_rpc_dp_size - src_rpc = self.__src_rpc - src_rpc_model_name = src_rpc.model_name - data_owner = self.__data_owner - buffer = self.__seqbuffer - stream = self.__stream - ctrl = self.__rpc_ctrl - - # fetch data from dataloader to fill the sequence buffer - blogger.info(f"Filling data into the buffer in a new epoch.") - fetch_data_start = time.perf_counter() - - # NOTE: PyTorch dataloader will shuffle data for us. - all_data: List[data_api.SequenceSample] = [] - received_ids = set() - - # NOTE: Currently we send dataloading requests until iterating - # over the entire dataset. This may lead to a huge memory waste - # with super-large datasets. Empirically, it's fine. - is_final_batch = [False for _ in range(src_rpc_dp_size)] - is_first_batch = [True for _ in range(src_rpc_dp_size)] - while not all(is_final_batch): - # Send request to model workers to get the specification of data. - # Data itself is not transferred to the master worker. - data_batches: List[data_api.DataBatchMeta | None] = [ - None for _ in range(src_rpc_dp_size) - ] - for i in range(src_rpc_dp_size): - if is_final_batch[i]: - data_batches[i] = None - continue - data_batches[i] = stream.call( - handlers=[f"__data{i}__"], - handle_type="fetch", - datas=[dict(first_batch=is_first_batch[i], first_poll=first_poll)], - verbose=False, - )[0] - is_final_batch[i] = data_batches[i].is_final_batch - is_first_batch[i] = False - - # Unpack batched sequences into individual sequences. - for dp_rank, x in enumerate(data_batches): - if x is None: - continue - if x.meta_sample is None: - continue - for xx in x.meta_sample.unpack(): - if xx.ids[0] in received_ids: - raise ValueError( - f"Duplicate data id {xx.ids[0]}. Is the final batch? {is_final_batch}." - ) - received_ids.add(xx.ids[0]) - # Store the owner information of the data. - # RPCs corountines will use this information to - # determine the src and dst of data transfer. - for k in xx.keys: - data_owner[(xx.ids[0], k)] = (src_rpc_model_name, dp_rank) - all_data += x.meta_sample.unpack() - - filtered_data = [] - for x in all_data: - if x.ids[0] in ctrl.hash_vals_to_ignore_in_recover: - ctrl.hash_vals_to_ignore_in_recover.remove(x.ids[0]) - ctrl.ids_to_clear.add(x.ids[0]) - else: - filtered_data.append(x) - all_data = filtered_data - - # We load data in a round-robin manner across different DP ranks, - # so we also need to shuffle the data to fuse different dataset splits. - random.shuffle(all_data) - - blogger.info( - f"Master worker loaded {len(all_data)} pieces of data. " - f"Training epoch {self.__rpc_ctrl.step_info.epoch + 1} approximately has {self._steps_per_epoch} steps. " - f"Each batch has {self._avg_tokens_per_batch:.2f} tokens in average. " - f"Remaining number of data to ignore: {len(self.__rpc_ctrl.hash_vals_to_ignore_in_recover)}." - ) - - # Store into buffer! - buffer_indices = buffer.put_batch_synced(all_data) - assert len(buffer_indices) == len(all_data) - - blogger.info( - f"Filling data finished. Time consumption: " - f"{time.perf_counter() - fetch_data_start:.3f}s." - ) - - # We should let model workers clear the data ignored during recover. - if ctrl.ids_to_clear: - self._clear_gpu_cache() - def _log_training_stats(self, e2e_time: float, time_since_configure: float): # calculate flops ######################################### @@ -1262,24 +438,7 @@ def _log_training_stats(self, e2e_time: float, time_since_configure: float): f"Time taken so far across all configurations: {time.perf_counter() - self.global_exp_tik:.2f}s" ) - def _clear_gpu_cache(self): - self.__stream.request( - handlers=self.__all_mw_handlers, - handle_type="clear_data_cache", - datas=[self.__rpc_ctrl.ids_to_clear for _ in self.__all_mw_handlers], - no_syn=True, - ) - self.__rpc_ctrl.ids_to_clear.clear() - def experiment_complete_exit(self): - self.__rpc_ctrl.stop.set() - for task in self.__asyncio_tasks: - task.cancel() - self.__asyncio_ctx.future.set_result(None) - # NOTE: stopping the loop immediately after cancelling tasks may - # raise warnings sometimes, but it doesn't matter. - self.__asyncio_ctx.loop.stop() - teardown_run_util_complete(self.__asyncio_ctx) logger.info( colorama.Style.RESET_ALL + colorama.Fore.YELLOW @@ -1292,9 +451,9 @@ def experiment_complete_exit(self): # Send requests to pause model workers. # Model workers will not respond to this message. self.__stream.request( - handlers=self.__all_mw_handlers, + handlers=list(range(self.config.n_model_workers)), handle_type="reset", - datas=[None for _ in self.__all_mw_handlers], + datas=[None for _ in list(range(self.config.n_model_workers))], ) self.__stream.close() constants.reset_run() @@ -1310,6 +469,8 @@ def experiment_complete_exit(self): ) wandb.finish() + if self.__summary_writer is not None: + self.__summary_writer.close() gc.collect() self.__initialized = False self.pause() @@ -1336,29 +497,3 @@ def __recover_save(self): logger.info( f"Number of data used in this epoch: {len(recover_info.hash_vals_to_ignore)}" ) - - # def _exit_hook(self, exit_status: worker_base.WorkerServerStatus): - # logger.info(f"Master worker exits with {exit_status}.") - # if os.getenv("REAL_SAVE_RECOVER_STATES", "0") != "1": - # return - # if exit_status == worker_base.WorkerServerStatus.ERROR: - # try: - # sleep_time = 600 - # current_sleep_time = 0 - # while current_sleep_time < sleep_time: - # logger.info( - # f"ERROR exit, waited {current_sleep_time} s for interruption ..." - # ) - # time.sleep(10) - # current_sleep_time += 10 - # except KeyboardInterrupt: - # logger.info("Received SIGINT, starting recover save") - - # self.__recover_save() - # ) - # time.sleep(10) - # current_sleep_time += 10 - # except KeyboardInterrupt: - # logger.info("Received SIGINT, starting recover save") - - # self.__recover_save() diff --git a/realhf/system/model_function_call.py b/realhf/system/model_function_call.py new file mode 100644 index 0000000000..0111e99d5a --- /dev/null +++ b/realhf/system/model_function_call.py @@ -0,0 +1,488 @@ +# Copyright 2025 Ant Group Inc. + +import asyncio +import dataclasses +import itertools +import json +import os +import time +import uuid +from collections import defaultdict +from typing import Dict, Hashable, List, Set, Tuple + +import wandb +from tensorboardX import SummaryWriter + +import realhf.api.core.config as config_api +import realhf.api.core.data_api as data_api +import realhf.api.core.dfg as dfg +import realhf.api.core.system_api as config_pkg +import realhf.base.recover as recover +import realhf.system.request_reply_stream as request_reply_stream +from realhf import ModelShardID +from realhf.api.core.config import ModelName +from realhf.api.core.model_api import ReaLModelConfig +from realhf.base import constants, logging, topology +from realhf.system.buffer import AsyncIOSequenceBuffer +from realhf.system.flops_counter import FlopsCounter +from realhf.system.redistributor import RedistribPlanner, RedistribStep + +logger = logging.getLogger(__name__, "system") +blogger = logging.getLogger("benchmark") + + +@dataclasses.dataclass +class RPCCorountineControl: + # for counting the number of finished training steps + # one training step corresponds to traversal of the whole DFG + train_count: asyncio.Queue + # For flushing requests + topo_level_count: asyncio.Queue + + # for training data management and data cleaning after each step + ids_to_clear: Set[Hashable] = dataclasses.field(default_factory=set) + flops_counter: FlopsCounter = dataclasses.field(default_factory=FlopsCounter) + + should_save: bool = False + should_eval: bool = False + should_ckpt: bool = False + step_info: recover.StepInfo = dataclasses.field(default_factory=recover.StepInfo) + + # recover information + used_hash_vals_this_epoch: List[int] = dataclasses.field(default_factory=list) + hash_vals_to_ignore_in_recover: List[int] = dataclasses.field(default_factory=list) + + +class ModelFunctionCall: + def __init__( + self, + rpc: dfg.MFCDef, + src_rpc: dfg.MFCDef, + stream: request_reply_stream.NameResolvingRequestClient, + msid2mwid: Dict[config_pkg.ModelShardID, int], + model_topos: Dict[str, topology.PipeModelDataParallelTopology], + model_configs: Dict[str, None | ReaLModelConfig], + ctrl: RPCCorountineControl, + buffer: AsyncIOSequenceBuffer, + redistrib_planner: RedistribPlanner, + summary_writer: SummaryWriter | None, + ): + + self.rpc = rpc + self.src_rpc = src_rpc + self.stream = stream + + self.n_model_workers = len(set(msid2mwid.values())) + + self.msid2mwid = msid2mwid + self.model_topos = model_topos + self.model_configs = model_configs + + self.mwid2msids = defaultdict(list) + for msid, mwid in msid2mwid.items(): + self.mwid2msids[mwid].append(msid) + + self.model_save_root = os.path.join( + constants.MODEL_SAVE_ROOT, + constants.experiment_name(), + constants.trial_name(), + ) + + self.rpc_ctrl = ctrl + self.buffer = buffer + self.redistrib_planner = redistrib_planner + + self.summary_writer = summary_writer + + @property + def dp_size(self): + return self.model_topos[self.rpc.model_name].get_dim("data") + + @property + def pp_size(self): + return self.model_topos[self.rpc.model_name].get_dim("pipe") + + def attach_payloads_with_hooks( + self, + payloads: Dict[config_api.ModelShardID, request_reply_stream.Payload], + mwids: List[int], + main_handlers: List[config_pkg.ModelShardID], + hook_type: str, + ) -> Tuple[Dict[config_api.ModelShardID, request_reply_stream.Payload], List[int]]: + assert hook_type in ["pre", "post"], hook_type + + rpc = self.rpc + model_topos = self.model_topos + model_configs = self.model_configs + + main_mwids = set([self.msid2mwid[h] for h in main_handlers]) + for hook in getattr(rpc, f"_{hook_type}_hooks"): + if isinstance(hook, dfg.ParamReallocHook): + assert (hook.source is None) != (hook.target is None), hook + if hook.source is None: + src_topo = model_topos[rpc.model_name] + dst_topo = model_topos[hook.target] + dst_config = model_configs[hook.target] + src_model_name, dst_model_name = rpc.model_name, hook.target + other_model_name = hook.target + other_topo = dst_topo + else: + src_topo = model_topos[hook.source] + dst_topo = model_topos[rpc.model_name] + dst_config = model_configs[rpc.model_name] + src_model_name, dst_model_name = hook.source, rpc.model_name + other_model_name = hook.source + other_topo = src_topo + + ps_data = { + "from_model_name": src_model_name, + "to_model_name": dst_model_name, + "from_topo": src_topo, + "to_topo": dst_topo, + "to_model_config": dst_config, + "eta": hook.eta, + } + for h in main_handlers: + getattr(payloads[h], f"{hook_type}_hooks").append("param_realloc") + getattr(payloads[h], f"{hook_type}_hook_data").append(ps_data) + other_handlers = [ + config_api.ModelShardID.from_parallelism_rank( + other_model_name, other_topo, j + ) + for j in range(other_topo.world_size()) + ] + for h in other_handlers: + if self.msid2mwid[h] not in mwids: + payloads[h] = request_reply_stream.Payload( + handler=h, + handle_name="empty", + ) + setattr(payloads[h], f"{hook_type}_hooks", ["param_realloc"]) + setattr(payloads[h], f"{hook_type}_hook_data", [ps_data]) + mwids.append(self.msid2mwid[h]) + elif self.msid2mwid[h] not in main_mwids: + hh = next( + hh + for hh in payloads + if self.msid2mwid[hh] == self.msid2mwid[h] + ) + getattr(payloads[hh], f"{hook_type}_hooks").append( + "param_realloc" + ) + getattr(payloads[hh], f"{hook_type}_hook_data").append(ps_data) + + elif isinstance(hook, dfg.OffloadHook): + for h in main_handlers: + getattr(payloads[h], f"{hook_type}_hooks").append("offload") + getattr(payloads[h], f"{hook_type}_hook_data").append( + dict(model_name=h.model_name) + ) + else: + raise NotImplementedError(f"Unknown hook type: {hook}") + return payloads, mwids + + def request( + self, + data_transfer_plan: List[RedistribStep], + partitioned_ids: List[Hashable], + meta_sample: data_api.SequenceSample, + handlers: List[config_pkg.ModelShardID], + ) -> Tuple[List[uuid.UUID], List[uuid.UUID]]: + + rpc = self.rpc + ctrl = self.rpc_ctrl + + dt_data = { + "target": rpc.model_name, + "plan": [json.dumps(dataclasses.asdict(x)) for x in data_transfer_plan], + "partitioned_ids": partitioned_ids, + "handle_name": rpc.interface_type.value, + "meta_sample": meta_sample, + "partitioned_ids": partitioned_ids, + } + + payloads = { + handler: request_reply_stream.Payload( + handler=handler, + handle_name=rpc.interface_type.value, + pre_hooks=["data_transfer"], + pre_hook_data=[dt_data], + data=rpc.name, + ) + for handler in handlers + } + if ctrl.should_eval: + for p in payloads.values(): + p.post_hooks.append("evaluate") + p.post_hook_data.append(dict(model_name=rpc.model_name)) + if ( + ctrl.should_save or ctrl.should_ckpt + ) and rpc.interface_type == dfg.ModelInterfaceType.TRAIN_STEP: + for p in payloads.values(): + p.post_hooks.append("save") + save_dir = os.path.join( + self.model_save_root, + rpc.model_name.role, + f"epoch{ctrl.step_info.epoch + 1}" + f"epochstep{ctrl.step_info.epoch_step + 1}" + f"globalstep{ctrl.step_info.global_step + 1}", + ) + p.post_hook_data.append( + dict( + model_name=rpc.model_name, + save_dir=save_dir, + recover_only=not ctrl.should_save, + ) + ) + mwids = [self.msid2mwid[h] for h in handlers] + assert len(mwids) == len(set(mwids)) + + for step in data_transfer_plan: + if step.root not in mwids: + handler = self.mwid2msids[step.root][0] + payloads[handler] = request_reply_stream.Payload( + handler=handler, + handle_name="empty", + pre_hooks=["data_transfer"], + pre_hook_data=[dt_data], + ) + mwids.append(step.root) + if step.comm_type == "gather": + for src in step.srcs: + if src not in mwids: + handler = self.mwid2msids[src][0] + payloads[handler] = request_reply_stream.Payload( + handler=handler, + handle_name="empty", + pre_hooks=["data_transfer"], + pre_hook_data=[dt_data], + ) + mwids.append(src) + + payloads, mwids = self.attach_payloads_with_hooks( + payloads, + mwids, + main_handlers=handlers, + hook_type="pre", + ) + payloads, mwids = self.attach_payloads_with_hooks( + payloads, + mwids, + main_handlers=handlers, + hook_type="post", + ) + main_payloads = [p for h, p in payloads.items() if h in handlers] + other_payloads = [p for h, p in payloads.items() if h not in handlers] + all_req_ids = self.stream.request( + payloads=main_payloads + other_payloads, + ) + return all_req_ids[: len(main_payloads)], all_req_ids[len(main_payloads) :] + + def data_parallel_dispatch( + self, buf_indices: List[int], sample: data_api.SequenceSample + ) -> Tuple[List[int], data_api.SequenceSample, List[Tuple[int, int]]]: + # Dispatch data to different data parallel ranks. + if self.rpc.is_generate(): + # The workload of generation is decided by batch size, instead of the generated length. + lens = [1 for _ in range(sample.bs)] + samples, forward_indices, _ = sample.split_with_lengths( + mb_spec=data_api.MicroBatchSpec(n_mbs=self.dp_size), + lens=lens, + ) + else: + samples, forward_indices, _ = sample.split( + data_api.MicroBatchSpec(n_mbs=self.dp_size) + ) + blogger.info( + f"DP split (DP size {self.dp_size}) for RPC {self.rpc.name}: " + f"#seqs: {[s.bs for s in samples]}, " + f"#tokens: {[sum([sum(lens) for lens in s.seqlens[s._get_split_key()]]) for s in samples]}" + ) + sample = data_api.SequenceSample.gather(samples) + buf_indices = [buf_indices[i] for i in forward_indices] + + partitions = data_api.SequenceSplitSpec( + sizes=[s.bs for s in samples] + ).partitions + return buf_indices, sample, partitions + + async def run_step(self, buf_indices, sample): + rpc = self.rpc + topo = self.model_topos[rpc.model_name] + ctrl = self.rpc_ctrl + + handlers = [ + config_pkg.ModelShardID.from_parallelism_rank(rpc.model_name, topo, j) + for j in range(topo.world_size()) + ] + + dp_head_indices = [ + topo.get_rank(data=i, pipe=topo.get_dim("pipe") - 1, model=0) + for i in range(self.dp_size) + ] + + ctrl.flops_counter.add_rpc(rpc, sample, self.model_configs[rpc.model_name]) + + # logger.info(f"Model rpc {rpc.name} requesting.") + + # Sample may be reordered here. + buf_indices, sample, partitions = self.data_parallel_dispatch( + buf_indices, sample + ) + + # Build data destinations: GPU id -> List[data ids] + partitioned_ids = [] + dests = {} + for dp_rank, (st, ed) in enumerate(partitions): + ranks = topo.filter_match(data=dp_rank) + for rank in ranks: + h = config_pkg.ModelShardID.from_parallelism_rank( + model_name=rpc.model_name, topo=topo, parallelism_rank=rank + ) + gpu_id = self.msid2mwid[h] + assert gpu_id not in dests + dests[gpu_id] = sample.ids[st:ed] + partitioned_ids.append(sample.ids[st:ed]) + for i in range(self.n_model_workers): + if i not in dests: + dests[i] = [] + + # NOTE: The data loaded from the dataset may be unevenly distributed across DP ranks. + # Only bcast works in this case. + if rpc.is_src: + pattern = "bcast" + else: + pattern = "gather-scatter" + data_transfer_plan = self.redistrib_planner.derive_plan( + dests, + keys=rpc.input_keys, + pattern=pattern, + ) + blogger.info(f"Data tranfer plan for `{rpc.name}`: {data_transfer_plan}.") + + # Update storage tracker for transferred data. + if rpc.is_src: + # NOTE: since the data we loaded may be unevenly distributed across DP ranks, + # we should change the owner of the data to the src RPC. + for i in range(topo.world_size()): + h = ModelShardID.from_parallelism_rank( + model_name=rpc.model_name, topo=topo, parallelism_rank=i + ) + is_dp_head = h.mp_rank == 0 and h.pp_rank == topo.get_dim("pipe") - 1 + gpu_id = self.msid2mwid[h] + for key in rpc.input_keys: + self.redistrib_planner.storage_tracker.add_data( + gpu_id, partitioned_ids[h.dp_rank], key=key, is_owner=is_dp_head + ) + else: + for step in data_transfer_plan: + if step.comm_type == "scatter": + for gpu_id, ids in zip(step.dsts, step.ids): + for key in step.keys: + self.redistrib_planner.storage_tracker.add_data( + gpu_id, ids, key=key, is_owner=False + ) + elif step.comm_type == "gather": + for key in step.keys: + self.redistrib_planner.storage_tracker.add_data( + step.root, + list(itertools.chain.from_iterable(step.ids)), + key=key, + is_owner=False, + ) + + # send partitioned data to model workers + req_ids, other_req_ids = self.request( + data_transfer_plan=data_transfer_plan, + partitioned_ids=partitioned_ids, + meta_sample=sample, + handlers=handlers, + ) + tik = time.perf_counter() + + await ctrl.topo_level_count.put(1) + logger.info(f"Model rpc {rpc.name} requested.") + + # Then, wait for all main requests to finish. + responses = await self.stream.gather_async(request_ids=req_ids) + # logger.info(f"rpc {rpc.name} received responses {req_ids}") + + # Filter out responses other than DP heads. + # Other repsonses are duplicated or None. + responses = [responses[i] for i in dp_head_indices] + + # If the returned data is a SequenceSample, it is the data returned by + # model function calls. The data shoulbe be amended into buffer. + # Otherwise, it's the train statistics and should be reduced and logged. + if isinstance(responses[-1], data_api.SequenceSample): + # Update storage tracker for generated data. + for dp_rank, x in enumerate(responses): + pp_size = topo.get_dim("pipe") + ranks = topo.filter_match(data=dp_rank, pipe=pp_size - 1, model=0) + for rank in ranks: + h = config_pkg.ModelShardID.from_parallelism_rank( + model_name=rpc.model_name, topo=topo, parallelism_rank=rank + ) + gpu_id = self.msid2mwid[h] + for k in rpc.output_keys: + self.redistrib_planner.storage_tracker.add_data( + gpu_id, + x.ids, + key=k, + is_owner=True, + ) + res = data_api.SequenceSample.gather(responses) + else: + res = data_api.gather_stat(responses) + + if rpc.log_return_value: + logger.info(f"RPC name {rpc.name} returns {res}") + + if isinstance(res, Dict): + wandb.log(res, step=ctrl.step_info.global_step) + if self.summary_writer is not None: + for key, val in res.items(): + self.summary_writer.add_scalar( + f"{key}", val, ctrl.step_info.global_step + ) + + logger.info( + f"Model rpc {rpc.name} finished. " + f"Run time {time.perf_counter() - tik:.4f}s." + ) + + # If this RPC is the final node in the dataflow graph, + # update the train counter. + # Otherwise, amend data in the buffer. + if rpc.is_dst: + ctrl.ids_to_clear = ctrl.ids_to_clear.union(sample.ids) + await ctrl.train_count.put(1) + else: + logger.info(f"Amending RPC {rpc.name} output keys: {res.keys}") + await self.buffer.amend_batch(buf_indices, res.unpack()) + + # Wait for all side-effect requests to finish. + # Side-effect or empty requests are required for data transfer + # and parameter synchronization. + # Wait them after the main request to log the oorrect MFC time. + await self.stream.gather_async(other_req_ids) + + async def run(self): + rpc = self.rpc + topo = self.model_topos[rpc.model_name] + + logger.info( + f"Running Model RPC, interface_type=#{rpc.interface_type}# " + f"(dp,mp,pp) = *({topo.get_dim('data')},{topo.get_dim('model')},{topo.get_dim('pipe')})*" + ) + + consumed = 0 + while True: + buf_indices, sample = await self.buffer.get_batch_for_rpc(rpc) + + await self.run_step(buf_indices, sample) + consumed += sample.bs + + # Ensure that parent RPCs will not be over-consumed. + if all(consumed >= c.n_seqs for c in rpc.all_successors()): + break diff --git a/realhf/system/model_worker.py b/realhf/system/model_worker.py index d2f4a78cc1..25700348fd 100644 --- a/realhf/system/model_worker.py +++ b/realhf/system/model_worker.py @@ -30,7 +30,6 @@ import realhf.api.core.dfg as dfg import realhf.api.core.system_api as system_api -import realhf.impl.model.comm.data_transfer as data_transfer_comm import realhf.impl.model.comm.global_comm as global_comm import realhf.impl.model.comm.param_realloc as param_realloc_comm from realhf.api.core.config import ModelName @@ -49,11 +48,12 @@ cuda_tmark, cuda_tmarked, dump_tmark_db, - gpu_utilization_monitor, ) from realhf.impl.model.nn.real_llm_api import ReaLModel from realhf.impl.model.utils import cuda_graph from realhf.system import request_reply_stream, worker_base +from realhf.system.data_manager import DataManager +from realhf.system.redistributor import RedistribStep # NOTE: Register all implemented datasets and models. import realhf.api.core.data_api as data_api # isort:skip @@ -257,11 +257,12 @@ def __lazy_setup(self): msid2mwid=self.config.msid2mwid, ) - self.__data_transfer_info = data_transfer_comm.setup_data_transfer( + self.data_manager = DataManager( model_topos=self.config.model_topos, msid2mwid=self.config.msid2mwid, data_transfer_pairs=self.config.data_transfer_pairs, ) + self.data_manager.setup_process_groups() self.__param_realloc_info = param_realloc_comm.setup_param_realloc( model_topos=self.config.model_topos, @@ -340,6 +341,8 @@ def __lazy_setup(self): for tmp_sample in self.__dataloader: self.__raw_samples += tmp_sample.meta().unpack() + self.__data_generator = enumerate(self.__dataloader) + self.__models: Dict[ModelName, model_api.Model] = dict() self.__model_is_handle: Dict[ModelName, bool] = dict() self.__interfaces: Dict[ModelName, model_api.ModelInterface] = dict() @@ -464,17 +467,6 @@ def __lazy_setup(self): self.__reply_queue = queue.Queue(maxsize=8) self.__request_sample_size = dict() - # Storing data loaded from the dataset and outputs of the - # model function call. - self.__data_storage: Dict[int, data_api.SequenceSample] = {} - - self.__data_sent_worker_indices: Dict[int, Dict[str, Set]] = ( - collections.defaultdict(lambda: collections.defaultdict(set)) - ) - self.__data_received_worker_indices: Dict[int, Dict[str, Set]] = ( - collections.defaultdict(lambda: collections.defaultdict(set)) - ) - self.__compute_input_queues = { model_name: dict( train_step=queue.Queue(4), @@ -581,7 +573,10 @@ def handle_non_blocking_request(self, request: request_reply_stream.Payload): elif request.handle_name == "fetch": dp_rank = int(re.search(r"__data(\d+)__", request.handler).group(1)) assert self.__has_dataset - if request.data["first_batch"]: + # Fetch. + try: + self.__dataset_batch_counter, cur_sample = next(self.__data_generator) + except StopIteration: # Upon the first fetch request, filter dataset and create dataloader. eval_scores_path = os.path.join( constants.MODEL_SAVE_ROOT, @@ -595,10 +590,8 @@ def handle_non_blocking_request(self, request: request_reply_stream.Payload): constants.trial_name(), f"dataset_indices_{dp_rank}.npy", ) - if ( - hasattr(self.__dataset, "filter") - and not request.data["first_poll"] - and os.path.exists(eval_scores_path) + if hasattr(self.__dataset, "filter") and os.path.exists( + eval_scores_path ): # Don't filter dataset on the first poll after recover. with open(eval_scores_path, "r", encoding="utf-8") as f: @@ -621,17 +614,15 @@ def handle_non_blocking_request(self, request: request_reply_stream.Payload): generator=g, ) self.__data_generator = enumerate(self.__dataloader) - - # Fetch. - self.__dataset_batch_counter, cur_sample = next(self.__data_generator) + self.__dataset_batch_counter, cur_sample = next(self.__data_generator) # Defer data that has not been used in the previous epoch. data_loaded = [] for x in cur_sample.unpack(): - if x.ids[0] in self.__data_storage: + if self.data_manager.has_data(x.ids[0]): continue data_loaded.append(x) - self.__data_storage[x.ids[0]] = x + self.data_manager.store(x) assert len(set([x.ids[0] for x in data_loaded])) == len(data_loaded) if len(data_loaded) > 0: @@ -652,13 +643,7 @@ def handle_non_blocking_request(self, request: request_reply_stream.Payload): elif request.handle_name == "clear_data_cache": with cuda_tmarked("clear_data_cache", CUDATimeMarkType.misc): ids = request.data - for _id in ids: - if _id in self.__data_storage: - del self.__data_storage[_id] - if _id in self.__data_sent_worker_indices: - del self.__data_sent_worker_indices[_id] - if _id in self.__data_received_worker_indices: - del self.__data_received_worker_indices[_id] + self.data_manager.remove(ids) gc.collect() if ( self.config.cuda_cache_cleanliness @@ -672,7 +657,7 @@ def handle_non_blocking_request(self, request: request_reply_stream.Payload): ) logger.info( "Get clear_data_cache, dump cuda tmark. " - f"Remaining data in local storage: {len(self.__data_storage)}. " + f"Remaining data in local storage: {self.data_manager.storage_size()}. " ) dump_tmark_db(self.__worker_index) res = request_reply_stream.NoResponse() @@ -707,9 +692,6 @@ def handle_blocking_request( # e.g., data transfer, parameter syncrhonization. pass elif request.handle_name == "initialize": - assert not self.__model_is_handle[ - request.handler.model_name - ], request.handler.model_name self.__models[request.handler.model_name] = self._backend.initialize( self._model, data ) @@ -942,7 +924,7 @@ def __handle_model_function_calls( for x in res.unpack(): # The input data must exist in the storage, otherwise # the model function call will not run. - self.__data_storage[x.ids[0]].update_(x) + self.data_manager.update(x) # Only return meta data back to the master worker. if isinstance(res, data_api.SequenceSample): @@ -969,32 +951,18 @@ def __handle_model_function_calls( @cuda_tmark("data_transfer", CUDATimeMarkType.comm) def __data_transfer_among_workers(self, hook_data: Dict[str, Any]): meta_sample = hook_data["meta_sample"] - comm_plan = data_transfer_comm.derive_data_transfer_plan( - keys=hook_data["keys"], - global_ids=meta_sample.ids, - consumer_name=hook_data["target"], - consumer_mapping=hook_data["target_mapping"], - producer_names=hook_data["producer_names"], - producer_mappings=hook_data["producer_mappings"], - data_transfer_info=self.__data_transfer_info, - ) - data_transfer_comm.run_data_transfer( - comm_plan=comm_plan, - meta_samples={x.ids[0]: x for x in meta_sample.unpack()}, - storage=self.__data_storage, - sent_worker_idx_table=self.__data_sent_worker_indices, - received_worker_idx_table=self.__data_received_worker_indices, - ) + plan = [RedistribStep(**json.loads(x)) for x in hook_data["plan"]] + self.data_manager.redistribute(meta_sample, plan=plan) if hook_data["target"] in self.__models: with constants.model_scope(hook_data["target"]): - local_ids = [ - meta_sample.ids[i] - for i in hook_data["target_mapping"][self._dp_rank] - ] + local_ids = hook_data["partitioned_ids"][self._dp_rank] r = data_api.SequenceSample.gather( - [self.__data_storage[_id] for _id in local_ids], + [ + self.data_manager.get(_id).to_device(constants.current_device()) + for _id in local_ids + ], keys=meta_sample.keys, ) self.__compute_input_queues[hook_data["target"]][ @@ -1351,7 +1319,6 @@ def __experiment_complete_exit(self): self.__models.clear() self.__backends.clear() self.__interfaces.clear() - self.__data_storage.clear() # Reset model worker states. self.__dist_env_resolved = False diff --git a/realhf/system/redistributor.py b/realhf/system/redistributor.py new file mode 100644 index 0000000000..caeeaad5be --- /dev/null +++ b/realhf/system/redistributor.py @@ -0,0 +1,343 @@ +# Copyright 2025 Ant Group Inc. + +import dataclasses +import itertools +import os +from collections import defaultdict +from typing import * + +from realhf.base.cluster import spec as cluster_spec + + +class GlobalStorageTracker: + def __init__(self, world_size: int): + self.storages: List[Dict[Hashable, List[str]]] + self.storages = [{} for _ in range(world_size)] + self.data_owner: Dict[Tuple[Hashable, str], int] + self.data_owner = {} + + def add_data(self, rank: int, ids: List[Hashable], key: str, is_owner: bool): + for data_id in ids: + if data_id not in self.storages[rank]: + self.storages[rank][data_id] = [key] + else: + if key not in self.storages[rank][data_id]: + self.storages[rank][data_id].append(key) + if is_owner: + self.data_owner[(data_id, key)] = rank + + def clear_data(self, ids: List[Hashable]): + for storage in self.storages: + for i in ids: + if i in storage: + storage.pop(i) + keys = list(self.data_owner.keys()) + for i, k in keys: + if i in ids: + self.data_owner.pop((i, k)) + + +@dataclasses.dataclass +class RedistribStep: + comm_type: str + root: int | None + srcs: List[int] | None + dsts: List[int] | None + ids: List[List[Hashable]] + keys: List[str] + + def __repr__(self) -> str: + if self.comm_type == "gather": + return f"Gather {self.keys} to {self.root} from {self.srcs}." + if self.comm_type == "scatter": + return f"Scatter {self.keys} from {self.root} to {self.dsts}." + if self.comm_type == "bcast": + return f"Bcast {self.keys} from {self.root} to {self.dsts}." + raise NotImplementedError() + + +class RedistribPlanner: + def __init__(self, storage_tracker: GlobalStorageTracker): + self.storage_tracker = storage_tracker + + def derive_plan( + self, + dests: Dict[int, List[Hashable]], + keys: List[str], + pattern: str = "gather-scatter", + ) -> List[RedistribStep]: + if pattern == "gather-scatter": + return self.derive_plan_gather_scatter(dests, keys) + elif pattern == "bcast": + return self.derive_plan_bcast(dests, keys) + raise NotImplementedError(f"Unknown data redistribution pattern: {pattern}") + + def derive_plan_gather_scatter( + self, dests: Dict[int, List[Hashable]], keys: List[str] + ) -> List[RedistribStep]: + self.dests = dests + + all_data_ids = set() + for all_samples in dests.values(): + for data_id in all_samples: + all_data_ids.add(data_id) + + transfer_plan = [] + for key in keys: + owners = sorted( + list( + set( + [ + self.storage_tracker.data_owner[(i, key)] + for i in all_data_ids + ] + ) + ) + ) + gather_ids = [] + for owner in owners: + this_owner_ids = [] + for i in all_data_ids: + if ( + i in self.storage_tracker.storages[owner] + and key in self.storage_tracker.storages[owner][i] + ): + this_owner_ids.append(i) + gather_ids.append(sorted(this_owner_ids)) + gather_step = RedistribStep( + comm_type="gather", + root=owners[0], + srcs=owners, + dsts=None, + ids=gather_ids, + keys=[key], + ) + + scatter_dsts = sorted([i for i in dests if dests[i]]) + scatter_ids = [sorted(dests[i]) for i in scatter_dsts] + scatter_step = RedistribStep( + comm_type="scatter", + root=owners[0], + dsts=scatter_dsts, + srcs=None, + ids=scatter_ids, + keys=[key], + ) + transfer_plan += [gather_step, scatter_step] + + # Prune the plan. + pop_indices = [] + for idx, step in enumerate(transfer_plan): + # 1. Omit the gather step if data has already been gathered before. + if step.comm_type == "gather": + all_gather_ids = list(itertools.chain.from_iterable(step.ids)) + key = step.keys[0] + if any( + i not in self.storage_tracker.storages[step.root] + for i in all_gather_ids + ): + continue + if any( + key not in self.storage_tracker.storages[step.root][i] + for i in all_gather_ids + ): + continue + pop_indices.append(idx) + # 2. Omit the gather + scatter step if data has already exists in all dst GPUs. + if step.comm_type == "scatter": + key = step.keys[0] + all_exists = True + for dst, ids in zip(step.dsts, step.ids): + if any(i not in self.storage_tracker.storages[dst] for i in ids): + all_exists = False + break + if any( + key not in self.storage_tracker.storages[dst][i] for i in ids + ): + all_exists = False + break + if all_exists: + pop_indices.append(idx) + pop_indices.append(idx - 1) + for pop_idx in reversed(sorted(set(pop_indices))): + transfer_plan.pop(pop_idx) + + # Merging the gather/scatter of different keys + gather_plan = {} + scatter_plan = {} + for step in transfer_plan: + if step.comm_type == "gather": + plan_key = ( + step.root, + tuple(sorted(step.srcs)), + tuple([tuple(sorted(ids)) for ids in step.ids]), + ) + if plan_key not in gather_plan: + gather_plan[plan_key] = step + else: + assert all( + k not in gather_plan[plan_key].keys for k in step.keys + ), ( + gather_plan[plan_key], + step, + plan_key, + ) + gather_plan[plan_key].keys += step.keys + if step.comm_type == "scatter": + plan_key = ( + step.root, + tuple(sorted(step.dsts)), + tuple([tuple(sorted(ids)) for ids in step.ids]), + ) + if plan_key not in scatter_plan: + scatter_plan[plan_key] = step + else: + assert all( + k not in scatter_plan[plan_key].keys for k in step.keys + ), ( + scatter_plan[plan_key], + step, + plan_key, + ) + scatter_plan[plan_key].keys += step.keys + + # Prioritize gather over scatter + return list(gather_plan.values()) + list(scatter_plan.values()) + + def derive_plan_bcast( + self, dests: Dict[int, List[Hashable]], keys: List[str] + ) -> List[RedistribStep]: + assert isinstance(keys, list), type(keys) + self.dests = dests + + # Get all requried data IDs. + all_data_ids = set() + for all_samples in self.dests.values(): + for data_id in all_samples: + all_data_ids.add(data_id) + + # The producers for each required data. + id2gpu_src = {} + for data_id in all_data_ids: + for key in keys: + id2gpu_src[(data_id, key)] = [] + for gpu_id, ids2keys in enumerate(self.storage_tracker.storages): + if data_id in ids2keys and key in ids2keys[data_id]: + id2gpu_src[(data_id, key)].append(gpu_id) + + # The consumers for each requried data. + id2gpu_dst = {} + for data_id in all_data_ids: + for key in keys: + id2gpu_dst[(data_id, key)] = [] + for gpu_id, ids in self.dests.items(): + if data_id in ids: + id2gpu_dst[(data_id, key)].append(gpu_id) + + self.transfer_plan = {} + + for data_id, key in itertools.product(all_data_ids, keys): + source_gpus = id2gpu_src[(data_id, key)] + target_gpus = id2gpu_dst[(data_id, key)] + + assert len(source_gpus) > 0, (data_id, key, id2gpu_src, id2gpu_dst) + + # Omit data transfer if it exists in the target GPU + target_gpus = [gpu for gpu in target_gpus if gpu not in source_gpus] + if not target_gpus: + continue + + # Find the "nearest" GPU for data transfer. + best_src = self._select_best_bcast_source(source_gpus, target_gpus) + + self.transfer_plan[(data_id, key)] = {"src": best_src, "dsts": target_gpus} + + return self._group_bcast_transfers() + + def _on_same_node(self, i, j) -> bool: + return (i // cluster_spec.n_gpus_per_node) == ( + j // cluster_spec.n_gpus_per_node + ) + + def _select_best_bcast_source(self, source_gpus, target_gpus): + same_node_counts = {} + for src in source_gpus: + same_node_count = sum( + 1 for dst in target_gpus if self._on_same_node(src, dst) + ) + same_node_counts[src] = same_node_count + + # Find the source that maximizes locality. + max_same_node = max(same_node_counts.values()) + best_sources = [ + src for src, count in same_node_counts.items() if count == max_same_node + ] + + # Find the source with the smallest workload. + src_load = defaultdict(int) + for plan in self.transfer_plan.values(): + src_gpu = plan["src"] + src_load[src_gpu] += len(plan["dsts"]) + return min(best_sources, key=lambda src: src_load[src]) + + def _group_bcast_transfers(self) -> List[RedistribStep]: + # Group data ids that should be transferred from "src" to "dsts" + src_to_transfers = defaultdict(lambda: defaultdict(list)) + for (data_id, key), plan in self.transfer_plan.items(): + src_to_transfers[(plan["src"], key)][tuple(sorted(plan["dsts"]))].append( + data_id + ) + + stages = [] + while any(src_to_transfers.values()): + stage = [] + used_dsts = set() + used_srcs = set() + + for (src, key), transfers in list(src_to_transfers.items()): + if src in used_srcs: + continue + if not transfers: + continue + + # Find a transfer that can be concurrent executed. + pop_key = None + for i, dsts in enumerate(transfers): + if not any(dst in used_dsts for dst in dsts): + pop_key = dsts + break + + if pop_key is not None: + data_ids = transfers.pop(pop_key) + stage.append( + RedistribStep( + comm_type="bcast", + root=src, + srcs=[src], + keys=[key], + dsts=pop_key, + ids=data_ids, + ) + ) + used_dsts.update(dsts) + used_srcs.add(src) + + if stage: + stages += stage + else: + for (src, key), transfers in list(src_to_transfers.items()): + if transfers: + dsts, data_ids = transfers.pop(0) + stages.append( + RedistribStep( + comm_type="bcast", + srcs=[src], + root=src, + dsts=dsts, + ids=data_ids, + keys=[key], + ) + ) + break + + return stages diff --git a/realhf/system/request_reply_stream.py b/realhf/system/request_reply_stream.py index 9147143b58..7ac238db23 100644 --- a/realhf/system/request_reply_stream.py +++ b/realhf/system/request_reply_stream.py @@ -138,6 +138,9 @@ def __init__( f"subscribers: {name_resolve.get_subtree(names.request_reply_stream(experiment_name, trial_name, PUBSUB_BARRIER_NAME))}." ) + def route_to(self, handler) -> int: + return self._handler_routing[handler] + def close(self): self.recv_socket.close() for send_socket in self.send_sockets: diff --git a/realhf/system/worker_base.py b/realhf/system/worker_base.py index f6661956e3..c63e8765db 100644 --- a/realhf/system/worker_base.py +++ b/realhf/system/worker_base.py @@ -4,7 +4,6 @@ import dataclasses import enum -import getpass import os import queue import re @@ -14,15 +13,7 @@ from typing import Any, Dict, List, Optional, Tuple import realhf.api.core.system_api as system_api -from realhf.base import ( - cluster, - logging, - monitor, - name_resolve, - names, - network, - timeutil, -) +from realhf.base import logging, name_resolve, names from realhf.base.gpu_utils import set_cuda_device logger = logging.getLogger("worker") @@ -580,6 +571,7 @@ def configure( expr_config.lazy_init() self.wandb_config = expr_config.wandb os.environ["WANDB_MODE"] = self.wandb_config.mode + self.tensorboard_config = expr_config.tensorboard config = expr_config.resolve_worker_config( self.__worker_type, self.__worker_index ) diff --git a/requirements.txt b/requirements.txt index bc0ca01de8..b7ebf85c35 100644 --- a/requirements.txt +++ b/requirements.txt @@ -51,4 +51,6 @@ paramiko torch>2.0.0 black>=25.1.0 cookiecutter>2.1.1 +asyncio +aiohttp httpx>=0.28.1 diff --git a/tests/comm/test_data_transfer.py b/tests/comm/test_data_transfer.py new file mode 100644 index 0000000000..db76060bd0 --- /dev/null +++ b/tests/comm/test_data_transfer.py @@ -0,0 +1,288 @@ +# Copyright 2025 Ant Group Inc. +# Copyright 2024 Wei Fu & Zhiyu Mei +# Licensed under the Apache License, Version 2.0 (the "License"). + +import dataclasses +import os +import pathlib +import pickle +import uuid +from typing import * + +import numpy as np +import pytest +import torch +import torch.distributed as dist + +from realhf.api.core.config import ModelName, ModelShardID +from realhf.api.core.data_api import SequenceSample +from realhf.base import constants, testing, topology +from realhf.base.testing import ( + LocalMultiProcessTest, + PipeModelDataParallelTopology, + init_global_constants, +) +from realhf.system.data_manager import DataManager +from realhf.system.redistributor import GlobalStorageTracker, RedistribPlanner + + +def get_data_manager( + from_model_name, + to_model_name, + from_pp_dp_mp, + to_pp_dp_mp, +): + + from_num_pp, from_num_dp, from_num_mp = from_pp_dp_mp + to_num_pp, to_num_dp, to_num_mp = to_pp_dp_mp + + from_world_size = from_num_dp * from_num_mp * from_num_pp + to_world_size = to_num_dp * to_num_mp * to_num_pp + + from_topo = topology.PipeModelDataParallelTopology( + num_dp=from_num_dp, + num_mp=from_num_mp, + num_pp=from_num_pp, + sequence_parallel=False, + gradient_checkpointing=False, + max_prompt_len=None, + gradient_accumulation_fusion=False, + ) + to_topo = topology.PipeModelDataParallelTopology( + num_dp=to_num_dp, + num_mp=to_num_mp, + num_pp=to_num_pp, + sequence_parallel=False, + gradient_checkpointing=False, + max_prompt_len=None, + gradient_accumulation_fusion=False, + ) + + model_topos = {from_model_name: from_topo, to_model_name: to_topo} + + msid2mwid = {} + for i in range(dist.get_world_size()): + # We assume the `from_model` occupies the first serveral GPUs, + # while the `to_model` occupies GPUs from the last one. + # For example, when the world size of `from_model` is 6 and + # the world size of `to_model` is 4, the GPU layout is: + # GPU 0-3: from_model (shard 0-3) + # GPU 4-5: from_model (shard 4-5) + to_model (shard 0-1) + # GPU 6-7: to_model (shard 2-3) + _model_names = [] + if i < from_world_size: + _model_names.append(from_model_name) + if i >= dist.get_world_size() - to_world_size: + _model_names.append(to_model_name) + for _model_name in _model_names: + if _model_name == from_model_name: + coord = model_topos[_model_name].get_coord(i) + else: + coord = model_topos[_model_name].get_coord( + i + to_world_size - dist.get_world_size() + ) + k = ModelShardID( + _model_name, + dp_rank=coord.data, + mp_rank=coord.model, + pp_rank=coord.pipe, + topo=model_topos[_model_name], + ) + msid2mwid[k] = i + + init_global_constants( + num_dp=from_num_dp, + num_mp=from_num_mp, + num_pp=from_num_pp, + topo=from_topo, + model_name=from_model_name, + sequence_parallel=False, + msid2mwid=msid2mwid, + ) + + init_global_constants( + num_dp=to_num_dp, + num_mp=to_num_mp, + num_pp=to_num_pp, + model_name=to_model_name, + sequence_parallel=False, + msid2mwid=msid2mwid, + ) + + return DataManager( + model_topos=model_topos, + msid2mwid=msid2mwid, + data_transfer_pairs=[(from_model_name, to_model_name)], + ) + + +def recursive_assert_equal(x1, x2): + if type(x1) != type(x2): + raise AssertionError(f"{type(x1)} != {type(x2)}") + if isinstance(x1, dict): + assert set(x1.keys()) == set(x2.keys()) + for k in x1.keys(): + recursive_assert_equal(x1[k], x2[k]) + elif dataclasses.is_dataclass(x1): + for f in dataclasses.fields(x1): + recursive_assert_equal(getattr(x1, f.name), getattr(x2, f.name)) + elif isinstance(x1, torch.Tensor): + assert torch.allclose(x1, x2), (x1, x2) + elif isinstance(x1, list): + assert len(x1) == len(x2) + for a, b in zip(x1, x2): + recursive_assert_equal(a, b) + else: + assert x1 == x2 + + +def _test_data_transfer( + tmp_path, + from_pp_dp_mp: Tuple, + to_pp_dp_mp: Tuple, +): + + from_model_name = ModelName("data_transfer_test", 0) + from_topo = PipeModelDataParallelTopology( + num_pp=from_pp_dp_mp[0], + num_mp=from_pp_dp_mp[-1], + num_dp=from_pp_dp_mp[1], + sequence_parallel=True, + gradient_checkpointing=True, + gradient_accumulation_fusion=True, + ) + to_model_name = ModelName("data_transfer_test", 1) + to_topo = PipeModelDataParallelTopology( + num_pp=to_pp_dp_mp[0], + num_mp=to_pp_dp_mp[-1], + num_dp=to_pp_dp_mp[1], + sequence_parallel=True, + gradient_checkpointing=True, + gradient_accumulation_fusion=True, + ) + + data_manager = get_data_manager( + from_model_name, + to_model_name, + from_pp_dp_mp, + to_pp_dp_mp, + ) + data_manager.setup_process_groups() + + storage_tracker = GlobalStorageTracker(dist.get_world_size()) + planner = RedistribPlanner(storage_tracker) + + key = "input_ids" + + world_size = dist.get_world_size() + samples = [] + for dp_rank in range(from_pp_dp_mp[1]): + gpu_id = data_manager.msid2mwid[ + ModelShardID( + from_model_name, + dp_rank=dp_rank, + mp_rank=0, + pp_rank=from_pp_dp_mp[0] - 1, + topo=from_topo, + ) + ] + storage_tracker.add_data( + gpu_id, + ids=[i + dp_rank * world_size for i in range(world_size)], + key=key, + is_owner=True, + ) + + seqlens = torch.randint(10, 1000, size=(world_size,)) + dist.all_reduce(seqlens) + input_ids = torch.randint( + 0, + 10000, + size=(int(sum(seqlens)),), + ) + dist.all_reduce(input_ids) + + s = SequenceSample.from_default( + ids=[i + dp_rank * world_size for i in range(world_size)], + seqlens=seqlens.numpy().tolist(), + data=dict(input_ids=input_ids), + ) + + if dist.get_rank() == 0: + for ss in s.unpack(): + with open(os.path.join(tmp_path, f"{ss.ids[0]}.pkl"), "wb") as f: + pickle.dump(ss, f) + + samples.append(s) + if dist.get_rank() == gpu_id: + for ss in s.unpack(): + data_manager.store(ss) + + dist.barrier() + + all_ids = list(range(world_size * from_topo.get_dim("data"))) + np.random.shuffle(all_ids) + _all_ids = [all_ids] + dist.broadcast_object_list(_all_ids, src=0) + all_ids = _all_ids[0] + + dests = {} + for rank in range(to_topo.world_size()): + coord = to_topo.get_coord(rank) + dp_size = to_topo.get_dim("data") + gpu_id = data_manager.msid2mwid[ + ModelShardID( + to_model_name, + dp_rank=coord.data, + mp_rank=coord.model, + pp_rank=coord.pipe, + topo=to_topo, + ) + ] + size_per_dp = len(all_ids) // dp_size + dests[gpu_id] = [coord.data * size_per_dp + i for i in range(size_per_dp)] + + for gpu_id in range(world_size): + if gpu_id not in dests: + dests[gpu_id] = [] + + plan = planner.derive_plan(dests, keys=[key]) + data_manager.redistribute(SequenceSample.gather(samples), plan) + dist.barrier() + + for i, s in data_manager.storage.items(): + with open(os.path.join(tmp_path, f"{i}.pkl"), "rb") as f: + ss = pickle.load(f) + recursive_assert_equal(ss, s) + print("success") + + +parallelism = [(1, 4, 2), (1, 8, 1)] + + +@pytest.mark.skipif( + os.cpu_count() < 32 or testing.get_free_mem_gb() < 50, + reason="The parameter reallocation test requires at least 32 CPUs and 50GB memory.", +) +@pytest.mark.parametrize("from_pp_dp_mp", [(1, 4, 2)]) +@pytest.mark.parametrize("to_pp_dp_mp", [(1, 8, 1)]) +@pytest.mark.distributed +def test_data_transfer( + tmp_path, + from_pp_dp_mp: Tuple, + to_pp_dp_mp: Tuple, +): + expr_name = uuid.uuid4() + trial_name = uuid.uuid4() + constants.set_force_cpu(True) + test_impl = LocalMultiProcessTest( + world_size=16, + func=_test_data_transfer, + expr_name=expr_name, + trial_name=trial_name, + timeout_secs=300, + tmp_path=tmp_path, + from_pp_dp_mp=from_pp_dp_mp, + to_pp_dp_mp=to_pp_dp_mp, + ) + test_impl.launch() diff --git a/tests/experiments/test_math_ppo.py b/tests/experiments/test_math_ppo.py index f9e6a72d22..d653af735b 100644 --- a/tests/experiments/test_math_ppo.py +++ b/tests/experiments/test_math_ppo.py @@ -51,7 +51,7 @@ def math_dataset(request, save_path): "dp,pp,mp", [ (1, 1, 1), - (2, 1, 1), + (2, 1, 2), (1, 2, 1), (1, 1, 2), ], @@ -133,6 +133,7 @@ def test_ppo_symm( def test_ppo_global_reshard( tmp_path_factory, tokenizer, + math_dataset, save_path, cpu_hf_model, mconfig, @@ -238,17 +239,17 @@ def test_ppo_global_reshard( ), ), ) - run_test_exp(exp_cfg) # Actor/critic train and ref_inf/rew_inf are on disjoint # device meshes and executed concurrently. -@pytest.mark.parametrize("actor_gen", [(1, 2, 1)]) -@pytest.mark.parametrize("critic_inf", [(1, 1, 2)]) +@pytest.mark.parametrize("actor_gen", [(2, 2, 1)]) +@pytest.mark.parametrize("critic_inf", [(2, 1, 2)]) def test_ppo_param_realloc_sub_device_mesh( tmp_path_factory, tokenizer, + math_dataset, save_path, cpu_hf_model, mconfig, @@ -269,7 +270,7 @@ def test_ppo_param_realloc_sub_device_mesh( mode="local", allocation_mode="manual", n_nodes=1, - n_gpus_per_node=2, + n_gpus_per_node=8, actor=ModelTrainEvalConfig( path=str(save_path), init_from_scratch=True, @@ -305,49 +306,51 @@ def test_ppo_param_realloc_sub_device_mesh( ), ), actor_gen=MFCConfig( + device_mesh="NODE01:0,1,2,3", parallel=ParallelismConfig( data_parallel_size=actor_gen[0], model_parallel_size=actor_gen[1], pipeline_parallel_size=actor_gen[2], - ) + ), ), actor_train=MFCConfig( - device_mesh="NODE01:0", + device_mesh="NODE01:4,5,6,7", parallel=ParallelismConfig( - data_parallel_size=1, + data_parallel_size=4, model_parallel_size=1, pipeline_parallel_size=1, ), ), critic_inf=MFCConfig( + device_mesh="NODE01:4,5,6,7", parallel=ParallelismConfig( data_parallel_size=critic_inf[0], model_parallel_size=critic_inf[1], pipeline_parallel_size=critic_inf[2], - ) + ), ), rew_inf=MFCConfig( - device_mesh="NODE01:1", + device_mesh="NODE01:4,5,6,7", parallel=ParallelismConfig( - data_parallel_size=1, + data_parallel_size=4, model_parallel_size=1, pipeline_parallel_size=1, ), ), ref_inf=MFCConfig( - device_mesh="NODE01:0", + device_mesh="NODE01:4,5,6,7", parallel=ParallelismConfig( data_parallel_size=1, - model_parallel_size=1, - pipeline_parallel_size=1, + model_parallel_size=2, + pipeline_parallel_size=2, ), ), critic_train=MFCConfig( - device_mesh="NODE01:1", + device_mesh="NODE01:4,5,6,7", parallel=ParallelismConfig( - data_parallel_size=1, + data_parallel_size=2, model_parallel_size=1, - pipeline_parallel_size=1, + pipeline_parallel_size=2, ), ), ) @@ -470,6 +473,8 @@ def test_ppo_save( ), ), ) + exp_cfg.actor.vllm.hybrid_train = True + exp_cfg.actor.vllm.enforce_eager = True run_test_exp(exp_cfg) diff --git a/tests/experiments/test_sft.py b/tests/experiments/test_sft.py index 8407673079..a7df7b9a43 100644 --- a/tests/experiments/test_sft.py +++ b/tests/experiments/test_sft.py @@ -35,7 +35,15 @@ def model_class(request): ], ) def test_sft_xl(tmp_path_factory, tokenizer, save_path, cpu_hf_model, dp, pp, tp): - test_sft(tmp_path_factory, tokenizer, save_path, cpu_hf_model, dp, pp, tp) + test_sft( + tmp_path_factory, + tokenizer, + save_path, + cpu_hf_model, + dp, + pp, + tp, + ) @pytest.mark.parametrize( diff --git a/tests/experiments/utils.py b/tests/experiments/utils.py index b1fdd1b3a3..579b09e2d6 100644 --- a/tests/experiments/utils.py +++ b/tests/experiments/utils.py @@ -18,7 +18,11 @@ def model_class(request): return request.param -def run_model_worker(cfg, mw, barrier): +def run_model_worker( + cfg, + mw, + barrier, +): constants.set_force_cpu(True) # Register all datasets and models import realhf.impl.dataset # isort: skip @@ -44,7 +48,11 @@ def run_model_worker(cfg, mw, barrier): initd = True -def run_test_exp(exp_cfg: Experiment, expr_name=None, trial_name=None): +def run_test_exp( + exp_cfg: Experiment, + expr_name=None, + trial_name=None, +): constants.set_force_cpu(True) # Register all datasets and models import realhf.impl.dataset # isort: skip @@ -74,7 +82,12 @@ def run_test_exp(exp_cfg: Experiment, expr_name=None, trial_name=None): testcase = testing.LocalMultiProcessTest( world_size=len(exp_setup.model_worker), func=[ - functools.partial(run_model_worker, cfg=exp_cfg, mw=mw, barrier=barrier) + functools.partial( + run_model_worker, + cfg=exp_cfg, + mw=mw, + barrier=barrier, + ) for mw in exp_setup.model_worker ], expr_name=expr_name or testing._DEFAULT_EXPR_NAME,