diff --git a/Dockerfile b/Dockerfile index 3ee8c8ef0e..e01ac04566 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,42 +1,4 @@ -ARG REAL_CPU_BASE_IMAGE -ARG REAL_GPU_BASE_IMAGE - -# >>>>>> CPU image -FROM ${REAL_CPU_BASE_IMAGE} as cpu - -ENV DEBIAN_FRONTEND=noninteractive -RUN apt update -RUN apt install -y ca-certificates -RUN sed -i "s@http://.*archive.ubuntu.com@https://mirrors.tuna.tsinghua.edu.cn@g" /etc/apt/sources.list -RUN sed -i "s@http://.*security.ubuntu.com@https://mirrors.tuna.tsinghua.edu.cn@g" /etc/apt/sources.list -RUN apt update -RUN apt install -y net-tools python3-pip pkg-config libopenblas-base libopenmpi-dev git - -RUN pip3 install -U pip -RUN pip3 config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple -# Install PyTorch in advance to prevent rebuilding this large Docker layer. -RUN pip3 install torch==2.3.1 - -RUN pip3 install deepspeed==0.14.0 megatron==0.6.0 - -COPY ./requirements.txt /requirements.txt -RUN pip3 install -r /requirements.txt && rm /requirements.txt - -COPY . /realhf -RUN REAL_NO_EXT=1 pip3 install -e /realhf --no-build-isolation -WORKDIR /realhf - -# >>>>>> Documentation images -# FROM cpu AS docs-builder -# RUN pip install -U sphinx sphinx-nefertiti -i https://pypi.tuna.tsinghua.edu.cn/simple -# RUN sphinx-build -M html /realhf/docs/source/ /realhf/docs/build/ -FROM nginx:alpine AS docs -COPY ./docs/build/html /usr/share/nginx/html -EXPOSE 80 -CMD ["nginx", "-g", "daemon off;"] - -# >>>>>> GPU image -FROM ${REAL_GPU_BASE_IMAGE} AS gpu +FROM nvcr.io/nvidia/pytorch:24.07-py3 AS gpu ENV DEBIAN_FRONTEND=noninteractive RUN apt update @@ -67,9 +29,9 @@ RUN pip3 install flash-attn==2.4.2 --no-build-isolation # Install grouped_gemm for MoE acceleration RUN pip3 install git+https://github.com/tgale96/grouped_gemm.git@v0.1.4 --no-build-isolation --no-deps -COPY . /realhf -RUN REAL_CUDA=1 pip3 install -e /realhf --no-build-isolation -WORKDIR /realhf +COPY . /AReaL +RUN REAL_CUDA=1 pip3 install -e /AReaL --no-build-isolation +WORKDIR /AReaL RUN git clone --depth=1 -b v0.6.3.post1 https://github.com/vllm-project/vllm.git /vllm RUN apt install kmod ccache -y @@ -85,4 +47,4 @@ RUN apt-get update && apt-get install -y python3.10-venv RUN git clone --depth=1 https://github.com/QwenLM/Qwen2.5-Math /qwen2_5-math && mv /qwen2_5-math/evaluation/latex2sympy /latex2sympy RUN python3 -m venv /sympy RUN /sympy/bin/pip install /latex2sympy -RUN /sympy/bin/pip install regex numpy tqdm datasets python_dateutil sympy==1.12 antlr4-python3-runtime==4.11.1 word2number Pebble timeout-decorator prettytable \ No newline at end of file +RUN /sympy/bin/pip install regex numpy tqdm datasets python_dateutil sympy==1.12 antlr4-python3-runtime==4.11.1 word2number Pebble timeout-decorator prettytable diff --git a/README.md b/README.md index 4b61b9cd63..d0c350fffb 100644 --- a/README.md +++ b/README.md @@ -103,20 +103,30 @@ We have listed detailed hardware requirements and the approach to setting up env *Figure 2. Total time consumption for RL training under varying computational resource settings for 10 epochs.* -## RL Training LRM directly from the Base Model +## Qwen2.5-7B-Zero RL Training -We start RL training from the base model Qwen2.5-7B with DeepSeek-R1-Zero style training. The initial training curves and results are presented below. As the training progresses, both the rewards and the response lengths gradually grow. A similar trend is also revealed [Open-Reasoner-Zero](https://github.com/Open-Reasoner-Zero/Open-Reasoner-Zero). **The simultaneous growth of average response lengths and rewards** can be considered as a critical sign of **the emergent deep thinking capabilities of LRM** to solve challenging reasoning problems. +We start RL training from the base model Qwen2.5-7B with DeepSeek-R1-Zero style training. The initial training curves and results are presented below. We conducted experiments on the data released by [Open-Reasoner-Zero](https://github.com/Open-Reasoner-Zero/Open-Reasoner-Zero). As the training progresses, both the rewards and the response lengths gradually grow. A similar trend is also revealed [Open-Reasoner-Zero](https://github.com/Open-Reasoner-Zero/Open-Reasoner-Zero). **The simultaneous growth of average response lengths and rewards** can be considered as a critical sign of **the emergent deep thinking capabilities of LRM** to solve challenging reasoning problems. -![undefined](/assets/7b_zero_curve.png) +![undefined](/assets/7b_zero_training_curve.png) *Figure 3. The training curve of Qwen2.5-7B-Zero during RL training* -| | MATH500 | AIME 2024 | AMC 2023 | + +We evaluate the performances of the intermediate checkpoints on the MATH500 and AIME24 datasets in the following figure. It appears that both the accuracy and the generated length continue to show an upward trend. + +![undefined](/assets/7b_zero_eval_acc.png) +*Figure 4. The test accuracy and response length evaluated on the MATH500 and AIME24 datasets.* + +We also conduct a experiment on the [DeepScalseR](https://github.com/agentica-project/deepscaler) dataset, which demonstrates a similar training trend. The evaluation results are presented in the following table. + +| | MATH-500 | AIME 2024 | AMC 2023 | | -------- | -------- | --------- | -------- | | Qwen2.5-7B | 34.0 | 2.0 | 17.8 | -| Open-Reasoner-Zero-7B Step = 150 | ~77 | ~15 | - | -| Qwen2.5-7B-Zero Step = 150 KL = 0 | 73.9 | 16.8 | 51.1 | -| Qwen2.5-7B-Zero Step = 150 KL = 0.001 | 76.3 | 13.3 | 53.7 | +| Open-Reasoner-Zero-7B, Step=150 | ~77 | ~15 | - | +| Qwen2.5-7B-Zero, Step=150 Dataset=DeepScaleR | 75.9 | 13.9 | 56.1 | +| Qwen2.5-7B-Zero, Step=150 Dataset=ORZ | 77.3 | 13.9 | 57.1 | +| Qwen2.5-7B-Zero, Step=200 Dataset=ORZ | 78.0 | 14.5 | 58.7 | +Table 2. Evaluation on AIME 2024, AMC 2023, and MATH-500. The results are reported using Pass@1 accuracy, which are averaged over 32 samples for each problem and evaluated with a temperature of 1.0. - **Training:** @@ -128,7 +138,7 @@ wget https://huggingface.co/datasets/inclusionAI/AReaL-RL-Data/resolve/main/data wget https://huggingface.co/datasets/inclusionAI/AReaL-RL-Data/resolve/main/data/id2info.json?download=true MODEL_PATH=${path_to_Qwen2.5-7B} -bash ./examples/train_7B_zero_n16_on_ray.sh $MODEL_PATH $DATA_PATH/prompts_for_zero.jsonl $DATA_PATH/id2info.json 16384 +bash ./examples/train_7B_zero_n16_on_ray.sh $MODEL_PATH $DATA_PATH/prompts_for_zero.jsonl $DATA_PATH/id2info.json 24000 ``` diff --git a/assets/7b_zero_curve.png b/assets/7b_zero_curve.png deleted file mode 100644 index 598f5bc688..0000000000 Binary files a/assets/7b_zero_curve.png and /dev/null differ diff --git a/assets/7b_zero_eval_acc.png b/assets/7b_zero_eval_acc.png new file mode 100644 index 0000000000..7080f84474 Binary files /dev/null and b/assets/7b_zero_eval_acc.png differ diff --git a/assets/7b_zero_training_curve.png b/assets/7b_zero_training_curve.png new file mode 100644 index 0000000000..3b1e4a15a6 Binary files /dev/null and b/assets/7b_zero_training_curve.png differ diff --git a/docker-compose.yml b/docker-compose.yml deleted file mode 100644 index e36526f7bd..0000000000 --- a/docker-compose.yml +++ /dev/null @@ -1,12 +0,0 @@ -version: '3' -services: - web: - build: - context: . - dockerfile: Dockerfile - target: docs - args: - REAL_CPU_BASE_IMAGE: ubuntu:22.04 - REAL_GPU_BASE_IMAGE: nvcr.io/nvidia/pytorch:23.10-py3 - ports: - - "7780:80" \ No newline at end of file diff --git a/examples/README.md b/examples/README.md index 3efbf71055..2f4c63fe7c 100644 --- a/examples/README.md +++ b/examples/README.md @@ -29,18 +29,52 @@ This tutorial provides a Docker image. Below are the tested software versions: | | Version | |---|:---:| | OS | CentOS 7 / Ubuntu 22.04 or any other system that meets the software requirements below | -| Nvidia Driver | 550.127.08 | +| NVIDIA Driver | 550.127.08 | | CUDA | 12.5 | | Git LFS | Refer to: https://docs.github.com/en/repositories/working-with-files/managing-large-files/installing-git-large-file-storage. Mainly used for downloading models, datasets, and AReaL project code. | | Docker | 27.5.1 | |NVIDIA Container Toolkit|[Installing the NVIDIA Container Toolkit](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html)| | AReaL Image | `ghcr.io/inclusionai/areal-runtime:v0.1.0`. This image includes AReaL's runtime dependencies and Ray components. | +Since the installation of NVIDIA Drivers and CUDA, as well as the mounting of shared storage, depends on node configurations and system versions, please complete these installations independently. This tutorial does not cover their setup. + +For multi-node training, ensure that the shared storage is mounted to the `/storage` directory on every node. All subsequent downloads and resources will be stored in this directory. The AReaL container will also mount this directory to `/storage` within the container, enabling seamless access during training. + + +# One-Click Environment Setup and Training Launch + +This section provides a one-click setup script to automatically configure the node environment: + +1. Install Docker, Git LFS, and NVIDIA Container Toolkit +2. Pull the AReaL image on each node +3. Download AReaL code, models, and datasets +4. Set up a Ray cluster +5. [Optional] Launch a training task within the Ray cluster + +Please perform the following operations on any chosen node: + +```bash +mkdir -p /storage/codes +cd /storage/codes/ +git clone https://github.com/inclusionAI/AReaL.git +cd /storage/codes/AReaL + +python ./examples/env/setup_env_and_start_train.py setup --private_key_file /path/to/ssh_key --ssh_port 22 --username root --hostnames NODE_IP_1 NODE_IP_2 NODE_IP_3 NODE_IP_4 --train_param 1.5B_n1 +``` + +`setup_env_and_start_train.py setup` arguments： + +- `private_key_file`: SSH secret key. Using by connecting nodes. +- `ssh_port`: SSH port +- `username`: SSH username +- `hostnames`: IP list. Split with space. Can be 1, 4, or 16 node IPs +- `train_param`: [Optional] Training parameters used to launch a training task immediately after environment setup. Valid options are: `1.5B_n1`, `1.5B_n4`, `1.5B_n16`, `7B_n4`, `7B_n16` + +If the script in this section fails to execute or encounters errors due to environmental discrepancies, you may manually configure the environment and launch training by following the instructions in the subsequent sections of this tutorial. + # Environment Setup -After preparing the nodes and the OS, follow this section to download the AReaL project code, models, and datasets, then start the Ray cluster. -- For multi-node training, mount the shared storage to the `/storage` directory on each node. All downloaded content will be stored here and later mounted into the AReaL environment image container. -- Since shared storage is used, downloading only needs to be done on one node. +Since shared storage is used, downloading only needs to be done on one node. ## Code and Cluster Configuration Clone the AReaL project code to `/storage/codes`: diff --git a/examples/README_zh.md b/examples/README_zh.md index 8fd81f5800..7a7c6c866e 100644 --- a/examples/README_zh.md +++ b/examples/README_zh.md @@ -37,7 +37,7 @@ ||版本说明| |---|---| |OS|CentOS 7 / Ubuntu 22.04 或其他满足下方软件运行的系统| -|Nvidia Driver|版本：550.127.08| +|NVIDIA Driver|版本：550.127.08| |CUDA|版本：12.5| |Git LFS|参考：[Git LFS 安装指南](https://docs.github.com/en/repositories/working-with-files/managing-large-files/installing-git-large-file-storage) 主要用于下载模型，数据集，AReaL 工程代码| |Docker|版本：27.5.1| @@ -45,11 +45,44 @@ |镜像|ghcr.io/inclusionai/areal-runtime:v0.1.0 这个镜像中包含运行依赖和 Ray 的相关组件| -# 运行环境配置 -在准备好节点和系统环境之后，按照本节的说明下载 AReaL 工程代码，模型，数据集，然后启动 Ray 集群。 +由于 NVIDIA Driver 和 CUDA 的安装以及共享存储的挂载与节点和系统版本有关，请自行完成安装，本教程不进行介绍。 -- 如果是多节点训练，请先将共享存储挂载到每个节点的 `/storage` 目录上，后续下载的内容都将放在这个目录下，并且最后挂载到 AReaL 环境镜像的容器里。 -- 由于使用了共享存储，下载操作只需要在一个节点上进行。 +如果是多节点训练，请先将共享存储挂载到每个节点的 `/storage` 目录上，后续下载的内容都将放在这个目录下，并且 AReaL 容器也会将该目录挂载到容器的 `/storage`，以便训练时访问。 + + +# 一键搭建环境并启动训练 + +本节提供一个一键安装脚本，自动完成节点的环境配置工作： +1. 安装 Docker，Git LFS，NVIDIA Container Toolkit +2. 在每个节点上拉取 AReaL 镜像 +3. 下载 AReaL 代码，模型，数据集 +4. 搭建 Ray 集群 +5. 【可选】在 Ray 集群中启动一个训练任务 + +请选择任意一个节点执行如下操作： + +```bash +mkdir -p /storage/codes +cd /storage/codes/ +git clone https://github.com/inclusionAI/AReaL.git +cd /storage/codes/AReaL + +python ./examples/env/setup_env_and_start_train.py setup --private_key_file /path/to/ssh_key --ssh_port 22 --username root --hostnames NODE_IP_1 NODE_IP_2 NODE_IP_3 NODE_IP_4 --train_param 1.5B_n1 +``` + +`setup_env_and_start_train.py setup` 参数说明： + +- `private_key_file`：SSH 私钥文件，用于连接节点 +- `ssh_port`：SSH 端口 +- `username`：SSH 用户名 +- `hostnames`：IP 列表，用空格分割。可以是 1/4/16 个节点 IP +- `train_param`：【可选】训练参数，用于在完成环境搭建后直接启动一个训练任务。可选值为 `1.5B_n1`，`1.5B_n4`，`1.5B_n16`，`7B_n4`，`7B_n16` + +如果因为环境差异，无法运行本节中的脚本或运行出现错误，也可以按照本教程后续章节的内容手动完成环境配置和启动训练。 + +# 环境配置 + +由于使用了共享存储，下载操作只需要在一个节点上完成。 ## 代码和集群配置 将 AReaL 项目代码克隆到 `/storage/codes` 中： diff --git a/examples/env/scripts/download-dataset.sh b/examples/env/scripts/download-dataset.sh new file mode 100644 index 0000000000..3addabe84b --- /dev/null +++ b/examples/env/scripts/download-dataset.sh @@ -0,0 +1,14 @@ +#!/bin/sh +set -e + +mkdir -p /storage/datasets/ +cd /storage/datasets/ +if [ $REGION == 'CN' ]; then + wget https://www.modelscope.cn/datasets/inclusionAI/AReaL-RL-Data/resolve/master/data/prompts_for_r1_distilled.jsonl + wget https://www.modelscope.cn/datasets/inclusionAI/AReaL-RL-Data/resolve/master/data/prompts_for_zero.jsonl + wget https://www.modelscope.cn/datasets/inclusionAI/AReaL-RL-Data/resolve/master/data/id2info.json +else + wget https://huggingface.co/datasets/inclusionAI/AReaL-RL-Data/resolve/main/data/prompts_for_r1_distilled.jsonl?download=true + wget https://huggingface.co/datasets/inclusionAI/AReaL-RL-Data/resolve/main/data/prompts_for_zero.jsonl?download=true + wget https://huggingface.co/datasets/inclusionAI/AReaL-RL-Data/resolve/main/data/id2info.json?download=true +fi \ No newline at end of file diff --git a/examples/env/scripts/download-model.sh b/examples/env/scripts/download-model.sh new file mode 100644 index 0000000000..24b8edcfca --- /dev/null +++ b/examples/env/scripts/download-model.sh @@ -0,0 +1,12 @@ +#!/bin/sh +set -e + +mkdir -p /storage/models +cd /storage/models +if [ $REGION == 'CN' ]; then + git clone https://www.modelscope.cn/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B.git + git clone https://www.modelscope.cn/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B.git +else + GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B + GIT_LFS_SKIP_SMUDGE=1 git clone https://huggingface.co/deepseek-ai/DeepSeek-R1-Distill-Qwen-7B +fi \ No newline at end of file diff --git a/examples/env/scripts/install-dependency.sh b/examples/env/scripts/install-dependency.sh new file mode 100644 index 0000000000..af012d8eda --- /dev/null +++ b/examples/env/scripts/install-dependency.sh @@ -0,0 +1,176 @@ +#!/bin/sh +OS=$(uname -s | tr '[:upper:]' '[:lower:]') +ARCH=$(uname -m | tr '[:upper:]' '[:lower:]') +DISTRO="" +GIT_LFS_VERSION="3.4.0" + +# get distro +if [ "$OS" = "linux" ]; then + if [ -f /etc/os-release ]; then + . /etc/os-release + DISTRO=$ID + else + echo "unknown linux distribution" + exit 1 + fi +else + echo "Unsupport os: $OS" + exit 1 +fi + +command_exists() { + command -v "$@" > /dev/null 2>&1 +} + +install_git() { + # check git exists + if command_exists git; then + return 0 + fi + + case "$DISTRO" in + ubuntu) + sudo apt install -y git + ;; + centos|alinux|amzn) + sudo yum install -y git || sudo dnf install -y git + ;; + *) + echo "Unsupported operating system $DISTRO" + exit 1 + ;; + esac +} + +install_git_lfs() { + # check git exists + if command_exists "git lfs"; then + return 0 + fi + + case "$ARCH" in + x86_64) + GIT_LFS_URL="https://github.com/git-lfs/git-lfs/releases/download/v3.4.0/git-lfs-linux-amd64-v${GIT_LFS_VERSION}.tar.gz" + ;; + aarch64) + GIT_LFS_URL="https://github.com/git-lfs/git-lfs/releases/download/v3.4.0/git-lfs-linux-arm64-v${GIT_LFS_VERSION}.tar.gz" + ;; + *) + echo "Unsupported arch $ARCH" + exit 1 + ;; + esac + + TEMP_DIR=$(mktemp -d) + echo "Start download to ${TEMP_DIR}" + curl -sL "$GIT_LFS_URL" -o "$TEMP_DIR/git-lfs.tar.gz" + if [ $? -ne 0 ]; then + echo "fail to download" + rm -rf "$TEMP_DIR" + exit 1 + fi + + echo "Start unzip" + tar -xzf "$TEMP_DIR/git-lfs.tar.gz" -C "$TEMP_DIR" + if [ $? -ne 0 ]; then + echo "fail to unzip" + rm -rf "$TEMP_DIR" + exit 1 + fi + + echo "Start install" + cd "$TEMP_DIR/git-lfs-${GIT_LFS_VERSION}" || exit 1 + sudo ./install.sh + if [ $? -ne 0 ]; then + echo "install fail" + rm -rf "$TEMP_DIR" + exit 1 + fi + + echo "Start config" + git lfs install + if [ $? -ne 0 ]; then + echo "git-lfs config error" + rm -rf "$TEMP_DIR" + exit 1 + fi + + rm -rf "$TEMP_DIR" +} + +install_docker() { + # check docker exists + if command_exists docker; then + return 0 + fi + + case "$DISTRO" in + ubuntu) + sudo apt-get update + sudo apt-get -y install apt-transport-https ca-certificates curl software-properties-common + sudo curl -fsSL http://mirrors.cloud.aliyuncs.com/docker-ce/linux/ubuntu/gpg | sudo apt-key add - + sudo add-apt-repository -y "deb [arch=$(dpkg --print-architecture)] http://mirrors.cloud.aliyuncs.com/docker-ce/linux/ubuntu $(lsb_release -cs) stable" + sudo apt-get -y install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin + ;; + centos) + sudo wget -O /etc/yum.repos.d/docker-ce.repo http://mirrors.cloud.aliyuncs.com/docker-ce/linux/centos/docker-ce.repo + sudo sed -i 's|https://mirrors.aliyun.com|http://mirrors.cloud.aliyuncs.com|g' /etc/yum.repos.d/docker-ce.repo + sudo dnf -y install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin + ;; + alinux) + sudo rm -f /etc/yum.repos.d/docker*.repo + sudo dnf -y remove docker-ce containerd.io docker-ce-rootless-extras docker-buildx-plugin docker-ce-cli docker-compose-plugin + sudo wget -O /etc/yum.repos.d/docker-ce.repo http://mirrors.cloud.aliyuncs.com/docker-ce/linux/centos/docker-ce.repo + sudo sed -i 's|https://mirrors.aliyun.com|http://mirrors.cloud.aliyuncs.com|g' /etc/yum.repos.d/docker-ce.repo + sudo dnf -y install dnf-plugin-releasever-adapter --repo alinux3-plus + sudo dnf -y install docker-ce docker-ce-cli containerd.io docker-buildx-plugin docker-compose-plugin + ;; + amzn) + sudo yum update -y + sudo dnf -y install docker + ;; + *) + echo "Unsupported operating system $DISTRO" + exit 1 + ;; + esac + + sudo systemctl enable docker + sudo systemctl start docker +} + +install_nvidia_container_toolkit() { + # check nvidia-container-toolkit exits + if command_exists nvidia-container-toolkit; then + return 0 + fi + + case "$DISTRO" in + ubuntu) + curl -fsSL https://nvidia.github.io/libnvidia-container/gpgkey | sudo gpg --dearmor -o /usr/share/keyrings/nvidia-container-toolkit-keyring.gpg \ + && curl -s -L https://nvidia.github.io/libnvidia-container/stable/deb/nvidia-container-toolkit.list | \ + sed 's#deb https://#deb [signed-by=/usr/share/keyrings/nvidia-container-toolkit-keyring.gpg] https://#g' | \ + sudo tee /etc/apt/sources.list.d/nvidia-container-toolkit.list + sudo apt-get update + sudo apt-get install -y nvidia-container-toolkit + ;; + centos|alinux|amzn) + curl -s -L https://nvidia.github.io/libnvidia-container/stable/rpm/nvidia-container-toolkit.repo | sudo tee /etc/yum.repos.d/nvidia-container-toolkit.repo + sudo dnf install -y nvidia-container-toolkit + ;; + *) + echo "Unsupported operating system $DISTRO" + exit 1 + ;; + esac +} + +restart_docker() { + sudo systemctl restart docker +} + +install_git +install_git_lfs +install_docker +install_nvidia_container_toolkit +restart_docker \ No newline at end of file diff --git a/examples/env/setup_env_and_start_train.py b/examples/env/setup_env_and_start_train.py new file mode 100644 index 0000000000..d3be83fdcc --- /dev/null +++ b/examples/env/setup_env_and_start_train.py @@ -0,0 +1,285 @@ +import argparse +import re + +import paramiko + + +def run_command_by_ssh( + private_key_file: str, + hostname: str, + port: int, + username: str, + command: str, +): + private_key = paramiko.RSAKey.from_private_key_file(private_key_file) + ssh_client = paramiko.SSHClient() + ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy) + ssh_client.connect( + hostname=hostname, + port=port, + username=username, + pkey=private_key, + ) + stdin, stdout, stderr = ssh_client.exec_command(command) + exit_status = stdout.channel.recv_exit_status() + if exit_status == 0: + res = stdout.read().decode() + err = "" + else: + res = "" + err = stderr.read().decode() + ssh_client.close() + return res,err,exit_status + +def setup_mount_nas( + private_key_file: str, + hostname: str, + port: int, + username: str, + nas_url: str, +): + cmd = 'mkdir -p /storage;mount -t nfs -o vers=4,minorversion=0,rsize=1048576,wsize=1048576,hard,timeo=600,retrans=2,noresvport {}:/ /storage'.format(nas_url) + res, err, exit_status = run_command_by_ssh(private_key_file,hostname,port,username,cmd) + if exit_status != 0: + print(err) + raise Exception('fail to mount_nas') + +def setup_download_codes( + private_key_file: str, + hostname: str, + port: int, + username: str, +): + cmd = 'apt install -y git || yum install -y git || dnf install -y git;mkdir -p /storage/codes;cd /storage/codes/;test -d AReaL || git clone https://github.com/inclusionAI/AReaL' + res, err, exit_status = run_command_by_ssh(private_key_file,hostname,port,username,cmd) + if exit_status != 0: + print(err) + raise Exception('fail to download codes') + +def setup_install_dependency( + private_key_file: str, + hostname: str, + port: int, + username: str, +): + cmd = 'bash /storage/codes/AReaL/examples/env/scripts/install-dependency.sh' + res, err, exit_status = run_command_by_ssh(private_key_file,hostname,port,username,cmd) + if exit_status != 0: + print(err) + raise Exception('fail to install dependency') + + +def setup_download_dataset( + private_key_file: str, + hostname: str, + port: int, + username: str, +): + cmd = 'bash /storage/codes/AReaL/examples/env/scripts/download-dataset.sh' + res, err, exit_status = run_command_by_ssh(private_key_file,hostname,port,username,cmd) + if exit_status != 0: + print(err) + raise Exception('fail to download dataset') + +def setup_download_model( + private_key_file: str, + hostname: str, + port: int, + username: str, +): + cmd = 'bash /storage/codes/AReaL/examples/env/scripts/download-model.sh' + res, err, exit_status = run_command_by_ssh(private_key_file,hostname,port,username,cmd) + if exit_status != 0: + print(err) + raise Exception('fail to download model') + +def setup_start_ray_header( + private_key_file: str, + hostname: str, + port: int, + username: str, +): + cmd = 'mkdir -p /storage/ray;cp /storage/codes/AReaL/examples/cluster_config_on_ray.json /storage/ray/cluster_config_on_ray.json; docker run -d --name r1-ray-head --privileged --gpus all --network host --shm-size 700g -v /storage:/storage ghcr.io/inclusionai/areal-runtime:v0.1.0 /bin/bash -c "ray start --head --port=6379 && tail -f /dev/null"' + res, err, exit_status = run_command_by_ssh(private_key_file,hostname,port,username,cmd) + if exit_status != 0: + print(err) + raise Exception('fail to start ray header') + +def setup_start_ray_worker( + private_key_file: str, + hostname: str, + port: int, + username: str, + header_hostname: str, +): + cmd = 'RAY_HEAD_IP={};docker run -d --name r1-ray-worker --privileged --gpus all --network host --shm-size 700g -v /storage:/storage ghcr.io/inclusionai/areal-runtime:v0.1.0 /bin/bash -c "ray start --address=$RAY_HEAD_IP:6379 && tail -f /dev/null"'.format(header_hostname) + res, err, exit_status = run_command_by_ssh(private_key_file,hostname,port,username,cmd) + if exit_status != 0: + print(err) + raise Exception('fail to start ray worker') + +def setup_start_traning( + train_param: str, + private_key_file: str, + hostname: str, + port: int, + username: str, +): + cmd = 'docker exec r1-ray-head /bin/bash -c "cd /storage/codes/AReaL;mkdir -p /storage/ray/train_batch_logs/;nohup bash ./examples/train_batch_{}.sh &> /storage/ray/train_batch_logs/{}.log &"'.format(train_param,train_param) + res, err, exit_status = run_command_by_ssh(private_key_file,hostname,port,username,cmd) + if exit_status != 0: + print(err) + raise Exception('fail to start train') + +# setup running environment +def setup(args): + total_nodes = len(args.hostnames) + # param check + if args.train_param != "": + node_required = int(re.findall(r"\d+", args.train_param)[-1]) + if node_required > total_nodes: + print("training needs {} node but {} found".format(node_required,total_nodes)) + return + # mount nas if needed + if args.nas_url != "": + print("starting mount nas") + for index in range(0,total_nodes): + private_key_file = args.private_key_file + hostname = args.hostnames[index] + port = args.ssh_port + username = args.username + setup_mount_nas( + private_key_file, + hostname, + port, + username, + args.nas_url, + ) + # download codes + print("starting download codes") + setup_download_codes( + args.private_key_file, + args.hostnames[0], + args.ssh_port, + args.username, + ) + # install dependency + print("starting install dependency") + for index in range(0,total_nodes): + private_key_file = args.private_key_file + hostname = args.hostnames[index] + port = args.ssh_port + username = args.username + setup_install_dependency( + private_key_file, + hostname, + port, + username, + ) + # download datasets + print("starting download datasets") + setup_download_dataset( + args.private_key_file, + args.hostnames[0], + args.ssh_port, + args.username, + ) + # download model + print("starting download model") + setup_download_model( + args.private_key_file, + args.hostnames[0], + args.ssh_port, + args.username, + ) + # start ray header + print("starting ray header") + setup_start_ray_header( + args.private_key_file, + args.hostnames[0], + args.ssh_port, + args.username, + ) + # start ray worker if needed + if total_nodes > 1: + print("starting ray worker") + for index in range(1,total_nodes): + private_key_file = args.private_key_file + hostname = args.hostnames[index] + port = args.ssh_port + username = args.username + header_hostname = args.hostnames[0] + setup_start_ray_worker( + private_key_file, + hostname, + port, + username, + header_hostname, + ) + print("setup success") + # start training if needed + if args.train_param != "": + print("starting training process") + setup_start_traning( + args.train_param, + args.private_key_file, + args.hostnames[0], + args.ssh_port, + args.username + ) + print("training process started") + +def main(): + parser = argparse.ArgumentParser() + subparsers = parser.add_subparsers() + + parser_setup = subparsers.add_parser( + name="setup", + help="setup training environment", + ) + parser_setup.add_argument( + "--private_key_file", + required=True, + help="", + type=str, + ) + parser_setup.add_argument( + "--hostnames", + required=True, + help="", + nargs="+", + type=str, + ) + parser_setup.add_argument( + "--ssh_port", + help="", + default=22, + type=int, + ) + parser_setup.add_argument( + "--username", + help="", + default="root", + type=str, + ) + parser_setup.add_argument( + "--nas_url", + help="", + default="", + type=str, + ) + parser_setup.add_argument( + "--train_param", + help="", + choices=["","1.5B_n1","1.5B_n4","1.5B_n16","7B_n4","7B_n16"], + default="", + type=str, + ) + parser_setup.set_defaults(func=setup) + + args = parser.parse_args() + args.func(args) + + +if __name__ == "__main__": + main() diff --git a/examples/train_1.5B_n16_on_ray.sh b/examples/train_1.5B_n16_on_ray.sh index 33f051cbfb..1e35883639 100644 --- a/examples/train_1.5B_n16_on_ray.sh +++ b/examples/train_1.5B_n16_on_ray.sh @@ -8,7 +8,7 @@ NODES="16" ALLOCATION_MODE="vllm.d64p1m1+d32p2m1" MAX_NEW_TOKENS=$4 MAX_NUM_SEQS=128 -PPO_MBS=1 +PPO_MBS=4 KL_CTL=0.001 MAX_TOKEN_PER_MB=$(expr 2048 + ${MAX_NEW_TOKENS} + 1024) @@ -105,9 +105,11 @@ python3 -m realhf.apps.quickstart ppo-math \ 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=1e-5 \ actor.optimizer.lr_scheduler_type=constant \ + actor.optimizer.initial_loss_scale=262144.0 \ + actor.optimizer.loss_scale_window=5 \ + actor.optimizer.hysteresis=2 \ 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} \ diff --git a/math_verify_utils_qwen.py b/math_verify_utils_qwen.py index bc5bb9f6b0..b345c3e4a0 100644 --- a/math_verify_utils_qwen.py +++ b/math_verify_utils_qwen.py @@ -4,24 +4,34 @@ import json from parser import extract_answer -from grader import math_equal +from grader import call_with_timeout, 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) - # if extract_answer.strip() == "": - # print (answer) - # raise - 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 + try: + extracted_answer = extract_answer(answer, "math", use_last_number=False) + extracted_solution = extract_answer(solution, "math", use_last_number=True) - return retval, (extracted_answer, extracted_solution) + # if extract_answer.strip() == "": + # print (answer) + # raise + if extracted_answer is None or extracted_answer.strip() in ["None", "none", ""]: + retval = 0 + elif math_equal(extracted_answer, extracted_solution, timeout=False): + # elif call_with_timeout(math_equal, extracted_answer, extracted_solution): + retval = 1 + else: + retval = 0 + + return retval, (extracted_answer, extracted_solution) + except: + return 0, ("None", "None") + + +def process_results_process(a, b, output_queue): + result = process_results(a, b) + output_queue.put(result) if __name__ == "__main__": @@ -36,9 +46,17 @@ def process_results(answer, solution): with open(f"/tmp/{args.tmp_id}-output.jsonl", "w", encoding="utf-8") as temp_file: for input_data in all_input_data: - r, (ans, sol) = process_results( - input_data["answer"], input_data["solution"] + # r, (ans, sol) = process_results( + # input_data["answer"], input_data["solution"] + # ) + tmp = call_with_timeout( + process_results_process, input_data["answer"], input_data["solution"] ) + if isinstance(tmp, bool): + r, (ans, sol) = 0, ("None", "None") + else: + r, (ans, sol) = tmp + res = {"retval": r, "ans": ans, "sol": sol} temp_file.write(json.dumps(res) + "\n") diff --git a/realhf/api/core/config.py b/realhf/api/core/config.py index 36aaff1ff7..b57aea01ca 100644 --- a/realhf/api/core/config.py +++ b/realhf/api/core/config.py @@ -16,12 +16,6 @@ class DatasetAbstraction: args: Dict[str, Any] = dataclasses.field(default_factory=dict) -@dataclasses.dataclass -class DataLoaderAbstraction: - type_: str = "default" - args: Dict[str, Any] = dataclasses.field(default_factory=dict) - - @dataclasses.dataclass class ModelWrapperAbstraction: type_: str @@ -187,10 +181,6 @@ class StandaloneModelShardAbstraction: model: ModelAbstraction backend: ModelBackendAbstraction # evaluation - eval_datasets: Optional[List[DatasetAbstraction]] = None - eval_dataloader: Optional[DataLoaderAbstraction] = dataclasses.field( - default_factory=lambda: DataLoaderAbstraction( - "packed_eval", args=dict(batch_size=128) - ) - ) + eval_dataset: Optional[DatasetAbstraction] = None + eval_bs: int = 128 should_instantiate: bool = True diff --git a/realhf/api/core/data_api.py b/realhf/api/core/data_api.py index 9556e88d29..3cb7c2c89a 100644 --- a/realhf/api/core/data_api.py +++ b/realhf/api/core/data_api.py @@ -109,14 +109,10 @@ class MicroBatchSpec: :param max_tokens_per_mb: The maximum number of tokens per micro- batch. :type max_tokens_per_mb: Optional[int] - :param balanced_seqs: Whether to balance the number of sequences per - micro-batch. Only effective when max_tokens_per_mb is None. - :type balanced_seqs: bool, optional """ n_mbs: int = 1 - max_tokens_per_mb: int | None = None - balanced_seqs: bool = False + max_tokens_per_mb: int = int(1e12) @classmethod def new(cls, mb_spec: "MicroBatchSpec", **kwargs): @@ -124,7 +120,6 @@ def new(cls, mb_spec: "MicroBatchSpec", **kwargs): fields = dict( n_mbs=mb_spec.n_mbs, max_tokens_per_mb=mb_spec.max_tokens_per_mb, - balanced_seqs=mb_spec.balanced_seqs, ) fields.update(kwargs) return cls(**fields) @@ -350,30 +345,6 @@ def _get_split_key(self) -> str: acc_seqlen = {k: sum(sum(l) for l in lens) for k, lens in self.seqlens.items()} return max(acc_seqlen, key=acc_seqlen.get) - def get_split_spec( - self, k: int, key: Optional[str] = None, min_size: int = 1 - ) -> SequenceSplitSpec: - """Get the partition specification for splitting the data into `k` - parts using a dynamic programming algorithm to achieve the most - balanced partitioning. - - :param k: The number of parts to split the data into. - :type k: int - :param key: The key to be used for splitting. If None, the key - with the largest total sequence length will be used. - :type key: Optional[str] - :param min_size: The minimum size of each partition. - :type min_size: int - :return: A SequenceSplitSpec object representing the - partitioning specification. - :rtype: SequenceSplitSpec - """ - if key is None: - key = self._get_split_key() - lens = [sum(lens) for lens in self.seqlens[key]] - partitions = datapack.min_abs_diff_partition(lens, k, min_size) - return SequenceSplitSpec(partitions=partitions) - def split_with_spec(self, spec: SequenceSplitSpec) -> List["SequenceSample"]: """Split the data according to the given spec.""" samples = [] @@ -419,47 +390,9 @@ def split_with_spec(self, spec: SequenceSplitSpec) -> List["SequenceSample"]: ) return samples - def split( - self, - k: int, - key: Optional[str] = None, - min_size: int = 1, - ) -> List["SequenceSample"]: - """Split the data into `k` parts. - - This method uses the specified key or the key with the largest total sequence length - to split the data into `k` parts. The partitioning ensures that each part meets the - minimum size requirement. - - :param k: The number of parts to split the data into. - :type k: int - :param key: The key to use for splitting. If None, the key with the largest - total sequence length will be used. - :type key: Optional[str] - :param min_size: The minimum size of each partition. - :type min_size: int - :return: A list of `SequenceSample` objects, each representing a part of the split data. - :rtype: List[SequenceSample] - """ - spec = self.get_split_spec(k, key, min_size) - return self.split_with_spec(spec) - - def divide_into_mbs( - self, mb_spec: MicroBatchSpec + def split_with_lengths( + self, mb_spec: MicroBatchSpec, lens: List[int] ) -> Tuple[List["SequenceSample"], List[int] | np.ndarray, List[int] | np.ndarray]: - if mb_spec.max_tokens_per_mb is None: - return ( - self.split( - mb_spec.n_mbs, - min_size=( - 1 if not mb_spec.balanced_seqs else self.bs // mb_spec.n_mbs - ), - ), - np.arange(self.bs), - np.arange(self.bs), - ) - - lens = [sum(lens) for lens in self.seqlens[self._get_split_key()]] group_indices = datapack.ffd_allocate( lens, mb_spec.max_tokens_per_mb, min_groups=mb_spec.n_mbs ) @@ -474,10 +407,24 @@ def divide_into_mbs( return sample.split_with_spec(spec), forward_indices, backward_indices - def divide_into_mbs_balanced( + def split( + self, mb_spec: MicroBatchSpec + ) -> Tuple[List["SequenceSample"], List[int] | np.ndarray, List[int] | np.ndarray]: + """Split the data into `n_mbs` parts. + + :param mb_spec: The configuration to split the data into. + `n_mbs` is the minimum number of micro-batches, + `max_tokens_per_mb` is the maximum number of tokens in each micro-batch. + If `max_tokens_per_mb` is a large value, defaults to balanced split. + :type mb_spec: MicroBatchSpec + """ + lens = [sum(lens) for lens in self.seqlens[self._get_split_key()]] + return self.split_with_lengths(mb_spec, lens) + + def synced_data_parallel_split( self, mb_spec: MicroBatchSpec ) -> List["SequenceSample"]: - mb_inputs, *_ = self.divide_into_mbs(mb_spec) + mb_inputs, *_ = self.split(mb_spec) all_n_mbs = [None for _ in range(constants.data_parallel_world_size())] dist.all_gather_object( all_n_mbs, len(mb_inputs), group=constants.data_parallel_group() @@ -487,7 +434,7 @@ def divide_into_mbs_balanced( # This method is called when max_tokens_per_mb is given and during training. # In this case, we evenly partition sequences across DP ranks, # so the recursion will always terminate when n_mbs = bs // dp_size - return self.divide_into_mbs_balanced( + return self.synced_data_parallel_split( MicroBatchSpec.new(mb_spec, n_mbs=max(all_n_mbs)) ) @@ -881,61 +828,3 @@ def make_dataset( logger.info(f"Dataset creation/loading time: {time.perf_counter() - tik:.3f}s") return dataset - - -ALL_DATALOADER_CLASSES = {} - - -def register_dataloader(name, dataloader_cls): - assert name not in ALL_DATALOADER_CLASSES - ALL_DATALOADER_CLASSES[name] = dataloader_cls - - -def make_dataloader( - cfg: Union[str, config_api.DataLoaderAbstraction], dataset: torch.utils.data.Dataset -) -> torch.utils.data.DataLoader: - if isinstance(cfg, str): - cfg = config_api.DataLoaderAbstraction(type_=cfg) - dataloader_cls = ALL_DATALOADER_CLASSES[cfg.type_] - return dataloader_cls(dataset, **cfg.args) - - -def PackedDataLoader(dataset, *args, **kwargs): - if not isinstance(getattr(dataset, "util", None), DatasetUtility): - raise ValueError("Dataset must have a `util` attribute of type DatasetUtility.") - g = torch.Generator() - g.manual_seed(dataset.util.seed) - - def seed_worker(worker_id): - worker_seed = torch.initial_seed() % 2**32 - np.random.seed(worker_seed) - random.seed(worker_seed) - - return torch.utils.data.DataLoader( - dataset, - *args, - collate_fn=SequenceSample.gather, - # NOTE: This is *NOT* the actual batch size for training. - # It is just a proper size to load data to workers. - batch_size=10240, - shuffle=True, - generator=g, - worker_init_fn=seed_worker, - **kwargs, - ) - - -def PackedEvalDataLoader(dataset, *args, **kwargs): - if not isinstance(getattr(dataset, "util", None), DatasetUtility): - raise ValueError("Dataset must have a `util` attribute of type DatasetUtility.") - return torch.utils.data.DataLoader( - dataset, - *args, - collate_fn=SequenceSample.gather, - shuffle=False, - **kwargs, - ) - - -register_dataloader("packed", PackedDataLoader) -register_dataloader("packed_eval", PackedEvalDataLoader) diff --git a/realhf/api/core/model_api.py b/realhf/api/core/model_api.py index de6c231485..210e65fc91 100644 --- a/realhf/api/core/model_api.py +++ b/realhf/api/core/model_api.py @@ -27,6 +27,10 @@ logger = logging.getLogger("model_api") +class ZeroTotalLossWeightException(Exception): + pass + + @dataclasses.dataclass class GenerationHyperparameters: """Generation hyperparameters. @@ -334,7 +338,9 @@ def train_batch( input_: SequenceSample, mb_spec: MicroBatchSpec, loss_fn: Callable[[torch.Tensor, SequenceSample], Tuple[torch.Tensor, Dict]], + loss_weight_fn: Callable[[torch.Tensor, SequenceSample], float], version_steps: int, + token_normalize_scope: Literal["global", "dp"] = "global", ) -> Tuple[torch.Tensor, Dict] | None: """Update the model with a batch of data and a loss function. @@ -345,6 +351,10 @@ def train_batch( :param loss_fn: The loss function. It takes the output of the forward pass and the input data, returning the loss and a dictionary of statistics. :type loss_fn: Callable[[torch.Tensor, SequenceSample], Tuple[torch.Tensor, Dict]] + :param loss_weight_fn: This function is used to calculate the number of valid tokens + when normalizing loss across micro batches and DP ranks. Can be `lambda: 1` + if just taking the average over batches. + :type loss_weight_fn: Callable[[torch.Tensor, SequenceSample], float] :param version_steps: The global step counter for this experiment, used by the backend to determine the learning rate schedule. :type version_steps: int @@ -354,6 +364,11 @@ def train_batch( which automatically schedules the forward and backward passes. For non-pipelined training, forward and backward passes are executed iteratively over mini-batches to accumulate gradients. If None, the batch will not be split. + :param global_normalize_scope: The scope of token-wise loss normalization. Choices: + global: average across all micro batches across DP ranks. + dp: average across micro batches in current DP rank. + Default to "global". + :type global_normalize_scope: Literal["global", "dp"] """ raise NotImplementedError() diff --git a/realhf/api/core/system_api.py b/realhf/api/core/system_api.py index 0cbf06a097..29dbc6b5d3 100644 --- a/realhf/api/core/system_api.py +++ b/realhf/api/core/system_api.py @@ -15,7 +15,6 @@ import realhf.api.core.dfg as dfg from realhf.api.core.config import ( - DataLoaderAbstraction, DatasetAbstraction, ModelAbstraction, ModelName, @@ -114,17 +113,13 @@ def system_setup( @dataclasses.dataclass class ModelWorker: - seed: int + base_seed: int shards: List[StandaloneModelShardAbstraction] # dataset, for source model workers tokenizer_name_or_path: Optional[str] = None datasets: Optional[List[Union[str, DatasetAbstraction]]] = None - dataloader: Union[str, DataLoaderAbstraction] = "packed" use_dataset_cache: bool = False dataset_cahce_root: str = constants.DATASET_CACHE_PATH - # cuda & cudnn config - cudnn_benchmark: bool = False - cudnn_deterministic: bool = False cuda_cache_cleanliness: bool = True cuda_cache_clear_freq: int = 10 torch_cache_mysophobia: bool = False @@ -210,12 +205,13 @@ class ExperimentSaveEvalControl: @dataclasses.dataclass class MasterWorker: + base_seed: int exp_ctrl: ExperimentSaveEvalControl # main components n_model_workers: int model_rpcs: List[dfg.MFCDef] = None model_topos: Dict[ModelName, topology.PipeModelDataParallelTopology] = None - msid2mwid: Dict[ModelShardID, int] = None + msid2mwid: Dict[ModelShardID | str, int] = None data_transfer_pairs: List[Tuple[str, str]] = None sync_param_pairs: List[Tuple[str, str]] = None worker_info: Optional[WorkerInformation] = None @@ -263,7 +259,11 @@ class ExperimentConfig: def __post_init__(self): self.master_worker = [ - MasterWorker(exp_ctrl=self.exp_ctrl, n_model_workers=len(self.model_worker)) + MasterWorker( + base_seed=self.model_worker[0].base_seed, + exp_ctrl=self.exp_ctrl, + n_model_workers=len(self.model_worker), + ) ] def lazy_init(self): diff --git a/realhf/api/quickstart/model.py b/realhf/api/quickstart/model.py index f329e931cb..c1621824f8 100644 --- a/realhf/api/quickstart/model.py +++ b/realhf/api/quickstart/model.py @@ -106,7 +106,10 @@ class OptimizerConfig: ) warmup_steps_proportion: float = 0.02 offload: bool = False - initial_loss_scale: float = 4096.0 + initial_loss_scale: float = 2**32 + min_loss_scale: float = 1.0 + loss_scale_window: float = 5 + hysteresis: int = 2 gradient_clipping: float = 1.0 diff --git a/realhf/base/datapack.py b/realhf/base/datapack.py index 2b0d248175..3e4ece2cb3 100644 --- a/realhf/base/datapack.py +++ b/realhf/base/datapack.py @@ -2,6 +2,7 @@ # Copyright 2024 Wei Fu & Zhiyu Mei # Licensed under the Apache License, Version 2.0 (the "License"). +import heapq import itertools from typing import Any, List, Tuple, Union @@ -148,28 +149,36 @@ def reorder_to_balanced_batches( return np.array(reordered_indices), max_diff -@numba.njit +# @numba.njit def _ffd_allocate( values: np.ndarray, capacity: int, min_groups: int ) -> List[List[int]]: + """A greedy allocation algorithm that partitions a list of numbers + into k groups, where the summation of each group is less than capacity + and k >= min_groups. We want to minimize k and make partitions as balanced + as possible. + + 1. Sort the numbers in reverse order. + 2. If the number of groups is less than, create a new group. + 3. If the new number fits into the smallest group, add it into the group. + 4. Otherwise, create a new group. + """ value_indices = np.argsort(-values) - group_indices = [] - group_values = [] + group_indices: List[List[int]] = [] + group_values: List[Tuple[float, int]] = [] + group_cnt = 0 for idx in value_indices: - if len(group_values) < min_groups: - group_values.append(values[idx]) - group_indices.append([idx]) - continue - placed = False - for i in range(len(group_values)): - if group_values[i] + values[idx] <= capacity: - group_values[i] += values[idx] - group_indices[i].append(idx) - placed = True - break - if not placed: - group_values.append(values[idx]) + if ( + len(group_values) < min_groups + or group_values[0][0] + values[idx] > capacity + ): + heapq.heappush(group_values, (float(values[idx]), group_cnt)) group_indices.append([idx]) + group_cnt += 1 + else: + v, group_idx = heapq.heappop(group_values) + heapq.heappush(group_values, (float(v + values[idx]), group_idx)) + group_indices[group_idx].append(idx) return group_indices @@ -188,15 +197,15 @@ def ffd_allocate(values: List[int], capacity: int, min_groups: int) -> List[List for i in range(100): st = time.monotonic() - nums = np.random.randint(512, 4000, size=(32768,)) + nums = np.random.randint(1024, 8192, size=(100,)) # k = np.random.randint(2, 20) # min_size = np.random.randint(1, len(nums) // k) # res = min_abs_diff_partition(nums, k, min_size) # assert all(y - x >= min_size for x, y in res) - max_tokens_per_mb = 655360 - n_groups = np.random.randint(10, 20) - groups = ffd_allocate(nums, max_tokens_per_mb, n_groups) - assert len(groups) >= n_groups + max_tokens_per_mb = 163840 + min_n_groups = np.random.randint(1, 8) + groups = ffd_allocate(nums, max_tokens_per_mb, min_n_groups) + assert len(groups) >= min_n_groups import itertools indices = list(itertools.chain(*groups)) @@ -207,6 +216,8 @@ def ffd_allocate(values: List[int], capacity: int, min_groups: int) -> List[List print( len(groups), + min_n_groups, + [sum(nums[i] for i in group) for group in groups], max(group_percent), min(group_percent), np.mean(group_percent), diff --git a/realhf/base/seeding.py b/realhf/base/seeding.py index b10a924e18..5e8cd48130 100644 --- a/realhf/base/seeding.py +++ b/realhf/base/seeding.py @@ -2,17 +2,29 @@ # Copyright 2024 Wei Fu & Zhiyu Mei # Licensed under the Apache License, Version 2.0 (the "License"). +import os import random import numpy as np import torch import transformers +_SEED = None + def set_random_seed(seed): + global _SEED + _SEED = seed + os.environ["PYTHONHASHSEED"] = str(seed) transformers.set_seed(seed) random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) if torch.cuda.is_available(): + torch.cuda.manual_seed(seed) torch.cuda.manual_seed_all(seed) + + +def get_seed() -> int: + global _SEED + return _SEED diff --git a/realhf/experiments/common/check.py b/realhf/experiments/common/check.py index 2418866e29..3094dd33f3 100644 --- a/realhf/experiments/common/check.py +++ b/realhf/experiments/common/check.py @@ -82,11 +82,6 @@ def check_valid_parallel_batch_size(rpc_alloc: RPCAllocation): factor = 1 if rpc.is_train() and rpc_alloc.parallel.pipeline_parallel_size > 1: factor = 2 - if mb_spec.balanced_seqs or ( - mb_spec.max_tokens_per_mb is not None and rpc.is_train() - ): - assert rpc.n_seqs % dp_size == 0, (rpc.n_seqs, dp_size) - return assert ( rpc.n_seqs diff --git a/realhf/experiments/common/common.py b/realhf/experiments/common/common.py index 5d08108ddf..2347e09728 100644 --- a/realhf/experiments/common/common.py +++ b/realhf/experiments/common/common.py @@ -15,10 +15,10 @@ 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 ( - DataLoaderAbstraction, DatasetAbstraction, ModelAbstraction, ModelBackendAbstraction, @@ -59,9 +59,7 @@ check_valid_vllm, ) from realhf.experiments.common.utils import ( - extract_decoupled_vllm_train_allocation, - extract_key_value_allocation, - extract_symmetric_allocation, + AllocationMode, get_topo, make_inf_backend_config, make_train_backend_config, @@ -112,10 +110,6 @@ class CommonExperimentConfig(Experiment): - ``heuristic``\: Allocate resources and configure parallel strategies using heuristic strategies obtained from a search. - - ``pipe_data``\: Identical parallelization (like DSChat) with pipe+data parallelism. For a world size under 8, only data parallelism will be used. - - - ``pipe_model``\: Identical parallelization (like DSChat) with pipe+model parallelism. For a world size under 8, only tensor-model parallelism will be used. - - A regex pattern like ``d${DP}p${PP}m${TP}``\: Identical parallelization for all MFCs with ${DP}-way data parallelism, ${PP}-way pipeline parallelism, and ${TP}-way model parallelism. - A regex pattern like ``vllm.{IdentPara}+{IdentPara}``\: Decoupled generation (vLLM) and training allocations with correspnding identical parallelization strategies. Note that the pipeline parallel degree of vLLM can only be 1. @@ -205,7 +199,7 @@ class CommonExperimentConfig(Experiment): recover_retries: int = 1 recover_after: int = 10 ignore_worker_error: bool = False - allocation_mode: str = "pipe_model" + allocation_mode: str = "" allocation_use_cache: bool = False n_nodes: int = 1 n_gpus_per_node: int = cluster_spec.n_gpus_per_node @@ -256,22 +250,17 @@ def datasets(self) -> List[DatasetAbstraction]: return NotImplementedError(f"datasets is not implemented in {self.__class__}") @property - def eval_datasets(self) -> List[DatasetAbstraction]: - """A list of dataset configurations used for evaluation. + def eval_dataset(self) -> DatasetAbstraction | None: + """The dataset configuration used for evaluation. Can be None if runtime evaluation is not needed. """ return None @property - def eval_dataloader(self) -> DataLoaderAbstraction: - """The dataloader configuration used for evaluation. - - Reserved to changed the evaluation batch size. Training does not - require this property because the batch size is handled in MFC - definitions. - """ - return DataLoaderAbstraction("packed_eval", args=dict(batch_size=128)) + def eval_bs(self) -> int: + """The batch size for runtime evaluation.""" + return 128 @property def tokenizer_name_or_path(self) -> str: @@ -364,6 +353,8 @@ def _get_rpc_allocations(self) -> List[RPCAllocation]: self.__check_legal_allocation_options() + self._allocation_mode = AllocationMode.from_str(self.allocation_mode) + rpcs = self.rpcs if self.allocation_mode == "search": # assert self.mode == "slurm" @@ -378,90 +369,72 @@ def _get_rpc_allocations(self) -> List[RPCAllocation]: break else: raise ValueError(f"RPC {rpc_alloc.rpc} not found in rpcs.") - elif ( - self.allocation_mode == "pipe_data" - or self.allocation_mode == "pipe_model" - or extract_symmetric_allocation(self.allocation_mode) - ): - if self.allocation_mode == "pipe_data": - dp, pp, mp = self.n_gpus_per_node, self.n_nodes, 1 - elif self.allocation_mode == "pipe_model": - dp, pp, mp = 1, self.n_nodes, self.n_gpus_per_node - else: - para = extract_symmetric_allocation(self.allocation_mode) - dp, pp, mp = para["d"], para["p"], para["m"] - if dp * pp * mp != self.n_nodes * self.n_gpus_per_node: - raise ValueError( - "The multiplication of 3D parallel degrees " - "does not equal to the number of gpus. " - f"dp={dp}, pp={pp}, mp={mp}, " - f"n_nodes={self.n_nodes}, n_gpus_per_node={self.n_gpus_per_node}" - ) - rpc_allocs: List[RPCAllocation] = [ - RPCAllocation( - rpc=rpc, - device_mesh=self.global_device_mesh, - parallel=ParallelismConfig( - data_parallel_size=dp, - pipeline_parallel_size=pp, - model_parallel_size=mp, - use_sequence_parallel=( - rpc.interface_type == ModelInterfaceType.TRAIN_STEP - and mp > 1 - ), - ), - ) - for rpc in rpcs.values() - ] - elif extract_decoupled_vllm_train_allocation(self.allocation_mode): - para = extract_decoupled_vllm_train_allocation(self.allocation_mode) - dp, pp, mp = para["d"], para["p"], para["m"] - vdp, vpp, vmp = para["vllm.d"], para["vllm.p"], para["vllm.m"] - vllm_world_size = vdp * vpp * vmp - if dp * pp * mp + vdp * vpp * vmp != self.n_nodes * self.n_gpus_per_node: - raise ValueError( - "The multiplication of 3D parallel degrees " - "does not equal to the number of gpus. " - "Note that the device mesh of vLLM should be disjoint from the device mesh of other MFCs, " - "so their summation should be equal to the total number of gpus. " - f"dp={dp}, pp={pp}, mp={mp}, vllm.dp={vdp}, vllm.pp={vpp}, vllm.mp={vmp}, " - f"n_nodes={self.n_nodes}, n_gpus_per_node={self.n_gpus_per_node}" - ) + elif self._allocation_mode.is_decoupled(): + paras = self._allocation_mode.parallel_strat + + gdp, gpp, gmp = paras["gen"]["d"], paras["gen"]["p"], paras["gen"]["m"] + gen_world_size = gdp * gpp * gmp assert ( - vllm_world_size < self.n_gpus_per_node - or vllm_world_size % self.n_gpus_per_node == 0 + gen_world_size < self.n_gpus_per_node + or gen_world_size % self.n_gpus_per_node == 0 ) - vllm_device_mesh, train_device_mesh = self.global_device_mesh.split( - vllm_world_size + gen_device_mesh, train_device_mesh = self.global_device_mesh.split( + gen_world_size ) - self.vllm_device_mesh = vllm_device_mesh + self.gen_device_mesh = gen_device_mesh self.train_device_mesh = train_device_mesh rpc_allocs = [] flag = False for rpc in rpcs.values(): - if rpc.interface_type == ModelInterfaceType.GENERATE: - if vpp != 1: + if rpc.is_generate(): + if gpp != 1: raise NotImplementedError( - "vllm pipeline parallel is not supported yet." + "vllm/sglang pipeline parallel is not supported yet." ) if flag: raise NotImplementedError( - "vllm does not support two generation RPCs for now." + "vllm/sglang does not support two generation RPCs for now." ) alloc = RPCAllocation( rpc=rpc, - device_mesh=vllm_device_mesh, + device_mesh=gen_device_mesh, parallel=ParallelismConfig( - data_parallel_size=vdp, - pipeline_parallel_size=vpp, - model_parallel_size=vmp, + data_parallel_size=gdp, + pipeline_parallel_size=gpp, + model_parallel_size=gmp, use_sequence_parallel=False, ), ) flag = True else: + rpc_name = rpc.name + if rpc_name in paras: + dp, pp, mp = ( + paras[rpc_name]["d"], + paras[rpc_name]["p"], + paras[rpc_name]["m"], + ) + else: + if "*" not in paras: + raise ValueError( + f"RPC {rpc_name} parallel strategy not given, " + "expect a `*` to specify the default parallel strategy." + ) + dp, pp, mp = paras["*"]["d"], paras["*"]["p"], paras["*"]["m"] + if ( + dp * pp * mp + gdp * gpp * gmp + != self.n_nodes * self.n_gpus_per_node + ): + raise ValueError( + "The multiplication of 3D parallel degrees " + "does not equal to the number of gpus. " + "Note that the device mesh of vLLM should be disjoint from the device mesh of other MFCs, " + "so their summation should be equal to the total number of gpus. " + f"dp={dp}, pp={pp}, mp={mp}, vllm.dp={gdp}, vllm.pp={gpp}, vllm.mp={gmp}, " + f"n_nodes={self.n_nodes}, n_gpus_per_node={self.n_gpus_per_node}" + ) alloc = RPCAllocation( rpc=rpc, device_mesh=train_device_mesh, @@ -478,10 +451,10 @@ def _get_rpc_allocations(self) -> List[RPCAllocation]: rpc_allocs.append(alloc) if not flag: raise ValueError( - "No generation RPC found. Please use the allocation mode without vllm." + "No generation RPC found. Please use the hybrid train allocation mode." ) - elif extract_key_value_allocation(self.allocation_mode): - paras = extract_key_value_allocation(self.allocation_mode) + elif self._allocation_mode.is_global_hybrid(): + paras = self._allocation_mode.parallel_strat rpc_allocs = [] for rpc_name, rpc in self.rpcs.items(): if rpc_name in paras: @@ -553,7 +526,7 @@ def _get_model_worker_configs( for i, j in itertools.product(range(self.n_nodes), range(self.n_gpus_per_node)): mw = ModelWorker( - seed=self.seed, + base_seed=self.seed, shards=[], datasets=self.datasets, torch_cache_mysophobia=self.torch_cache_mysophobia, @@ -564,8 +537,8 @@ def _get_model_worker_configs( # vLLM enabled model worker, shortcut case if ( - extract_decoupled_vllm_train_allocation(self.allocation_mode) - and self.vllm_device_mesh.mapping[i, j] + self._allocation_mode.is_decoupled() + and self.gen_device_mesh.mapping[i, j] ): gen_rpc_alloc = next( alloc @@ -611,7 +584,6 @@ def _get_model_worker_configs( backend=ModelBackendAbstraction( "vllm", args=dict( - seed=self.seed, model_path=model_cfg.path, **vllm_dict_args, ), @@ -691,7 +663,6 @@ def _get_model_worker_configs( backend = ModelBackendAbstraction( "vllm", args=dict( - seed=self.seed, model_path=model_cfg.path, **vllm_dict_args, ), @@ -713,8 +684,8 @@ def _get_model_worker_configs( ), model=model, backend=backend, - eval_datasets=self.eval_datasets, - eval_dataloader=self.eval_dataloader, + eval_dataset=self.eval_dataset, + eval_bs=self.eval_bs, ) ) shard_counter[model_name] += 1 diff --git a/realhf/experiments/common/sft_exp.py b/realhf/experiments/common/sft_exp.py index ef3a7ba296..06622e1793 100644 --- a/realhf/experiments/common/sft_exp.py +++ b/realhf/experiments/common/sft_exp.py @@ -5,7 +5,6 @@ import dataclasses from realhf.api.core.config import ( - DataLoaderAbstraction, DatasetAbstraction, ModelInterfaceAbstraction, ModelInterfaceType, @@ -82,22 +81,18 @@ def datasets(self): ] @property - def eval_datasets(self): - return [ - DatasetAbstraction( - "prompt_answer", - args=dict( - max_length=self.dataset.max_seqlen, - dataset_path=self.dataset.valid_path, - ), - ) - ] + def eval_dataset(self): + return DatasetAbstraction( + "prompt_answer", + args=dict( + max_length=self.dataset.max_seqlen, + dataset_path=self.dataset.valid_path, + ), + ) @property - def eval_dataloader(self): - return DataLoaderAbstraction( - "packed_eval", args=dict(batch_size=self.dataset.valid_bs_n_seqs) - ) + def eval_bs(self) -> int: + return self.dataset.valid_bs_n_seqs @property def tokenizer_name_or_path(self): diff --git a/realhf/experiments/common/utils.py b/realhf/experiments/common/utils.py index 4bcae1e354..8de0dca0ee 100644 --- a/realhf/experiments/common/utils.py +++ b/realhf/experiments/common/utils.py @@ -3,6 +3,8 @@ # Licensed under the Apache License, Version 2.0 (the "License"). import collections +import dataclasses +import enum import itertools import re from typing import * @@ -224,58 +226,94 @@ def resolve_rpc_hooks( logger.info(f"Add offload hook for rpc {rpc.name} for role {rpc.role}") -def extract_symmetric_allocation(allocation_mode: str) -> Dict | None: - for x, y, z in itertools.permutations(["d", "m", "p"]): - pattern = rf"{x}(\d+){y}(\d+){z}(\d+)" - m = re.match(pattern, allocation_mode) - if not m: - continue - a, b, c = map(int, m.groups()) - return { - x: a, - y: b, - z: c, - } +class AllocationType(enum.Enum): + DECOUPLED = 1 + GLOBAL_HYBRID = 2 -def extract_decoupled_vllm_train_allocation(allocation_mode: str) -> Dict | None: - pattern = re.compile(r"(?:vllm\.(.+?)\+(.+))|(?:(.+?)\+vllm\.(.+))") - m = pattern.match(allocation_mode) - if not m: - return - if m.group(1): - vllm_alloc = m.group(1) - other_alloc = m.group(2) - else: - vllm_alloc = m.group(4) - other_alloc = m.group(3) - vllm_alloc = extract_symmetric_allocation(vllm_alloc) - other_alloc = extract_symmetric_allocation(other_alloc) - if not vllm_alloc: - return - if not other_alloc: - return - other_alloc.update({"vllm." + k: v for k, v in vllm_alloc.items()}) - return other_alloc +@dataclasses.dataclass +class AllocationMode: + type_: AllocationType + parallel_strat: Dict[str, Dict[str, int]] + + def is_decoupled(self): + return self.type_ == AllocationType.DECOUPLED + + def is_global_hybrid(self): + return self.type_ == AllocationType.GLOBAL_HYBRID + + @classmethod + def from_str(cls, allocation_mode: str): + 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) + if alloc_decoupled: + return cls(AllocationType.DECOUPLED, alloc_decoupled) + if alloc_3d: + return cls(AllocationType.GLOBAL_HYBRID, alloc_3d) + if alloc_hybrid: + return cls(AllocationType.GLOBAL_HYBRID, alloc_hybrid) + raise NotImplementedError(f"Failed to parse allocation: {allocation_mode}") + @staticmethod + def extract_3d_alloc(allocation_mode: str) -> Dict | None: + for x, y, z in itertools.permutations(["d", "m", "p"]): + pattern = rf"{x}(\d+){y}(\d+){z}(\d+)" + m = re.match(pattern, allocation_mode) + if not m: + continue + a, b, c = map(int, m.groups()) + # to be consistent with the key-value pattern + return { + "*": { + x: a, + y: b, + z: c, + } + } -def parse_key_value_pairs(s: str): - pattern = re.compile(r"([^:,]+):([^:,]+)") - matches = pattern.findall(s) - if not matches: - return None - return {key: value for key, value in matches} + @staticmethod + def extract_decoupled_alloc(allocation_mode: str) -> Dict | None: + pattern = re.compile( + r"(?:(?:vllm|sglang)\.(.+?)\+(.+))|(?:(.+?)\+(?:vllm|sglang)\.(.+))" + ) + m = pattern.match(allocation_mode) + if not m: + return + if m.group(1): + gen_alloc = m.group(1) + other_alloc = m.group(2) + else: + gen_alloc = m.group(4) + other_alloc = m.group(3) + gen_alloc = AllocationMode.extract_3d_alloc(gen_alloc) + if not gen_alloc: + return + other_alloc = AllocationMode.extract_3d_alloc( + other_alloc + ) or AllocationMode.extract_key_value_alloc(other_alloc) + if not other_alloc: + return + other_alloc.update({"gen": gen_alloc["*"]}) + return other_alloc + @staticmethod + def extract_key_value_alloc( + allocation_mode: str, + ) -> Dict[str, Dict[str, int]] | None: + def parse_key_value_pairs(s: str): + pattern = re.compile(r"([^:,]+):([^:,]+)") + matches = pattern.findall(s) + if not matches: + return None + return {key: value for key, value in matches} -def extract_key_value_allocation( - allocation_mode: str, -) -> Dict[str, Dict[str, int]] | None: - allocs = parse_key_value_pairs(allocation_mode) - if not allocs: - return - for k, v in allocs.items(): - v = extract_symmetric_allocation(v) - if not v: + allocs = parse_key_value_pairs(allocation_mode) + if not allocs: return - allocs[k] = v - return allocs + for k, v in allocs.items(): + v = AllocationMode.extract_3d_alloc(v) + if not v: + return + allocs[k] = v["*"] + return allocs diff --git a/realhf/impl/model/backend/inference.py b/realhf/impl/model/backend/inference.py index 8649701076..5a323121f2 100644 --- a/realhf/impl/model/backend/inference.py +++ b/realhf/impl/model/backend/inference.py @@ -103,7 +103,7 @@ def forward( post_hook=post_hook, aggregate_fn=aggregate_fn, ) - mb_inputs, fwd_indices, bwd_indices = input_.divide_into_mbs(mb_spec) + mb_inputs, fwd_indices, bwd_indices = input_.split(mb_spec) if constants.parallelism_rank() == 0: logger.info( f"MB spec: {mb_spec}, #mbs={len(mb_inputs)}, " @@ -161,14 +161,7 @@ def generate( # NOTE: Interleave mini-batches in the pipeline results will not decrease # the memory usage, because we need to hold all KV-caches for different # mini-batches, so we split mini-batches in the outer loop. - if mb_spec.max_tokens_per_mb is not None and constants.parallelism_rank() == 0: - logger.warning( - "Generation will ignore max_tokens_per_mb because the length is not predictable." - ) - mb_spec = MicroBatchSpec.new( - mb_spec, max_tokens_per_mb=None, balanced_seqs=True - ) - mb_inputs, *_ = input_.divide_into_mbs(mb_spec) + mb_inputs, *_ = input_.split(mb_spec) if constants.parallelism_rank() == 0: logger.info( f"MB spec: {mb_spec}, #mbs={len(mb_inputs)}, " diff --git a/realhf/impl/model/backend/megatron.py b/realhf/impl/model/backend/megatron.py index b2ab7d962b..e913d20de1 100644 --- a/realhf/impl/model/backend/megatron.py +++ b/realhf/impl/model/backend/megatron.py @@ -664,7 +664,6 @@ def _exec_backward_pass( loss: torch.Tensor = tensor_buffer.get( "losses", micro_batch_id, remove=True ) - loss = self.engine.optim.scale_loss(loss) loss.backward() tensor_buffer.put("losses", micro_batch_id, loss.detach().clone()) return @@ -745,12 +744,14 @@ def train_batch( input_: SequenceSample, mb_spec: MicroBatchSpec, loss_fn: Callable, + loss_weight_fn: Callable, + token_normalize_scope: str, version_steps: int, ): with megatron_ctx(): self.engine.zero_grad() if constants.pipe_parallel_world_size() > 1: - mb_inputs = input_.divide_into_mbs_balanced( + mb_inputs = input_.synced_data_parallel_split( MicroBatchSpec.new( mb_spec, n_mbs=mb_spec.n_mbs * self.pipe_runner.default_train_mbs, @@ -765,10 +766,24 @@ def train_batch( input_=input_, mb_spec=mb_spec, loss_fn=loss_fn, + loss_weight_fn=loss_weight_fn, + token_normalize_scope=token_normalize_scope, version_steps=version_steps, ) - mb_inputs = input_.divide_into_mbs_balanced(mb_spec) + mb_inputs = input_.synced_data_parallel_split(mb_spec) + total_loss_weight = torch.tensor( + sum([loss_weight_fn(mb) for mb in mb_inputs]), dtype=torch.float32 + ) + if token_normalize_scope == "global": + dist.all_reduce( + total_loss_weight, group=constants.data_parallel_group() + ) + if total_loss_weight == 0: + raise model_api.ZeroTotalLossWeightException( + "The sum of loss weights of all micro batches is zero." + ) + if constants.parallelism_rank() == 0: logger.info( f"MB spec: {mb_spec}, #mbs={len(mb_inputs)}, " @@ -795,8 +810,17 @@ def train_batch( max_seqlen=max_seqlen, ).logits loss, _stat = loss_fn(model_output, mb_input) + loss_scale = loss_weight_fn(mb_inputs[i]) / total_loss_weight + if token_normalize_scope == "global": + # Megatron will average gradients across DP ranks. + # If we normalize loss across micro batches of all DP ranks, + # we should revert the effect of gradient averaging in megatron + # to make sure loss from each token is scaled properly. + loss_scale *= constants.data_parallel_world_size() + loss_scale *= self.engine.optim.get_loss_scale().item() + loss *= loss_scale with cuda_tmarked("bwd", CUDATimeMarkType.backward): - self.engine.optim.scale_loss(loss).backward() + loss.backward() for k, v in _stat.items(): stat[k] += v @@ -937,6 +961,9 @@ def _initialize( overlap_grad_reduce=self.overlap_grad_reduce, overlap_param_gather=self.overlap_param_gather, clip_grad=self.optimizer.gradient_clipping, + min_loss_scale=self.optimizer.min_loss_scale, + loss_scale_window=self.optimizer.loss_scale_window, + hysteresis=self.optimizer.hysteresis, ) with megatron_ctx(): diff --git a/realhf/impl/model/backend/mock_train.py b/realhf/impl/model/backend/mock_train.py index c176e1d3f4..8e6ff36ed7 100644 --- a/realhf/impl/model/backend/mock_train.py +++ b/realhf/impl/model/backend/mock_train.py @@ -107,6 +107,8 @@ def train_batch( input_: SequenceSample, mb_spec: MicroBatchSpec, loss_fn: Callable, + loss_weight_fn: Callable, + token_normalize_scope: str, version_steps: int, ): self.optimizer.zero_grad() @@ -120,10 +122,21 @@ def train_batch( input_=input_, mb_spec=mb_spec, loss_fn=loss_fn, + loss_weight_fn=loss_weight_fn, + token_normalize_scope=token_normalize_scope, version_steps=version_steps, ) - mb_inputs = input_.divide_into_mbs_balanced(mb_spec) + mb_inputs = input_.synced_data_parallel_split(mb_spec) + total_loss_weight = torch.tensor( + sum([loss_weight_fn(mb) for mb in mb_inputs]), dtype=torch.float32 + ) + if token_normalize_scope == "global": + dist.all_reduce(total_loss_weight, group=constants.data_parallel_group()) + if total_loss_weight == 0: + raise model_api.ZeroTotalLossWeightException( + "The sum of loss weights of all micro batches is zero." + ) if constants.parallelism_rank() == 0: logger.info( @@ -147,6 +160,10 @@ def train_batch( max_seqlen=max_seqlen, ).logits loss, _stat = loss_fn(model_output, mb_input) + loss_scale = loss_weight_fn(mb_inputs[i]) / total_loss_weight + if token_normalize_scope == "global": + loss_scale *= constants.data_parallel_world_size() + loss *= loss_scale for k, v in _stat.items(): stat[k] += v diff --git a/realhf/impl/model/backend/pipe_runner.py b/realhf/impl/model/backend/pipe_runner.py index 3b026883c2..e9daaf182a 100644 --- a/realhf/impl/model/backend/pipe_runner.py +++ b/realhf/impl/model/backend/pipe_runner.py @@ -17,7 +17,10 @@ import realhf.impl.model.parallelism.pipeline_parallel.static_schedule as schedule import realhf.impl.model.utils.cuda_graph as cuda_graph from realhf.api.core.data_api import MicroBatchSpec, SequenceSample -from realhf.api.core.model_api import GenerationHyperparameters +from realhf.api.core.model_api import ( + GenerationHyperparameters, + ZeroTotalLossWeightException, +) from realhf.base.datapack import flat2d from realhf.impl.model.nn.real_llm_api import ReaLModel from realhf.impl.model.nn.real_llm_base import PipeCacheData, PipeTransferData @@ -668,7 +671,7 @@ def _exec_forward_pass( "input_cache", micro_batch_id, remove=True ) loss, stats = loss_fn(model_output, input_cache) - loss = loss / tensor_buffer.get("n_pp_mbs", micro_batch_id) + loss = loss * tensor_buffer.get("loss_scale", micro_batch_id) tensor_buffer.put("losses", micro_batch_id, loss) tensor_buffer.put("stats", micro_batch_id, stats) @@ -751,6 +754,7 @@ def _exec_recv_grads( @dataclasses.dataclass class PipeTrainInstrSet: + engine: Any def _exec_optimizer_step(self, *args, **kwargs): raise NotImplementedError() @@ -804,7 +808,7 @@ def forward( mb_spec = MicroBatchSpec.new( mb_spec, n_mbs=self.default_inf_mbs * mb_spec.n_mbs ) - mb_inputs, fwd_indices, bwd_indices = input_.divide_into_mbs(mb_spec) + mb_inputs, fwd_indices, bwd_indices = input_.split(mb_spec) if constants.parallelism_rank() == 0: logger.info( f"MB spec: {mb_spec}, #mbs={len(mb_inputs)}, " @@ -880,11 +884,8 @@ def generate( # When the global batch is fixed, not matter how many micro-batches we # split, the all-together KV-cache memory usage will not be changed, # so it's useless to split micro-batches here. - mb_spec = MicroBatchSpec( - n_mbs=self.default_inf_mbs, - balanced_seqs=True, - ) - mb_inputs, *_ = input_.divide_into_mbs(mb_spec) + mb_spec = MicroBatchSpec(n_mbs=self.default_inf_mbs) + mb_inputs, *_ = input_.split(mb_spec) if constants.parallelism_rank() == 0: logger.info( f"MB spec: {mb_spec}, #mbs={len(mb_inputs)}, " @@ -992,6 +993,8 @@ def train_batch( input_: SequenceSample, mb_spec: MicroBatchSpec, loss_fn: Callable, + loss_weight_fn: Callable, + token_normalize_scope: str, version_steps: int, ): # TODO: return whether update success @@ -1003,7 +1006,13 @@ def train_batch( mb_spec = MicroBatchSpec.new( mb_spec, n_mbs=mb_spec.n_mbs * self.default_train_mbs ) - mb_inputs = input_.divide_into_mbs_balanced(mb_spec) + mb_inputs = input_.synced_data_parallel_split(mb_spec) + total_loss_weight = torch.tensor( + sum([loss_weight_fn(mb) for mb in mb_inputs]), dtype=torch.float32 + ) + if token_normalize_scope == "global": + dist.all_reduce(total_loss_weight, group=constants.data_parallel_group()) + if constants.parallelism_rank() == 0: logger.info( f"MB spec: {mb_spec}, #mbs={len(mb_inputs)}, " @@ -1016,6 +1025,15 @@ def train_batch( tensor_buffer = TensorBuffer() for i in range(n_pp_mbs): tensor_buffer.put("n_pp_mbs", i, n_pp_mbs) + loss_scale = loss_weight_fn(mb_inputs[i]) / total_loss_weight + if token_normalize_scope == "global": + # Megatron will average gradients across DP ranks. + # If we normalize loss across micro batches of all DP ranks, + # we should revert the effect of gradient averaging in megatron + # to make sure loss from each token is scaled properly. + loss_scale *= constants.data_parallel_world_size() + loss_scale *= instr_set.engine.optim.get_loss_scale().item() + tensor_buffer.put("loss_scale", i, loss_scale) tensor_buffer.put("version_steps", i, version_steps) tensor_buffer.put("loss_fn", i, loss_fn) diff --git a/realhf/impl/model/backend/vllm.py b/realhf/impl/model/backend/vllm.py index f20ccd5e98..7bb60913aa 100644 --- a/realhf/impl/model/backend/vllm.py +++ b/realhf/impl/model/backend/vllm.py @@ -31,7 +31,7 @@ class LLMEngine_: from realhf.api.core import data_api, model_api from realhf.api.quickstart.model import vLLMConfig -from realhf.base import constants, logging +from realhf.base import constants, logging, seeding logger = logging.getLogger("vLLM backend") @@ -165,7 +165,6 @@ def generate( @dataclasses.dataclass class vLLMGenerationBackend(vLLMConfig, model_api.ModelBackend): - seed: int = 0 model_path: str = "" def _initialize( @@ -187,7 +186,7 @@ def _initialize( skip_tokenizer_init=False, trust_remote_code=True, max_model_len=self.max_model_len, - seed=self.seed, + seed=seeding.get_seed(), dtype=torch.float16, kv_cache_dtype=self.kv_cache_type, device=constants.current_device(), diff --git a/realhf/impl/model/interface/math_parser.py b/realhf/impl/model/interface/math_parser.py index e94811688e..d9af254963 100644 --- a/realhf/impl/model/interface/math_parser.py +++ b/realhf/impl/model/interface/math_parser.py @@ -84,7 +84,7 @@ def parse_line(prompt_str, generated, query_id): preexec_fn=os.setsid, stdout=subprocess.DEVNULL, ) - pro.wait(timeout=int(os.getenv("REAL_MATH_EVAL_TIMEOUT_SECS", "10"))) + pro.wait() try: os.killpg(os.getpgid(pro.pid), signal.SIGTERM) except ProcessLookupError: @@ -175,7 +175,7 @@ def parse_lines_in_parallel( procs.append(pro) for pro in procs: try: - pro.wait(timeout=int(os.getenv("REAL_MATH_EVAL_TIMEOUT_SECS", "10"))) + pro.wait() except Exception as e: pass try: diff --git a/realhf/impl/model/interface/ppo_interface.py b/realhf/impl/model/interface/ppo_interface.py index ca63267b45..760f2a0ae3 100644 --- a/realhf/impl/model/interface/ppo_interface.py +++ b/realhf/impl/model/interface/ppo_interface.py @@ -7,7 +7,7 @@ import functools import itertools import time -from typing import Dict, Optional, Tuple +from typing import Dict, Literal, Optional, Tuple import torch import torch.distributed as dist @@ -186,6 +186,7 @@ class PPOActorInterface(model_api.ModelInterface): mask_too_long: bool = False use_dense_reward: bool = False reward_delta: bool = True + token_normalize_scope: Literal["global", "dp"] = "global" def __post_init__(self): if self.adaptive_kl_ctl: @@ -614,9 +615,11 @@ def train_step( ) # NOTE: We cannot randomly shuffle data here because # data must have the same shape across different pipeline stages. - datas = input_.split( - self.n_minibatches, - min_size=input_.bs // self.n_minibatches, + datas, *_ = input_.split(MicroBatchSpec(n_mbs=self.n_minibatches)) + logger.info( + f"PPO minibatch split (size {self.n_minibatches}): " + f"#seqs: {[s.bs for s in datas]}, " + f"#tokens: {[sum([sum(lens) for lens in s.seqlens[s._get_split_key()]]) for s in datas]}" ) if self.use_dense_reward: @@ -672,6 +675,7 @@ def train_step( # Run mini-batched PPO training! train_stats = collections.defaultdict(lambda: 0) + for data in datas: stats = module.train_batch( input_=data, @@ -685,6 +689,8 @@ def train_step( early_stop_kl=self.early_stop_kl, temperature=self.gconfig.temperature, ), + loss_weight_fn=lambda x: x.data["ppo_loss_mask"].count_nonzero(), + token_normalize_scope=self.token_normalize_scope, ) if stats: @@ -888,6 +894,7 @@ class PPOCriticInterface(model_api.ModelInterface): mask_too_long: bool = False use_dense_reward: bool = False reward_delta: bool = True + token_normalize_scope: Literal["global", "dp"] = "global" def __post_init__(self): if self.adaptive_kl_ctl: @@ -1086,9 +1093,11 @@ def train_step( ) # NOTE: We cannot randomly shuffle data here because # data must have the same shape across different pipeline stages. - datas = input_.split( - self.n_minibatches, - min_size=input_.bs // self.n_minibatches, + datas, *_ = input_.split(MicroBatchSpec(n_mbs=self.n_minibatches)) + logger.info( + f"PPO minibatch split (size {self.n_minibatches}): " + f"#seqs: {[s.bs for s in datas]}, " + f"#tokens: {[sum([sum(lens) for lens in s.seqlens[s._get_split_key()]]) for s in datas]}" ) # Logging. @@ -1111,6 +1120,8 @@ def train_step( kl_adapter=self.kl_adapter, rms=None if not self.value_norm else self.rms, ), + loss_weight_fn=lambda x: x.data["ppo_loss_mask"].count_nonzero(), + token_normalize_scope=self.token_normalize_scope, ) if stats: diff --git a/realhf/impl/model/interface/sft_interface.py b/realhf/impl/model/interface/sft_interface.py index 5dba66895b..545750d30e 100644 --- a/realhf/impl/model/interface/sft_interface.py +++ b/realhf/impl/model/interface/sft_interface.py @@ -2,7 +2,7 @@ # Copyright 2024 Wei Fu & Zhiyu Mei # Licensed under the Apache License, Version 2.0 (the "License"). -from typing import Dict +from typing import Dict, Literal import torch import torch.distributed as dist @@ -89,6 +89,7 @@ def compute_packed_sft_loss( class SFTInterface(model_api.ModelInterface): + token_normalize_scope: Literal["global", "dp"] = "global" def train_step( self, model: model_api.Model, data: SequenceSample, mb_spec: MicroBatchSpec @@ -100,6 +101,10 @@ def train_step( stat = module.train_batch( input_=data, loss_fn=compute_packed_sft_loss, + loss_weight_fn=lambda x: x.data["prompt_mask"] + .logical_not() + .count_nonzero(), + token_normalize_scope=self.token_normalize_scope, mb_spec=mb_spec, version_steps=model.version.global_step, ) @@ -147,10 +152,7 @@ def evaluate( res = module.eval_batch( input_=x.to_device(device), loss_fn=compute_packed_sft_loss, - mb_spec=MicroBatchSpec( - n_mbs=constants.pipe_parallel_world_size(), - balanced_seqs=True, - ), + mb_spec=MicroBatchSpec(), ) if res is not None: diff --git a/realhf/system/flops_counter.py b/realhf/system/flops_counter.py new file mode 100644 index 0000000000..5e1a0740d3 --- /dev/null +++ b/realhf/system/flops_counter.py @@ -0,0 +1,117 @@ +import dataclasses +from typing import * + +from realhf.api.core.data_api import SequenceSample +from realhf.api.core.dfg import MFCDef, ModelInterfaceType +from realhf.api.core.model_api import ReaLModelConfig +from realhf.base.monitor import ( + caculuate_llama_forward_flops, + calculate_llama_gen_flops, + calculate_llama_train_flops, +) + + +@dataclasses.dataclass +class FlopsCounter: + train_configs: List[ReaLModelConfig] = dataclasses.field(default_factory=list) + train_bs: List[int] = dataclasses.field(default_factory=list) + train_seqlens: List[List[int]] = dataclasses.field(default_factory=list) + + inf_configs: List[ReaLModelConfig] = dataclasses.field(default_factory=list) + inf_bs: List[int] = dataclasses.field(default_factory=list) + inf_seqlens: List[List[int]] = dataclasses.field(default_factory=list) + + gen_configs: List[ReaLModelConfig] = dataclasses.field(default_factory=list) + gen_bs: List[int] = dataclasses.field(default_factory=list) + prompt_lens: List[List[int]] = dataclasses.field(default_factory=list) + gen_len: List[int] = dataclasses.field(default_factory=list) + + def clear(self): + self.train_bs.clear() + self.train_seqlens.clear() + + self.inf_bs.clear() + self.inf_seqlens.clear() + + self.gen_bs.clear() + self.prompt_lens.clear() + self.gen_len.clear() + + self.train_configs.clear() + self.inf_configs.clear() + self.gen_configs.clear() + + def add_rpc( + self, rpc: MFCDef, sample: SequenceSample, model_config: ReaLModelConfig + ): + # Record the data amount for each interface to compute FLOPs. + # Since the user may arbitrarily specify input/output keys, + # we can only try to find the most probable key name for computing FLOPs. + # If such keys do not exist, we will use the key with the longest + # sequence length in this model function call. + acc_seqlens = { + k: sum(sum(x) for x in slens) for k, slens in sample.seqlens.items() + } + seqlen_key = max(sample.seqlens, key=acc_seqlens.get) + flops_seqlens = [sum(x) for x in sample.seqlens[seqlen_key]] + if rpc.interface_type == ModelInterfaceType.GENERATE: + self.gen_configs.append(model_config) + self.gen_bs.append(sample.bs) + self.gen_len.append( + rpc.interface_impl.args["generation_config"]["min_new_tokens"] + ) + self.prompt_lens.append(flops_seqlens) + elif rpc.interface_type == ModelInterfaceType.TRAIN_STEP: + self.train_configs.append(model_config) + self.train_bs.append(sample.bs) + self.train_seqlens.append(flops_seqlens) + elif rpc.interface_type == ModelInterfaceType.INFERENCE: + self.inf_configs.append(model_config) + self.inf_bs.append(sample.bs) + self.inf_seqlens.append(flops_seqlens) + + def get_flops(self) -> int: + flops = 0 + for train_bs, train_seqlens, real_config in zip( + self.train_bs, + self.train_seqlens, + self.train_configs, + ): + flops += calculate_llama_train_flops( + checkpoint_activations_factor=4, + batch_size=train_bs, + seqlens=train_seqlens, + num_layers=real_config.n_layers, + hidden_size=real_config.hidden_dim, + intermediate_size=real_config.intermediate_dim, + vocab_size=real_config.vocab_size, + ) + for inf_bs, inf_seqlens, real_config in zip( + self.inf_bs, + self.inf_seqlens, + self.inf_configs, + ): + flops += caculuate_llama_forward_flops( + batch_size=inf_bs, + seqlens=inf_seqlens, + num_layers=real_config.n_layers, + hidden_size=real_config.hidden_dim, + intermediate_size=real_config.intermediate_dim, + vocab_size=real_config.vocab_size, + ) + for gen_bs, prompt_lens, gen_len, real_config in zip( + self.gen_bs, + self.prompt_lens, + self.gen_len, + self.gen_configs, + ): + flops += calculate_llama_gen_flops( + batch_size=gen_bs, + prompt_lens=prompt_lens, + gen_len=gen_len, + num_layers=real_config.n_layers, + hidden_size=real_config.hidden_dim, + intermediate_size=real_config.intermediate_dim, + vocab_size=real_config.vocab_size, + ) + return flops diff --git a/realhf/system/master_worker.py b/realhf/system/master_worker.py index d6ac2ddbc4..445ff3effa 100644 --- a/realhf/system/master_worker.py +++ b/realhf/system/master_worker.py @@ -11,6 +11,7 @@ import itertools import os import pprint +import random import re import time import uuid @@ -40,6 +41,7 @@ logging, name_resolve, names, + seeding, timeutil, topology, ) @@ -48,12 +50,8 @@ setup_run_until_complete, teardown_run_util_complete, ) -from realhf.base.monitor import ( - caculuate_llama_forward_flops, - calculate_llama_gen_flops, - calculate_llama_train_flops, -) from realhf.system.buffer import AsyncIOSequenceBuffer +from realhf.system.flops_counter import FlopsCounter logger = logging.getLogger("master worker", "system") blogger = logging.getLogger("benchmark") @@ -107,37 +105,6 @@ def _attach_param_realloc_hooks( return payload -@dataclasses.dataclass -class InterfaceDataAmount: - train_configs: List[ReaLModelConfig] = dataclasses.field(default_factory=list) - train_bs: List[int] = dataclasses.field(default_factory=list) - train_seqlens: List[List[int]] = dataclasses.field(default_factory=list) - - inf_configs: List[ReaLModelConfig] = dataclasses.field(default_factory=list) - inf_bs: List[int] = dataclasses.field(default_factory=list) - inf_seqlens: List[List[int]] = dataclasses.field(default_factory=list) - - gen_configs: List[ReaLModelConfig] = dataclasses.field(default_factory=list) - gen_bs: List[int] = dataclasses.field(default_factory=list) - prompt_lens: List[List[int]] = dataclasses.field(default_factory=list) - gen_len: List[int] = dataclasses.field(default_factory=list) - - def clear(self): - self.train_bs.clear() - self.train_seqlens.clear() - - self.inf_bs.clear() - self.inf_seqlens.clear() - - self.gen_bs.clear() - self.prompt_lens.clear() - self.gen_len.clear() - - self.train_configs.clear() - self.inf_configs.clear() - self.gen_configs.clear() - - @dataclasses.dataclass class RPCCorountineControl: ## Shared resources ## @@ -157,9 +124,7 @@ class RPCCorountineControl: # for training data management and data cleaning after each step ids_to_clear: Set[int] = dataclasses.field(default_factory=set) - data_amount: InterfaceDataAmount = dataclasses.field( - default_factory=InterfaceDataAmount - ) + flops_counter: FlopsCounter = dataclasses.field(default_factory=FlopsCounter) should_save: bool = False should_eval: bool = False @@ -408,31 +373,7 @@ async def model_rpc_request_func( buf_indices, sample = await buffer.get_batch_for_rpc(rpc) - # Record the data amount for each interface to compute FLOPs. - # Since the user may arbitrarily specify input/output keys, - # we can only try to find the most probable key name for computing FLOPs. - # If such keys do not exist, we will use the key with the longest - # sequence length in this model function call. - acc_seqlens = { - k: sum(sum(x) for x in slens) for k, slens in sample.seqlens.items() - } - seqlen_key = max(sample.seqlens, key=acc_seqlens.get) - flops_seqlens = [sum(x) for x in sample.seqlens[seqlen_key]] - if rpc.interface_type == dfg.ModelInterfaceType.GENERATE: - ctrl.data_amount.gen_configs.append(model_configs[rpc.model_name]) - ctrl.data_amount.gen_bs.append(sample.bs) - ctrl.data_amount.gen_len.append( - rpc.interface_impl.args["generation_config"]["min_new_tokens"] - ) - ctrl.data_amount.prompt_lens.append(flops_seqlens) - elif rpc.interface_type == dfg.ModelInterfaceType.TRAIN_STEP: - ctrl.data_amount.train_configs.append(model_configs[rpc.model_name]) - ctrl.data_amount.train_bs.append(sample.bs) - ctrl.data_amount.train_seqlens.append(flops_seqlens) - elif rpc.interface_type == dfg.ModelInterfaceType.INFERENCE: - ctrl.data_amount.inf_configs.append(model_configs[rpc.model_name]) - ctrl.data_amount.inf_bs.append(sample.bs) - ctrl.data_amount.inf_seqlens.append(flops_seqlens) + ctrl.flops_counter.add_rpc(rpc, sample, model_configs[rpc.model_name]) this_rpc_consumed_seqs += sample.bs @@ -440,21 +381,27 @@ async def model_rpc_request_func( # Dispatch data to different data parallel ranks. dp_size = topo.get_dim("data") - pp_size = topo.get_dim("pipe") - if rpc.mb_spec.balanced_seqs or ( - rpc.mb_spec.max_tokens_per_mb is not None - and rpc.interface_type == dfg.ModelInterfaceType.TRAIN_STEP - ): - # For a train RPC, we must assure that all DP ranks have the same number - # of micro-batches, so we must evenly distribute the sequences. - assert sample.bs % dp_size == 0 - min_n_seqs_per_dp = sample.bs // dp_size + 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: - min_n_seqs_per_dp = pp_size * rpc.min_n_seqs_per_pass * rpc.mb_spec.n_mbs - if rpc.interface_type == dfg.ModelInterfaceType.TRAIN_STEP and pp_size > 1: - min_n_seqs_per_dp *= 2 - split_spec = sample.get_split_spec(dp_size, min_size=int(min_n_seqs_per_dp)) - partitions = split_spec.partitions + 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 @@ -599,6 +546,8 @@ class MasterWorker(worker_base.Worker): def _configure(self, config: config_pkg.MasterWorker): self.config = config + seeding.set_random_seed(self.config.base_seed + self.config.n_model_workers) + self.__model_topos: Dict[ModelName, topology.PipeModelDataParallelTopology] = ( config.model_topos ) @@ -1248,12 +1197,9 @@ def _maybe_request_load_data(self, first_poll: bool): filtered_data.append(x) all_data = filtered_data - # Reorder loaded (meta-)data and store them into the buffer. - # NOTE: The reordered indices prioritize longer sequences for detecting OOM errors early. - # reorder_indices, _ = datapack.reorder_to_balanced_batches( - # np.array(seqlens), src_rpc.n_seqs - # ) - # all_data: List[data_api.SequenceSample] = [all_data[i] for i in reorder_indices] + # 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. " @@ -1287,52 +1233,10 @@ def _log_training_stats(self, e2e_time: float, time_since_configure: float): flops = None tflops_per_gpu = float("inf") else: - flops = 0 - for train_bs, train_seqlens, real_config in zip( - self.__rpc_ctrl.data_amount.train_bs, - self.__rpc_ctrl.data_amount.train_seqlens, - self.__rpc_ctrl.data_amount.train_configs, - ): - flops += calculate_llama_train_flops( - checkpoint_activations_factor=4, - batch_size=train_bs, - seqlens=train_seqlens, - num_layers=real_config.n_layers, - hidden_size=real_config.hidden_dim, - intermediate_size=real_config.intermediate_dim, - vocab_size=real_config.vocab_size, - ) - for inf_bs, inf_seqlens, real_config in zip( - self.__rpc_ctrl.data_amount.inf_bs, - self.__rpc_ctrl.data_amount.inf_seqlens, - self.__rpc_ctrl.data_amount.inf_configs, - ): - flops += caculuate_llama_forward_flops( - batch_size=inf_bs, - seqlens=inf_seqlens, - num_layers=real_config.n_layers, - hidden_size=real_config.hidden_dim, - intermediate_size=real_config.intermediate_dim, - vocab_size=real_config.vocab_size, - ) - for gen_bs, prompt_lens, gen_len, real_config in zip( - self.__rpc_ctrl.data_amount.gen_bs, - self.__rpc_ctrl.data_amount.prompt_lens, - self.__rpc_ctrl.data_amount.gen_len, - self.__rpc_ctrl.data_amount.gen_configs, - ): - flops += calculate_llama_gen_flops( - batch_size=gen_bs, - prompt_lens=prompt_lens, - gen_len=gen_len, - num_layers=real_config.n_layers, - hidden_size=real_config.hidden_dim, - intermediate_size=real_config.intermediate_dim, - vocab_size=real_config.vocab_size, - ) + flops = self.__rpc_ctrl.flops_counter.get_flops() tflops = flops / (e2e_time * (10**12)) tflops_per_gpu = flops / (e2e_time * self.config.n_model_workers * (10**12)) - self.__rpc_ctrl.data_amount.clear() + self.__rpc_ctrl.flops_counter.clear() ######################################### epoch = self.__rpc_ctrl.step_info.epoch + 1 diff --git a/realhf/system/model_worker.py b/realhf/system/model_worker.py index 2c4fa588c9..d2f4a78cc1 100644 --- a/realhf/system/model_worker.py +++ b/realhf/system/model_worker.py @@ -125,10 +125,7 @@ def _configure(self, cfg: system_api.ModelWorker): self.__worker_index = cfg.worker_info.worker_index - torch.backends.cudnn.benchmark = cfg.cudnn_benchmark - torch.backends.cudnn.deterministic = cfg.cudnn_deterministic - - seeding.set_random_seed(cfg.seed) + seeding.set_random_seed(cfg.base_seed + self.__worker_index) # Reveal process group identity of this worker to world. gpu_utils.reveal_pg_identity( @@ -307,7 +304,8 @@ def __lazy_setup(self): datasets = [ data_api.make_dataset( d, - self.config.seed, + # NOTE: we must use the same seed to ensure the same dataset split + self.config.base_seed, self.__dataset_dp_rank, self.__dataset_dp_size, self.config.tokenizer_name_or_path, @@ -326,10 +324,20 @@ def __lazy_setup(self): else: self.__dataset = torch.utils.data.ConcatDataset(datasets) + g = torch.Generator() + g.manual_seed(seeding.get_seed()) + self.__dataloader = torch.utils.data.DataLoader( + self.__dataset, + collate_fn=data_api.SequenceSample.gather, + # NOTE: This is *NOT* the actual batch size for training. + # It is just a proper size to load data to workers. + batch_size=10240, + shuffle=True, + generator=g, + ) + self.__raw_samples = [] - for tmp_sample in data_api.make_dataloader( - self.config.dataloader, self.__dataset - ): + for tmp_sample in self.__dataloader: self.__raw_samples += tmp_sample.meta().unpack() self.__models: Dict[ModelName, model_api.Model] = dict() @@ -362,7 +370,11 @@ def __lazy_setup(self): ) # Recover indices for dynamic dataset - if self.__has_dataset and hasattr(self.__dataset, "filter"): + if ( + s.id.model_name == src_rpc.model_name + and self.__has_dataset + and hasattr(self.__dataset, "filter") + ): dataset_indices_path = os.path.join( constants.MODEL_SAVE_ROOT, constants.experiment_name(), @@ -405,30 +417,27 @@ def __lazy_setup(self): interface_impl[0] ) - if s.eval_datasets is not None and s.eval_dataloader is not None: - eval_datasets = [ - data_api.make_dataset( - d, - self.config.seed, - s.id.dp_rank, - s.id.topo.get_dim("data"), - self.__models[s.id.model_name].tokenizer, - self.config.worker_info.experiment_name, - self.config.worker_info.trial_name, - cache_root=( - None - if not self.config.use_dataset_cache - else self.config.dataset_cahce_root - ), - ) - for d in s.eval_datasets - ] - if len(eval_datasets) > 1: - eval_dataset = torch.utils.data.ConcatDataset(eval_datasets) - else: - eval_dataset = eval_datasets[0] - eval_dataloader = data_api.make_dataloader( - s.eval_dataloader, eval_dataset + if s.eval_dataset is not None: + eval_dataset = data_api.make_dataset( + s.eval_dataset, + # NOTE: we must use the same seed to ensure the same dataset split + self.config.base_seed, + s.id.dp_rank, + s.id.topo.get_dim("data"), + self.__models[s.id.model_name].tokenizer, + self.config.worker_info.experiment_name, + self.config.worker_info.trial_name, + cache_root=( + None + if not self.config.use_dataset_cache + else self.config.dataset_cahce_root + ), + ) + eval_dataloader = torch.utils.data.DataLoader( + eval_dataset, + batch_size=s.eval_bs, + collate_fn=data_api.SequenceSample.gather, + shuffle=False, ) else: eval_dataloader = None @@ -600,8 +609,16 @@ def handle_non_blocking_request(self, request: request_reply_stream.Payload): dataset_indices_path, self.__dataset.active_indices, ) - self.__dataloader = data_api.make_dataloader( - self.config.dataloader, self.__dataset + g = torch.Generator() + g = g.set_state(self.__dataloader.generator.get_state()) + self.__dataloader = torch.utils.data.DataLoader( + self.__dataset, + collate_fn=data_api.SequenceSample.gather, + # NOTE: This is *NOT* the actual batch size for training. + # It is just a proper size to load data to workers. + batch_size=10240, + shuffle=True, + generator=g, ) self.__data_generator = enumerate(self.__dataloader) diff --git a/realhf/system/request_reply_stream.py b/realhf/system/request_reply_stream.py index 5375bc6b9d..9147143b58 100644 --- a/realhf/system/request_reply_stream.py +++ b/realhf/system/request_reply_stream.py @@ -156,7 +156,7 @@ def post(self, payload: Payload) -> uuid.UUID: def request( self, - handlers: List[str] | None = None, + handlers: List[str | int] | None = None, handle_type: str | None = None, datas: List[Any] | None = None, payloads: List[Payload] | None = None, @@ -222,7 +222,7 @@ def request( def call( self, - handlers: List[str] | None = None, + handlers: List[str | int] | None = None, handle_type: str | None = None, datas: List[Any] | None = None, payloads: List[Payload] | None = None, diff --git a/requirements.txt b/requirements.txt index 3fe62118cb..bc0ca01de8 100644 --- a/requirements.txt +++ b/requirements.txt @@ -38,18 +38,17 @@ scipy seaborn setuptools>=61.0 tqdm -transformers==4.42.3 networkx==3.3 matplotlib tabulate aiofiles pydantic -black -isort +isort==5.13.2 clang-format +ninja +paramiko # To eliminate security risks torch>2.0.0 black>=25.1.0 cookiecutter>2.1.1 -tensorflow>2.13.0 -tensorflow-cpu>2.12.1 \ No newline at end of file +httpx>=0.28.1 diff --git a/tests/comm/test_param_realloc.py b/tests/comm/test_param_realloc.py index e99b4479dd..78c97718cf 100644 --- a/tests/comm/test_param_realloc.py +++ b/tests/comm/test_param_realloc.py @@ -496,6 +496,8 @@ def _test_para_realloc( if not is_critic else compute_critic_loss ), + loss_weight_fn=lambda: 1, + token_normalize_scope="dp", version_steps=i, ) diff --git a/tests/data/test_sequence_gather_split.py b/tests/data/test_sequence_gather_split.py index 67f1f372c3..cee68e9def 100644 --- a/tests/data/test_sequence_gather_split.py +++ b/tests/data/test_sequence_gather_split.py @@ -243,19 +243,6 @@ def test_gather_split(sample_type: str, dp: int): x = SequenceSample.gather(samples) - # Test gather-split-gather cosistency - for k in x.keys: - y = SequenceSample.gather(x.split(dp, key=k, min_size=1)) - recursive_assert_equal(x, y) - - # Test balanced split - balanced_size = sum(batch_sizes) // dp - for k in x.keys: - splitted = x.split(dp, key=k, min_size=balanced_size) - assert all(len(s.ids) >= balanced_size for s in splitted) - y = SequenceSample.gather(splitted) - recursive_assert_equal(x, y) - # Test split to original samples spec = SequenceSplitSpec(sizes=batch_sizes) ss = x.split_with_spec(spec) @@ -264,11 +251,10 @@ def test_gather_split(sample_type: str, dp: int): # Test split to the finest granularity total_bs = sum(batch_sizes) - for k in x.keys: - ss = x.split(total_bs, key=k, min_size=1) - assert len(ss) == total_bs - y = SequenceSample.gather(ss) - recursive_assert_equal(x, y) + ss, _, backward_indices = x.split(MicroBatchSpec(n_mbs=x.bs)) + assert len(ss) == total_bs + y = SequenceSample.reorder(SequenceSample.gather(ss), backward_indices) + recursive_assert_equal(x, y) # Test divide micro batch and merge back for seqlens in [ @@ -281,7 +267,7 @@ def test_gather_split(sample_type: str, dp: int): n_mbs=np.random.randint(1, 10), max_tokens_per_mb=np.random.randint(800, 1000), ) - mb_data, fwd_indices, bwd_indices = x.divide_into_mbs(mb_spec) + mb_data, fwd_indices, bwd_indices = x.split(mb_spec) for id_x, id_y in zip( [x.ids[i] for i in fwd_indices],