diff --git a/.github/workflows/test-areal.yml b/.github/workflows/test-areal.yml new file mode 100644 index 0000000000..1e939ed10c --- /dev/null +++ b/.github/workflows/test-areal.yml @@ -0,0 +1,20 @@ +name: Test AReaL + +on: + pull_request: + branches: [ main ] + paths: + - .github/workflows/test-areal.yml + - areal/** + - ci/** + workflow_dispatch: + +jobs: + test-areal: + environment: + name: AReaL-unittests + runs-on: gpu + steps: + - uses: actions/checkout@v4 + + - run: bash ./ci/test_areal.sh diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml new file mode 100644 index 0000000000..e3897d5e6b --- /dev/null +++ b/.pre-commit-config.yaml @@ -0,0 +1,60 @@ +repos: + - repo: https://github.com/pre-commit/mirrors-clang-format + rev: v19.1.7 + hooks: + - id: clang-format + files: \.(c|cc|cxx|cpp|h|hpp|hxx|cu|cuh)$ + args: + - --style=file + - --fallback-style=Google + + - repo: https://github.com/pre-commit/pre-commit-hooks + rev: v4.4.0 + hooks: + - id: check-yaml + files: \.(yaml|yml)$ + - id: end-of-file-fixer + files: \.(py|md|yaml|yml|c|cc|cxx|cpp|h|hpp|hxx|cu|cuh)$ + - id: trailing-whitespace + files: \.(py|md|yaml|yml|c|cc|cxx|cpp|h|hpp|hxx|cu|cuh)$ + + - repo: https://github.com/executablebooks/mdformat + rev: 0.7.17 + hooks: + - id: mdformat + args: + - --wrap=88 + additional_dependencies: + - mdformat-gfm + - mdformat-tables + - mdformat-frontmatter + + - repo: https://github.com/pycqa/autoflake + rev: v2.3.1 + hooks: + - id: autoflake + args: + - -i + - -r + - --remove-all-unused-imports + - --remove-unused-variables + files: ^(areal/|examples/lite/).*\.py$ + + - repo: https://github.com/pycqa/isort + rev: 5.13.2 + hooks: + - id: isort + args: + - --profile=black + - --multi-line=3 + - --line-length=88 + files: \.py$ + + - repo: https://github.com/psf/black/ + rev: 25.1.0 + hooks: + - id: black + language_version: python3 + args: + - --line-length=88 + files: \.py$ diff --git a/README.md b/README.md index 58658430c3..1ed2e83d1d 100644 --- a/README.md +++ b/README.md @@ -7,165 +7,186 @@ WeChat Group |

-ReaL +ReaL -AReaL (Ant Reasoning RL) is an open-source **fully asynchronous reinforcement learning training system** for large reasoning models developed at **the RL Lab, Ant Research**. Built upon the open-source project [RealHF](https://github.com/openpsi-project/ReaLHF), we are fully committed to open-source by providing training details, data, and infrastructure required to reproduce results along with the model itself. AReaL aims to help everyone build their own AI agents easily and affordably. Our team loves milk tea because it's delicious, customizable, and affordable. We hope you enjoy our project just like how you enjoy real-world milk tea (cheers). +AReaL (Ant Reasoning RL) is an open-source **fully asynchronous reinforcement learning +training system** for large reasoning models developed at **the RL Lab, Ant Research**. +Built upon the open-source project [ReaLHF](https://github.com/openpsi-project/ReaLHF), +we are fully committed to open-source by providing training details, data, and +infrastructure required to reproduce results along with the model itself. AReaL aims to +help everyone build their own AI agents easily and affordably. Our team loves milk tea +because it's delicious, customizable, and affordable. We hope you enjoy our project just +like how you enjoy real-world milk tea (cheers). **AReaL Highlights** -+ 🔥 **[NEW] Asynchronous RL:** With algorithm-system co-design, AReaL supports fully asynchronous RL for **the fastest training**! Experimental support for multi-turn agentic RL is also provided. -+ 🛠️ **Open & Reproducible**: We continuously release _all code, datasets, and training recipes_ for RL training of LLMs. -+ 🚀 **Scalability**: AReaL can seamlessly adapt to different computational resource settings, ranging from a single node to 1K GPUs. -+ 🔪 **Cutting-Edge Performance:** AReaL can produce models with cutting-edge reasoning capabilities in math and coding. We are also actively working on agentic tasks. +- ⚡ **\[NEW\] AReaL-lite:** Our new + release AReaL-lite is a **light-weight** and **algorithm-first** codebase that + prioritizes better development experiences for AI researchers. As a result, AReaL-lite + delivers most AReaL functionalities while maintains its high performance with much + fewer lines of code. This allows users to build their own **agentic** training + workflows with minimal efforts. +- 🔥 **Asynchronous RL**: With algorithm-system co-design, AReaL supports fully + asynchronous RL for **the fastest training speed**! Experimental support for + multi-turn agentic RL is also provided. +- 🛠️ **Open & Reproducible**: We continuously release _all code, datasets, and training + recipes_ for RL training of LLMs. +- 🚀 **Scalability**: AReaL can seamlessly adapt to different computational resource + settings, ranging from a single node to 1K GPUs. +- 🔪 **Cutting-Edge Performance:** AReaL can produce models with cutting-edge reasoning + capabilities in math and coding. We are also actively working on agentic tasks. ## News -**[2025/06/03] (v0.3, boba²)** We release **boba²** (double-boba) for fully asynchronous RL training, which achieves a **2.77x speedup while obtaining on-par or even better training performance** compared to synchronous systems. Moreover, asynchronous RL makes it extremely easy to set up multi-turn agentic RL training! Check out [our v0.3 overview blog](/blog/AReaL_v0_3.md) and the [research paper](https://arxiv.org/pdf/2505.24298). +**\[2025/07/31\] (AReaL-lite)** We introduce AReaL-lite, a **light-weight** version of +AReaL designed specifically for AI researchers and rapid prototyping. AReaL-lite +features an **algorithm-first** API design that prioritizes ease of use and algorithm +development, while inherently supporting **fully asynchronous agentic RL**. With 80% +fewer lines of code, AReaL-lite maintains 90% of AReaL's high performance and core +functionality. Check out [our AReaL-lite design doc](/areal/README.md) and +[the quickstart guide](https://inclusionai.github.io/AReaL/tutorial/quickstart.html) to +begin your journey with **AReaL-lite**! -**[2025/03/31] (v0.2, boba)** Here comes our next milestone release - boba! Please call it A-ReaL-boba! This release includes much faster training with SGLang support and SOTA 7B and 32B models on math reasoning. Check our [v0.2 technical blog](/blog/AReaL_v0_2.md). +**\[2025/06/03\] (v0.3, boba²)** We release **boba²** (double-boba) for fully +asynchronous RL training, which achieves a **2.77x speedup while obtaining on-par or +even better training performance** compared to synchronous systems. Moreover, +asynchronous RL makes it extremely easy to set up multi-turn agentic RL training! Check +out [our v0.3 overview blog](/blog/AReaL_v0_3.md) and the +[research paper](https://arxiv.org/pdf/2505.24298). -**[2025/02/24] (v0.1)** Our initial release includes reproducible results for 1.5B and 7B LRMs. Check our [v0.1 technical blog](/blog/AReaL_v0_1.md). +**\[2025/03/31\] (v0.2, boba)** Here comes our next milestone release - boba! Please +call it A-ReaL-boba! This release includes much faster training with SGLang support and +SOTA 7B and 32B models on math reasoning. Check our +[v0.2 technical blog](/blog/AReaL_v0_2.md). -## Release Highlights +**\[2025/02/24\] (v0.1)** Our initial release includes reproducible results for 1.5B and +7B LRMs. Check our [v0.1 technical blog](/blog/AReaL_v0_1.md). -In our AReaL-boba² (A-ReaL-double-boba) release, we highlight the top 3 most important features: +## AReaL-lite Release Highlights -+ A fully asynchronous RL training pipeline with **system and RL algorithm co-design**, achieving over 2.77x speedup without any performance drop. Check the [benchmark scripts and instructions here](https://github.com/inclusionAI/AReaL/tree/main/benchmark/verl_v0_3_0_post1_76084d3). +New highlights in AReaL-lite: -+ SOTA coding models, i.e., a 14B model with a **69.1 score on LCB-v5**. To reproduce, check the [configs](https://github.com/inclusionAI/AReaL/tree/main/examples/configs/v0.3-qwen3-code) and [instructions](https://inclusionai.github.io/AReaL/references/reproduce.html). +- Instead of the *system-first* architecture in old AReaL, AReaL-lite follows an + **algorithm-first** API design that aims to provide the following key features: -+ Experimental support for **multi-turn** agentic RL training. Check our [complete example](https://inclusionai.github.io/AReaL/customization/agent.html). + - **Light-weight** & **easy-to-write** algorithm and training workflow customization. + - **Easy to scale up** without knowing system and infrastructure details. + - **Adaptable and plugable:** Smooth to integrate with other modern AI applications. -For the complete system design and more training details, please check [our v0.3 blog](/blog/AReaL_v0_3.md) and our [research paper](https://arxiv.org/pdf/2505.24298). + These features make AReaL-lite easy for AI researchers to adopt, understand, and + develop effectively and efficiently. To learn more about the design principles of + AReaL, please read the [AReaL-lite design doc](/areal/README.md)! -**Jump to the [quickstart section](https://github.com/inclusionAI/AReaL?tab=readme-ov-file#getting-started) if you want to quickly run an experiment and get your hands dirty!** 😈 +- A much more *light-weight* codebase compared to old AReaL codebase with only **20%** # + lines of code, with a detailed + [code walkthrough](https://inclusionai.github.io/AReaL/areal/gsm8k_grpo.html) on an + GRPO-on-GSM8K example. Save your time & efforts for code reading! -### Overview of Asynchronous RL Training +- Smoother customization for your own **algorithms** and **agentic & RLVR rollout** RL + within a single file! Check + [here](https://inclusionai.github.io/AReaL/customization/agent.html) for agent & RLVR + customization and [here](https://inclusionai.github.io/AReaL/tutorial/algorithm.html) + for algorithm customization. -During the synchronous RL training process, a generation step must wait until the longest sequence completes within the batch of LLM outputs. Due to the varying output lengths for LRMs, a synchronous RL system suffers from massive GPU idle time, leading to training inefficiency. Some recent works ([DeepCoder](https://pretty-radio-b75.notion.site/DeepCoder-A-Fully-Open-Source-14B-Coder-at-O3-mini-Level-1cf81902c14680b3bee5eb349a512a51), [Intellect](https://www.primeintellect.ai/blog/intellect-2)) propose overlapping a single training step with a single generation step to accelerate training. However, the largest bottleneck remains unchanged: the samples within a batch are still from the same model version, leading to waiting and GPU idle time. +Good old stuff from AReaL: -![Synchronous vs One-step Overlap RL](/assets/sync_one_step_gen.png) +- High performance and scalability with fully asynchronous RL training. Check our + [boba² (v0.3) blog](/blog/AReaL_v0_3.md) for details. -*Fig.1. Left: Execution timeline of synchronous RL training. Right: Execution timeline of one-step overlap RL system.* +- A single command line to launch an experiment, no matter on a single node or a + large-scale distributed cluster. -AReaL adopts a fully asynchronous RL training framework that completely decouples generation from training. In AReaL, LLM generation runs in a streaming manner, with each rollout worker continuously producing outputs without waiting. Meanwhile, trainer workers perform parallel model updates upon receiving training batches. +Now, let us run an example experiment with AReaL-lite following the quickstart guide +below! -![Asynchronous RL Training](/assets/async_timeline.png) +## Getting Started with AReaL-lite -*Fig 2. Execution timeline of our fully asynchronous RL system.* +Our training scripts will automatically download the dataset (openai/gsm8k) and model +(Qwen/Qwen2-1.5B-Instruct). On a single node, runs: -AReaL follows a system-algorithm co-design principle: on the system side, AReaL efficiently syncs model parameters and carefully controls the staleness of each training sample; on the algorithm side, AReaL improves the objective of PPO to make async-RL stable. - -We compare the scalability of **asynchronous RL** training based on our AReaL-boba² system with **classical synchronous RL** training (we adopt the fastest open-source system veRL, main branch on 05/07/2025) across different model sizes and different numbers of H800 GPUs. AReaL demonstrates much improved scaling capabilities with respect to training throughput. This is also partially due to AReaL decoupling training and generation, leading to much fewer GPU memory fragments. - -![Scaling Comparison](/assets/async_scaling_vs_verl.png) - -*Fig.3 The scaling trend of asynchronous RL (based on AReaL-boba2) and classical synchronous RL (based on veRL) with different model sizes. Dotted lines indicate ideal linear scaling.* - -### SOTA Code Generation Model by AReaL-boba² - -We use **Qwen3** as our base model. After asynchronous RL training, we achieve SOTA results on LiveCodeBench, Codeforces, and CodeContests benchmarks. - -| **Model (8B)** | **LiveCodeBench v5**
**(2024.10-2025.2)** | **Codeforces** | **CodeContests** | -| :---: | :---: | :---: | :---: | -| Qwen3-8B | 58.8 | 1879/96.7% | 31.4 | -| DeepSeek-R1-0528-Qwen3-8B | 58.4 | 1945/97.3% | 31.0 | -| [🤗 AReaL-boba²-8B-Open](https://huggingface.co/inclusionAI/AReaL-boba-2-8B-subset) | 62.0 | 1933/97.2% | **41.4** | -| [🤗 AReaL-boba²-8B](https://huggingface.co/inclusionAI/AReaL-boba-2-8B) | **63.0** | **1962/97.5%** | 40.8 | - -| **Model (14B)** | **LiveCodeBench v5**
**(2024.10-2025.2)** | **Codeforces** | **CodeContests** | -| :---: | :---: | :---: | :---: | -| Qwen3-14B | 65.4 | 1978/97.7% | 38.3 | -| DeepCoder-14B-Preview | 60.6 | 1936/95.3% | 40.1 | -| [🤗 AReaL-boba²-14B-Open](https://huggingface.co/inclusionAI/AReaL-boba-2-14B-subset) | 67.3 | 1990/97.8% | **46.2** | -| [🤗 AReal-boba²-14B](https://huggingface.co/inclusionAI/AReaL-boba-2-14B) | **69.1** | **2044/98.2%** | 46.1 | +```bash +python3 -m arealite.launcher.local examples/arealite/gsm8k_grpo.py --config examples/arealite/configs/gsm8k_grpo.yaml +``` -| **Larger Models** | **LiveCodeBench v5**
**(2024.10-2025.2)** | **Codeforces** | **CodeContests** | -| :---: | :---: | :---: | :---: | -| Qwen3-235B | 70.7 | 2056 | - | -| DeepSeek-R1 | 64.3 | 2029 | - | -| OpenAI-o3-mini (Medium) | 66.3 | 2036 | - | +On a Ray cluster with 2 nodes & 8 GPUs each node, runs (Remember to change paths in the +YAML file to your own shared storage): -*Table 1: Coding Task Performance Comparison. AReaL-boba²-8B/14B-Open denotes training results on open-source data. AReaL-boba²-8B/14B models are trained with an additional small amount of internal data and achieve SOTA performance on LiveCodeBench, Codeforces & CodeContests.* +```bash +python3 -m arealite.launcher.ray examples/arealite/gsm8k_grpo.py --config examples/arealite/configs/gsm8k_grpo.yaml \ + cluster.n_nodes=2 \ + cluster.n_gpus_per_node=8 +``` -We highlight the [tutorials](https://inclusionai.github.io/AReaL/customization/dataset.html) and [code walkthroughs](https://inclusionai.github.io/AReaL/developer/overview.html) about the following key features for asynchronous training: + -+ [Streaming generation and reward computation](https://inclusionai.github.io/AReaL/developer/rollout/rollout_worker.html) -+ [Interruptible rollout](https://inclusionai.github.io/AReaL/developer/rollout/gserver.html) -+ [Data staleness control with the rollout controller](https://inclusionai.github.io/AReaL/developer/rollout/gserver.html) -+ [The adoption of decoupled PPO loss](https://inclusionai.github.io/AReaL/customization/algorithm.html) + -AReaL-boba² allows you to independently customize the [dataset](https://inclusionai.github.io/AReaL/customization/dataset.html), [rollout behavior](https://inclusionai.github.io/AReaL/customization/agent.html), and the [training algorithm](https://inclusionai.github.io/AReaL/customization/algorithm.html), without needing to modify the heavy system-level code. +For more detailed guide on how to run experiments in AReaL-lite, please check out +[our quickstart guide](https://inclusionai.github.io/AReaL/tutorial/quickstart.html)! -In particular, we show a simple example to develop a multi-turn math agent for RL training. Please see the learning curve below and reference the [step-by-step guide](https://inclusionai.github.io/AReaL/customization/agent.html) if you want to implement your own agentic RL project. +## Switching from legacy AReaL to AReaL-lite -## Getting Started -Obtain the training data: -- [Math](https://huggingface.co/datasets/inclusionAI/AReaL-boba-Data) -- [Code](https://huggingface.co/datasets/inclusionAI/AReaL-boba-2-RL-Code) +We also provide a convenient script to convert your AReaL YAML config into AReaL-lite +config in one command line. First you need to locate your AReaL config either modified +from files from `examples` folder, or generated when you run your experiments in +`//` folder. Runs: -For code training data, a simple preprocessing script was provided in `examples/data_preprocess/preprocess_training_data.py`: -```bash -python3 preprocess_training_data.py --data_path $original_data_path --output_path $training_data_path ``` - -Train Qwen3 1.7B locally (Remember to modify `dataset.path` in the script below): -```bash -bash examples/run_async_ppo.sh +python arealite/utils/convert.py --convert_src AReaL --src_config_path --template_path examples/arealite/configs/gsm8k_grpo.yaml --output_path ``` -Evaluation: +Then you should be able to run experiments with your old settings on AReaL-lite! -```bash -cd evaluation -# Evaluate the model -python eval_and_aggregate.py \ - --model_path ${MODEL_PATH} \ - --output_path ${OUTPUT_PATH} \ - --data_names aime24,aime25 \ - --max_gen_tokens 32768 \ - --data_names codeforces,lcb_v5 \ - --prompt_type qwen3-think-pure \ - --temperature 1.0 -``` +## AReaL-lite vs legacy AReaL + +AReaL-lite is an initiative to fully refactor AReaL, addressing historical issues such +as redundant code and unnecessary system-level abstractions. Currently, AReaL-lite +provides a lightweight codebase that enables fast prototyping for new RL training +workflows and algorithms on a relatively small scale. For large-scale experiments (1K+ +GPUs), we recommend using the battle-tested legacy AReaL to ensure stability. In the +future, we will continue developing AReaL-lite by expanding its APIs, migrating legacy +features, introducing new functionality, and validating the system through large-scale +experiments. ## Resources -+ [Documentation](https://inclusionai.github.io/AReaL/) -+ [Contributing](https://inclusionai.github.io/AReaL/contrib.html) +- [Documentation](https://inclusionai.github.io/AReaL/) +- [Contributing](https://inclusionai.github.io/AReaL/contrib.html) ### Quickstart -+ [Installation](https://inclusionai.github.io/AReaL/tutorial/installation.html) -+ [Example: Improving the math capability of Qwen3 with PPO](https://inclusionai.github.io/AReaL/tutorial/quickstart.html) +- [Installation](https://inclusionai.github.io/AReaL/tutorial/installation.html) +- [AReaL-lite Quickstart](https://inclusionai.github.io/AReaL/tutorial/quickstart.html) -### Benchmark and Reproduction +### Code Walkthrough -+ **Reproduce boba² Code Models** - - 🤗 **Model weights**: [8B-code](https://huggingface.co/inclusionAI/AReaL-boba-2-8B), [14B-code](https://huggingface.co/inclusionAI/AReaL-boba-2-14B), [8B-code-open](https://huggingface.co/inclusionAI/AReaL-boba-2-8B-subset), [14B-code-open](https://huggingface.co/inclusionAI/AReaL-boba-2-14B-subset) - - [Evaluation Guide](https://inclusionai.github.io/AReaL/tutorial/eval.html) - - [Training configs](https://github.com/inclusionAI/AReaL/tree/main/examples/configs/v0.3-qwen3-code) and [instructions](https://inclusionai.github.io/AReaL/references/reproduce.html) -+ [Scripts for Benchmark Training Throughput](https://github.com/inclusionAI/AReaL/tree/main/benchmark/verl_v0_3_0_post1_76084d3) +- [Running GRPO on GSM8K dataset with AReaL-lite](https://inclusionai.github.io/AReaL/arealite/gsm8k_grpo.html) -### Customization Guide +### Customization -- [Use your own dataset](https://inclusionai.github.io/AReaL/customization/dataset.html) -- [Modifying the reward function and rollout behavior (multi-turn agentic RL)](https://inclusionai.github.io/AReaL/customization/agent.html) -- [Modifying PPO to GRPO](https://inclusionai.github.io/AReaL/customization/algorithm.html#grouped-advantage-normalization) -- [Developing the decoupled PPO loss](https://inclusionai.github.io/AReaL/customization/algorithm.html#the-decoupled-ppo-loss) +- [Customize dataset with AReaL-lite](https://inclusionai.github.io/AReaL/customization/dataset.html) +- [Customize Agentic/RVLR rollout workflows with AReaL-lite](https://inclusionai.github.io/AReaL/customization/agent.html) +- [Customize algorithms with AReaL-lite](https://inclusionai.github.io/AReaL/customization/algorithm.html) -### System Code Walkthrough +### AReaL Legacy -+ [Trainer](https://inclusionai.github.io/AReaL/developer/trainer/model_worker.html) -+ [Model Backend and Algorithm Interface](https://inclusionai.github.io/AReaL/developer/trainer/algo_interface.html) -+ [Rollout Controller](https://inclusionai.github.io/AReaL/developer/rollout/gserver.html) -+ [Streaming generation and reward computation](https://inclusionai.github.io/AReaL/developer/rollout/rollout_worker.html) +For old AReaL documentation, check the legacy sections in our +[documentation](https://inclusionai.github.io/AReaL/). To reproduce AReaL boba & boba² +results, check our +[reproduction guide with legacy AReaL](https://inclusionai.github.io/AReaL/references/reproduce.html). ## Future Plan -AReaL is under active development. We plan to have minor releases weekly and major releases monthly. Community engagement and contributions are extremely welcome. We are also **hiring interns and full-time employees** with open positions in both the US and China. +AReaL is under active development. We plan to have minor releases weekly and major +releases monthly. Community engagement and contributions are extremely welcome. We are +also **hiring interns and full-time employees** with open positions in both the US and +China. For the research and development plan already in place, please see the following list: @@ -174,26 +195,46 @@ For the research and development plan already in place, please see the following - [x] Support for SGLang - [x] RL training with coding problems - [x] Asynchronous generation and RL training -- [ ] Optimizations for distributed training: expert parallel for MOE and zero-bubble pipelining -- [ ] RL for vision-language models (VLM) +- [ ] Optimizations for distributed training: expert parallel for MOE and zero-bubble + pipelining +- [x] RL for vision-language models (VLM) - [x] Multi-turn agentic RL -- [ ] Function calling and tool use +- [x] Function calling and tool use ### Algorithm Development - [x] RL training recipes for 1.5B and 7B models - [x] A complete RL training recipe for 32B models - [ ] Sample-efficient multi-task RL algorithms -- [ ] Agentic capabilities with end-to-end RL +- [x] Agentic capabilities with end-to-end RL - [ ] Stable RL training for larger MOE models ## Acknowledgement -We would like to note that major contributors are from the RL Lab at Ant Research and the Institute for Interdisciplinary Information Sciences, Tsinghua University. - -Our team has also received invaluable assistance from the Data Intelligence Lab at Ant Research for data support and from the Super Computing Technology (SCT) team at Ant Group, particularly in the realm of large-scale cluster operations and maintenance. - -We also appreciate all the pioneering works from the community, particularly the [ReaLHF](https://github.com/openpsi-project/ReaLHF) project from OpenPsi Inc. and other projects, including but not limited to [DeepScaleR](https://github.com/agentica-project/deepscaler), [Open-Reasoner-Zero](https://github.com/Open-Reasoner-Zero/Open-Reasoner-Zero/tree/main), [OpenRLHF](https://github.com/OpenRLHF/OpenRLHF), [VeRL](https://github.com/volcengine/verl), [SGLang](https://github.com/sgl-project/sglang), [QwQ](https://github.com/QwenLM/QwQ), [Light-R1](https://github.com/Qihoo360/Light-R1) and [DAPO](https://github.com/BytedTsinghua-SIA/DAPO). +We would like to note that major contributors are from the RL Lab at Ant Research and +the Institute for Interdisciplinary Information Sciences, Tsinghua University. + +Our team has also received invaluable assistance from the following groups (listed in +alphabetical order): + +• The Data Intelligence Lab at Ant Research for their data support • The +[Relaxed System Lab](https://github.com/Relaxed-System-Lab) from HKUST for seamless +cooperation on many system-related aspects • The +[SGLang team](https://github.com/sgl-project/sglang) for supporting customized features +for asynchronous RL training and their contributions during the development of +AReaL-lite • The Super Computing Technology (SCT) team at Ant Group, particularly for +their expertise in large-scale cluster operations and maintenance + +We also appreciate all the pioneering works from the community, particularly the +[ReaLHF](https://github.com/openpsi-project/ReaLHF) project from OpenPsi Inc. and other +projects, including but not limited to +[DeepScaleR](https://github.com/agentica-project/deepscaler), +[Open-Reasoner-Zero](https://github.com/Open-Reasoner-Zero/Open-Reasoner-Zero/tree/main), +[OpenRLHF](https://github.com/OpenRLHF/OpenRLHF), +[VeRL](https://github.com/volcengine/verl), +[SGLang](https://github.com/sgl-project/sglang), [QwQ](https://github.com/QwenLM/QwQ), +[Light-R1](https://github.com/Qihoo360/Light-R1) and +[DAPO](https://github.com/BytedTsinghua-SIA/DAPO). ## Citation @@ -210,12 +251,12 @@ We also appreciate all the pioneering works from the community, particularly the ```bibtex @misc{fu2025areal, - title={AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning}, + title={AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning}, author={Wei Fu and Jiaxuan Gao and Xujie Shen and Chen Zhu and Zhiyu Mei and Chuyi He and Shusheng Xu and Guo Wei and Jun Mei and Jiashu Wang and Tongkai Yang and Binhang Yuan and Yi Wu}, year={2025}, eprint={2505.24298}, archivePrefix={arXiv}, primaryClass={cs.LG}, - url={https://arxiv.org/abs/2505.24298}, + url={https://arxiv.org/abs/2505.24298}, } ``` diff --git a/areal/README.md b/areal/README.md new file mode 100644 index 0000000000..3b179ef7ea --- /dev/null +++ b/areal/README.md @@ -0,0 +1,561 @@ +# AReaL-lite Design Doc + +## TL;DR + +Follow our +[step-by-step code walk-through](https://inclusionai.github.io/AReaL/lite/gsm8k_grpo.html) +to immediately get started with AReaL-lite! + +## Motivation + +AReaL presents several challenges that make it difficult for AI researchers to adopt, +understand, and develop with effectively. The primary issue stems from its +*system-first* rather than *algorithm-first* architecture and API design. An +*algorithm-first* design aims to provide three key features: + +- **Light-weight & easy-to-write customization:** Users can implement their algorithms + and training workflows with minimal and concentrated code, often in just a few files + or even a single file. +- **Easy to scale up:** Experiments can be scaled up seamlessly without requiring + knowledge of underlying system or infrastructure details. +- **Adaptable and plugable:** Users is free to integrate the system with code or APIs + from other AI libraries, or plug APIs from the system into other frameworks. + +We believe that AReaL, as well as other existing RL systems, falls short in fully +delivering these features. For example, the RL training workflow in AReaL is built +around multiple *workers* executing *model function calls* in a *DAG* (Directed Acyclic +Graph). To customize a training workflow, researchers first need to understand these +system-level concepts. Then they are forced to find code to modify, which is scattered +around in the codebase. It is also nearly impossible to exploit packages like `datasets` +since it is not compatible with the workers. This gap is the core motivation behind +AReaL-lite: rebuilding AReaL with an algorithm-first architecture and APIs. + +Beyond architectural concerns, AReaL suffers from accumulated technical debt. The +codebase contains substantial legacy code inherited from previous projects that no +longer serves a purpose but significantly increases complexity for both users and +developers. Even experienced core developers sometimes struggle with debugging due to +this accumulated complexity. + +The landscape of RL workflow development tools has matured considerably, making it +possible to achieve comparable efficiency with significantly fewer lines of code. This +presents an ideal opportunity to redesign the API and distill the massive codebase into +something clean and maintainable. Rather than pursuing maximum efficiency, our goal is +to deliver 90% of AReaL's functionality while dramatically reducing code complexity and +user burden. This philosophy drives AReaL-lite — the lightweight version of AReaL. + +AReaL-lite serves as the first phase in AReaL's broader refactoring initiative. It +functions both as a standalone training library with intuitive interfaces and as the +foundation for AReaL's future core API definitions. The plan is to transform AReaL's +current worker-based architecture into an algorithm-first architecture similar to +AReaL-lite, where AReaL will **extend** AReaL-lite's APIs and implementations to support +additional backends for efficient large-scale training. + +## Design Principles + +To achieve *algorithm-first* and *light-weight* while maintaining the efficiency, our +design is guided by seven core principles: + +1. **Native asynchronous RL training support** — Built from the ground up for + disentangled generation and training +1. **System-less design** — Minimize exposure to system concepts like "PlacementGroup" +1. **PyTorch-centric approach** — Use raw PyTorch types without unnecessary abstractions +1. **Transparent algorithm orchestration** — Make the flow of operations clear and + understandable +1. **Developer-friendly navigation** — Enable easy access to implementation details via + Ctrl+click in IDEs +1. **Ecosystem compatibility** — Integrate smoothly with existing ML/RL tools +1. **Single-file customization** — Allow RL pipeline modifications within a single file + +## Architecture + +### Core Directory Structure + +``` +areal/ +├── api/ # Abstract interfaces and dataclasses +├── engine/ # Training and inference engines +├── launcher/ # Launcher for different backends +├── tests/ # Standalone test scripts +└── workflow/ # Custom RL rollout workflows +``` + +### Component Overview + +The AReaL-lite codebase is structured into four distinct layers: the API layer, backend +layer, customization layer, and entry point layer. As illustrated in the figure below, +workflow and algorithm customization logic resides in separate layers above the backend. +We prioritize keeping the entry point and customization layers clean and intuitive, +isolating them from the complex backend implementation. With AReaL-lite, users can +define their custom training workflows and algorithms entirely within a single entry +point file. + +![arealite-layers](../assets/areal_lite_layers.png) + +#### 1. API Layer (`api/`) + +The API layer establishes clean contracts between components through abstract interfaces +and data classes: + +- **`engine_api.py`**: Defines `TrainEngine` for SPMD-based distributed training + backends and `InferenceEngine` for streaming LLM inference +- **`workflow.py`**: Defines `RolloutWorkflow` for RL data collection within a unified + method interface +- **`cli_args.py`**: Contains configuration dataclasses for all system components + +The workflow object invokes `InferenceEngine` to complete data collection following +customized patterns, providing flexibility while maintaining consistency. + +#### 2. Backend Layer (`engine/`) + +The backend layer provides adapters for third-party libraries, ensuring they conform to +the APIs defined in `engine_api.py`. These components deliver core inference and +training capabilities: + +- **`fsdp_engine.py`**: FSDP-based training engine utilizing PyTorch FSDP2 +- **`sglang_remote.py`**: Client interface for generation with remote SGLang servers + +#### 3. Customization Layer (`engine/ppo/`, `workflow/`) + +This layer applies backend capabilities to implement specific reinforcement learning +pipelines: + +- **`engine/ppo/actor.py`**: PPO algorithm implementation that leverages a `TrainEngine` +- **`workflow/rlvr.py`**: RLVR workflow that utilizes an `InferenceEngine` to sample + multiple responses for each prompt + +#### 4. Entry Point Layer (`examples/lite/`) + +The entry point layer composes customization layer implementations to create complete RL +training pipelines. While we provide several reference examples, users have complete +freedom to adapt these to their specific use cases. + +Entry points can be launched using launchers from `areal/launcher/`, similar to other +distributed launch tools like `torchrun`: + +```bash +python3 -m areal.launcher.ray entrypoint.py --config my-config.yaml +``` + +## Usage Examples + +### Basic RL Training + +Users must provide a YAML configuration file, though they can override configuration +parameters for hyperparameter searches or other experimental needs: + +```bash +# Launch with Ray launcher: 4 nodes (4 GPUs each), 3 nodes for generation, 1 node for training +python3 -m areal.launcher.ray examples/lite/gsm8k_grpo.py \ + --config examples/lite/configs/gsm8k_grpo.yaml \ + experiment_name= \ + trial_name= \ + allocation_mode=sglang.d12p1t1+d4p1t1 \ + cluster.n_nodes=4 \ + cluster.n_gpus_per_node=4 + +# Launch with Slurm launcher: 16 nodes (8 GPUs each), 12 nodes for generation, 4 nodes for training +python3 -m areal.launcher.slurm examples/lite/gsm8k_grpo.py \ + --config examples/lite/configs/gsm8k_grpo.yaml \ + experiment_name= \ + trial_name= \ + allocation_mode=sglang.d96p1t1+d32p1t1 \ + cluster.n_nodes=16 \ + cluster.n_gpus_per_node=8 +``` + +## Customization Guide + +For detailed customization instructions, please refer to our documentation: + +- [Adding new agents](https://inclusionai.github.io/AReaL/customization/agent.html) +- [Adding new datasets](https://inclusionai.github.io/AReaL/customization/dataset.html) +- [Adding new algorithms](https://inclusionai.github.io/AReaL/customization/algorithm.html) + +## Implementation Details + +### Entry Point Design Philosophy + +We considered two primary design patterns for entry points, each with distinct +trade-offs: + +#### Single-Controller Pattern + +The most modular approach uses a single-controller pattern where only one process in the +cluster executes the main coordination logic: + +```python +def my_reward_fn(prompt, completion, prompt_ids, completion_ids, **kwargs): + return len(completion_ids) + +class MyRolloutWorkflow: + async def arun_episode(self, engine: InferenceEngine, + data: Dict[str, Any]) -> Dict[str, Tensor]: + message = [ + {"role": "system", "message": ...}, + {"role": "user", "message": ...}, + ] + + for _ in range(self.config.num_turns): + text = tokenizer.apply_chat_template(message, tools=self.env.list_tools()) + req = LLMRequest(text=text, ...) + resp = await engine.agenerate(req) + tool_name, tool_args = parse_tool(resp) + cur_time = await self.env.aexecute(tool_name, tool_args) + message += [{"role": "user", "message": f"The current time is {cur_time}"}] + + reward = my_reward_fn(None, None, None, req.input_ids, **data) + return output + +def main_grpo(): + config, _ = load_expr_config(args, GRPOConfig) + + # Create rollout workflow + workflow = MyRolloutWorkflow() + + # Single-controller mode initialization + scheduler = SlurmScheduler() + rollout = RolloutController(SGLangEngine(config.rollout), scheduler) + actor = TrainController(MegatronGRPOActor(config.actor), scheduler) + + # Training loop + dataloader = StatefulDataloader(dataset) + for _ in range(max_steps): + # Collect trajectories using rollout workflow + batch = rollout.rollout_batch(next(dataloader), workflow=workflow) + batch: DistributedBatch + + # Prepare training inputs + adv_batch = actor.compute_advantages_and_returns(batch) + batch['advantages'] = adv_batch['advantages'] + + # Execute PPO update + stats = actor.ppo_update(batch) + + # Update inference engine weights (non-blocking to prevent NCCL blocking) + future = rollout.update_weights(wcfg) + actor.upload_weights(wcfg) + future.result() +``` + +**Advantages:** + +- Provides maximum flexibility for device allocation, scheduling, and data arrangement + +**Disadvantages:** + +- Introduces multiple abstractions (`TrainController`, `Scheduler`, `DistributedBatch`) + that complicate the script + +#### SPMD Pattern + +Given that AI researchers are more familiar with the SPMD (Single Program, Multiple +Data) pattern used in standard model training, we also provide an entry point following +this approach. With N GPUs dedicated to training, N processes execute the following +code: + +```python +def main_grpo(): + config, _ = load_expr_config(args, GRPOConfig) + + # Create rollout workflow + workflow = MyRolloutWorkflow() + + # SPMD mode initialization + rollout = RemoteSGLangEngine(config.rollout) + actor = MegatronGRPOActor(config.actor) + + # Training loop + dataloader = StatefulDataloader(dataset) + for _ in range(max_steps): + # Data collection (only on data parallel head) + if is_dp_head: + batch = rollout.rollout_batch(next(dataloader), workflow=workflow) + batch_list = [batch] + else: + batch_list = [None] + + # Broadcast data to all processes + batch = dist.broadcast(batch_list, src=0, group=model_parallel_group)[0] + batch: TensorDict + + # Prepare training inputs + adv_batch = actor.compute_advantages_and_returns(batch) + batch['advantages'] = adv_batch['advantages'] + + # Execute PPO update + stats = actor.ppo_update(batch) + + # Update weights (coordinated across processes) + if rank == 0: + future = rollout.update_weights(wcfg) + actor.upload_weights(wcfg) + if rank == 0: + future.result() +``` + +The SPMD pattern uses only concepts familiar to AI researchers, though it requires some +control flow branching based on parallelism strategy. + +### Training Engine Architecture + +The training engine operates at two abstraction levels to balance flexibility with ease +of use. + +#### Basic Level: Backend Adapters + +The foundational level provides unified interfaces for RL algorithms, handling +computation, parameter management, and weight updates for inference engines. Each RL +training experiment must use one of the implemented backends: + +```python +class TrainEngine(abc.ABC): + + def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): + """Initialize distributed training environment and load models.""" + raise NotImplementedError() + + def destroy(self): + """Clean up engine resources and release GPU memory.""" + pass + + def upload_weights(self, meta: WeightUpdateMeta): + """Upload weights to inference engine (blocking operation).""" + raise NotImplementedError() + + def save(self, meta: SaveLoadMeta): + """Save model weights and optimizer states for checkpointing.""" + raise NotImplementedError() + + def load(self, meta: SaveLoadMeta): + """Load model weights and optimizer states from checkpoint.""" + raise NotImplementedError() + + def train_batch( + self, + input_: TensorDict, + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], + ) -> Dict[str, float]: + """Update model parameters using provided batch and loss function.""" + raise NotImplementedError() + + @torch.no_grad() + def forward( + self, + input_: TensorDict, + output_seqlens: List[int] | None = None, + post_hook: Callable[[torch.Tensor, TensorDict], Any] | None = None, + aggregate_fn: Callable[[List[Any]], Any] = torch.cat, + ) -> Any | None: + """Execute gradient-free forward pass for inference.""" + raise NotImplementedError() +``` + +#### Algorithm Level: Extended Engines + +Extended engines like PPO Actor provide algorithm-specific organization and +computational interfaces. They leverage backend core methods (such as `forward`) to +generate algorithm-required tensors and execute specialized model updates. The produced +objects (e.g., `FSDPPPOActor`) are also instances of `TrainEngine`s, but affiliated with +methods that specially designed for the algorithm (e.g., `ppo_update`). + +```python +class PPOActor: + + def __init__(self, config: PPOActorConfig, engine: TrainEngine): + self.config = config + self.engine = engine + + @torch.no_grad() + def compute_logp( + self, + data: TensorDict, + temperature: Optional[float] = None, + ) -> torch.Tensor | None: + + def calc_logprobs(logits, input_data): + labels = torch.roll(input_data["input_ids"], shifts=-1, dims=-1) + logprobs = gather_logprobs(logits, labels, temperature or 1.0) + return logprobs + + self.engine.eval() + return self.engine.forward( + input_=data, + post_hook=calc_logprobs, + aggregate_fn=lambda xs: torch.cat(xs, dim=-1), + ) + + def compute_advantages(self, data: TensorDict) -> None: + """Compute advantages for PPO training.""" + # Implementation details... + pass + + def ppo_update(self, data: TensorDict) -> List[Dict[str, float]]: + """Execute PPO policy update.""" + # Implementation details... + pass + +class FSDPPPOActor(FSDPEngine): + """FSDP-backed PPO Actor implementation.""" + + def __init__(self, config: PPOActorConfig): + super().__init__(config) + self.actor = PPOActor(config, self) + + @torch.no_grad() + def compute_logp(self, *args, **kwargs) -> torch.Tensor | None: + return self.actor.compute_logp(*args, **kwargs) + + @torch.no_grad() + def compute_advantages(self, *args, **kwargs) -> None: + self.actor.compute_advantages(*args, **kwargs) + + def ppo_update(self, *args, **kwargs) -> List[Dict[str, float]]: + return self.actor.ppo_update(*args, **kwargs) +``` + +### Inference Engine Design + +The inference engine's core functionality revolves around `generate` and +`update_weights` methods. These methods can interface with HTTP server APIs or invoke +local LLM engines: + +```python +class InferenceEngine(abc.ABC): + + def initialize(self, addr: str | None, ft_spec): + """Initialize distributed inference environment and load models.""" + raise NotImplementedError() + + def destroy(self): + """Clean up engine resources and release GPU memory.""" + pass + + async def agenerate(self, req: LLMRequest) -> LLMResponse: + """Generate response asynchronously for the given request.""" + raise NotImplementedError() + + def update_weights(self, meta: WeightUpdateMeta) -> Future: + """Update inference engine weights asynchronously.""" + raise NotImplementedError() +``` + +#### Workflow Integration + +User-defined rollout workflows utilize the `InferenceEngine` to generate trajectories. +The workflow's `arun_episode` method produces one or more trajectories from a single +prompt. The generation process is streaming rather than batched, with each dataset item +processed independently. Here's a simplified RLVR example: + +```python +class RLVRWorkflow(RolloutWorkflow): + async def arun_episode(self, engine: InferenceEngine, data): + input_ids = self.tokenizer.apply_chat_template( + data["messages"], + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + + n_samples = self.gconfig.n_samples + req = LLMRequest( + rid=uuid.uuid4().hex, + input_ids=input_ids, + gconfig=self.gconfig.new(n_samples=1), + ) + + # Generate multiple responses concurrently + resps = await asyncio.gather(*[engine.agenerate(req) for _ in range(n_samples)]) + + results = [] + for resp in resps: + reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) + + results.append(TensorDict(res, batch_size=[1])) + + return concat_padded_tensors(results) +``` + +#### Batch Processing and Asynchronous Operations + +While individual trajectory collection is straightforward, batching and asynchronous +execution require additional infrastructure. The `InferenceEngine` provides high-level +methods: `submit`, `wait`, `rollout_batch`, and `prepare_batch`. + +The `rollout_batch` method submits multiple `workflow.arun_episode` jobs to an +asynchronous thread pool using `submit`, then waits for completion using `wait`. The +`prepare_batch` method separates submission and waiting to enable asynchronous rollout: + +```python +def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None: + try: + self.input_queue.put_nowait((data, workflow)) + except Full: + raise RuntimeError("Input queue full. Consider increasing queue_size.") + +def wait( + self, + count: int, + timeout: float | None = None, + should_accept: Callable | None = None, +) -> TensorDict: + """Wait for specified number of results with optional filtering.""" + # Implementation details... + pass + +def rollout_batch( + self, data: List[Dict[str, Any]], workflow: "RolloutWorkflow" +) -> TensorDict: + """Submit batch requests and wait for all results.""" + for item in data: + self.submit(item, workflow) + return self.wait(count=len(data)) + +def prepare_batch( + self, + dataloader: StatefulDataLoader, + workflow: "RolloutWorkflow", +): + """Prepare batch for asynchronous processing.""" + # Implementation details... + pass +``` + +### RolloutWorkflow Interface + +The `RolloutWorkflow` class provides the `arun_episode` method with a standardized +signature for collecting agent trajectories: + +```python +class MyRolloutWorkflow: + def __init__(self, config: Any): + self.config = config + self.tool_executor = ToolExecutor() + self.tool_executor.register_tool(get_current_time) + + async def arun_episode(self, engine: InferenceEngine, + data: Dict[str, Any]) -> Dict[str, Tensor]: + req = LLMRequest(input_ids=data['input_ids'], ...) + + for _ in range(self.config.num_turns): + resp = await engine.agenerate(req) + res = await self.tool_executor.aexecute_tool(resp.completion) + req.input_ids += res + + reward = my_reward_fn(None, None, None, req.input_ids, **data) + return output +``` + +### Controller Architecture + +`RolloutController` and `TrainController` mirror the APIs of `InferenceEngine` and +`TrainEngine` respectively. Controllers handle engine deployment across the cluster and +manage data distribution, invoking engine methods through remote procedure calls (RPCs). +This architecture enables distributed operation while maintaining familiar interfaces +for users. diff --git a/areal/api/cli_args.py b/areal/api/cli_args.py new file mode 100644 index 0000000000..2df2444676 --- /dev/null +++ b/areal/api/cli_args.py @@ -0,0 +1,809 @@ +import argparse +import os +from dataclasses import asdict, dataclass, field +from pathlib import Path +from typing import Dict, List, Optional + +import uvloop + +uvloop.install() +from hydra import compose as hydra_compose +from hydra import initialize as hydra_init +from omegaconf import MISSING, OmegaConf + +from areal.utils.fs import get_user_tmp + + +@dataclass +class MicroBatchSpec: + """Specification for splitting micro-batches during training.""" + + n_mbs: Optional[int] = field( + default=1, + metadata={ + "help": "Number of micro-batches (or minimum number if max_tokens_per_mb is set). Used when max_tokens_per_mb is None or as minimum count", + }, + ) + max_tokens_per_mb: Optional[int] = field( + default=None, + metadata={ + "help": "Maximum tokens per micro-batch. When set, n_mbs becomes the minimum number of micro-batches", + }, + ) + + @classmethod + def new(cls, mb_spec: "MicroBatchSpec", **kwargs): + """Create new spec with updated fields while maintaining Omegaconf compatibility.""" + fields = dict( + n_mbs=mb_spec.n_mbs, + max_tokens_per_mb=mb_spec.max_tokens_per_mb, + ) + fields.update(kwargs) + return cls(**fields) + + +@dataclass +class GenerationHyperparameters: + """Controls text generation behavior for RL training.""" + + n_samples: int = field( + default=1, metadata={"help": "Number of sequences to generate per prompt."} + ) + max_new_tokens: int = field( + default=16384, metadata={"help": "Maximum number of tokens to generate."} + ) + min_new_tokens: int = field( + default=0, metadata={"help": "Minimum number of tokens to generate."} + ) + greedy: bool = field( + default=False, + metadata={"help": "Whether to use greedy decoding (max probability)."}, + ) + top_p: float = field( + default=1.0, + metadata={"help": "Nucleus sampling probability threshold (0.0, 1.0]."}, + ) + top_k: int = field( + default=int(1e8), + metadata={"help": "Number of highest probability tokens to consider."}, + ) + temperature: float = field( + default=1.0, + metadata={"help": "Sampling temperature. Higher values increase diversity."}, + ) + stop_token_ids: List[int] = field( + default_factory=list, + metadata={"help": "Stop generation when encoutering these token ids."}, + ) + + def new(self, **kwargs): + args = asdict(self) + args.update(kwargs) + return GenerationHyperparameters(**args) + + +# Train Engine Configs + + +@dataclass +class OptimizerConfig: + """Configuration for model optimization during training. + Note: + Set type to "empty" for models that won't be trained. + """ + + type: str = field( + default="adam", + metadata={"help": "Optimizer type", "choices": ["adam", "empty"]}, + ) + lr: float = field(default=2e-5, metadata={"help": "Learning rate"}) + weight_decay: float = field(default=0.05, metadata={"help": "Weight decay"}) + beta1: float = field(default=0.9, metadata={"help": "Adam beta1 parameter"}) + beta2: float = field(default=0.95, metadata={"help": "Adam beta2 parameter"}) + eps: float = field(default=1e-5, metadata={"help": "Adam epsilon parameter"}) + min_lr_ratio: float = field( + default=0.0, + metadata={ + "help": "Minimum learning rate ratio after annealing", + }, + ) + lr_scheduler_type: str = field( + default="constant", + metadata={ + "help": "Learning rate scheduler type", + "choices": ["linear", "cosine", "constant"], + }, + ) + warmup_steps_proportion: float = field( + default=0.001, + metadata={ + "help": "Proportion of training steps for warmup", + }, + ) + offload: bool = field( + default=False, metadata={"help": "Enable optimizer state offloading"} + ) + initial_loss_scale: float = field( + default=2**32, metadata={"help": "Initial loss scaling factor"} + ) + min_loss_scale: float = field( + default=1.0, metadata={"help": "Minimum loss scaling factor"} + ) + loss_scale_window: float = field( + default=5, metadata={"help": "Window size for loss scaling adjustment"} + ) + hysteresis: int = field( + default=2, metadata={"help": "Hysteresis (scaling factor) for loss scaling"} + ) + gradient_clipping: float = field( + default=1.0, metadata={"help": "Gradient clipping threshold"} + ) + + +@dataclass +class FSDPWrapPolicy: + transformer_layer_cls_to_wrap: Optional[List[str]] = field( + default=None, + metadata={"help": "A list of transformer layer names for FSDP to wrap."}, + ) + + +@dataclass +class FSDPEngineConfig: + wrap_policy: Optional[FSDPWrapPolicy] = field( + default=None, + metadata={"help": "FSDP wrap policy, specifying model layers to wrap."}, + ) + offload_params: bool = field( + default=False, + metadata={"help": "Whether to offload FSDP parameters to CPU."}, + ) + + +@dataclass +class DeepSpeedAutoTPEngineConfig: + autotp_size: Optional[int] = field( + default=1, + metadata={"help": "DeepSpeed AutoTP size"}, + ) + + +@dataclass +class TrainEngineConfig: + experiment_name: str = MISSING + trial_name: str = MISSING + path: str = field(default="", metadata={"help": "Path to HuggingFace checkpoint"}) + attn_impl: str = field( + default="flash_attention_2", + metadata={ + "help": "Attention implementation for huggingface transformers model.", + "choices": ["flash_attention_2"], + }, + ) + init_from_scratch: bool = field( + default=False, metadata={"help": "Initialize model weights randomly"} + ) + init_critic_from_actor: bool = field( + default=False, + metadata={"help": "Initialize critic/reward model from LM checkpoint"}, + ) + # Runtime microbatch limit + mb_spec: MicroBatchSpec = field(default_factory=MicroBatchSpec) + pad_to_maximum: bool = field( + default=False, + metadata={ + "help": ( + "Whether to pad each microbatch to the length upper bound specified by mb_spec. " + "Can reduce memory fragmentation but slow down training." + ) + }, + ) + + # Training Backend Configuration + disable_dropout: bool = field(default=False) + gradient_checkpointing: bool = field( + default=True, metadata={"help": "Enable gradient checkpointing"} + ) + dtype: str = field(default="bfloat16", metadata={"help": "Parameter dtype."}) + grad_reduce_dtype: str = field( + default="float32", metadata={"help": "Gradient reduce dtype."} + ) + optimizer: Optional[OptimizerConfig] = field( + default=None, metadata={"help": "Optimizer configuration"} + ) + backend: str = "" + fsdp: FSDPEngineConfig = field(default_factory=FSDPEngineConfig) + ds_auto_tp: DeepSpeedAutoTPEngineConfig = field( + default_factory=DeepSpeedAutoTPEngineConfig + ) + + +@dataclass +class PPOActorConfig(TrainEngineConfig): + # Core PPO/GRPO Parameters + group_size: int = field( + default=1, metadata={"help": "Number of sequences in each group"} + ) + group_adv_norm: bool = field( + default=False, + metadata={ + "help": "Normalize advantages within each prompt group rather than globally" + }, + ) + ppo_n_minibatches: int = field( + default=4, metadata={"help": "Number of minibatches for each PPO update"} + ) + eps_clip: float = field( + default=0.2, metadata={"help": "Clipping factor for policy ratio"} + ) + c_clip: Optional[float] = field( + default=None, + metadata={ + "help": "Dual clipping factor for policy ratio, must > 1.0. None disables dual clipping." + }, + ) + temperature: float = field( + default=1.0, metadata={"help": "Temperature during generation."} + ) + # Reward + group_reward_norm: bool = field( + default=False, + metadata={ + "help": "Normalize final reward of each sequence (GRPO-style) to reduce length bias" + }, + ) + reward_scaling: float = field( + default=1.0, metadata={"help": "Reward scaling factor"} + ) + reward_bias: float = field(default=0.0, metadata={"help": "Reward bias"}) + reward_clip: float = field( + default=20.0, metadata={"help": "Maximum absolute value for reward clipping"} + ) + mask_no_eos_with_zero: bool = field( + default=False, + metadata={ + "help": "Mask truncated generations (no EOS token) and exclude from training" + }, + ) + + # Advantage Estimation + discount: float = field( + default=1.0, metadata={"help": "Discount factor for future rewards"} + ) + gae_lambda: float = field( + default=1.0, metadata={"help": "Lambda parameter for GAE"} + ) + adv_norm: bool = field( + default=True, metadata={"help": "Enable advantage normalization globally"} + ) + + # KL Control + kl_ctl: float = field(default=0.1, metadata={"help": "KL divergence coefficient"}) + + # Asynchronous RL + recompute_logprob: bool = field( + default=False, + metadata={"help": "Recompute logp and replace the logp returned by inference."}, + ) + use_decoupled_loss: bool = field( + default=False, + metadata={ + "help": "Use the decoupled loss. Implicitly enable recompute_logprob." + }, + ) + behav_imp_weight_cap: Optional[float] = field( + default=None, + metadata={ + "help": "We filter out the tokens where behav_imp_weight exceeds behav_imp_weight_cap when computing the loss, must be > 1.0, use_decoupled_loss must be true" + }, + ) + + +@dataclass +class SGLangConfig: + """Configuration for SGLang runtime. Refer to: + https://github.com/sgl-project/sglang for detailed documentation. + """ + + model_path: str = "" + random_seed: int = 1 + skip_tokenizer_init: bool = False + disable_cuda_graph: bool = False + disable_radix_cache: bool = False + disable_cuda_graph_padding: bool = False + enable_nccl_nvls: bool = False + disable_outlines_disk_cache: bool = False + disable_custom_all_reduce: bool = False + disable_overlap_schedule: bool = False + enable_mixed_chunk: bool = False + enable_dp_attention: bool = False + enable_ep_moe: bool = False + enable_torch_compile: bool = False + torch_compile_max_bs: int = 32 + cuda_graph_max_bs: Optional[int] = None + cuda_graph_bs: Optional[List[int]] = None + torchao_config: str = "" + enable_nan_detection: bool = False + enable_p2p_check: bool = False + triton_attention_reduce_in_fp32: bool = False + triton_attention_num_kv_splits: int = 8 + num_continuous_decode_steps: int = 1 + enable_memory_saver: bool = False + allow_auto_truncate: bool = False + # NOTE: to avoid the illegal memory access error + attention_backend: Optional[str] = "flashinfer" + sampling_backend: Optional[str] = None + context_length: Optional[int] = 32768 + mem_fraction_static: Optional[float] = 0.9 + max_running_requests: Optional[int] = None + # NOTE: chunked_prefill_size is by default 8192 on GPUs with 80GB mem in SGLang, + # but we disable it to avoid precision issues + chunked_prefill_size: Optional[int] = -1 + max_prefill_tokens: int = 32768 + schedule_policy: str = "lpm" + schedule_conservativeness: float = 1.0 + cpu_offload_gb: int = 0 + dtype: str = "bfloat16" + kv_cache_dtype: str = "auto" + # logging + log_level: str = "warning" + log_level_http: Optional[str] = "warning" + log_requests: bool = False + log_requests_level: int = 0 + show_time_cost: bool = False + enable_metrics: bool = True # Exports Prometheus-like metrics + # The interval (in decoding iterations) to log throughput + # and update prometheus metrics + decode_log_interval: int = 1 + + # Use staticmethod to make OmegaConf happy. + @staticmethod + def build_cmd( + sglang_config: "SGLangConfig", + tp_size, + base_gpu_id, + host, + port, + dist_init_addr: Optional[str] = None, + ): + args = SGLangConfig.build_args( + sglang_config=sglang_config, + tp_size=tp_size, + base_gpu_id=base_gpu_id, + host=host, + port=port, + dist_init_addr=dist_init_addr, + ) + + # convert to flags + flags = [] + for k, v in args.items(): + if v is None or v is False or v == "": + continue + if v is True: + flags.append(f"--{k.replace('_','-')}") + elif isinstance(v, list): + flags.append(f"--{k.replace('_','-')} {' '.join(map(str, v))}") + else: + flags.append(f"--{k.replace('_','-')} {v}") + return f"python3 -m sglang.launch_server {' '.join(flags)}" + + @staticmethod + def build_args( + sglang_config: "SGLangConfig", + tp_size, + base_gpu_id, + host, + port, + dist_init_addr: Optional[str] = None, + ): + from realhf.base import pkg_version + from realhf.experiments.common.utils import asdict as conf_as_dict + + args: Dict = conf_as_dict(sglang_config) + args = dict( + host=host, + port=port, + # Model and tokenizer + tokenizer_path=sglang_config.model_path, + tokenizer_mode="auto", + load_format="auto", + trust_remote_code=True, + device="cuda", + is_embedding=False, + # Other runtime options + tp_size=tp_size, + # Because we have set CUDA_VISIBLE_DEVICES to a single GPU in each process + base_gpu_id=base_gpu_id, + nnodes=1, + node_rank=0, + # initialization addresses and ports + dist_init_addr=dist_init_addr, + **args, + ) + if not pkg_version.is_version_greater_or_equal("sglang", "0.4.9.post2"): + raise RuntimeError("Needs sglang>=0.4.9.post2 to run the code.") + return args + + +@dataclass +class InferenceEngineConfig: + experiment_name: str = MISSING + trial_name: str = MISSING + max_concurrent_rollouts: None | int = field( + default=None, + metadata={ + "help": "Maximum number of concurrent rollouts to the inference engine. Defaults to consumer_batch_size." + }, + ) + queue_size: None | int = field( + default=None, + metadata={"help": "Input/Output queue size for async rollout."}, + ) + consumer_batch_size: int = field( + default=1, + metadata={"help": "Batch size for consuming rollouts from the queue."}, + ) + max_head_offpolicyness: int = field( + default=0, + metadata={ + "help": "Maximum off-policyness for the head. " + "If the current version is more than this many versions behind, " + "the request will not be accepted.", + }, + ) + enable_rollout_tracing: bool = field(default=False) + schedule_policy: str = field( + default="round_robin", + metadata={"help": "Request scheduling policy", "choices": ["round_robin"]}, + ) + setup_timeout: float = field(default=120.0) + request_timeout: float = field( + default=3600, metadata={"help": "Timeout for HTTP requests."} + ) + request_retries: int = field( + default=3, metadata={"help": "Number of retries for failed requests."} + ) + + +@dataclass +class _Timer: + experiment_name: str = MISSING + trial_name: str = MISSING + fileroot: str = MISSING + freq_epochs: Optional[int] = field( + default=None, + metadata={ + "help": "Trigger frequency in epochs. None disables epoch-based saving." + }, + ) + freq_steps: Optional[int] = field( + default=None, + metadata={ + "help": "Trigger frequency in steps. None disables step-based saving." + }, + ) + freq_secs: Optional[int] = field( + default=None, + metadata={ + "help": "Trigger frequency in seconds. None disables time-based saving." + }, + ) + + +@dataclass +class EvaluatorConfig(_Timer): + pass + + +@dataclass +class SaverConfig(_Timer): + pass + + +@dataclass +class WandBConfig: + mode: str = "disabled" + entity: Optional[str] = None + project: Optional[str] = None + name: Optional[str] = None + job_type: Optional[str] = None + group: Optional[str] = None + notes: Optional[str] = None + tags: Optional[List[str]] = None + config: Optional[Dict] = None + + +@dataclass +class SwanlabConfig: + project: Optional[str] = None + name: Optional[str] = None + config: Optional[Dict] = None + logdir: Optional[str] = None + mode: Optional[str] = "local" + api_key: Optional[str] = os.getenv("SWANLAB_API_KEY", None) + + +@dataclass +class TensorBoardConfig: + path: Optional[str] = None + + +@dataclass +class StatsLoggerConfig: + experiment_name: str = MISSING + trial_name: str = MISSING + fileroot: str = MISSING + wandb: WandBConfig = field( + default_factory=WandBConfig, + metadata={"help": "Weights & Biases configuration."}, + ) + swanlab: SwanlabConfig = field( + default_factory=SwanlabConfig, + metadata={"help": "SwanLab configuration."}, + ) + tensorboard: TensorBoardConfig = field( + default_factory=TensorBoardConfig, + metadata={"help": "TensorBoard configuration. Only 'path' field required."}, + ) + + +@dataclass +class NameResolveConfig: + type: str = field( + default="nfs", + metadata={ + "help": "Type of the distributed KV store for name resolving.", + "choices": ["nfs", "etcd3", "ray"], + }, + ) + nfs_record_root: str = field( + default="/tmp/areal/name_resolve", + metadata={ + "help": "Record root for NFS name resolving. Should be available in all nodes." + }, + ) + etcd3_addr: str = field( + default="localhost:2379", metadata={"help": "Address of the ETCD3 server."} + ) + ray_actor_name: str = field( + default="ray_kv_store", + metadata={"help": "Name of the distributed Ray KV store."}, + ) + + +@dataclass +class ClusterSpecConfig: + name_resolve: NameResolveConfig = field( + default_factory=NameResolveConfig, + metadata={"help": "Name resolving configuration."}, + ) + cluster_name: str = field( + default="local", + metadata={"help": "Name of the cluster. Used to set specific environs."}, + ) + fileroot: str = field( + default=get_user_tmp(), + metadata={ + "help": "Root for logs and checkpoints. Should be available to all nodes." + }, + ) + n_nodes: int = field( + default=32, + metadata={ + "help": "The size of the cluster. Used to decide slurm hostname suffix." + }, + ) + n_gpus_per_node: int = field( + default=8, + metadata={"help": "GPUs per node (physically)."}, + ) + + +@dataclass +class DatasetConfig: + path: str = field( + default=MISSING, + metadata={ + "help": "Path to the dataset. Can be a local path or a HuggingFace dataset name." + }, + ) + type: Optional[str] = field( + default=None, + metadata={"help": "Type of training method.e.g., 'sft', 'rl', etc."}, + ) + batch_size: int = field( + default=1, metadata={"help": "Batch size of the dataloader"} + ) + shuffle: bool = field( + default=True, metadata={"help": "Whether to shuffle the dataset"} + ) + pin_memory: bool = field( + default=False, + metadata={ + "help": "Pin memory for faster data loading (set True for GPU training)" + }, + ) + num_workers: int = field( + default=0, metadata={"help": "Number of worker processes for data loading"} + ) + drop_last: bool = field(default=True) + + +@dataclass +class SlurmLauncherConfig: + """Configuration for launching the SGLang server with Slurm.""" + + srun_additional_args: str = field( + default="--overlap --mpi=pmi2 -K --chdir $PWD", + metadata={"help": "Additional arguments to pass to the srun command."}, + ) + container_type: str = field( + default="apptainer", + metadata={ + "help": "Type of containers used in slurm", + "choices": ["apptainer", "none"], + }, + ) + mount: str = field( + default="/storage:/storage", metadata={"help": "Mount path for slurm."} + ) + trainer_image: str = field( + default="", metadata={"help": "slurm image for trainers."} + ) + inference_server_image: str = field( + default="", metadata={"help": "slurm image for LLM inference."} + ) + + +@dataclass +class LauncherConfig: + """Configuration for launching the SGLang server.""" + + inference_server_cpus_per_gpu: int = field( + default=4, + metadata={"help": "Number of CPUs allocated per GPU for inference server. "}, + ) + inference_server_mem_per_gpu: int = field( + default=32 * 1024, + metadata={"help": "Memory allocated per GPU for inference server in MB. "}, + ) + trainer_cpus_per_gpu: int = field( + default=4, + metadata={"help": "Number of CPUs allocated per GPU for training. "}, + ) + trainer_mem_per_gpu: int = field( + default=32 * 1024, + metadata={"help": "Memory allocated per GPU for training in MB. "}, + ) + inference_server_env_vars: str = field( + default="", + metadata={ + "help": "Environment variables for inference server, seperated by commas. " + "Example: 'ENV1=val1,ENV2=val2'. " + }, + ) + trainer_env_vars: str = field( + default="", + metadata={ + "help": "Environment variables for training, seperated by commas. " + "Example: 'ENV1=val1,ENV2=val2'. " + }, + ) + slurm: SlurmLauncherConfig = field( + default_factory=SlurmLauncherConfig, + metadata={"help": "Slurm launcher configuration."}, + ) + + +@dataclass +class BaseExperimentConfig: + # NOTE: we need this unified config class because different experiments + # have different config structures, e.g., GRPO has two engine configs, + # but SFT only has a single one. We use subclasses to represent these structures. + experiment_name: str = field( + default=MISSING, + metadata={"help": "Name of the experiment (no '_' or '/'). Required."}, + ) + trial_name: str = field( + default=MISSING, + metadata={"help": "Name of the trial (no '-' or '/'). Required."}, + ) + cluster: ClusterSpecConfig = field( + default_factory=ClusterSpecConfig, + metadata={"help": "Cluster specification. Mainly used by slurm."}, + ) + allocation_mode: str = field( + default="", + metadata={ + "help": "GPU parallel strategy allocation mode. " + "Options: manual/heuristic or pattern-based." + }, + ) + seed: int = field(default=1, metadata={"help": "Random seed for reproducibility."}) + total_train_epochs: int = field( + default=1, metadata={"help": "Total number of epochs to train the model."} + ) + total_train_steps: Optional[int] = field( + default=None, + metadata={ + "help": "Terminate training after this number of steps. " + "For benchmarking purposes only. None indicates normal training." + }, + ) + total_train_n_seqs: Optional[int] = field( + default=None, + metadata={ + "help": "Terminate training after consuming this number of samples. " + "For benchmarking purposes only. None indicates normal training." + }, + ) + tokenizer_path: str = field(default="") + + train_dataset: DatasetConfig = field(default_factory=DatasetConfig) + valid_dataset: Optional[DatasetConfig] = field(default=None) + + saver: SaverConfig = field(default_factory=SaverConfig) + checkpointer: SaverConfig = field(default_factory=SaverConfig) + evaluator: EvaluatorConfig = field(default_factory=EvaluatorConfig) + stats_logger: StatsLoggerConfig = field(default_factory=StatsLoggerConfig) + + server_only: bool = False + sglang: SGLangConfig = field(default_factory=SGLangConfig) + launcher: LauncherConfig = field(default_factory=LauncherConfig) + + +@dataclass +class SFTConfig(BaseExperimentConfig): + model: TrainEngineConfig = field(default_factory=TrainEngineConfig) + + +@dataclass +class GRPOConfig(BaseExperimentConfig): + async_training: bool = field(default=True) + gconfig: GenerationHyperparameters = field( + default_factory=GenerationHyperparameters + ) + rollout: InferenceEngineConfig = field(default_factory=InferenceEngineConfig) + actor: PPOActorConfig = field(default_factory=PPOActorConfig) + ref: PPOActorConfig = field(default_factory=PPOActorConfig) + + +def parse_cli_args(argv: List[str]): + parser = argparse.ArgumentParser() + parser.add_argument( + "--config", help="The path of the main configuration file", required=True + ) + args, overrides = parser.parse_known_args(argv) + # Initialize hydra config + config_file = Path(args.config).absolute() + assert config_file.exists() + # hydra only recognize relative paths + relpath = Path(os.path.relpath(str(config_file), Path(__file__).parent.absolute())) + hydra_init(config_path=str(relpath.parent), job_name="app", version_base=None) + cfg = hydra_compose( + config_name=str(relpath.name).split(".yaml")[0], + overrides=overrides, + ) + return cfg, config_file + + +def to_structured_cfg(cfg, config_cls): + # Merge with the default configuration. + # The yaml and commandline can omit some default values defined in python dataclasses. + default_cfg = OmegaConf.structured(config_cls) + cfg = OmegaConf.merge(default_cfg, cfg) + return cfg + + +def load_expr_config(argv: List[str], config_cls): + cfg, config_file = parse_cli_args(argv) + cfg = to_structured_cfg(cfg, config_cls=config_cls) + cfg = OmegaConf.to_object(cfg) + assert isinstance(cfg, BaseExperimentConfig) + # Setup environment + from realhf.base import constants, name_resolve + + constants.set_experiment_trial_names(cfg.experiment_name, cfg.trial_name) + name_resolve.reconfigure(cfg.cluster.name_resolve) + return cfg, str(config_file) diff --git a/areal/api/engine_api.py b/areal/api/engine_api.py new file mode 100644 index 0000000000..8728df9103 --- /dev/null +++ b/areal/api/engine_api.py @@ -0,0 +1,188 @@ +import abc +from concurrent.futures import Future +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional + +import torch +import torch.distributed as dist +from tensordict import TensorDict +from torchdata.stateful_dataloader import StatefulDataLoader + +from areal.api.io_struct import ( + FinetuneSpec, + LLMRequest, + LLMResponse, + ParamSpec, + SaveLoadMeta, + WeightUpdateMeta, +) + +if TYPE_CHECKING: + from areal.api.workflow_api import RolloutWorkflow + + +@dataclass +class Scheduling: + cpu: int + gpu: int + mem: int + nodelist: Optional[str] = None + exclude: Optional[str] = None + partition: Optional[str] = None + container_image: Optional[str] = None + env_vars: Dict[str, str] = field(default_factory=dict) + # time utils from "https://slurm.schedmd.com/sbatch.html" + time_limit: Optional[str] = None # see "--time" option for format + begin: Optional[str] = None # see "--begin" option for format + deadline: Optional[str] = None # see "--deadline" option for format + + +class TrainEngine(abc.ABC): + + def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): + """Initialize environments for distributed training and load models.""" + raise NotImplementedError() + + @property + def parallelism_group(self) -> dist.ProcessGroup: + """The global communication group of this engine.""" + raise NotImplementedError() + + def get_scheduling_config(self) -> Scheduling: + """Get the scheduling configuration for the engine, e.g., image, cpu/gpu/memory size.""" + raise NotImplementedError() + + def destroy(self): + """Destroy the engine and release GPU memory.""" + + def train(self, mode: bool = True): + """Set the engine to the train mode.""" + raise NotImplementedError() + + def eval(self): + """Set the engine to the eval mode.""" + return self.train(False) + + def upload_weights(self, meta: WeightUpdateMeta): + """Upload weights to the inference engine (in a blocking manner).""" + raise NotImplementedError() + + def get_param_specs(self) -> List[ParamSpec]: + """Get the parameter specifications for the model.""" + raise NotImplementedError() + + def set_version(self, version: int): + """Set the current weight version in the train engine.""" + raise NotImplementedError() + + def get_version(self) -> int: + """Get the current weight version in the train engine.""" + raise NotImplementedError() + + def save(self, meta: SaveLoadMeta): + """Save model weights (and optimizer states) for later use.""" + raise NotImplementedError() + + def load(self, meta: SaveLoadMeta): + """Load model weights and optimizer states from a file.""" + raise NotImplementedError() + + def step_lr_scheduler(self): + """Step learning rate scheduler. + + Since PPO uses minibatch updates, this method just need to be called once after a few train_batch calls. + It is separated from train_batch to allow for more flexible scheduling. + """ + raise NotImplementedError() + + def train_batch( + self, + input_: TensorDict, + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], + ) -> Dict[str, float]: + """Update the model with a batch of data and a loss function.""" + raise NotImplementedError() + + @torch.no_grad() + def eval_batch( + self, + input_: TensorDict, + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], + ) -> torch.Tensor | None: + """Evaluate the model using the forward pass and loss function.""" + raise NotImplementedError() + + @torch.no_grad() + def forward( + self, + input_: TensorDict, + output_seqlens: List[int] | None = None, + post_hook: Callable[[torch.Tensor, TensorDict], Any] | None = None, + aggregate_fn: Callable[[List[Any]], Any] = torch.cat, + ) -> Any | None: + """Run the forward pass or inference on the model. Note that it is gradient-free.""" + raise NotImplementedError() + + +class InferenceEngine(abc.ABC): + + def initialize(self, addr: str | None, ft_spec): + """Initialize environments for distributed inference and load models.""" + raise NotImplementedError() + + def destroy(self): + """Destroy the engine and release GPU memory.""" + + async def agenerate(self, req: LLMRequest) -> LLMResponse: + """Asynchronously generate a response for the given request.""" + raise NotImplementedError() + + def update_weights(self, meta: WeightUpdateMeta) -> Future: + """Update weights in the inference engine in a non-blocking manner.""" + raise NotImplementedError() + + def set_version(self, version: int) -> None: + """Set the current weight version in the inference engine.""" + raise NotImplementedError() + + def get_version(self) -> int: + """Get the current weight version in the inference engine.""" + raise NotImplementedError() + + def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None: + """Asynchronously submit a request to the inference engine. Exits immediately.""" + raise NotImplementedError() + + def wait( + self, + count: int, + timeout: float | None = None, + should_accept: Callable | None = None, + ) -> TensorDict: + """Wait for a specified number of requests to complete, with a timeout.""" + raise NotImplementedError() + + def rollout_batch( + self, data: List[Dict[str, Any]], workflow: "RolloutWorkflow" + ) -> TensorDict: + """Submit a batch of requests to the inference engine and wait for the results.""" + raise NotImplementedError() + + def prepare_batch( + self, + dataloader: StatefulDataLoader, + workflow: "RolloutWorkflow", + should_accept: Callable | None = None, + ): + """Asynchronously submit and wait until a full batch is ready.""" + raise NotImplementedError() + + def pause(self): + """Pause request submission for async rollout. Used during evaluation to prevent data over generation.""" + raise NotImplementedError() + + def resume(self): + """Resume request submission for async rollout.""" + raise NotImplementedError() diff --git a/areal/api/env_api.py b/areal/api/env_api.py new file mode 100644 index 0000000000..490d69d592 --- /dev/null +++ b/areal/api/env_api.py @@ -0,0 +1,28 @@ +import abc +from typing import Any, Dict, List + + +class Environment(abc.ABC): + + async def ainitialize(self): + """ + Performs the initialization logic for the environment asynchronously. + + For stateful environments, this is where resources are created and + prepared (e.g., launching a browser). + """ + + def list_tools(self) -> List[Dict[str, Any]]: + """Lists all available tools in the environment.""" + return [] + + async def aexecute(self, tool_name: str, tool_args: Dict[str, Any]) -> Any: + """Executes a tool in the environment asynchronously.""" + raise NotImplementedError() + + async def aclose(self): + """ + Destroys the environment asynchronously, releasing all held resources. + + This method is critical for stateful environments (e.g., a browser session). + """ diff --git a/areal/api/io_struct.py b/areal/api/io_struct.py new file mode 100644 index 0000000000..09204877d2 --- /dev/null +++ b/areal/api/io_struct.py @@ -0,0 +1,247 @@ +# Copyright 2025 Ant Group Inc. +# Licensed under the Apache License, Version 2.0 +import enum +import itertools +import os +import re +import uuid +from dataclasses import dataclass, field +from typing import TYPE_CHECKING, Any, Dict, List, Literal, Optional, Tuple + +from PIL.Image import Image as ImageObject +from transformers import AutoProcessor, PreTrainedTokenizerFast + +from areal.api.cli_args import GenerationHyperparameters, SaverConfig +from areal.utils.network import find_free_ports, gethostip + +if TYPE_CHECKING: + from areal.api.engine_api import TrainEngine + + +@dataclass +class LLMRequest: + rid: str = field(default_factory=lambda: str(uuid.uuid4())) + input_ids: List[int] = field(default_factory=list) + gconfig: GenerationHyperparameters = field( + default_factory=GenerationHyperparameters + ) + metadata: Dict[str, Any] = field(default_factory=dict) + model_id: Optional[str] = None + + +@dataclass +class LLMResponse: + # outputs + input_tokens: List[int] = field(default_factory=list) + output_tokens: List[int] = field(default_factory=list) + output_logprobs: List[float] = field(default_factory=list) + output_versions: List[int] = field(default_factory=list) + stop_reason: Literal["length", "stop", "interrupt"] = "stop" + + # statistics + latency: float = float("inf") + ttft: float = float("inf") # Time to first token + itl: List[float] = field(default_factory=list) # List of inter-token latencies + + @property + def input_len(self) -> int: + return len(self.input_tokens) + + @property + def output_len(self) -> int: + return len(self.output_tokens) + + +@dataclass +class VLMRequest(LLMRequest): + image_data: Optional[List[ImageObject | str]] = field(default_factory=list) + + +@dataclass +class VLMResponse(LLMResponse): + input_images: List[ImageObject | str] = field(default_factory=list) + + +@dataclass +class FinetuneSpec: + total_train_epochs: int + dataset_size: int + train_batch_size: int + + @property + def total_train_steps(self): + # assuming drop_last + return self.total_train_epochs * (self.dataset_size // self.train_batch_size) + + @property + def steps_per_epoch(self): + return self.dataset_size // self.train_batch_size + + +class AllocationType(enum.Enum): + COLOCATE = 0 + DECOUPLED_vLLM = 1 + DECOUPLED_SGLANG = 2 + + +@dataclass +class AllocationMode: + type_: AllocationType + parallel_strat: Dict[str, Dict[str, int]] + + @property + def gen_tp_size(self) -> int: + return self.parallel_strat["gen"]["t"] + + @property + def gen_pp_size(self) -> int: + return self.parallel_strat["gen"]["p"] + + @property + def gen_dp_size(self) -> int: + return self.parallel_strat["gen"]["d"] + + @property + def gen_world_size(self) -> int: + return self.gen_dp_size * self.gen_pp_size * self.gen_tp_size + + @property + def train_tp_size(self) -> int: + return self.parallel_strat["*"]["t"] + + @property + def train_pp_size(self) -> int: + return self.parallel_strat["*"]["p"] + + @property + def train_dp_size(self) -> int: + return self.parallel_strat["*"]["d"] + + @property + def train_world_size(self) -> int: + return self.train_dp_size * self.train_pp_size * self.train_tp_size + + @classmethod + def from_str(cls, allocation_mode: str): + if "vllm" in allocation_mode: + alloc_decoupled = AllocationMode.extract_decoupled_alloc(allocation_mode) + return cls(AllocationType.DECOUPLED_vLLM, alloc_decoupled) + elif "sglang" in allocation_mode: + alloc_decoupled = AllocationMode.extract_decoupled_alloc(allocation_mode) + return cls(AllocationType.DECOUPLED_SGLANG, alloc_decoupled) + else: + return cls( + AllocationType.COLOCATE, + AllocationMode.extract_parallelism_strategy(allocation_mode), + ) + raise NotImplementedError(f"Failed to parse allocation: {allocation_mode}") + + @staticmethod + def extract_parallelism_strategy(allocation_mode: str) -> Dict: + for x, y, z in itertools.permutations(["d", "t", "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, + } + } + raise ValueError( + f"Unknown how to resolve parallelism strategy: {allocation_mode}" + ) + + @staticmethod + def extract_decoupled_alloc(allocation_mode: str) -> Dict: + pattern = re.compile( + r"(?:(?:vllm|sglang)\.(.+?)\+(.+))|(?:(.+?)\+(?:vllm|sglang)\.(.+))" + ) + m = pattern.match(allocation_mode) + if not m: + raise ValueError( + f"Unknown how to resolve decoupled allocation: {allocation_mode}" + ) + 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_parallelism_strategy(gen_alloc) + other_alloc = AllocationMode.extract_parallelism_strategy(other_alloc) + other_alloc.update({"gen": gen_alloc["*"]}) + return other_alloc + + +@dataclass +class ParamSpec: + name: str + shape: Tuple + dtype: str + + +@dataclass +class WeightUpdateMeta: + type: Literal["disk", "nccl"] + path: str | None = None + alloc_mode: AllocationMode | None = None + + nccl_master_address: str = "127.0.0.1" + nccl_master_port: int = 29500 + nccl_param_specs: List[ParamSpec] = field(default_factory=list) + nccl_group_name: str = "update_weight_group" + + @classmethod + def from_disk( + cls, + saver_config: SaverConfig, + ) -> "WeightUpdateMeta": + from areal.utils.saver import Saver + + path = os.path.join( + Saver.get_save_checkpoint_root(saver_config), + "weight_update", + ) + return cls( + type="disk", + path=path, + ) + + @classmethod + def from_fsdp_nccl( + cls, + allocation_mode: AllocationMode, + fsdp_engine: "TrainEngine", + nccl_group_name: str = "update_weight_group", + ): + return cls( + type="nccl", + alloc_mode=allocation_mode, + nccl_master_address=gethostip(), + nccl_master_port=find_free_ports(1)[0], + nccl_param_specs=fsdp_engine.get_param_specs(), + nccl_group_name=nccl_group_name, + ) + + +@dataclass +class SaveLoadMeta: + path: str + weight_format: str + with_optim: bool + tokenizer: PreTrainedTokenizerFast | None + processor: AutoProcessor | None + base_model_path: str | None + naive_distributed: bool = False + + +@dataclass +class RolloutStat: + submitted: int = 0 + accepted: int = 0 + running: int = 0 diff --git a/areal/api/reward_api.py b/areal/api/reward_api.py new file mode 100644 index 0000000000..24b412ab1a --- /dev/null +++ b/areal/api/reward_api.py @@ -0,0 +1,23 @@ +from typing import List + + +def reward_fn( + prompt: str, + completions: str, + prompt_ids: List[int], + completion_ids: List[int], + **kwargs, +): + """This function is a placeholder for the reward function that will be used in the RLVR pipeline. + + In general, there's no restriction on the signature and implementation of this function in customized rollout workflows. + It would be convinent to follow this signature and directly use it in our predefined rollout workflows. + + :param prompt: The string representing the task to be completed. + :param completions: The string representing the trajectory generated by the model. + :param prompt_ids: The token IDs of the prompt. + :param completion_ids: The token IDs of the trajectory generated by the model. + :param kwargs: Other attributes of the data in the dataset, such as solutions, input_outputs, etc. + Any other attributes in the dataset will be passed as keyword arguments to this function. + :rtype: float + """ diff --git a/areal/api/workflow_api.py b/areal/api/workflow_api.py new file mode 100644 index 0000000000..d81ed478f7 --- /dev/null +++ b/areal/api/workflow_api.py @@ -0,0 +1,261 @@ +import asyncio +import itertools +import queue +import threading +import time +import traceback +from typing import TYPE_CHECKING, Any, Callable, Dict, List + +import torch.distributed as dist +import uvloop +from tensordict import TensorDict +from torchdata.stateful_dataloader import StatefulDataLoader + +from areal.api.cli_args import InferenceEngineConfig +from areal.api.engine_api import InferenceEngine +from areal.api.io_struct import RolloutStat +from areal.utils.data import concat_padded_tensors +from realhf.base import logging + +if TYPE_CHECKING: + from areal.api.engine_api import InferenceEngine + +logger = logging.getLogger("areal.workflow_api") + + +ROLLOUT_POLL_WAIT_TIME = 0.05 + + +class RolloutWorkflow: + + async def arun_episode( + self, engine: "InferenceEngine", data: Dict[str, Any] + ) -> TensorDict: + """Run a single episode of the workflow. + + See concrete example implementations under the `areal/workflow` directory. + """ + raise NotImplementedError() + + +class WorkflowExecutor: + + def __init__( + self, + config: InferenceEngineConfig, + inference_engine: "InferenceEngine", + ): + config.max_concurrent_rollouts = ( + config.max_concurrent_rollouts or config.consumer_batch_size + ) + self.config = config + self.exiting = threading.Event() + self.paused = threading.Event() + self.lock = threading.Lock() + + self.inference_engine = inference_engine + + qsize = config.queue_size or config.max_concurrent_rollouts * 16 + self.input_queue = queue.Queue(maxsize=qsize) + self.output_queue = queue.Queue(maxsize=qsize) + self.result_cache: List[TensorDict] = [] + + self.rollout_stat = RolloutStat() + + def initialize(self): + self.rollout_tasks: Dict[str, asyncio.Task] = {} + self.rollout_thread = threading.Thread(target=self._rollout_thread) + self.rollout_thread.start() + + def destroy(self): + self.exiting.set() + self.rollout_thread.join() + + def get_capacity(self): + if dist.is_initialized(): + world_size = dist.get_world_size() + else: + world_size = 1 + + with self.lock: + max_concurrent_rollouts = max( + 1, self.config.max_concurrent_rollouts // world_size + ) + capacity = max_concurrent_rollouts - len(self.rollout_tasks) + # Staleness control + version = self.inference_engine.get_version() + ofp = self.config.max_head_offpolicyness + sample_cnt = self.rollout_stat.accepted + self.rollout_stat.running + consumer_bs = max(1, self.config.consumer_batch_size // world_size) + capacity = min(capacity, (ofp + version + 1) * consumer_bs - sample_cnt) + return capacity + + def _rollout_thread(self): + """Thread that runs the rollout loop.""" + try: + uvloop.run(self._rollout_thread_async()) + except Exception: + traceback.print_exc() + + async def _rollout_thread_async(self): + rollout_tasks = self.rollout_tasks + rid = 0 + try: + while not self.exiting.is_set(): + # Check capacity + capacity = self.get_capacity() + # Create new rollout task + self.lock.acquire() + while ( + capacity > 0 + and not self.paused.is_set() + and self.input_queue.qsize() > 0 + ): + data, workflow = self.input_queue.get_nowait() + logger.debug(f"Get data from puller: {data}") + task = asyncio.create_task( + workflow.arun_episode(self.inference_engine, data), + name=str(rid), + ) + rollout_tasks[str(rid)] = task + self.rollout_stat.submitted += 1 + self.rollout_stat.running += 1 + if self.config.enable_rollout_tracing: + logger.info( + f"Submit rollout rid {rid}. " + f"Submit: {self.rollout_stat.submitted}, " + f"running: {self.rollout_stat.running}, " + f"accepted: {self.rollout_stat.accepted}." + ) + capacity -= 1 + rid += 1 + tasks = list(rollout_tasks.values()) + self.lock.release() + + # Wait for rollout completion + done = [] + if tasks: + done, _ = await asyncio.wait( + tasks, + timeout=ROLLOUT_POLL_WAIT_TIME, + return_when=asyncio.FIRST_COMPLETED, + ) + # Collect done results + for task in done: + traj = await task + traj: TensorDict + task_rid = task.get_name() + with self.lock: + rollout_tasks.pop(task_rid) + self.rollout_stat.accepted += 1 + + self.rollout_stat.running -= 1 + if self.config.enable_rollout_tracing: + logger.info( + f"Finish rollout {task_rid}. " + f"Submit: {self.rollout_stat.submitted}, " + f"running: {self.rollout_stat.running}, " + f"accepted: {self.rollout_stat.accepted}." + ) + try: + self.output_queue.put_nowait(traj) + except queue.Full: + raise RuntimeError( + "Output queue full. Please increase queue_size." + ) + + await asyncio.sleep(1) + except Exception: + traceback.print_exc() + finally: + # Cancel remaining tasks + with self.lock: + for task in rollout_tasks.values(): + if not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None: + try: + self.input_queue.put_nowait((data, workflow)) + except queue.Full: + raise RuntimeError("Input queue full. Please increase queue_size.") + + def wait( + self, + count: int, + timeout: float | None = None, + should_accept: Callable | None = None, + ) -> TensorDict: + tik = time.perf_counter() + accepted = len(self.result_cache) + timeout = timeout or float(7 * 24 * 3600) + while ( + accepted < count + and not self.exiting.is_set() + and time.perf_counter() - tik < timeout + ): + try: + result = self.output_queue.get(timeout=ROLLOUT_POLL_WAIT_TIME) + if should_accept is None or should_accept(result): + self.result_cache.append(result) + accepted += 1 + else: + with self.lock: + self.rollout_stat.accepted -= 1 + except queue.Empty: + pass + if self.exiting.is_set(): + raise RuntimeError("Rollout engine is exiting, cannot wait for results.") + if accepted < count: + raise TimeoutError( + f"Timed out waiting for {count} rollouts, " f"only received {accepted}." + ) + results, self.result_cache = ( + self.result_cache[:count], + self.result_cache[count:], + ) + return concat_padded_tensors(results) + + def rollout_batch( + self, data: List[Dict[str, Any]], workflow: "RolloutWorkflow" + ) -> TensorDict: + """Submit a batch of requests to the inference engine and wait for the results.""" + for item in data: + self.submit(item, workflow) + return self.wait(count=len(data)) + + def prepare_batch( + self, + dataloader: StatefulDataLoader, + workflow: "RolloutWorkflow", + should_accept: Callable | None = None, + ): + if not hasattr(self, "data_generator"): + self.data_generator = itertools.cycle(dataloader) + assert dataloader.batch_size is not None + while True: + # Submit at least two batches to allow maximum overlap + if ( + self.get_capacity() + dataloader.batch_size > 0 + and self.input_queue.qsize() + dataloader.batch_size + < self.input_queue.maxsize + ): + data = next(self.data_generator) + for item in data: + self.submit(item, workflow=workflow) + try: + return self.wait( + dataloader.batch_size, timeout=1, should_accept=should_accept + ) + except TimeoutError: + pass + + def pause(self): + self.paused.set() + + def resume(self): + self.paused.clear() diff --git a/areal/dataset/__init__.py b/areal/dataset/__init__.py new file mode 100644 index 0000000000..c70de81c90 --- /dev/null +++ b/areal/dataset/__init__.py @@ -0,0 +1,43 @@ +from typing import Optional + +import transformers + +VALID_DATASETS = ["gsm8k", "clevr_count_70k"] + + +def get_custom_dataset( + path: str, + rank: int, + world_size: int, + type: str = "sft", + split: Optional[str] = None, + tokenizer: Optional[transformers.PreTrainedTokenizerFast] = None, + processor: Optional[transformers.AutoProcessor] = None, + **kwargs, +): + + if "gsm8k" in path and type == "sft": + from areal.dataset.gsm8k import get_gsm8k_sft_dataset + + return get_gsm8k_sft_dataset(path, split, tokenizer, rank, world_size, **kwargs) + elif "gsm8k" in path and type == "rl": + from areal.dataset.gsm8k import get_gsm8k_rl_dataset + + return get_gsm8k_rl_dataset(path, split, rank, world_size, **kwargs) + elif "clevr_count_70k" in path and type == "sft": + from areal.dataset.clevr_count_70k import get_clevr_count_70k_sft_dataset + + return get_clevr_count_70k_sft_dataset( + path, split, processor, rank, world_size, **kwargs + ) + elif "clevr_count_70k" in path and type == "rl": + from areal.dataset.clevr_count_70k import get_clevr_count_70k_rl_dataset + + return get_clevr_count_70k_rl_dataset( + path, split, processor, rank, world_size, **kwargs + ) + else: + raise ValueError( + f"Dataset {path} with split {split} and training type {type} is not supported. " + f"Supported datasets are: {VALID_DATASETS}. " + ) diff --git a/areal/dataset/clevr_count_70k.py b/areal/dataset/clevr_count_70k.py new file mode 100644 index 0000000000..c268f91a78 --- /dev/null +++ b/areal/dataset/clevr_count_70k.py @@ -0,0 +1,125 @@ +import math +from io import BytesIO +from typing import Any, Dict, Optional, Union + +from datasets import load_dataset +from datasets.distributed import split_dataset_by_node +from PIL.Image import Image as ImageObject + + +def convert_image( + image: Union[Dict[str, Any], ImageObject, str], + max_pixels: Optional[int], +) -> ImageObject: + if max_pixels is not None and (image.width * image.height) > max_pixels: + resize_factor = math.sqrt(max_pixels / (image.width * image.height)) + width, height = int(image.width * resize_factor), int( + image.height * resize_factor + ) + image = image.resize((width, height)) + + if image.mode != "RGB": + image = image.convert("RGB") + with BytesIO() as output: + image.save(output, format="JPEG") + return output.getvalue() + + +def get_clevr_count_70k_sft_dataset(path, split, processor, rank, world_size): + """ + "clevr_count_70k": { + "image_key": "images", + "question_key": "problem", + "answer_key": "answer" + }, + """ + dataset = load_dataset(path=path, split=split) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + + tokenizer = processor.tokenizer + + def process_example(example, idx): + # Add query_id column + images = example["images"] + if "qwen" in processor.image_processor.image_processor_type.lower(): + image_token = "<|vision_start|><|image_pad|><|vision_end|>" + else: + image_token = processor.image_token if processor is not None else "" + example["problem"] = ( + example["problem"].replace("", image_token).replace("different", "") + ) + processed_images = [] + for image in images: + processed_images.append(convert_image(image, 336 * 336)) + example["images"] = processed_images + example["seq"] = example["problem"] + example["answer"] + tokenizer.eos_token + + return example + + dataset = dataset.map( + lambda example, idx: process_example(example, idx), + with_indices=True, + ) + + def _process(example): + text = example["seq"] + processed_input = processor( + text=[text], + images=example["images"], + padding=False, + return_tensors="pt", + return_length=True, + return_attention_mask=False, + ) + + example["input_ids"] = processed_input["input_ids"].squeeze(0) + example["pixel_values"] = processed_input["pixel_values"] + example["image_grid_thw"] = processed_input["image_grid_thw"] + answer_token = tokenizer.encode(example["answer"]) + loss_mask = [0] * (len(example["input_ids"]) - len(answer_token)) + [1] * len( + answer_token + ) + example["loss_mask"] = loss_mask + return example + + dataset = dataset.map( + lambda x: _process(x), remove_columns=["images", "seq", "problem", "answer"] + ) + return dataset + + +def get_clevr_count_70k_rl_dataset(path, split, processor, rank, world_size): + dataset = load_dataset(path=path, split=split) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + + def process(sample): + processed_images = [ + convert_image(image, 336 * 336) for image in sample["images"] + ] + if "qwen" in processor.image_processor.image_processor_type.lower(): + image_token = "<|vision_start|><|image_pad|><|vision_end|>" + else: + image_token = processor.image_token if processor is not None else "" + system_prompt = { + "role": "system", + "content": ( + "Solve the following question: count the number of items in the image and provide the final answer in [ ] format, ensuring that only the number is inside the brackets without any additional text or explanations. " + ), + } + + messages = [ + { + "role": "user", + "content": sample["problem"] + .replace("", image_token) + .replace("different", ""), + } + ] + messages.insert(0, system_prompt) + messages = processor.tokenizer.apply_chat_template( + messages, add_generation_prompt=True, tokenize=False + ) + return {"messages": messages, "images": processed_images} + + dataset = dataset.map(process).remove_columns(["problem"]) + return dataset diff --git a/areal/dataset/gsm8k.py b/areal/dataset/gsm8k.py new file mode 100644 index 0000000000..7c7744722f --- /dev/null +++ b/areal/dataset/gsm8k.py @@ -0,0 +1,30 @@ +from datasets import load_dataset +from datasets.distributed import split_dataset_by_node + + +def get_gsm8k_sft_dataset(path, split, tokenizer, rank, world_size): + dataset = load_dataset(path=path, name="main", split=split) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + + def process(sample): + seq_token = tokenizer.encode( + sample["question"] + sample["answer"] + tokenizer.eos_token + ) + prompt_token = tokenizer.encode(sample["question"]) + loss_mask = [0] * len(prompt_token) + [1] * (len(seq_token) - len(prompt_token)) + return {"input_ids": seq_token, "loss_mask": loss_mask} + + dataset = dataset.map(process).remove_columns(["question", "answer"]) + return dataset + + +def get_gsm8k_rl_dataset(path, split, rank, world_size): + dataset = load_dataset(path=path, name="main", split=split) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + + def process(sample): + messages = [{"role": "user", "content": sample["question"]}] + return {"messages": messages} + + dataset = dataset.map(process).remove_columns(["question"]) + return dataset diff --git a/areal/engine/__init__.py b/areal/engine/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/areal/engine/base_hf_engine.py b/areal/engine/base_hf_engine.py new file mode 100644 index 0000000000..748d58940c --- /dev/null +++ b/areal/engine/base_hf_engine.py @@ -0,0 +1,416 @@ +import gc +import os +import time +from typing import Any, Callable, Dict, List + +import torch +import torch.distributed as dist +from tensordict import TensorDict +from transformers import ( + AutoConfig, + AutoModelForCausalLM, + AutoModelForImageTextToText, + AutoProcessor, + PretrainedConfig, + PreTrainedTokenizerFast, + get_constant_schedule_with_warmup, + get_linear_schedule_with_warmup, +) + +from areal.api.cli_args import TrainEngineConfig +from areal.api.engine_api import FinetuneSpec, TrainEngine +from areal.utils.data import ( + MicroBatchList, + amend_position_ids, + pack_tensor_dict, + pad_and_stack_tensors_along_first_dim, + pad_mb_list, + reorder_list, + split_padded_tensor_dict_into_mb_list, + unpack_sequence, + unsqueeze_mb_list, +) +from areal.utils.fsdp import get_cosine_schedule_with_warmup +from areal.utils.model import VALID_VISION_MODELS, disable_dropout_in_model +from realhf.api.core.data_api import load_hf_processor_and_tokenizer, load_hf_tokenizer +from realhf.base import constants, logging + +logger = logging.getLogger("Base HF Engine") + + +class BaseHFEngine(TrainEngine): + def __init__(self, config: TrainEngineConfig): + self.config = config + self.optimizer_config = config.optimizer + + self.model: torch.nn.Module + self.optimizer: torch.optim.Optimizer + self.tokenizer: PreTrainedTokenizerFast + self.processor: AutoProcessor | None = None + # huggingface model config + self.model_config: PretrainedConfig + self._version: int = 0 + + # initialization + self.initialized = False + self.own_global_group = False + self._parallelism_group: dist.ProcessGroup + self.weight_update_group_initialized = False + + self.model_config = AutoConfig.from_pretrained( + pretrained_model_name_or_path=self.config.path, + trust_remote_code=True, + ) + self.is_vision_model = self.model_config.model_type in VALID_VISION_MODELS + + self.world_size = int(os.environ["WORLD_SIZE"]) + + def set_version(self, version: int): + self._version = version + + def get_version(self) -> int: + return self._version + + def train(self, mode: bool = True): + assert self.model is not None + self.model.train(mode=mode) + return self + + @property + def parallelism_group(self) -> dist.ProcessGroup: + assert self.initialized + return self._parallelism_group + + def create_process_group(self): + # Required by NCCL weight update group for SGLang + os.environ["NCCL_CUMEM_ENABLE"] = "0" + os.environ["NCCL_NVLS_ENABLE"] = "0" + if not dist.is_initialized(): + # TODO: Handle the condition when WORLD_SIZE and RANK is not set in launcher + # NOTE: device_id **SHOULD NOT** be passed into init_process_group, + # otherwise initializing the NCCL weight update group will be wrong! + dist.init_process_group( + backend="nccl", + timeout=constants.NCCL_DEFAULT_TIMEOUT, + ) + self.own_global_group = True + self._parallelism_group = dist.new_group() + + def create_device_model(self): + torch.cuda.set_device(int(os.environ["LOCAL_RANK"])) + self.device = torch.device(int(os.environ["LOCAL_RANK"])) + + dtype = getattr(torch, self.config.dtype) + + if self.is_vision_model: + if dtype == torch.float16: + raise ValueError( + "Vision models do not support float16 dtype. Please use bfloat16." + ) + if self.config.init_from_scratch: + raise ValueError( + "Vision models do not support initialization from scratch. Please use a pretrained model." + ) + self.processor, self.tokenizer = load_hf_processor_and_tokenizer( + self.config.path + ) + + tik = time.perf_counter() + with torch.device("cuda"): + model = AutoModelForImageTextToText.from_pretrained( + pretrained_model_name_or_path=self.config.path, + trust_remote_code=True, + torch_dtype=dtype, + attn_implementation=self.config.attn_impl, + ) + if self.config.disable_dropout: + disable_dropout_in_model(model) + else: + self.tokenizer = load_hf_tokenizer(self.config.path) + tik = time.perf_counter() + with torch.device("cuda"): + if self.config.init_from_scratch: + # initialize scratch model from config + # NOTE: VLM cannot directly load state dict using this + # random initialized model, so otherwise we call + # from_pretrained rather than loading weights into this random model. + model = AutoModelForCausalLM.from_config( + self.model_config, + torch_dtype=dtype, + attn_implementation=self.config.attn_impl, + ) + else: + model = AutoModelForCausalLM.from_pretrained( + pretrained_model_name_or_path=self.config.path, + trust_remote_code=True, + torch_dtype=dtype, + attn_implementation=self.config.attn_impl, + ) + if self.config.disable_dropout: + disable_dropout_in_model(model) + + if self.config.gradient_checkpointing: + model.gradient_checkpointing_enable( + gradient_checkpointing_kwargs={"use_reentrant": False} + ) + logger.info(f"Model creation and loading time: {time.perf_counter() - tik}") + self.model = model + + def create_optimizer(self, ft_spec: FinetuneSpec): + if self.optimizer_config is None: + return + assert self.model is not None + # Set up optimizer + tik = time.perf_counter() + assert ( + self.optimizer_config.type == "adam" + ), "Only AdamW optimizer is supported in this engine." + lr = self.optimizer_config.lr + weight_decay = self.optimizer_config.weight_decay + beta1 = self.optimizer_config.beta1 + beta2 = self.optimizer_config.beta2 + eps = self.optimizer_config.eps + + self.optimizer = torch.optim.AdamW( + self.model.parameters(), + lr=lr, + weight_decay=weight_decay, + betas=(beta1, beta2), + eps=eps, + ) + total_train_steps = ft_spec.total_train_steps + num_warmup_steps = int( + self.optimizer_config.warmup_steps_proportion * total_train_steps + ) + + if self.optimizer_config.lr_scheduler_type == "cosine": + self.lr_scheduler = get_cosine_schedule_with_warmup( + self.optimizer, + num_warmup_steps, + total_train_steps, + min_lr_ratio=self.optimizer_config.min_lr_ratio, + ) + elif self.optimizer_config.lr_scheduler_type == "linear": + self.lr_scheduler = get_linear_schedule_with_warmup( + self.optimizer, + num_warmup_steps, + total_train_steps, + ) + elif self.optimizer_config.lr_scheduler_type == "constant": + self.lr_scheduler = get_constant_schedule_with_warmup( + self.optimizer, + num_warmup_steps, + ) + else: + raise ValueError( + f"Unknown lr scheduler type {self.optimizer_config.lr_scheduler_type}" + ) + logger.info(f"Create optimizer time: {time.perf_counter() - tik}") + + def destroy(self): + """Destroy the engine and release GPU memory.""" + del self.optimizer + del self.model + gc.collect() + torch.cuda.empty_cache() + gc.collect() + dist.destroy_process_group(self.parallelism_group) + if self.own_global_group: + dist.destroy_process_group() + self.initialized = False + + def save_optimizer_state(self, path: str): + # Save FSDP sharded state dict on each rank + assert self.optimizer is not None + assert dist.is_initialized() + rank = dist.get_rank() + shard_path = os.path.join( + path, f"optim_world_size_{self.world_size}_rank_{rank}.pt" + ) + state_dict = self.optimizer.state_dict() + torch.save(state_dict, shard_path) + dist.barrier(device_ids=[self.device.index]) + + def load_optimizer_state(self, path: str): + # Load FSDP sharded state dict + assert self.optimizer is not None + assert dist.is_initialized() + rank = dist.get_rank() + shard_path = os.path.join( + path, f"optim_world_size_{self.world_size}_rank_{rank}.pt" + ) + optimizer_state_dict = torch.load(shard_path, weights_only=False) + self.optimizer.load_state_dict(optimizer_state_dict) + dist.barrier(device_ids=[self.device.index]) + + def step_lr_scheduler(self): + assert self.lr_scheduler is not None + self.lr_scheduler.step() + + def prepare_mb_list(self, input_: TensorDict) -> MicroBatchList: + assert "attention_mask" in input_ and "input_ids" in input_ + if self.is_vision_model: + assert ( + "pixel_values" in input_ and "image_grid_thw" in input_ + ), "For vision-language models, pixel_values and image_grid_thw must be present in input_" + + if isinstance(input_, dict): + input_ = TensorDict(input_, batch_size=[input_["input_ids"].shape[0]]) + input_ = amend_position_ids(input_) + + mb_list = split_padded_tensor_dict_into_mb_list(input_, self.config.mb_spec) + mb_list.mbs = [pack_tensor_dict(mb) for mb in mb_list.mbs] + mb_list = pad_mb_list( + mb_list, + pad_value=0.0, + pad_to_maximum=self.config.pad_to_maximum, + ) + logger.info( + f"Microbatch #tokens (rank {dist.get_rank()}): {mb_list.group_lens}, " + f"padded to: {mb_list.padded_to_lengths}, padding lengths: {mb_list.padding_lengths}" + ) + # NOTE: We unsqueeze here because huggingface transformer models requires + # packed input to be of shape [1, total_seqlen]. + mb_list = unsqueeze_mb_list(mb_list) + + # FIXME: the resulting max_seqlen is a tensor rather than an integer + for mb in mb_list.mbs: + mb["max_seqlen"] = int(mb["max_seqlen"]) + mb["use_cache"] = False + for mb in mb_list.padded_mbs: + mb["max_seqlen"] = int(mb["max_seqlen"]) + mb["use_cache"] = False + + return mb_list + + def train_batch( + self, + input_: TensorDict, + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], + ) -> Dict[str, float]: + """Train on a batch using gradient accumulation.""" + input_ = input_.to(self.device) + assert self.optimizer is not None + assert self.optimizer_config is not None + assert self.lr_scheduler is not None + + self.optimizer.zero_grad() + mb_list = self.prepare_mb_list(input_) + + total_loss_weight = torch.tensor( + sum([loss_weight_fn(mb) for mb in mb_list.mbs]), dtype=torch.float32 + ) + assert total_loss_weight != 0 + dist.all_reduce(total_loss_weight) + + # Process microbatches with gradient accumulation + for i, (pad_length, padded_mb_input, mb_input) in enumerate( + zip(mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs) + ): + + outputs = self.model(**padded_mb_input) + + logits = outputs.logits.squeeze(0) + logits = logits[:-pad_length] if pad_length > 0 else logits + loss = loss_fn(logits, mb_input) + + loss_scale = loss_weight_fn(mb_input) / total_loss_weight + + # Scale loss for accumulation + # Revert gradient averaging across dp ranks + # FIXME: should be DP size + loss_scale *= self.world_size + + loss *= loss_scale + loss.backward() + + grad_norm = torch.nn.utils.clip_grad_norm_( + self.model.parameters(), + self.optimizer_config.gradient_clipping, + norm_type=2.0, + error_if_nonfinite=False, + foreach=None, + ) + if not torch.isfinite(grad_norm): + self.optimizer.zero_grad() + update_successful = False + else: + self.optimizer.step() + update_successful = True + + current_lr = self.lr_scheduler.get_last_lr()[0] + return dict( + update_successful=float(update_successful), + grad_norm=float(grad_norm) if grad_norm is not None else float("nan"), + lr=current_lr, + ) + + @torch.no_grad() + def eval_batch( + self, + input_: TensorDict, + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], + ) -> torch.Tensor | None: + """Evaluate on a batch.""" + input_ = input_.to(self.device) + mb_list = self.prepare_mb_list(input_) + total_loss_weight = torch.tensor( + sum([loss_weight_fn(mb) for mb in mb_list.mbs]), dtype=torch.float32 + ) + assert total_loss_weight != 0 + + total_loss = 0.0 + total_weight = 0.0 + + for pad_length, padded_mb_input, mb_input in zip( + mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs + ): + outputs = self.model(**padded_mb_input) + logits = outputs.logits.squeeze(0) + logits = logits[:-pad_length] if pad_length > 0 else logits + loss = loss_fn(logits, mb_input) + + # Simple weight calculation (could be improved) + loss_scale = loss_weight_fn(mb_input) / total_loss_weight + total_loss += loss.item() * loss_scale + total_weight += loss_scale + + return torch.tensor(total_loss / total_weight) + + @torch.no_grad() + def forward( + self, + input_: TensorDict, + output_seqlens: List[int] | None = None, + post_hook: Callable[[torch.Tensor, TensorDict], Any] | None = None, + aggregate_fn: Callable[[List[Any]], Any] = torch.cat, + ) -> Any | None: + """Forward pass with optional post-processing.""" + input_ = input_.to(self.device) + cu_seqlens = pack_tensor_dict(input_)["cu_seqlens"] + mb_list = self.prepare_mb_list(input_) + + if output_seqlens is None: + output_seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).cpu().numpy().tolist() + + results = [] + for pad_length, padded_mb_input, mb_input in zip( + mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs + ): + outputs = self.model(**padded_mb_input) + logits = outputs.logits.squeeze(0) + logits = logits[:-pad_length] if pad_length > 0 else logits + + if post_hook: + result = post_hook(logits, mb_input) + results.append(result) + else: + results.append(logits) + + res = aggregate_fn(results) + output_seqlens = [output_seqlens[i] for i in mb_list.forward_indices] + unpacked = unpack_sequence(res, lens=output_seqlens, dim=0) + reordered = reorder_list(unpacked, mb_list.backward_indices) + return pad_and_stack_tensors_along_first_dim(reordered) diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py new file mode 100644 index 0000000000..a2d76b7f37 --- /dev/null +++ b/areal/engine/fsdp_engine.py @@ -0,0 +1,272 @@ +import os +import time +from datetime import datetime +from typing import Callable, Dict, List, Optional + +import torch +import torch.distributed as dist +from tensordict import TensorDict +from torch.distributed._tensor import DTensor +from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + get_model_state_dict, +) +from transformers import AutoProcessor, PreTrainedTokenizerFast + +from areal.api.cli_args import TrainEngineConfig +from areal.api.engine_api import FinetuneSpec +from areal.api.io_struct import ParamSpec, SaveLoadMeta, WeightUpdateMeta +from areal.engine.base_hf_engine import BaseHFEngine +from areal.utils.distributed import init_custom_process_group +from areal.utils.fsdp import ( + CPUOffloadPolicy, + MixedPrecisionPolicy, + apply_fsdp2, + create_fsdp_device_mesh, + fsdp2_clip_grad_norm_, + fsdp2_load_full_state_dict, +) +from areal.utils.save_load import get_state_dict_from_repo_id_or_path +from realhf.base import logging, name_resolve, names, pkg_version + +logger = logging.getLogger("FSDPEngine") + + +class FSDPEngine(BaseHFEngine): + def __init__(self, config: TrainEngineConfig): + super().__init__(config) + # FSDP options + self.mixed_precision_policy = None + self.device_mesh = None + self.cpu_offload = None + + def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): + # Initialize distributed enviroments and load model. + assert addr is None, "FSDPEngine does not support remote initialization." + assert pkg_version.is_version_greater_or_equal( + "torch", "2.4.0" + ), f"areal only supports FSDP2, which requires torch>=2.4.0" + + self.create_process_group() + self.create_device_model() + + # Wrap with FSDP2 + # Simple auto wrap policy + self.mixed_precision_policy = MixedPrecisionPolicy( + param_dtype=getattr(torch, self.config.dtype), + reduce_dtype=getattr(torch, self.config.grad_reduce_dtype), + cast_forward_inputs=True, + ) + self.device_mesh = create_fsdp_device_mesh(self.world_size, self.world_size) + # sharding_strategy = ShardingStrategy.FULL_SHARD + self.cpu_offload = ( + CPUOffloadPolicy() if self.config.fsdp.offload_params else None + ) + fsdp_kwargs = { + "mesh": self.device_mesh, + "mp_policy": self.mixed_precision_policy, + "offload_policy": self.cpu_offload, + "reshard_after_forward": True, + } + tik = time.perf_counter() + apply_fsdp2(self.model, fsdp_kwargs, self.config.fsdp.wrap_policy) + logger.info(f"Applying FSDP2 time: {time.perf_counter() - tik}") + + self.create_optimizer(ft_spec) + self.initialized = True + + def save(self, meta: SaveLoadMeta): + if meta.weight_format == "hf": + self._save_model_to_hf(meta.path, meta.tokenizer, meta.processor) + elif meta.weight_format == "dcp": + # TODO: implement DCP save/load for FSDP + raise NotImplementedError("DCP format saving is not implemented yet. ") + else: + raise ValueError(f"Unknown weight format {meta.weight_format}. ") + + if meta.with_optim: + self.save_optimizer_state(meta.path) + + def load(self, meta: SaveLoadMeta): + if meta.weight_format == "hf": + self._load_model_from_hf(meta.path) + elif meta.weight_format == "dcp": + # TODO: implement DCP save/load for FSDP + raise NotImplementedError("DCP format loading is not implemented yet. ") + else: + raise ValueError(f"Unknown weight format {meta.weight_format}. ") + + if meta.with_optim: + self.load_optimizer_state(meta.path) + + def _save_model_to_hf( + self, + path: str, + tokenizer: Optional[PreTrainedTokenizerFast], + processor: Optional[AutoProcessor], + ): + """Save model in HuggingFace format.""" + if self.model is None: + raise RuntimeError("Model not initialized") + os.makedirs(path, exist_ok=True) + + # FSDP2 checkpoint saving + # Get full state dict with FSDP2 + options = StateDictOptions(full_state_dict=True, cpu_offload=True) + state_dict = get_model_state_dict(self.model, options=options) + + # save huggingface model on rank 0 + if dist.get_rank() == 0: + os.makedirs(path, exist_ok=True) + self.model.save_pretrained(path, state_dict=state_dict) + self.model_config.save_pretrained(path) + if tokenizer is not None: + tokenizer.save_pretrained(path) + if processor is not None: + processor.save_pretrained(path) + + dist.barrier(device_ids=[self.device.index]) + + def _load_model_from_hf(self, path: str): + """Load model from HuggingFace format.""" + if dist.get_rank() == 0: + full_state = get_state_dict_from_repo_id_or_path(path) + else: + full_state = {} + + fsdp2_load_full_state_dict( + self.model, + full_state, + self.cpu_offload, + tie_word_embeddings=self.model_config.tie_word_embeddings, + ) + + def upload_weights(self, meta: WeightUpdateMeta): + if meta.type == "nccl": + if not self.weight_update_group_initialized: + self._init_distributed_weight_update(meta) + self._update_weights_from_distributed() + dist.barrier(device_ids=[self.device.index]) + torch.cuda.synchronize() + elif meta.type == "disk": + self._save_model_to_hf(meta.path, self.tokenizer, self.processor) + # dist.barrier() are called when _save_model_to_hf finished + if dist.get_rank() == 0: + update_name = names.update_weights_from_disk( + self.config.experiment_name, + self.config.trial_name, + self.get_version(), + ) + name_resolve.add( + update_name, str(datetime.now().timestamp()), keepalive_ttl=120 + ) + else: + raise ValueError(f"Unknown weight update type {meta.type}") + + def _init_distributed_weight_update(self, meta: WeightUpdateMeta): + # NOTE: Processes launched with torchrun will set the following env var to True, + # which blocks creating another TCP store for weight update. + os.environ["TORCHELASTIC_USE_AGENT_STORE"] = str(False) + if dist.get_rank() == 0: + self.weight_update_group = init_custom_process_group( + backend="nccl", + world_size=meta.alloc_mode.gen_world_size + 1, + init_method=f"tcp://{meta.nccl_master_address}:{meta.nccl_master_port}", + rank=0, + group_name=meta.nccl_group_name, + ) + # NOTE: sglang v0.4.9.post2 or later does not have the barrier call + self.weight_update_group_initialized = True + + def _update_weights_from_distributed(self): + """Broadcast parameters from rank 0 (FSDP2 compatible).""" + + for name, param in self.model.named_parameters(): + if isinstance(param.data, DTensor): + tensor = param.data.full_tensor() + else: + tensor = param.data + if dist.get_rank() == 0: + logger.debug( + f"Broadcasting {name} with shape {tensor.shape}", flush=True + ) + dist.broadcast(tensor, src=0, group=self.weight_update_group) + del tensor # optional, for memory hygiene + torch.cuda.empty_cache() + + def get_param_specs(self) -> List[ParamSpec]: + param_specs = [] + for name, param in self.model.named_parameters(): + if isinstance(param.data, DTensor): + tensor = param.data.full_tensor() + else: + tensor = param.data + param_specs.append( + ParamSpec( + name=name, + shape=tuple(tensor.shape), + dtype=str(tensor.dtype).split("torch.")[1], + ) + ) + del tensor # free memory if full_tensor was created + return param_specs + + def train_batch( + self, + input_: TensorDict, + loss_fn: Callable[[torch.Tensor, TensorDict], torch.Tensor], + loss_weight_fn: Callable[[TensorDict], float], + ) -> Dict[str, float]: + """Train on a batch using gradient accumulation.""" + input_ = input_.to(self.device) + assert self.optimizer is not None + assert self.optimizer_config is not None + assert self.lr_scheduler is not None + + self.optimizer.zero_grad() + mb_list = self.prepare_mb_list(input_) + + total_loss_weight = torch.tensor( + sum([loss_weight_fn(mb) for mb in mb_list.mbs]), dtype=torch.float32 + ) + assert total_loss_weight != 0 + dist.all_reduce(total_loss_weight) + + # Process microbatches with gradient accumulation + for i, (pad_length, padded_mb_input, mb_input) in enumerate( + zip(mb_list.padding_lengths, mb_list.padded_mbs, mb_list.mbs) + ): + self.model.set_requires_gradient_sync(i == len(mb_list.mbs) - 1) + outputs = self.model(**padded_mb_input) + + logits = outputs.logits.squeeze(0) + logits = logits[:-pad_length] if pad_length > 0 else logits + loss = loss_fn(logits, mb_input) + loss_scale = loss_weight_fn(mb_input) / total_loss_weight + + # Scale loss for accumulation + # Revert gradient averaging across dp ranks + loss_scale *= self.world_size + + loss *= loss_scale + loss.backward() + + # NOTE: grad norm clip function is different + + grad_norm = fsdp2_clip_grad_norm_( + self.model.parameters(), max_norm=self.optimizer_config.gradient_clipping + ) + + if not torch.isfinite(grad_norm): + self.optimizer.zero_grad() + update_successful = False + else: + self.optimizer.step() + update_successful = True + + current_lr = self.lr_scheduler.get_last_lr()[0] + return dict( + update_successful=float(update_successful), + grad_norm=float(grad_norm) if grad_norm is not None else float("nan"), + lr=current_lr, + ) diff --git a/areal/engine/ppo/actor.py b/areal/engine/ppo/actor.py new file mode 100644 index 0000000000..0f283f8267 --- /dev/null +++ b/areal/engine/ppo/actor.py @@ -0,0 +1,333 @@ +import functools +from typing import Dict, List, Optional + +import torch +from tensordict import TensorDict + +from areal.api.cli_args import MicroBatchSpec, PPOActorConfig +from areal.api.engine_api import TrainEngine +from areal.engine.fsdp_engine import FSDPEngine +from areal.utils.data import split_padded_tensor_dict_into_mb_list +from areal.utils.functional import ( + gather_logprobs, + gather_logprobs_entropy, + masked_normalization, + ppo_actor_loss_fn, +) +from realhf.base import stats_tracker + + +class PPOActor: + + def __init__(self, config: PPOActorConfig, engine: TrainEngine): + self.config = config + self.engine = engine + + self.reward_bias = config.reward_bias + self.reward_scaling = config.reward_scaling + self.reward_clip = config.reward_clip + + self.group_reward_norm = config.group_reward_norm + self.group_adv_norm = config.group_adv_norm + self.group_size = config.group_size + + self.kl_ctl = config.kl_ctl + + self.adv_norm = config.adv_norm + self.discount = config.discount + self.gae_lambda = config.gae_lambda + self.mask_no_eos_with_zero = config.mask_no_eos_with_zero + + self.temperature = config.temperature + + @torch.no_grad() + def compute_logp( + self, + data: TensorDict, + temperature: Optional[float] = None, + ) -> torch.Tensor | None: + + def calc_logprobs(logits, input_data): + labels = torch.roll(input_data["input_ids"], shifts=-1, dims=-1) + logprobs = gather_logprobs(logits, labels, temperature or 1.0) + return logprobs + + self.engine.eval() + return self.engine.forward( + input_=data, + post_hook=calc_logprobs, + aggregate_fn=lambda xs: torch.cat(xs, dim=-1), + ) + + def compute_advantages(self, data: TensorDict) -> None: + bs = data["input_ids"].shape[0] + max_seqlen = data["input_ids"].shape[1] + batch_indices = torch.arange( + bs, device=data["input_ids"].device, dtype=torch.long + ) + + # Compute rewards using the reward function in synchronous RLVR pipeline. + reward_score = data["rewards"] + reward_score = (reward_score + self.reward_bias) * self.reward_scaling + reward_score = torch.clip( + reward_score, max=self.reward_clip, min=-self.reward_clip + ) + if self.group_reward_norm: + for i in range(bs // self.group_size): + s = slice(i * self.group_size, (i + 1) * self.group_size) + r = reward_score[s] + reward_score[s] = (r - r.mean()) / (r.std() + 1e-9) + + loss_mask = data["loss_mask"].float() + loss_mask = torch.roll(loss_mask, shifts=-1, dims=-1) + # Apply the mask to log probabilities. + if not self.config.use_decoupled_loss and self.config.recompute_logprob: + # Overwrite logprobs produced by the inference engine + old_logp = data["logprobs"] = data["prox_logp"] + else: + old_logp = torch.roll(data["logprobs"], shifts=-1, dims=-1) + if not self.config.use_decoupled_loss: + # prox logp not available, use inferenced logp + data["prox_logp"] = old_logp + ref_logp = data.get("ref_logp", torch.zeros_like(old_logp)) + ref_logp *= loss_mask + old_logp *= loss_mask + + # Compute KL-regularized rewards. + attn_mask = data["attention_mask"] + seqlens = attn_mask.sum(-1).long() + seq_no_eos_mask = seqlens == attn_mask.shape[1] + rewards = -self.kl_ctl * (old_logp - ref_logp) + kl_rewards = rewards.clone() + # KL rewards at the next token after eos is zero. + rewards[batch_indices, seqlens - 1] = 0 + indices = torch.clip(seqlens - 2, min=0) + if self.mask_no_eos_with_zero: + rewards[batch_indices, indices] += torch.where( + seq_no_eos_mask, 0, reward_score + ) + else: + rewards[batch_indices, indices] += reward_score + + # Compute GAE. + if "values" not in data: + values = torch.zeros_like(rewards) + else: + values = data["values"] + advantages_reversed = [ + torch.zeros(bs, dtype=torch.float32, device=values.device) + ] + lastgaelam = 0 + for t in reversed(range(max_seqlen - 1)): + nextvalues = values[:, t + 1] + if t == max_seqlen - 2: + nextvalues *= seq_no_eos_mask + delta = rewards[:, t] + self.discount * nextvalues - values[:, t] + lastgaelam = delta + self.discount * self.gae_lambda * lastgaelam + advantages_reversed.append(lastgaelam) + advantages = torch.stack(advantages_reversed[::-1], dim=1) + + # Optionally perform advantage normalization. + if self.adv_norm or self.group_adv_norm: + if self.group_adv_norm: + adv_list = [] + for i in range(0, bs, self.group_size): + s = slice(i * self.group_size, (i + 1) * self.group_size) + adv = advantages[s] + m = loss_mask[s] + adv_list.append(masked_normalization(adv, m, all_reduce=False)) + advantages = torch.cat(adv_list, 0) + else: + advantages = masked_normalization(advantages, loss_mask) + + # Store data in the dict. + data["advantages"] = advantages + data["kl_rewards"] = kl_rewards + data["tot_rewards"] = rewards + data["loss_mask"] = loss_mask + # because we have rolled old_logp by -1 + data["logprobs"] = old_logp + + def ppo_update(self, data: TensorDict) -> List[Dict[str, float]]: + attn_mask = data["attention_mask"] + loss_mask = data["loss_mask"] + reward_score = data["rewards"] + seqlens = attn_mask.sum(-1) + + all_stats = [] + ########## Logging code starts ########## + result_denominators = { + "correct_n_seqs": (reward_score > 0).bool(), + "incorrect_n_seqs": (reward_score <= 0).bool(), + } + global_denominators = dict( + n_seqs=torch.ones_like(reward_score, dtype=torch.bool), + n_tokens=torch.ones_like(loss_mask, dtype=torch.bool), + n_valid_tokens=loss_mask.bool(), + **result_denominators, + ) + stats_tracker.denominator(**global_denominators) + stats_tracker.stat( + correct_seq_len=seqlens.float(), denominator="correct_n_seqs" + ) + stats_tracker.stat( + incorrect_seq_len=seqlens.float(), denominator="incorrect_n_seqs" + ) + + stats = dict( + advantages=data["advantages"], + kl_rewards=data["kl_rewards"], + final_reward=data["tot_rewards"], + ) + stats_tracker.stat(**stats, denominator="n_valid_tokens") + + prompt_lens = [] + prompt_lens = data["attention_mask"].sum(-1) - data["loss_mask"].sum(-1) + seq_stats = dict( + no_eos_ratios=(seqlens == attn_mask.shape[-1]).float(), + task_reward=reward_score.float(), + prompt_len=prompt_lens.float(), + seq_len=seqlens.float(), + ) + stats_tracker.stat(**seq_stats, denominator="n_seqs") + scalars = dict( + mask_no_eos_with_zero=self.config.mask_no_eos_with_zero, + eps_clip=self.config.eps_clip, + ) + if self.config.c_clip is not None: + scalars["c_clip"] = self.config.c_clip + scalars["use_dual_clip"] = 1 + else: + scalars["use_dual_clip"] = 0 + if self.config.behav_imp_weight_cap is not None: + scalars["behav_imp_weight_cap"] = self.config.behav_imp_weight_cap + stats_tracker.scalar(**scalars) + + global_stats = stats_tracker.export(reduce_group=self.engine.parallelism_group) + for k in global_denominators: + keys = list(global_stats.keys()) + for k2 in keys: + if k2.endswith(k): + global_stats.pop(k2) + ########## Logging code ends ########## + + for key in ["rewards", "tot_rewards", "kl_rewards", "versions"]: + data.pop(key, None) + # NOTE: calling engine.train() is critical to enabling gradient checkpointing + self.engine.train() + mb_inputs = split_padded_tensor_dict_into_mb_list( + data, + mb_spec=MicroBatchSpec(n_mbs=self.config.ppo_n_minibatches), + ) + for mb in mb_inputs.mbs: + train_stat = self.engine.train_batch( + mb, + loss_fn=functools.partial( + grpo_loss_fn, + temperature=self.temperature, + eps_clip=self.config.eps_clip, + c_clip=self.config.c_clip, + behav_imp_weight_cap=self.config.behav_imp_weight_cap, + ), + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), + ) + stats_tracker.scalar(**train_stat) + all_stats.append( + stats_tracker.export(reduce_group=self.engine.parallelism_group) + ) + all_stats[0].update(global_stats) + return all_stats + + +class FSDPPPOActor(FSDPEngine): + + def __init__(self, config: PPOActorConfig): + super().__init__(config) + self.actor = PPOActor(config, self) + + @torch.no_grad() + def compute_logp(self, *args, **kwargs) -> torch.Tensor | None: + return self.actor.compute_logp(*args, **kwargs) + + @torch.no_grad() + def compute_advantages(self, *args, **kwargs) -> None: + self.actor.compute_advantages(*args, **kwargs) + + def ppo_update(self, *args, **kwargs) -> List[Dict[str, float]]: + return self.actor.ppo_update(*args, **kwargs) + + +def grpo_loss_fn( + logits: torch.Tensor, + input_data: Dict, + temperature: float, + eps_clip: float, + c_clip: float | None, + behav_imp_weight_cap: float | None, +): + """Loss function for actor step, all inputs should be splitted into + pipeline micro batches, returns loss and logging stats.""" + input_ids = input_data["input_ids"] + old_logp = input_data["logprobs"] + advantages = input_data["advantages"] + loss_mask = input_data["loss_mask"].bool() + prox_logp = input_data["prox_logp"] + + logprobs, entropy = gather_logprobs_entropy( + logits, torch.roll(input_ids, shifts=-1, dims=-1), temperature + ) + entropy = entropy.detach() + loss, stat = ppo_actor_loss_fn( + logprobs=logprobs, + old_logprobs=old_logp, + advantages=advantages, + eps_clip=eps_clip, + loss_mask=loss_mask, + c_clip=c_clip, + proximal_logprobs=prox_logp, + behav_imp_weight_cap=behav_imp_weight_cap, + ) + + # Log training statistics + stats_tracker.denominator( + n_tokens=torch.ones(logits.shape[0], dtype=torch.bool, device=logits.device), + n_valid_tokens=loss_mask.bool(), + clipped_tokens=stat["clip_mask"], + dual_clipped_tokens=stat["dual_clip_mask"], + ) + + stats_tracker.stat( + importance_weight=stat["importance_weight"], + approx_kl=stat["approx_kl"], + new_logp=logprobs.detach(), + old_logp=old_logp, + entropy=entropy.float(), + actor_loss=stat["loss"], + clip_ratio=stat["clip_mask"].float(), + dual_clip_ratio=stat["dual_clip_mask"].float(), + denominator="n_valid_tokens", + ) + if "behave_imp_weight" in stat: + stats_tracker.denominator(unclipped_behave_tokens=stat["behave_mask"]) + stats_tracker.stat( + behave_imp_weight=stat["behave_imp_weight"], + behave_approx_kl=stat["behave_approx_kl"], + denominator="unclipped_behave_tokens", + ) + vocab_min_logits = logits.detach().min(-1).values.float() + vocab_max_logits = logits.detach().max(-1).values.float() + stats_tracker.stat( + vocab_min_logits=vocab_min_logits, + vocab_max_logits=vocab_max_logits, + denominator="n_tokens", + ) + + clip_mask = stat["clip_mask"] + clipped_new_logp = torch.where(clip_mask, logprobs.detach(), 0.0) + clipped_old_logp = torch.where(clip_mask, old_logp, 0.0) + stats_tracker.stat( + clipped_new_logp=clipped_new_logp, + clipped_old_logp=clipped_old_logp, + denominator="clipped_tokens", + ) + return loss diff --git a/areal/engine/sft/lm_engine.py b/areal/engine/sft/lm_engine.py new file mode 100644 index 0000000000..8e2647cc69 --- /dev/null +++ b/areal/engine/sft/lm_engine.py @@ -0,0 +1,83 @@ +import torch +import torch.utils.data +from tensordict import TensorDict + +from areal.api.cli_args import TrainEngineConfig +from areal.api.engine_api import TrainEngine +from areal.engine.fsdp_engine import FSDPEngine +from areal.utils.functional import gather_logprobs +from realhf.base import stats_tracker + + +class LMEngine: + def __init__(self, engine: TrainEngine): + self.engine = engine + + def train_lm(self, data: TensorDict): + self.engine.train() + return self.engine.train_batch( + input_=data, + loss_fn=compute_packed_sft_loss, + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), + ) + + def evaluate_lm(self, data): + self.engine.eval() + self.engine.eval_batch( + input_=data, + loss_fn=compute_packed_sft_loss, + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), + ) + + +class FSDPLMEngine(FSDPEngine): + def __init__(self, config: TrainEngineConfig): + super().__init__(config) + self.lm_engine = LMEngine(self) + + def train_lm(self, data): + return self.lm_engine.train_lm(data) + + def evaluate_lm(self, data): + return self.lm_engine.evaluate_lm(data) + + +def compute_packed_sft_loss(logits: torch.Tensor, input_: TensorDict) -> torch.Tensor: + packed_input_ids: torch.Tensor = input_["input_ids"] + cu_seqlens: torch.Tensor = input_["cu_seqlens"] + loss_mask = input_["loss_mask"].bool() + + logprobs = gather_logprobs(logits, torch.roll(packed_input_ids, shifts=-1, dims=-1)) + loss_mask = torch.roll(loss_mask, shifts=-1, dims=-1) + logprobs = torch.where(loss_mask, logprobs, 0) + + loss = -logprobs.sum() / loss_mask.count_nonzero() + with torch.no_grad(): + seqlogp = torch.zeros( + cu_seqlens.shape[0] - 1, device=logits.device, dtype=torch.float64 + ) + for i in range(cu_seqlens.shape[0] - 1): + m = loss_mask[cu_seqlens[i] : cu_seqlens[i + 1]] + logp = logprobs[cu_seqlens[i] : cu_seqlens[i + 1]] + seqlogp[i] = torch.where(m, logp.detach(), 0.0).sum() / (m.count_nonzero()) + + ## Loggin stats + stats_tracker.denominator( + n_seqs=torch.ones( + cu_seqlens.shape[0] - 1, dtype=torch.bool, device=logprobs.device + ), + n_tokens=torch.ones(logits.shape[0], dtype=torch.bool, device=logits.device), + n_valid_tokens=loss_mask, + prompt_tokens=loss_mask.logical_not(), + ) + stats_tracker.stat(ppl=(-seqlogp).exp().float(), denominator="n_seqs") + stats_tracker.stat(loss=-logprobs.detach(), denominator="n_valid_tokens") + vocab_min_logits = logits.detach().min(-1).values.float() + vocab_max_logits = logits.detach().max(-1).values.float() + stats_tracker.stat( + vocab_min_logits=vocab_min_logits, + vocab_max_logits=vocab_max_logits, + denominator="n_tokens", + ) + + return loss diff --git a/areal/engine/sglang_remote.py b/areal/engine/sglang_remote.py new file mode 100644 index 0000000000..76ba18ec84 --- /dev/null +++ b/areal/engine/sglang_remote.py @@ -0,0 +1,440 @@ +import asyncio +import os +import random +import shutil +import time +from concurrent.futures import Future, ProcessPoolExecutor +from datetime import datetime +from typing import Any, Callable, Dict, List + +import aiohttp +import requests +import uvloop +from tensordict import TensorDict +from torchdata.stateful_dataloader import StatefulDataLoader + +from areal.api.cli_args import InferenceEngineConfig +from areal.api.engine_api import InferenceEngine +from areal.api.io_struct import ( + FinetuneSpec, + LLMRequest, + LLMResponse, + VLMRequest, + VLMResponse, + WeightUpdateMeta, +) +from areal.api.workflow_api import RolloutWorkflow, WorkflowExecutor +from areal.utils.http import arequest_with_retry, get_default_connector +from realhf.base import logging, name_resolve, names + +logger = logging.getLogger(__name__) + + +RID_CACHE_SIZE = 128 + + +class RemoteSGLangEngine(InferenceEngine): + + def __init__(self, config: InferenceEngineConfig): + self.config = config + + self.rid_to_address = {} + # Maintain the addresses for the recent 128 requests + self.rid_queue = [] + + self.addresses = os.getenv("AREAL_LLM_SERVER_ADDRS").split(",") + self.session = None + + if not self.addresses: + raise RuntimeError("No configured SGLang servers.") + + self.server_idx = random.randint(0, len(self.addresses) - 1) + self.distributed_weight_update_initialized = False + self._version = 0 + + self.workflow_executor = WorkflowExecutor( + config=config, + inference_engine=self, + ) + + def _wait_for_server(self, address): + base_url = f"http://{address}" + tik = time.time() + while time.time() - tik < self.config.setup_timeout: + if self.check_health(base_url): + return + time.sleep(1) + raise RuntimeError("server launch failed") + + def check_health(self, base_url): + # Check server endpoint + try: + response = requests.get(f"{base_url}/health", timeout=30) + return response.status_code == 200 + except requests.exceptions.RequestException as e: + return False + + def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None = None): + logger.info("Waiting for server ready...") + for addr_ in self.addresses: + self._wait_for_server(addr_) + logger.info("Servers are all ready!") + self.executor = ProcessPoolExecutor(max_workers=1) + self.workflow_executor.initialize() + + def destroy(self): + self.executor.shutdown() + + def set_version(self, version): + self._version = version + + def get_version(self): + return self._version + + def choose_server(self) -> str: + if self.config.schedule_policy == "round_robin": + server = self.addresses[self.server_idx] + self.server_idx = (self.server_idx + 1) % len(self.addresses) + return server + raise NotImplementedError("Only round-robin scheduling is implemented.") + + async def agenerate( + self, req: LLMRequest | VLMRequest + ) -> LLMResponse | VLMResponse: + """Async version of generate using aiohttp.""" + if self.session is None: + # NOTE: Lazily initialize aiohttp.ClientSession since it needs to be initialized + # inside asyncio loop in WorkflowExecutor + self.session = aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout( + total=self.config.request_timeout, + sock_connect=self.config.request_timeout, + connect=self.config.request_timeout, + ), + read_bufsize=1024 * 1024 * 10, + connector=get_default_connector(), + ) + # Prepare request payload + gconfig = req.gconfig + stop_token_ids = gconfig.stop_token_ids + + if gconfig.n_samples != 1: + raise ValueError( + "RemoteSGLangEngine does not support n_samples > 1. " + "Please call generate for multiple times with n_samples = 1." + ) + sample_params = { + "top_p": gconfig.top_p, + "top_k": gconfig.top_k, + "max_new_tokens": gconfig.max_new_tokens, + "temperature": 0.0 if gconfig.greedy else gconfig.temperature, + "stop_token_ids": stop_token_ids, + } + if isinstance(req, VLMRequest): + # VLMRequest has image_data + payload = { + "input_ids": req.input_ids.copy(), + "image_data": req.image_data, # ImageObject or str + "sampling_params": sample_params, + "return_logprob": True, + "stream": False, + } + else: + # NOTE: rid should NOT be passed in payload + payload = { + "input_ids": req.input_ids.copy(), + "sampling_params": sample_params, + "return_logprob": True, + "stream": False, + } + + # Make request + start_time = time.perf_counter() + accumulated_output_tokens = [] + accumulated_output_logprobs = [] + accumulated_versions = [] + + # A single "rid" shares the same sever to allow KV cache reuse + if req.rid in self.rid_to_address: + server_addr = self.rid_to_address[req.rid] + else: + server_addr = self.choose_server() + if len(self.rid_queue) >= RID_CACHE_SIZE: + # Remove the oldest entry if cache is full + oldest_rid = self.rid_queue.pop(0) + self.rid_to_address.pop(oldest_rid, None) + self.rid_to_address[req.rid] = server_addr + self.rid_queue.append(req.rid) + + # Deal with rollout interruption + # "abort" is the stop reason for later v0.4.9.post2 after + # we call the pause_generation endpoint + stop_reason = None + while ( + stop_reason != "stop" + and len(accumulated_output_tokens) < gconfig.max_new_tokens + ): + # Request is interrupted, wait for some time to avoid interfering + # with update weights requests + if stop_reason is not None: + await asyncio.sleep(0.5) + + # loop until the generation is complete + result = await arequest_with_retry( + session=self.session, + addr=server_addr, + endpoint="/generate", + payload=payload, + method="POST", + max_retries=self.config.request_retries, + timeout=self.config.request_timeout, + ) + + meta_info = result["meta_info"] + # Check if generation is complete + finish_reason = meta_info["finish_reason"] + stop_reason = finish_reason["type"] + if ( + stop_reason == "abort" + and finish_reason.get("message") == "Abort before prefill" + ): + continue + + # Parse response + output_tokens = [x[1] for x in meta_info["output_token_logprobs"]] + output_logprobs = [x[0] for x in meta_info["output_token_logprobs"]] + + # Update accumulated outputs + accumulated_output_tokens.extend(output_tokens) + accumulated_output_logprobs.extend(output_logprobs) + # FIXME: Update with actual server versions + accumulated_versions.extend([-1] * len(output_tokens)) + + payload["input_ids"] += result["output_ids"] + sample_params["max_new_tokens"] -= len(output_tokens) + + latency = time.perf_counter() - start_time + + if isinstance(req, VLMRequest): + response = VLMResponse( + input_tokens=req.input_ids, + input_images=req.image_data, + output_tokens=accumulated_output_tokens, + output_logprobs=accumulated_output_logprobs, + output_versions=accumulated_versions, + stop_reason=stop_reason, + latency=latency, + ttft=latency, # Simplified for non-streaming + ) + else: + response = LLMResponse( + input_tokens=req.input_ids, + output_tokens=accumulated_output_tokens, + output_logprobs=accumulated_output_logprobs, + output_versions=accumulated_versions, + stop_reason=stop_reason, + latency=latency, + ttft=latency, # Simplified for non-streaming + ) + return response + + def update_weights(self, meta: WeightUpdateMeta): + for addr in self.addresses: + res = requests.post(f"http://{addr}/pause_generation") + res.raise_for_status() + fut = Future() + if meta.type == "nccl": + fut = self.executor.submit( + update_weights_from_distributed, + meta, + self.addresses, + self.config.request_timeout, + not self.distributed_weight_update_initialized, + ) + + def callback(fut): + self.distributed_weight_update_initialized = True + + fut.add_done_callback(callback) + elif meta.type == "disk": + # Update weights from disk + # Use ProcessPool to bypass python GIL for running async coroutines + fut = self.executor.submit( + update_weights_from_disk, + self.config.experiment_name, + self.config.trial_name, + self.get_version(), + self.addresses, + meta.path, + self.config.request_retries, + self.config.request_timeout, + ) + + def callback(fut): + shutil.rmtree(meta.path, ignore_errors=True) + + fut.add_done_callback(callback) + else: + raise NotImplementedError(f"Unsupported weight update type: {meta.type}") + + def callback(fut): + for addr in self.addresses: + res = requests.post(f"http://{addr}/continue_generation") + res.raise_for_status() + + fut.add_done_callback(callback) + return fut + + def submit(self, data: Dict[str, Any], workflow: RolloutWorkflow) -> None: + return self.workflow_executor.submit(data, workflow) + + def wait( + self, + count: int, + timeout: float | None = None, + should_accept: Callable | None = None, + ) -> TensorDict: + return self.workflow_executor.wait( + count, + timeout=timeout, + should_accept=should_accept, + ) + + def rollout_batch( + self, data: List[Dict[str, Any]], workflow: "RolloutWorkflow" + ) -> TensorDict: + return self.workflow_executor.rollout_batch(data, workflow) + + def prepare_batch( + self, + dataloader: StatefulDataLoader, + workflow: RolloutWorkflow, + ): + return self.workflow_executor.prepare_batch(dataloader, workflow) + + def pause(self): + """Pause request submission for async rollout. Used during evaluation to prevent data over generation.""" + return self.workflow_executor.pause() + + def resume(self): + """Resume request submission for async rollout.""" + return self.workflow_executor.resume() + + +def update_weights_from_disk( + experiment_name, + trial_name, + model_version, + addresses, + path, + request_retries, + request_timeout, +): + async def _fn(): + # Wait for model checkpoints of meta.version + update_name = names.update_weights_from_disk( + experiment_name, trial_name, model_version + ) + save_timestamp = float(name_resolve.wait(update_name, timeout=120)) + load_timestamp = datetime.now().timestamp() + logger.info( + f"Begin update weights from {path}, responded in {(load_timestamp - save_timestamp):.2f}s" + ) + session = aiohttp.ClientSession( + timeout=aiohttp.ClientTimeout( + total=request_timeout, + sock_connect=request_timeout, + connect=request_timeout, + ), + read_bufsize=1024 * 1024 * 10, + connector=get_default_connector(), + ) + jobs = [ + arequest_with_retry( + addr=addr, + session=session, + endpoint="/update_weights_from_disk", + payload=dict(model_path=str(path)), + method="POST", + max_retries=request_retries, + timeout=request_timeout, + ) + for addr in addresses + ] + await asyncio.gather(*jobs) + await session.close() + logger.info( + f"Loading weights done in {(datetime.now().timestamp() - load_timestamp):.2f}s" + ) + + return uvloop.run(_fn()) + + +def update_weights_from_distributed( + meta: WeightUpdateMeta, + addresses: List[str], + request_timeout, + init_group: bool, +): + async def _fn(): + tik = time.perf_counter() + if init_group: + await asyncio.gather( + *[ + ainit_weights_update_group(addr, i, meta, request_timeout) + for i, addr in enumerate(addresses) + ] + ) + await asyncio.gather( + *[ + arequest_with_retry( + addr=addr, + endpoint="/update_weights_from_distributed", + payload={ + "names": [pspec.name for pspec in meta.nccl_param_specs], + "dtypes": [pspec.dtype for pspec in meta.nccl_param_specs], + "shapes": [pspec.shape for pspec in meta.nccl_param_specs], + "group_name": meta.nccl_group_name, + }, + method="POST", + max_retries=1, + timeout=request_timeout, + ) + for addr in addresses + ] + ) + + logger.info(f"Distributed update weights done in {time.perf_counter() - tik}s") + + return uvloop.run(_fn()) + + +async def ainit_weights_update_group( + addr: str, + server_idx: int, + meta: WeightUpdateMeta, + request_timeout: float, +): + assert meta.alloc_mode is not None + if meta.alloc_mode.gen_pp_size != 1: + raise NotImplementedError( + "NCCL weight update with PP size > 1 is not implemented yet." + ) + rank_offset = 1 + server_idx * meta.alloc_mode.gen_tp_size + payload = { + "master_address": meta.nccl_master_address, + "master_port": str(meta.nccl_master_port), + "rank_offset": rank_offset, + "world_size": meta.alloc_mode.gen_world_size + 1, + "backend": "nccl", + "group_name": meta.nccl_group_name, + } + res = await arequest_with_retry( + addr=addr, + endpoint="/init_weights_update_group", + payload=payload, + method="POST", + max_retries=1, + timeout=request_timeout, + ) + assert res["success"] diff --git a/areal/experimental/autotp_engine.py b/areal/experimental/autotp_engine.py new file mode 100644 index 0000000000..f5ec055765 --- /dev/null +++ b/areal/experimental/autotp_engine.py @@ -0,0 +1,128 @@ +import os + +import torch +import torch.distributed as dist +from safetensors.torch import save_file + +from areal.api.cli_args import TrainEngineConfig +from areal.api.engine_api import FinetuneSpec, SaveLoadMeta, WeightUpdateMeta +from areal.engine.base_hf_engine import BaseHFEngine +from areal.utils.save_load import ( + get_state_dict_from_repo_id_or_path, + is_existing_local_path, +) +from realhf.base import constants, logging + +logger = logging.getLogger("DeepSpeedAutoTPEngine") + + +class DeepSpeedAutoTPEngine(BaseHFEngine): + def __init__(self, config: TrainEngineConfig): + super().__init__(config) + + def initialize(self, addr: str | None, ft_spec: FinetuneSpec | None): + """Initialize distributed communication and model.""" + assert ( + addr is None + ), "DeepSpeedAutoTPEngine does not support remote initialization." + import deepspeed + + self.create_process_group() + + world_size = int(os.environ.get("WORLD_SIZE")) + deepspeed.init_distributed( + dist_backend="nccl", + world_size=world_size, + timeout=constants.NCCL_DEFAULT_TIMEOUT, + ) + self.create_device_model() + # NOTE: the device context manager does not work here. + self.model = deepspeed.tp_model_init( + self.model, + tp_size=self.config.ds_auto_tp.autotp_size, + dtype=getattr(torch, self.config.dtype), + ).to(self.device) + self.create_optimizer(ft_spec) + self.initialized = True + + def _check_autotp(self): + tp_size = self.config.ds_auto_tp.autotp_size + config = self.model_config + num_attention_heads = config.num_attention_heads + num_key_value_heads = config.num_key_value_heads + hidden_size = config.hidden_size + intermediate_size = config.intermediate_size + + return ( + num_attention_heads % tp_size == 0 + and num_key_value_heads % tp_size == 0 + and hidden_size % tp_size == 0 + and intermediate_size % tp_size == 0 + ) + + def save(self, meta: SaveLoadMeta): + if meta.weight_format != "naive_distributed": + raise ValueError(f"Unknown weight format {meta.weight_format}. ") + if self.model is None: + raise RuntimeError("Model not initialized") + + rank = dist.get_rank() + world_size = dist.get_world_size() + if rank == 0: + os.makedirs(meta.path, exist_ok=True) + self.model_config.save_pretrained( + meta.path, + ) + if meta.tokenizer is not None: + meta.tokenizer.save_pretrained( + meta.path, + ) + + state_dict = self.model.state_dict() + if hasattr(self.model, "module"): + state_dict = { + k.replace("module.", "", 1) if k.startswith("module.") else k: v.cpu() + for k, v in state_dict.items() + } + else: + state_dict = {k: v.cpu() for k, v in state_dict.items()} + + # Only support store parameters from model partitions respectively + gathered_state_dicts = None + if rank == 0: + gathered_state_dicts = [None for _ in range(world_size)] + dist.gather_object( + obj=state_dict, object_gather_list=gathered_state_dicts, dst=0 + ) + if rank == 0: + for i, state_dict in enumerate(gathered_state_dicts): + save_file(state_dict, f"{meta.path}/rank_{i:02d}_model.safetensors") + if meta.with_optim: + self.save_optimizer_state(meta.path) + + def load(self, meta: SaveLoadMeta): + if meta.weight_format != "naive_distributed": + raise ValueError(f"Unknown weight format {meta.weight_format}. ") + rank = dist.get_rank() + # Only support load full model parameters from huggingface + # and load model partition locally + if rank == 0 or is_existing_local_path(meta.path): + path = f"{meta.path}/rank_{rank:02d}_model.safetensors" + full_state = get_state_dict_from_repo_id_or_path(meta.path) + + if hasattr(self.model, "module") and not hasattr(full_state): + full_state = { + f"module.{k}" if not k.startswith("module.") else k: v + for k, v in full_state.items() + } + self.model.load_state_dict( + full_state, strict=not self.model_config.tie_word_embeddings + ) + if self.model_config.tie_word_embeddings: + self.model.tie_weights() + + if meta.with_optim: + self.load_optimizer_state(meta.path) + + def upload_weights(self, meta: WeightUpdateMeta): + raise ValueError(f"update weight not implemented {meta.type}") diff --git a/areal/experimental/sglang_engine.py b/areal/experimental/sglang_engine.py new file mode 100644 index 0000000000..d8425ee4bb --- /dev/null +++ b/areal/experimental/sglang_engine.py @@ -0,0 +1,371 @@ +import asyncio +import threading +import time +import traceback +from concurrent.futures import ThreadPoolExecutor +from queue import Empty, Full, Queue +from typing import TYPE_CHECKING, Any, Callable, Dict, List, Optional + +import sglang as sgl +import torch.distributed as dist +from tensordict import TensorDict + +from areal.api.cli_args import InferenceEngineConfig +from areal.api.engine_api import InferenceEngine +from areal.api.io_struct import LLMRequest, LLMResponse, RolloutStat, WeightUpdateMeta +from realhf.base import logging, name_resolve, names, pkg_version + +if TYPE_CHECKING: + from areal.api.workflow_api import RolloutWorkflow +logger = logging.getLogger(__name__) + +if pkg_version.is_available("sglang"): + if pkg_version.is_version_greater_or_equal("sglang", "0.4.4"): + SGLANG_TOKEN_OUTPUT_IDENTIFIER = "output_ids" + else: + SGLANG_TOKEN_OUTPUT_IDENTIFIER = "token_ids" + +ROLLOUT_POLL_WAIT_TIME = 0.4 +RID_CACHE_SIZE = 128 + +""" +Local SGLang Inference Engine +SGLangEngine currently only supports single-controller. Cannot be used in SPMD +""" + + +class SGLangEngine(InferenceEngine): + + def __init__( + self, + config: InferenceEngineConfig, + engine_args: Optional[Dict[str, Any]] = None, + ): + config.max_concurrent_rollouts = ( + config.max_concurrent_rollouts or config.consumer_batch_size + ) + self.config = config + self.engine_args = engine_args or {} + + qsize = config.queue_size or config.max_concurrent_rollouts * 10 + self.input_queue = Queue(maxsize=qsize) + self.output_queue = Queue(maxsize=qsize) + self.result_cache = [] + + self.exiting = threading.Event() + self.lock = threading.Lock() + + self.rollout_stat = RolloutStat() + + self._version = 0 + + def initialize(self, addr: str | None, ft_spec: Optional[Dict[str, Any]] = None): + self.engine = sgl.Engine(**self.engine_args) + + self.rollout_thread = threading.Thread(target=self._rollout_thread) + self.rollout_thread.start() + + def destroy(self): + self.exiting.set() + self.rollout_thread.join() + + if hasattr(self, "engine") and self.engine is not None: + try: + self.engine.shutdown() + except Exception as e: + logger.warning(f"Error shutting down engine: {e}") + + def set_version(self, version): + with self.lock: + self._version = version + + def get_version(self): + with self.lock: + return self._version + + def _rollout_thread(self): + """Thread that runs the rollout loop.""" + try: + asyncio.run(self._rollout_thread_async()) + except Exception as e: + traceback.print_exc() + raise e + + async def _rollout_thread_async(self): + data = None + + rollout_tasks: Dict[str, asyncio.Task] = {} + rid = 0 + + try: + while not self.exiting.is_set(): + # Load next data from controller + if data is None: + try: + data, workflow = self.input_queue.get_nowait() + logger.info(f"Get data from puller: {data}") + except Empty: + logger.debug(f"No data from puller stream.") + + # Check capacity + if dist.is_initialized(): + world_size = dist.get_world_size() + else: + world_size = 1 + + cannot_rollout_reason = [] + capacity = max(1, self.config.max_concurrent_rollouts // world_size) + can_rollout = len(rollout_tasks) < capacity + if not can_rollout: + cannot_rollout_reason.append( + f"Exceeding capacity: # running tasks {len(rollout_tasks)} >= capacity {capacity}" + ) + + # Staleness control + version = self.get_version() + ofp = self.config.max_head_offpolicyness + with self.lock: + sample_cnt = self.rollout_stat.accepted + self.rollout_stat.running + expected_version = sample_cnt // self.config.consumer_batch_size + not_staled = expected_version <= ofp + version + can_rollout &= not_staled + if not not_staled: + cannot_rollout_reason.append( + f"Staled: expected version ({expected_version}) = " + f"global sample cnt ({sample_cnt}) // batch size ({self.config.consumer_batch_size}), " + f"current latest version {version}, " + f"offpolicyness {self.config.max_head_offpolicyness}." + ) + + if not can_rollout: + logger.debug( + f"Cannot submit new rollouts. " + + "\n".join(cannot_rollout_reason) + ) + + # Create new rollout task + if can_rollout and data is not None: + task = asyncio.create_task( + workflow.arun_episode(self, data), name=str(rid) + ) + rollout_tasks[str(rid)] = task + + with self.lock: + self.rollout_stat.submitted += 1 + self.rollout_stat.running += 1 + logger.info( + f"Submit rollout rid {rid}. " + f"Submit: {self.rollout_stat.submitted}, " + f"running: {self.rollout_stat.running}, " + f"accepted: {self.rollout_stat.accepted}." + ) + + rid += 1 + data = None + + # Wait for rollout completion + tasks = list(rollout_tasks.values()) + done = [] + if tasks: + done, _ = await asyncio.wait( + tasks, + timeout=ROLLOUT_POLL_WAIT_TIME, + return_when=asyncio.FIRST_COMPLETED, + ) + else: + await asyncio.sleep(ROLLOUT_POLL_WAIT_TIME) + + # Collect done results + for task in done: + traj = await task + traj: TensorDict + task_rid = task.get_name() + rollout_tasks.pop(task_rid) + self.rollout_stat.accepted += 1 + + try: + self.output_queue.put_nowait(traj) + except Full: + raise RuntimeError( + "Output queue full. Please increase queue_size." + ) + + with self.lock: + self.rollout_stat.running -= 1 + logger.info( + f"Finish rollout {task_rid}. " + f"Submit: {self.rollout_stat.submitted}, " + f"running: {self.rollout_stat.running}, " + f"accepted: {self.rollout_stat.accepted}." + ) + finally: + # Cancel remaining tasks + for task in rollout_tasks.values(): + if not task.done(): + task.cancel() + try: + await task + except asyncio.CancelledError: + pass + + async def agenerate(self, req: LLMRequest) -> LLMResponse: + """Async version of generate using local sglang engine.""" + if not hasattr(self, "engine") or self.engine is None: + raise RuntimeError( + "Local SGLang engine is not initialized, cannot generate." + ) + + # Prepare request payload + gconfig = req.gconfig + stop_token_ids = gconfig.stop_token_ids + + if gconfig.n_samples != 1: + raise ValueError( + "LocalSGLangEngine does not support n_samples > 1. " + "Please call generate for multiple times with n_samples = 1." + ) + sample_params = { + "top_p": gconfig.top_p, + "top_k": gconfig.top_k, + "max_new_tokens": gconfig.max_new_tokens, + "temperature": 0.0 if gconfig.greedy else gconfig.temperature, + "stop_token_ids": stop_token_ids, + } + + completions = "" + prompt = req.text if req.text else None + input_ids = req.input_ids if req.input_ids else None + + # Make request + start_time = time.perf_counter() + accumulated_output_tokens = [] + accumulated_output_logprobs = [] + accumulated_versions = [] + stop_reason = "length" + while ( + stop_reason != "stop" + and len(accumulated_output_tokens) < gconfig.max_new_tokens + ): + + try: + outputs = await self.engine.async_generate( + prompt=prompt, + input_ids=input_ids, + sampling_params=sample_params, + return_logprob=True, + ) + + completions += outputs["text"] + if prompt is None: + prompt = outputs["text"] + else: + prompt += outputs["text"] + + meta_info = outputs["meta_info"] + output_tokens = [x[1] for x in meta_info["output_token_logprobs"]] + output_logprobs = [x[0] for x in meta_info["output_token_logprobs"]] + + finish_reason = meta_info.get("finish_reason", {}) + stop_reason = finish_reason.get("type", "length") + + accumulated_output_tokens.extend(output_tokens) + accumulated_output_logprobs.extend(output_logprobs) + accumulated_versions.extend([-1] * len(output_tokens)) + + except Exception as e: + raise RuntimeError(f"Local SGLang engine generation failed: {e}") + + latency = time.perf_counter() - start_time + + return LLMResponse( + completions=completions, + input_tokens=req.input_ids if req.input_ids else [], + output_tokens=accumulated_output_tokens, + output_logprobs=accumulated_output_logprobs, + output_versions=accumulated_versions, + stop_reason=stop_reason, + latency=latency, + ttft=latency, + ) + + def update_weights(self, meta): + executor = ThreadPoolExecutor(max_workers=1) + return executor.submit(self._update_weights, meta) + + def _update_weights(self, meta: WeightUpdateMeta): + if not hasattr(self, "engine") or self.engine is None: + raise RuntimeError( + "Local SGLang engine is not initialized, cannot update weights." + ) + if meta.type == "disk": + try: + update_name = names.update_weights_from_disk( + self.config.experiment_name, + self.config.trial_name, + meta.model_version, + ) + save_timestamp = int(name_resolve.wait(update_name, timeout=120)) + load_timestamp = time.time_ns() + logger.info( + f"Begin update weights from {meta.path}, responded in {(load_timestamp - save_timestamp)/1e6:.2f} ms" + ) + # Update weights from disk, + self.engine.update_weights_from_disk(model_path=meta.path) + + logger.info( + f"Loading weights done in {(time.time_ns() - load_timestamp)/1e6:.2f} ms" + ) + self.set_version(meta.model_version) + except Exception as e: + logger.error(f"Failed to update weights: {e}") + raise + else: + raise NotImplementedError(f"Unsupported weight update type: {meta.type}") + + def submit(self, data: Dict[str, Any], workflow: "RolloutWorkflow") -> None: + try: + self.input_queue.put_nowait((data, workflow)) + except Full: + raise RuntimeError("Input queue full. Please increase queue_size.") + + def wait(self, count: int, timeout: float, should_accept: Callable) -> TensorDict: + tik = time.perf_counter() + accepted = len(self.result_cache) + while ( + accepted < count + and not self.exiting.is_set() + and time.perf_counter() - tik < timeout + ): + try: + result = self.output_queue.get(timeout=ROLLOUT_POLL_WAIT_TIME) + if should_accept(result): + self.result_cache.append(result) + accepted += 1 + else: + with self.lock: + self.rollout_stat.accepted -= 1 + except Empty: + time.sleep(ROLLOUT_POLL_WAIT_TIME) + if self.exiting.is_set(): + raise RuntimeError("Rollout engine is exiting, cannot wait for results.") + if accepted < count: + raise TimeoutError( + f"Timed out waiting for {count} rollouts, " f"only received {accepted}." + ) + results, self.result_cache = ( + self.result_cache[:count], + self.result_cache[count:], + ) + return TensorDict.cat(results, dim=0) + + def rollout( + self, data: List[Dict[str, Any]], workflow: "RolloutWorkflow" + ) -> TensorDict: + """Submit a batch of requests to the inference engine and wait for the results.""" + for item in data: + self.submit(item, workflow) + return self.wait( + count=len(data), + timeout=self.config.request_timeout, + should_accept=lambda x: True, + ) diff --git a/areal/experimental/tests/test_sglang_local_engine.py b/areal/experimental/tests/test_sglang_local_engine.py new file mode 100644 index 0000000000..f716d2b97a --- /dev/null +++ b/areal/experimental/tests/test_sglang_local_engine.py @@ -0,0 +1,140 @@ +# ================================================================ +# NOTE: This test file is dedicated to LocalSGLangEngine testing. +# +# Unlike remote engine setup which is managed via a pytest fixture, +# the LocalSGLangEngine requires explicit `initialize()` to construct +# the engine instance at runtime. Therefore, each test must manually +# create and destroy the engine. +# +# Because of this lifecycle difference, tests for local and remote +# engines cannot be placed in the same test file. +# ================================================================ +import os +import time +import uuid + +import pytest +import torch +from tensordict import TensorDict + +from areal.api.cli_args import ( + GenerationHyperparameters, + InferenceEngineConfig, + SGLangConfig, +) +from areal.api.io_struct import LLMRequest, LLMResponse +from areal.experimental.sglang_engine import SGLangEngine +from areal.workflow.rlvr import RLVRWorkflow +from realhf.api.core.data_api import load_hf_tokenizer +from realhf.base import seeding + +EXPR_NAME = "test_sglang_local_engine" +TRIAL_NAME = "trial_0" +MODEL_PATH = "/storage/testing/models/Qwen__Qwen3-1.7B/" +if not os.path.exists(MODEL_PATH): + MODEL_PATH = "Qwen/Qwen2-0.5B" + + +def build_engine_config(**kwargs): + return InferenceEngineConfig( + experiment_name=EXPR_NAME, + trial_name=TRIAL_NAME, + **kwargs, + ) + + +def build_engine_args(): + return SGLangConfig.build_args( + sglang_config=SGLangConfig(mem_fraction_static=0.3, enable_metrics=False), + model_path=MODEL_PATH, + tp_size=1, + base_gpu_id=0, + served_model_name=MODEL_PATH, + skip_tokenizer_init=False, + ) + + +@pytest.mark.asyncio +async def test_local_sglang_generate(): + seeding.set_random_seed(1, EXPR_NAME) + config = build_engine_config() + engine = SGLangEngine(config, engine_args=build_engine_args()) + engine.initialize(None, None) + + req = LLMRequest( + rid=str(uuid.uuid4()), + text="hello! how are you today", + gconfig=GenerationHyperparameters(max_new_tokens=16), + ) + resp = await engine.agenerate(req) + + assert isinstance(resp, LLMResponse) + assert resp.input_tokens == req.input_ids + assert ( + len(resp.output_logprobs) + == len(resp.output_tokens) + == len(resp.output_versions) + ) + engine.destroy() + + +@pytest.mark.parametrize("n_samples", [1, 2, 4]) +def test_local_sglang_rollout(n_samples): + seeding.set_random_seed(1, EXPR_NAME) + config = build_engine_config(max_concurrent_rollouts=2, consumer_batch_size=2) + engine = SGLangEngine(config, engine_args=build_engine_args()) + engine.initialize(None, None) + + gconfig = GenerationHyperparameters( + max_new_tokens=16, greedy=False, n_samples=n_samples + ) + tokenizer = load_hf_tokenizer(MODEL_PATH) + + workflow = RLVRWorkflow( + reward_fn=lambda **kwargs: 1.0, + gconfig=gconfig, + tokenizer=tokenizer, + ) + + data = {"messages": [{"role": "user", "content": "Hello, how are you?"}]} + result = engine.rollout([data] * 2, workflow=workflow) + + print("Here is the result ", result) + assert isinstance(result, TensorDict) + assert result.batch_size == torch.Size([2 * n_samples]) + engine.destroy() + + +@pytest.mark.parametrize("ofp", [1, 2, 4, 8, 16]) +@pytest.mark.parametrize("bs", [2, 4]) +@pytest.mark.parametrize("n_samples", [2, 1]) +def test_local_sglang_staleness_control(bs, ofp, n_samples): + seeding.set_random_seed(1, EXPR_NAME) + config = build_engine_config(consumer_batch_size=bs, max_head_offpolicyness=ofp) + engine = SGLangEngine(config, engine_args=build_engine_args()) + engine.initialize(None, None) + + gconfig = GenerationHyperparameters( + max_new_tokens=16, greedy=False, n_samples=n_samples + ) + tokenizer = load_hf_tokenizer(MODEL_PATH) + + workflow = RLVRWorkflow( + reward_fn=lambda **kwargs: 1.0, + gconfig=gconfig, + tokenizer=tokenizer, + ) + + data = {"messages": [{"role": "user", "content": "Hello, how are you?"}]} + for _ in range(bs * 2): + engine.submit(data, workflow=workflow) + time.sleep(5) + assert engine.output_queue.qsize() == min(bs * 2, bs * (ofp + 1)) + + engine.set_version(1) + for _ in range(bs * 2): + engine.submit(data, workflow=workflow) + time.sleep(5) + assert engine.output_queue.qsize() == min(bs * 4, bs * (ofp + 2)) + + engine.destroy() diff --git a/areal/launcher/local.py b/areal/launcher/local.py new file mode 100644 index 0000000000..c5ecc737f5 --- /dev/null +++ b/areal/launcher/local.py @@ -0,0 +1,304 @@ +import getpass +import os +import re +import signal as signal_module +import subprocess +import sys +import time +from collections import defaultdict +from typing import Dict, List, Optional, Tuple, Union + +import psutil + +from areal.api.cli_args import SGLangConfig, parse_cli_args, to_structured_cfg +from areal.api.io_struct import AllocationMode, AllocationType +from areal.utils.network import find_free_ports, gethostip +from realhf.base import gpu_utils, logging, name_resolve, names +from realhf.scheduler.client import JobException, JobInfo, JobState + +logger = logging.getLogger("Local Scheduler") +JOB_STATE_TO_PROCESS_STATUS = { + JobState.NOT_FOUND: [], + JobState.PENDING: [psutil.STATUS_PARKED], + JobState.RUNNING: [ + psutil.STATUS_RUNNING, + psutil.STATUS_SLEEPING, + psutil.STATUS_DISK_SLEEP, + psutil.STATUS_TRACING_STOP, + psutil.STATUS_WAKING, + psutil.STATUS_WAITING, + psutil.STATUS_LOCKED, + psutil.STATUS_IDLE, + ], + JobState.COMPLETED: [ + psutil.STATUS_DEAD, + psutil.STATUS_STOPPED, + psutil.STATUS_ZOMBIE, + ], + JobState.FAILED: [], + JobState.CANCELLED: [], +} + +PROCESS_STATUS_TO_JOB_STATE = {} +for job_state, process_statuses in JOB_STATE_TO_PROCESS_STATUS.items(): + for process_status in process_statuses: + PROCESS_STATUS_TO_JOB_STATE[process_status] = job_state + + +def terminate_process_and_children(pid: int, signal: Optional[Union[str, int]] = None): + if signal is None: + signal = signal_module.SIGKILL + if isinstance(signal, str): + signal = getattr(signal_module, signal) + try: + parent = psutil.Process(pid) + children = parent.children(recursive=True) + for child in children: + terminate_process_and_children(child.pid) + parent.send_signal(signal) + except psutil.NoSuchProcess: + pass + + +class LocalLauncher: + def __init__(self, experiment_name: str, trial_name: str, fileroot: str): + self.experiment_name = experiment_name + self.trial_name = trial_name + self.fileroot = fileroot + + self._jobs: Dict[str, subprocess.Popen] = {} + self._job_counter: Dict[str, int] = defaultdict(int) + self._job_states = {} + + self._gpu_counter = 0 + self._cuda_devices: List[str] = os.environ.get( + "CUDA_VISIBLE_DEVICES", ",".join(map(str, range(gpu_utils.gpu_count()))) + ).split(",") + if len(self._cuda_devices) < 1: + raise RuntimeError( + f"Local mode can only run when there is at least one GPU. " + f"CUDA_VISIBLE_DEVICES is currently set to {os.environ['CUDA_VISIBLE_DEVICES']}." + ) + + @property + def run_name(self): + return f"{self.experiment_name}_{self.trial_name}" + + def log_path_of(self, job_name: str) -> str: + log_path = f"{self.fileroot}/logs/{getpass.getuser()}/{self.experiment_name}/{self.trial_name}" + os.makedirs(log_path, exist_ok=True) + return os.path.join(log_path, f"{job_name}.log") + + def __del__(self): + self.wait() + + def submit_array( + self, + job_name: str, + cmd: str | List[str], + count: int = 1, + gpu: int = 0, + env_vars: Optional[Dict] = None, + ): + if env_vars is None: + env_vars = {} + if not isinstance(cmd, list): + cmd = [cmd] * count + offset = self._job_counter[job_name] + for i in range(count): + if gpu > 0: + # Allocate GPUs in a round-robin manner + visible_devices = [] + for _ in range(gpu): + available_device_id = self._gpu_counter % len(self._cuda_devices) + self._gpu_counter += 1 + visible_devices.append(available_device_id) + env_vars["CUDA_VISIBLE_DEVICES"] = ",".join( + str(self._cuda_devices[j]) for j in visible_devices + ) + c = ( + " ".join(str(k) + "=" + str(v) for k, v in env_vars.items()) + + " stdbuf -oL " + + cmd[i] + ) + c = f"{c} | tee -a {self.log_path_of(job_name)}" + logger.info("Starting local process with command: %s", c) + process = subprocess.Popen(c, shell=isinstance(c, str)) + self._jobs[f"{job_name}/{offset + i}"] = process + self._job_counter[job_name] += 1 + + def submit( + self, + job_name: str, + cmd: str | List[str], + gpu: int = 0, + env_vars: Optional[Dict] = None, + ): + self.submit_array(job_name=job_name, cmd=cmd, gpu=gpu, env_vars=env_vars) + + def stop(self, job_name, signal=None): + assert any(k.startswith(job_name) for k in self._jobs) + keys = [k for k, p in self._jobs.items() if k.startswith(job_name)] + procs = [p for k, p in self._jobs.items() if k.startswith(job_name)] + logger.info( + f"Stopping local process with signal {signal if signal else 'SIGKILL'}, " + f"pid: {[p.pid for p in procs]}" + ) + for p in procs: + terminate_process_and_children(p.pid, signal=signal) + for p in procs: + p.wait() + for k, p in zip(keys, procs): + self._jobs.pop(k) + del p + + def stop_all(self, signal=None): + # signal argument is ignored in local stop_all + for name in self._job_counter: + self.stop(name, signal=signal) + + def find(self, job_name): + if job_name in self._jobs: + return JobInfo(name=job_name, state=JobState.RUNNING, host="localhost") + else: + return JobInfo(name=job_name, state=JobState.NOT_FOUND) + + def find_all(self, job_name_regex=".*"): + rs = [] + for name in self._jobs: + if re.fullmatch(job_name_regex, name): + rs.append(self.find(name)) + return rs + + def wait( + self, + timeout=None, + check_status: Tuple[JobState, ...] = ( + JobState.CANCELLED, + JobState.FAILED, + JobState.NOT_FOUND, + ), + remove_status: Tuple[JobState, ...] = (JobState.COMPLETED,), + update=False, + ): + deadline = None if timeout is None else time.time() + timeout + logger.info( + "Waiting for %d local running processes, pids: %s", + len(self._jobs), + " ".join(str(job.pid) for job in self._jobs.values()), + ) + left = set(self._jobs.keys()) + num_jobs_left = len(left) + + while len(left) > 0: + to_remove = [] + if len(left) < num_jobs_left: + num_jobs_left = len(left) + logger.info(f"Waiting for {num_jobs_left} jobs.") + if deadline is not None and time.time() > deadline: + raise TimeoutError( + f"Timeout waiting for {self.run_name}: {', '.join(sorted(left))}" + ) + # update job states + for job_name in list(left): + job = self._jobs[job_name] + pid = job.pid + process = psutil.Process(pid) + self._job_states[job_name] = PROCESS_STATUS_TO_JOB_STATE.get( + process.status(), JobState.NOT_FOUND + ) + + for job_name in list(left): + state = self._job_states[job_name] + if state in check_status: + raise JobException( + run_name=self.run_name, + worker_type=job_name.split("/")[0], + host="local", + reason=state, + ) + if state in remove_status: + logger.info(f"Job {job_name} is {state}.(Removed)") + left.remove(job_name) + to_remove.append(job_name) + + if update: + for k in to_remove: + self._jobs.pop(k) + worker_type = k.split("/")[0] + assert worker_type in self._job_counter + self._job_counter[worker_type] -= 1 + if self._job_counter[worker_type] <= 0: + self._job_counter.pop(worker_type) + + time.sleep(2) + + +def main_local(): + cfg, _ = parse_cli_args(sys.argv[2:]) + name_resolve.reconfigure(cfg.cluster.name_resolve) + name_resolve.clear_subtree( + names.trial_root(experiment_name=cfg.experiment_name, trial_name=cfg.trial_name) + ) + alloc_mode = AllocationMode.from_str(cfg.allocation_mode) + + launcher = LocalLauncher(cfg.experiment_name, cfg.trial_name, cfg.cluster.fileroot) + + server_cmd = [] + server_addrs = [] + if alloc_mode.type_ == AllocationType.DECOUPLED_SGLANG: + base_seed = cfg.sglang.random_seed + cfg.sglang = to_structured_cfg(cfg.sglang, SGLangConfig) + ports = find_free_ports(alloc_mode.gen_dp_size * 2, port_range=(10000, 50000)) + host_ip = gethostip() + host = "localhost" if not cfg.sglang.enable_metrics else host_ip + for i in range(alloc_mode.gen_dp_size): + cfg.sglang.random_seed = base_seed + i + cmd = SGLangConfig.build_cmd( + cfg.sglang, + host=host, + tp_size=alloc_mode.gen_tp_size, + base_gpu_id=0, + port=ports[i * 2], + dist_init_addr=f"localhost:{ports[i*2+1]}", + ) + server_cmd.append(cmd) + server_addrs.append(f"{host}:{ports[i * 2]}") + + # Launch inference servers. + launcher.submit_array( + job_name="llm_server", + cmd=server_cmd, + count=alloc_mode.gen_dp_size, + gpu=alloc_mode.gen_pp_size * alloc_mode.gen_tp_size, + ) + logger.info( + f"LLM inference server launched at: AREAL_LLM_SERVER_ADDRS={','.join(server_addrs)}" + ) + + # Launch trainer entrypoint + if alloc_mode.type_ == AllocationType.COLOCATE or not cfg.server_only: + launcher.submit( + job_name="trainer", + cmd=f"torchrun --nnodes 1 --nproc-per-node {alloc_mode.train_world_size} --master-addr localhost --master-port {find_free_ports(1, (10000, 50000))[0]} {' '.join(sys.argv[1:])}", + gpu=alloc_mode.train_world_size, + env_vars=dict(AREAL_LLM_SERVER_ADDRS=",".join(server_addrs)), + ) + + try: + launcher.wait( + check_status=( + JobState.CANCELLED, + JobState.FAILED, + JobState.NOT_FOUND, + JobState.COMPLETED, + ), + remove_status=(), + ) + except (KeyboardInterrupt, JobException, TimeoutError) as e: + launcher.stop_all("SIGTERM") + raise e + + +if __name__ == "__main__": + main_local() diff --git a/areal/launcher/ray.py b/areal/launcher/ray.py new file mode 100644 index 0000000000..3df5b2861e --- /dev/null +++ b/areal/launcher/ray.py @@ -0,0 +1,408 @@ +import getpass +import importlib.util +import os +import pathlib +import sys +import time +from typing import Dict, List, Optional + +import ray +import ray.exceptions +from ray.runtime_env import RuntimeEnv +from ray.util.placement_group import PlacementGroup +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from areal.api.cli_args import ( + ClusterSpecConfig, + LauncherConfig, + SGLangConfig, + parse_cli_args, + to_structured_cfg, +) +from areal.api.io_struct import AllocationMode, AllocationType +from areal.utils.launcher import ( + get_env_vars, + validate_config_for_distributed_launcher, + wait_sglang_server_addrs, +) +from areal.utils.ray import get_placement_group_master_ip_and_port +from realhf.base import logging, name_resolve, names +from realhf.scheduler.client import JobException, JobState + +logger = logging.getLogger("RayLauncher") + +RAY_WAIT_CHECK_TIME_INTERVAL = 5 # seconds +DEFAULT_MAIN_FUNC_NAME = "main" + + +def run_func(file_path, function_name, *args, **kwargs): + # Convert the file path to a module name + module_name = file_path.replace("/", "_").replace(".", "_") + + # Load the module from file path + spec = importlib.util.spec_from_file_location(module_name, file_path) + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + + # Get the function and execute it + try: + function = getattr(module, function_name) + except AttributeError as e: + raise ValueError( + f"Function '{function_name}' not found in module '{module_name}'. " + f"Please ensure the name of the main function in your entry point " + f"is '{function_name}'." + ) from e + return function(*args, **kwargs) + + +class RayLauncher: + def __init__(self, experiment_name: str, trial_name: str, fileroot: str): + self.experiment_name = experiment_name + self.trial_name = trial_name + self.fileroot = fileroot + + # job_name to ray future + self.jobs = {} + + @property + def run_name(self): + return f"{self.experiment_name}_{self.trial_name}" + + def log_path_of(self, job_name: str) -> str: + log_path = f"{self.fileroot}/logs/{getpass.getuser()}/{self.experiment_name}/{self.trial_name}" + os.makedirs(log_path, exist_ok=True) + return os.path.join(log_path, f"{job_name}.log") + + def submit( + self, + job_name: str, + file_path: str, + func_name: str, + args: List[str], # arguments to pass to the function + gpus: int, + cpus: int, + mem: int, # MB + env_vars: Optional[Dict] = None, + placement_group: Optional[PlacementGroup] = None, + bundle_index: int = -1, + kwargs: Optional[ + Dict[str, str] + ] = None, # keyword arguments to pass to the function + ): + if kwargs is None: + kwargs = {} + runtime_env = RuntimeEnv( + env_vars=env_vars or dict(), + ) + scheduling_strategy = ( + PlacementGroupSchedulingStrategy( + placement_group=placement_group, + placement_group_bundle_index=bundle_index, + placement_group_capture_child_tasks=True, + ) + if placement_group is not None + else "DEFAULT" + ) + future = ray.remote( + num_cpus=cpus, + num_gpus=gpus, + memory=mem * 1024 * 1024, # Convert MB to bytes + runtime_env=runtime_env, + scheduling_strategy=scheduling_strategy, + )(run_func).remote(file_path, func_name, *args, **kwargs) + self.jobs[job_name] = future + return future + + def submit_array( + self, + job_name: str, + file_path: str, + func_name: str, + count: int, + nodes: int, + list_args: List[List], + gpus_per_task: int, + cpus_per_task: int, + mem_per_task: int, # MB + list_kwargs: List[Dict] | None = None, + env_vars: Optional[Dict] = None, + amend_torch_dist_env: bool = False, + ): + """Submit an array of jobs to Ray with ray placement groups. + + Note: Here we use `ray.remote` instead of `ray job submit` since `ray job submit` + does not support placement groups, and can not specify which node to run the job on. + Therefore we could not know the IP address of jobs for torch distributed initialization. + """ + + if count % nodes != 0: + raise ValueError( + f"Count {count} is not divisible by nodes {nodes}. " + "Please ensure that count is a multiple of nodes." + ) + assert ( + len(list_args) == count + ), f"Length of list_args {len(list_args)} does not match count {count}." + if list_kwargs is not None: + assert ( + len(list_kwargs) == count + ), f"Length of list_kwargs {len(list_kwargs)} does not match count {count}." + + tasks_per_node = count // nodes + gpus_per_node = gpus_per_task * tasks_per_node + cpus_per_node = cpus_per_task * tasks_per_node + mem_per_node = mem_per_task * tasks_per_node + + placement_group = ray.util.placement_group( + bundles=[ + { + "CPU": cpus_per_node, + "GPU": gpus_per_node, + "memory": mem_per_node * 1024 * 1024, # Convert MB to bytes + } + ] + * nodes, + strategy="STRICT_SPREAD", + ) + try: + ray.get(placement_group.ready(), timeout=30) + except ray.exceptions.GetTimeoutError as e: + logger.error( + "Ray placement group timeout, please check if the resource requirement " + "for your experiment exceeds the available resources in the cluster. \n" + f"ray.nodes(): {ray.nodes()} \n" + f"Placement Group bundles: " + f"cpus_per_node={cpus_per_node}, gpus_per_node={gpus_per_node}, " + f"mem_per_node={mem_per_node}MB, nodes={nodes}" + ) + raise e + + if amend_torch_dist_env: + host_ip, port = get_placement_group_master_ip_and_port(placement_group) + logger.info( + f"Amend torch distributed env vars: " + f"MASTER_ADDR={host_ip}, PORT={port}" + ) + + futures = [] + for i in range(count): + args = list_args[i] + kwargs = list_kwargs[i] if list_kwargs is not None else {} + + # manage environment variables + env_vars = env_vars or {} + if "CUDA_VISIBLE_DEVICES" in env_vars: + logger.warning( + "Setting CUDA_VISIBLE_DEVICES before running ray jobs may result in unexpected behavior." + ) + + node_id = i // tasks_per_node + _env_vars = { + **env_vars, + } + + if amend_torch_dist_env: + assert gpus_per_task == 1 + # NOTE: Here we only provide environment variables for torch distributed + # initialization, and LOCAL_RANK for torch.device. + # Other environment variables automatically set by torchrun are not set, and + # they should be never accessed in trainer code. + _env_vars.update( + { + "RANK": str(i), + "WORLD_SIZE": str(count), + # Ray will automatically isolate CUDA_VISIBLE_DEVICES for each GPU + "LOCAL_RANK": "0", + "MASTER_ADDR": str(host_ip), + "MASTER_PORT": str(port), + } + ) + future = self.submit( + job_name=f"{job_name}:{i}", + file_path=file_path, + func_name=func_name, + args=args, + gpus=gpus_per_task, + cpus=cpus_per_task, + mem=mem_per_task, + env_vars=_env_vars, + placement_group=placement_group, + bundle_index=node_id, + kwargs=kwargs, + ) + futures.append(future) + + return futures + + def stop(self, job_name: str, force: bool = False): + """Stop a job by name.""" + if job_name in self.jobs: + future = self.jobs[job_name] + try: + ray.cancel(future, force=force) + except Exception as e: + logger.error(f"Failed to cancel job {job_name}: {e}") + return + self.jobs.pop(job_name, None) + logger.info(f"Job {job_name} stopped.") + else: + logger.warning(f"Job {job_name} not found in running jobs.") + + def stop_all(self, force: bool = False): + """Stop all jobs.""" + for job_name in list(self.jobs.keys()): + self.stop(job_name, force=force) + logger.info("All jobs stopped.") + self.jobs.clear() + + def wait( + self, check_status=(JobState.FAILED,), remove_status=(JobState.COMPLETED,) + ): + """Check every RAY_WAIT_CHECK_TIME_INTERVAL seconds for the status of all jobs. + If a ray job returns, its status changes to JobState.COMPLETED. + If a ray job failed, its status changes to JobState.FAILED. + If any job is in check_status, stop all jobs at once. + If any job is in remove status, remove them from job list. + Return if all jobs are removed from job list, or some job is in check status. + """ + for status in list(check_status) + list(remove_status): + assert status in [ + JobState.COMPLETED, + JobState.FAILED, + ], "In RayLauncher.wait, we only check completed or failed jobs." + logger.info(f"Waiting for {len(self.jobs)} jobs.") + while self.jobs: + job_status = {} + for job_name, future in list(self.jobs.items()): + try: + r = ray.get(future, timeout=0.1) + logger.info(f"Job {job_name} completed with result: {r}") + job_status[job_name] = JobState.COMPLETED + except ray.exceptions.RayTaskError as e: + logger.error(f"Job {job_name} failed with error: {e}.") + job_status[job_name] = JobState.FAILED + except ray.exceptions.GetTimeoutError: + continue + + for job_name, status in job_status.items(): + if status in check_status: + logger.info(f"Job {job_name} is {status}, stopping all jobs.") + self.stop_all(force=True) + return + if status in remove_status: + logger.info(f"Job {job_name} is {status}, removed.") + self.jobs.pop(job_name) + + time.sleep(RAY_WAIT_CHECK_TIME_INTERVAL) + + +def ray_main(): + # usage: python -m areal.launcher.ray --config [] + ray.init() + config, config_file = parse_cli_args(sys.argv[2:]) + config.launcher = to_structured_cfg(config.launcher, LauncherConfig) + config.cluster = to_structured_cfg(config.cluster, ClusterSpecConfig) + config.sglang = to_structured_cfg(config.sglang, SGLangConfig) + validate_config_for_distributed_launcher(config) + + name_resolve.reconfigure(config.cluster.name_resolve) + name_resolve.clear_subtree( + names.trial_root( + experiment_name=config.experiment_name, trial_name=config.trial_name + ) + ) + + n_nodes = config.cluster.n_nodes + n_gpus_per_node = config.cluster.n_gpus_per_node + launcher = RayLauncher( + experiment_name=config.experiment_name, + trial_name=config.trial_name, + fileroot=config.cluster.fileroot, + ) + allocation_mode = config.allocation_mode + allocation_mode = AllocationMode.from_str(allocation_mode) + sglang_addrs = [] + n_sglang_nodes = 0 + if allocation_mode.type_ == AllocationType.DECOUPLED_SGLANG: + # Launcher should launch SGLang servers according to allocation mode. + sglang_tp_size = allocation_mode.gen_tp_size + n_sglang_servers = allocation_mode.gen_dp_size + n_sglang_nodes = allocation_mode.gen_world_size // n_gpus_per_node + + base_seed = config.sglang.random_seed + sglang_args_list = [ + [sys.argv[2:] + [f"sglang.random_seed={base_seed + i}"]] + for i in range(n_sglang_servers) + ] + sglang_entry_point = str( + pathlib.Path(__file__).resolve().parent.joinpath("sglang_server.py") + ) + launcher.submit_array( + job_name="llm_server", + file_path=sglang_entry_point, + func_name=DEFAULT_MAIN_FUNC_NAME, + count=n_sglang_servers, + nodes=n_sglang_nodes, + list_args=sglang_args_list, + gpus_per_task=sglang_tp_size, + cpus_per_task=config.launcher.inference_server_cpus_per_gpu + * sglang_tp_size, + mem_per_task=config.launcher.inference_server_mem_per_gpu * sglang_tp_size, + env_vars=get_env_vars( + config.cluster.cluster_name, + config.launcher.inference_server_env_vars, + ), + ) + # Get SGLang server addresses via name_resolve + try: + sglang_addrs = wait_sglang_server_addrs( + config.experiment_name, + config.trial_name, + n_sglang_servers, + ) + except TimeoutError as e: + launcher.stop_all(force=True) + raise e + + trainer_n_nodes = n_nodes - n_sglang_nodes + trainer_entry_point = sys.argv[1] + n_trainer_processes = trainer_n_nodes * config.cluster.n_gpus_per_node + trainer_args_list = [[sys.argv[2:]] for _ in range(n_trainer_processes)] + if not config.server_only: + # In ray, we launch trainer in the granularity of processes (1 GPU per process) + # We amend environment variable similar to torchrun to ensure correct initialization of + # torch distributed. + launcher.submit_array( + job_name="trainer", + file_path=trainer_entry_point, + func_name=DEFAULT_MAIN_FUNC_NAME, + count=trainer_n_nodes * config.cluster.n_gpus_per_node, + nodes=trainer_n_nodes, + list_args=trainer_args_list, + gpus_per_task=1, + cpus_per_task=config.launcher.trainer_cpus_per_gpu, + mem_per_task=config.launcher.trainer_mem_per_gpu, + env_vars=dict( + **get_env_vars( + config.cluster.cluster_name, + config.launcher.trainer_env_vars, + ), + AREAL_LLM_SERVER_ADDRS=",".join(sglang_addrs), + ), + amend_torch_dist_env=True, + ) + + try: + launcher.wait(check_status=(JobState.COMPLETED, JobState.FAILED)) + except (KeyboardInterrupt, JobException, TimeoutError) as e: + launcher.stop_all(force=True) + raise e + + +if __name__ == "__main__": + # usage: python -m areal.launcher.ray \ + # --config [] \ + # launcher.ray.main_func_name= + ray_main() diff --git a/areal/launcher/sglang_server.py b/areal/launcher/sglang_server.py new file mode 100644 index 0000000000..b5242a5d48 --- /dev/null +++ b/areal/launcher/sglang_server.py @@ -0,0 +1,155 @@ +import os +import subprocess +import sys +import time +import uuid +from typing import Optional + +import requests + +from areal.api.cli_args import ( + ClusterSpecConfig, + NameResolveConfig, + SGLangConfig, + parse_cli_args, + to_structured_cfg, +) +from areal.api.io_struct import AllocationMode, AllocationType +from areal.utils.launcher import TRITON_CACHE_PATH +from areal.utils.network import find_free_ports, gethostip +from realhf.base import logging, name_resolve, names + +logger = logging.getLogger("SGLangServer Wrapper") + + +def launch_server_cmd(command: str) -> subprocess.Popen: + """ + Execute a shell command and return its process handle. + """ + # Replace newline continuations and split the command string. + command = command.replace("\\\n", " ").replace("\\", " ") + parts = command.split() + _env = os.environ.copy() + # To avoid DirectoryNotEmpty error caused by triton + triton_cache_path = _env.get("TRITON_CACHE_PATH", TRITON_CACHE_PATH) + unique_triton_cache_path = os.path.join(triton_cache_path, str(uuid.uuid4())) + _env["TRITON_CACHE_PATH"] = unique_triton_cache_path + return subprocess.Popen( + parts, + text=True, + env=_env, + stdout=sys.stdout, + stderr=subprocess.STDOUT, + ) + + +def wait_for_server(base_url: str, timeout: Optional[int] = None) -> None: + """Wait for the server to be ready by polling the /v1/models endpoint. + + Args: + base_url: The base URL of the server + timeout: Maximum time to wait in seconds. None means wait forever. + """ + start_time = time.time() + while True: + try: + response = requests.get( + f"{base_url}/v1/models", + headers={"Authorization": "Bearer None"}, + ) + if response.status_code == 200: + time.sleep(5) + break + + if timeout and time.time() - start_time > timeout: + raise TimeoutError("Server did not become ready within timeout period") + except requests.exceptions.RequestException: + time.sleep(1) + + +class SGLangServerWrapper: + def __init__( + self, + experiment_name: str, + trial_name: str, + sglang_config: SGLangConfig, + tp_size: int, + n_gpus_per_node: int, + ): + self.experiment_name = experiment_name + self.trial_name = trial_name + self.config = sglang_config + self.tp_size = tp_size + self.server_process = None + self.n_gpus_per_node = n_gpus_per_node + + def run(self): + gpus_per_server = len(os.getenv("CUDA_VISIBLE_DEVICES").split(",")) + server_local_idx = ( + int(os.getenv("CUDA_VISIBLE_DEVICES").split(",")[0]) // gpus_per_server + ) + n_servers_per_node = max(1, self.n_gpus_per_node // gpus_per_server) + ports_per_server = 40000 // n_servers_per_node + port_range = ( + server_local_idx * ports_per_server + 10000, + (server_local_idx + 1) * ports_per_server + 10000, + ) + server_port, dist_init_port = find_free_ports(2, port_range) + + dist_init_addr = f"localhost:{dist_init_port}" + host_ip = gethostip() + + cmd = SGLangConfig.build_cmd( + self.config, + tp_size=self.tp_size, + base_gpu_id=0, + host=host_ip, + port=server_port, + dist_init_addr=dist_init_addr, + ) + self.server_process = launch_server_cmd(cmd) + wait_for_server(f"http://{host_ip}:{server_port}") + + name = names.gen_servers(self.experiment_name, self.trial_name) + name_resolve.add_subentry(name, f"{host_ip}:{server_port}") + + logger.info(f"SGLang server launched at: http://{host_ip}:{server_port}") + return_code = self.server_process.wait() + logger.info( + f"SGLang server at http://{host_ip}:{server_port} exits, returncode={return_code}" + ) + + def __del__(self): + if self.server_process and self.server_process.poll() is None: + logger.info("Terminating SGLang server process...") + self.server_process.terminate() + self.server_process.wait() + logger.info("SGLang server process terminated.") + + +def main(argv): + config, _ = parse_cli_args(argv) + config.sglang = to_structured_cfg(config.sglang, SGLangConfig) + config.cluster = to_structured_cfg(config.cluster, ClusterSpecConfig) + config.cluster.name_resolve = to_structured_cfg( + config.cluster.name_resolve, NameResolveConfig + ) + name_resolve.reconfigure(config.cluster.name_resolve) + + allocation_mode = config.allocation_mode + allocation_mode = AllocationMode.from_str(allocation_mode) + assert allocation_mode.type_ == AllocationType.DECOUPLED_SGLANG + tp_size = allocation_mode.gen_tp_size + + sglang_server = SGLangServerWrapper( + config.experiment_name, + config.trial_name, + config.sglang, + tp_size, + n_gpus_per_node=config.cluster.n_gpus_per_node, + ) + sglang_server.run() + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/areal/launcher/slurm.py b/areal/launcher/slurm.py new file mode 100644 index 0000000000..2cd4c43c20 --- /dev/null +++ b/areal/launcher/slurm.py @@ -0,0 +1,528 @@ +import getpass +import os +import re +import subprocess +import sys +import time +from typing import Dict, List, Optional, Tuple + +import realhf.base.logging as logging +from areal.api.cli_args import ( + ClusterSpecConfig, + LauncherConfig, + SGLangConfig, + parse_cli_args, + to_structured_cfg, +) +from areal.api.io_struct import AllocationMode, AllocationType +from areal.utils.launcher import ( + get_env_vars, + validate_config_for_distributed_launcher, + wait_sglang_server_addrs, +) +from areal.utils.slurm import ( + APPTAINER_CMD_TEMPLATE, + SBATCH_SCRIPT_TEMPLATE, + SRUN_CMD_TEMPLATE, + cancel_jobs, + query_jobs, +) +from realhf.base import logging, name_resolve, names +from realhf.scheduler.client import JobException, JobInfo, JobState + +logger = logging.getLogger("SlurmLauncher") + +SLURM_WAIT_CHECK_TIME_INTERVAL = 5 + + +class SlurmLauncher: + def __init__( + self, experiment_name: str, trial_name: str, fileroot: str, container_type: str + ): + self.experiment_name = experiment_name + self.trial_name = trial_name + self.fileroot = fileroot + self.container_type = container_type + + # slurm_job_id -> JobInfo + self.jobs: Dict[int, JobInfo] = {} + self.job_names = [] + + @property + def run_name(self) -> str: + """Returns the run name of this launcher.""" + return f"{self.experiment_name}_{self.trial_name}" + + def slurm_name(self, job_name: str) -> str: + """Returns the slurm name of a job.""" + return f"{self.experiment_name}_{self.trial_name}:{job_name}" + + def log_path_of(self, job_name: str) -> str: + log_path = f"{self.fileroot}/logs/{getpass.getuser()}/{self.experiment_name}/{self.trial_name}" + os.makedirs(log_path, exist_ok=True) + return os.path.join(log_path, f"{job_name}.log") + + def sbatch_path_of(self, job_name: str) -> str: + sbatch_path = f"{self.fileroot}/logs/{getpass.getuser()}/{self.experiment_name}/{self.trial_name}" + os.makedirs(sbatch_path, exist_ok=True) + return os.path.join(sbatch_path, f"{job_name}.sh") + + def submit(self, job_name, cmd, **kwargs): + """Submits and launch a job with SBATCH. + + Args: + cmd (str or List[str]): The core command to be executed. + """ + return self.submit_array(job_name, cmd, count=1, **kwargs) + + def find_job_id(self, job_name: str): + job_name = self.slurm_name(job_name) + for job_id, job_info in self.jobs.items(): + if job_info.name == job_name: + return job_id + return None + + def submit_array( + self, + job_name: str, + cmd: List[str] | str, + count: int, + nodes: int, + n_gpus_per_node: int, + cpus_per_task: int, + mem_per_task: int, # MB + container_image: str, + srun_additional_args: str = "", + container_mounts: Optional[str] = None, + env_vars: Optional[Dict] = None, + nodelist: Optional[str] = None, + exclude: Optional[str] = None, + ): + """Submits and launch a job array with SBATCH. + Note that a job array has one (unique) slurm name, and one (unique) slurm id. + + Args: + job_name (str): The job name of the job array. The actual slurm name will be + `_:`. + cmd (str or List[str]): The core command to be executed. + count (int): The number of jobs in the array. + """ + assert job_name not in self.job_names, ( + f"Job {job_name} is already submitted. " + "Please use a different job name or stop the existing job." + ) + if isinstance(cmd, str): + cmd = [cmd] + assert len(cmd) == count, ( + f"Command length {len(cmd)} does not match the job count {count}. " + "Please provide a command for each job in the array." + ) + assert count % nodes == 0, ( + f"Job count {count} must be divisible by the number of nodes {nodes}. " + "Please adjust the job count or the number of nodes." + ) + ntasks_per_node = count // nodes + assert n_gpus_per_node % ntasks_per_node == 0, ( + "GPUs must be evenly distributed across tasks. " + f"Current #GPUs per node {n_gpus_per_node}, #tasks per node {ntasks_per_node}." + ) + + mem_per_cpu = mem_per_task // cpus_per_task # MB per CPU + mem_per_node = ( + mem_per_task * count // nodes + 1024 * 10 + ) # make sure slurm does not run out of resources + + sbatch_options = [ + f"--job-name={self.slurm_name(job_name)}", + f"--output={self.log_path_of(job_name)}", + "--open-mode=append", + "--no-requeue", + f"--nodes={nodes}-{nodes}", + f"--ntasks-per-node={ntasks_per_node}", + f"--gres=gpu:{n_gpus_per_node}", + f"--cpus-per-task={cpus_per_task}", + f"--mem={mem_per_node}M", + ] + + if nodelist: + sbatch_options.append(f"--nodelist={nodelist}") + if exclude: + sbatch_options.append(f"--exclude={exclude}") + + sbatch_options_str = "\n".join([f"#SBATCH {opt}" for opt in sbatch_options]) + + if env_vars is None: + env_vars = dict() + n_gpus_per_task = n_gpus_per_node // ntasks_per_node + assert ( + "CUDA_VISIBLE_DEVICES" not in env_vars + ), "CUDA_VISIBLE_DEVICES should be automatically resolved by Launcher instead of manually assigned." + + srun_cmds = [] + for i in range(count): + # resolve CUDA_VISIBLE_DEVICES for each task + gpu_id_start = (i % ntasks_per_node) * n_gpus_per_task + gpu_id_end = ((i % ntasks_per_node) + 1) * n_gpus_per_task + node_id = i // ntasks_per_node + _env_vars = { + **env_vars, + "CUDA_VISIBLE_DEVICES": ",".join( + str(x) for x in range(gpu_id_start, gpu_id_end) + ), + } + # Prepare the command for each job in the array + job_cmd = cmd[i] + + if self.container_type == "apptainer": + env_string = " ".join( + "--env {}={}".format(k, v) for k, v in _env_vars.items() + ) + apptainer_cmd = APPTAINER_CMD_TEMPLATE.format( + container_mounts=container_mounts or "", + container_env_strings=env_string, + container_image=container_image, + cmd=job_cmd, + ) + srun_cmd = SRUN_CMD_TEMPLATE.format( + additional_args=srun_additional_args, + nodes=1, + ntasks=1, + node_id=node_id, + n_gpus_per_node=n_gpus_per_task, + cpus_per_task=cpus_per_task, + mem_per_cpu=mem_per_cpu, + cmd=apptainer_cmd, + ) + elif self.container_type == "none": + env_string = "--export=" + ",".join( + "{}={}".format(k, v) for k, v in _env_vars.items() + ) + srun_additional_args = srun_additional_args + " " + env_string + srun_cmd = SRUN_CMD_TEMPLATE.format( + additional_args=srun_additional_args, + nodes=1, + ntasks=1, + node_id=node_id, + n_gpus_per_node=n_gpus_per_task, + cpus_per_task=cpus_per_task, + mem_per_cpu=mem_per_cpu, + cmd=job_cmd, + ) + else: + raise ValueError( + f"Unsupported container type: {self.container_type}. " + "Supported types are 'apptainer' and 'none'." + ) + srun_cmds.append(srun_cmd) + + srun_cmds = "\n".join(srun_cmds) + sbatch_script = SBATCH_SCRIPT_TEMPLATE.format( + sbatch_options=sbatch_options_str, + srun_additional_args=srun_additional_args, + srun_cmds=srun_cmds, + ) + sbatch_file_path = self.sbatch_path_of(f"{job_name}") + with open(sbatch_file_path, "w") as f: + f.write(sbatch_script) + + # Submit the job + try: + output = ( + subprocess.check_output(["sbatch", sbatch_file_path]) + .decode("utf-8") + .strip() + ) + logger.info( + f"Submitted Slurm job {self.slurm_name(job_name)} to scheduler. To check the output, run \n\t`tail -f {self.log_path_of(job_name)}`." + ) + except subprocess.CalledProcessError as e: + logger.warning( + f"Failed to submit job {self.slurm_name(job_name)}. " + f"For debugging, please make sure your sbatch command works " + f"and check generated sbatch file on {sbatch_file_path}." + ) + logger.error(f"Error message: {e}") + return + + match = re.search(r"Submitted batch job (\d+)", output) + slurm_job_id = int(match.group(1)) if match else None + if slurm_job_id is None: + logger.warning( + f"Failed to obtain job id for job {self.slurm_name(job_name)}. " + f"sbatch output: {output}" + ) + return + + assert isinstance(slurm_job_id, int) + self.jobs[slurm_job_id] = JobInfo( + name=self.slurm_name(job_name), + state=JobState.PENDING, + slurm_id=slurm_job_id, + ) + self._update_all() + + def stop(self, job_name, force=False): + """Stops a running job. + + Raises exception if there is no such job, but passes if the job + has stopped either successfully or not. + + Args: + job_name: The job name of the job array to stop. + The actual slurm job name will be `_:`. + """ + signal = "SIGKILL" if force else "SIGTERM" + job_id = self.find_job_id(job_name) + if not job_id: + return + return cancel_jobs(slurm_ids=[job_id], signal=signal) + + def stop_all(self, force=False): + """Stops all running jobs.""" + signal = "SIGKILL" if force else "SIGTERM" + return cancel_jobs(slurm_ids=list(self.jobs.keys()), signal=signal) + + def find(self, job_name) -> JobInfo | None: + """Gets the status of a job of this job. + + Args: + job_name: The job name of the job array to find. + The actual slurm job name will be `_:`. + + Returns: + A JobInfo if the job is found, or None otherwise. + """ + self._update_all() + job_id = self.find_job_id(job_name) + return self.jobs[job_id] if job_id else None + + def find_all(self, job_name_regex=".*") -> List[JobInfo]: + """Finds jobs. + + Args: + job_name_regex: job name regex. + + Returns: + A list of found JobInfo. + """ + self._update_all() + infos = [] + for r in self.jobs.values(): + job_name = r.name.split(":")[-1] # Extract the job name from slurm name + if re.fullmatch(job_name_regex, job_name): + infos.append(r) + return infos + + def _find_job_with_status( + self, + status: List[JobState], + ) -> List[JobInfo]: + """Finds jobs with the given status. + + Args: + status: A list of JobState to filter jobs. + + Returns: + A list of JobInfo with the given status. + """ + self._update_all() + return [r for r in self.jobs.values() if r.state in status] + + def wait( + self, + timeout=None, + check_status: Tuple[JobState, ...] = ( + JobState.CANCELLED, + JobState.FAILED, + JobState.NOT_FOUND, + ), + remove_status: Tuple[JobState, ...] = (JobState.COMPLETED,), + update=False, + ): + """Waits until all jobs submitted via this client instance finish.""" + # begin wait + deadline = None if timeout is None else time.time() + timeout + + num_jobs_left = len(self.jobs) + left = list(self.jobs.keys()) + logger.info( + f"Waiting for {num_jobs_left} jobs. Jobs IDs: " + f"{','.join(sorted([str(x.slurm_id) for x in self.jobs.values()]))}." + ) + while len(left) > 0: + if len(left) < num_jobs_left: + num_jobs_left = len(left) + logger.info(f"Waiting for {num_jobs_left} jobs.") + if deadline is not None and time.time() > deadline: + raise TimeoutError( + f"Timeout waiting for {num_jobs_left} jobs. Job ID: " + f"{','.join(sorted([str(x.slurm_id) for x in self.jobs.values()]))}." + ) + self._update_all() + left = list(self.jobs.keys()) + for slurm_id in list(left): + slurm_info = self.jobs[slurm_id] + if slurm_info.slurm_id is None: + continue + if slurm_info.state in check_status: + raise JobException( + run_name=self.run_name, + worker_type=slurm_info.name, + host=slurm_info.host, + reason=slurm_info.state, + ) + if slurm_info.state in remove_status: + logger.info( + f"Job {slurm_info.name} is {slurm_info.state}. (Removed)" + ) + left.remove(slurm_id) + if update: + self.jobs.pop(slurm_info.slurm_id) + time.sleep(SLURM_WAIT_CHECK_TIME_INTERVAL) + + def _update_all(self): + """Updates the status of all jobs.""" + try: + slurm_infos = query_jobs(slurm_ids=list(self.jobs.keys())) + for slurm_info in slurm_infos: + assert slurm_info.slurm_id is not None + self.jobs[slurm_info.slurm_id] = slurm_info + except subprocess.CalledProcessError: + logger.warning( + "Calling squeue failed. Check slurm manually if you continue to see this warning." + ) + + +def slurm_main(): + config, _ = parse_cli_args(sys.argv[2:]) + config.launcher = to_structured_cfg(config.launcher, LauncherConfig) + config.cluster = to_structured_cfg(config.cluster, ClusterSpecConfig) + config.sglang = to_structured_cfg(config.sglang, SGLangConfig) + validate_config_for_distributed_launcher(config) + + name_resolve.reconfigure(config.cluster.name_resolve) + name_resolve.clear_subtree( + names.trial_root( + experiment_name=config.experiment_name, trial_name=config.trial_name + ) + ) + + n_nodes = config.cluster.n_nodes + n_gpus_per_node = config.cluster.n_gpus_per_node + + launcher = SlurmLauncher( + experiment_name=config.experiment_name, + trial_name=config.trial_name, + fileroot=config.cluster.fileroot, + container_type=config.launcher.slurm.container_type, + ) + allocation_mode = config.allocation_mode + allocation_mode = AllocationMode.from_str(allocation_mode) + sglang_cmds = [] + sglang_addrs = [] + n_sglang_nodes = 0 + if allocation_mode.type_ == AllocationType.DECOUPLED_SGLANG: + # Launcher should launch SGLang servers according to allocation mode. + sglang_tp_size = allocation_mode.gen_tp_size + n_sglang_servers = allocation_mode.gen_dp_size + n_sglang_nodes = allocation_mode.gen_world_size // n_gpus_per_node + + base_seed = config.sglang.random_seed + sglang_server_cmd_template = f"python3 -m areal.launcher.sglang_server {' '.join(sys.argv[2:])} sglang.random_seed={{seed}}" + for i in range(n_sglang_servers): + sglang_cmd = sglang_server_cmd_template.format( + seed=base_seed + i, + ) + sglang_cmds.append(sglang_cmd) + + launcher.submit_array( + job_name="llm_server", + cmd=sglang_cmds, + count=n_sglang_servers, + nodes=n_sglang_nodes, + n_gpus_per_node=config.cluster.n_gpus_per_node, + cpus_per_task=config.launcher.inference_server_cpus_per_gpu + * sglang_tp_size, + mem_per_task=config.launcher.inference_server_mem_per_gpu * sglang_tp_size, + srun_additional_args=config.launcher.slurm.srun_additional_args, + container_image=config.launcher.slurm.inference_server_image, + container_mounts=config.launcher.slurm.mount, + env_vars=get_env_vars( + config.cluster.cluster_name, + config.launcher.inference_server_env_vars, + ), + ) + # Get SGLang server addresses by name resolve + try: + sglang_addrs = wait_sglang_server_addrs( + config.experiment_name, + config.trial_name, + n_sglang_servers, + ) + except TimeoutError as e: + launcher.stop_all(force=True) + raise e + + trainer_n_nodes = n_nodes - n_sglang_nodes + # Here $head_node_ip is the IP address of the first node in the job array. + # $trainer_port is a free port on the head node. + # Both of them are obtained in by the SBATCH script. + trainer_cmd_template = ( + f"torchrun --nnodes={{nnodes}} --nproc-per-node={{nproc_per_node}} --node-rank {{node_rank}} " + f"--master-addr $head_node_ip --master-port $trainer_port {' '.join(sys.argv[1:])}" + ) + + trainer_cmds = [] + for i in range(trainer_n_nodes): + # In slurm, we launch trainer in the granularity of nodes with torchrun command. + trainer_cmds.append( + trainer_cmd_template.format( + nnodes=trainer_n_nodes, + nproc_per_node=config.cluster.n_gpus_per_node, + node_rank=i, + ) + ) + + if not config.server_only: + # launch trainers + launcher.submit_array( + job_name="trainer", + cmd=trainer_cmds, + count=trainer_n_nodes, + nodes=trainer_n_nodes, + n_gpus_per_node=config.cluster.n_gpus_per_node, + cpus_per_task=config.launcher.trainer_cpus_per_gpu + * config.cluster.n_gpus_per_node, + mem_per_task=config.launcher.trainer_mem_per_gpu + * config.cluster.n_gpus_per_node, + container_image=config.launcher.slurm.trainer_image, + srun_additional_args=config.launcher.slurm.srun_additional_args, + container_mounts=config.launcher.slurm.mount, + env_vars=dict( + **get_env_vars( + config.cluster.cluster_name, + config.launcher.trainer_env_vars, + ), + AREAL_LLM_SERVER_ADDRS=",".join(sglang_addrs), + ), + ) + + try: + launcher.wait( + check_status=( + JobState.CANCELLED, + JobState.FAILED, + JobState.NOT_FOUND, + JobState.COMPLETED, + ), + remove_status=(), + ) + except (KeyboardInterrupt, JobException, TimeoutError) as e: + launcher.stop_all(force=True) + raise e + + +if __name__ == "__main__": + # usage: python -m areal.launcher.slurm \ + # --config [] + slurm_main() diff --git a/areal/tests/test_fsdp_engine_nccl.py b/areal/tests/test_fsdp_engine_nccl.py new file mode 100644 index 0000000000..25c643423b --- /dev/null +++ b/areal/tests/test_fsdp_engine_nccl.py @@ -0,0 +1,124 @@ +import os +import subprocess +import sys +import time + +import pytest +import requests + +from areal.api.cli_args import ( + InferenceEngineConfig, + OptimizerConfig, + SGLangConfig, + TrainEngineConfig, +) +from areal.api.io_struct import AllocationMode, FinetuneSpec, WeightUpdateMeta +from areal.engine.fsdp_engine import FSDPEngine +from areal.engine.sglang_remote import RemoteSGLangEngine +from realhf.base import network + +EXPR_NAME = "test_fsdp_engine_nccl" +TRIAL_NAME = "trial_nccl" +MODEL_PATH = "/storage/testing/models/Qwen__Qwen3-1.7B/" +if not os.path.exists(MODEL_PATH): + MODEL_PATH = "Qwen/Qwen2-0.5B" +PORT = 13998 +DIST_PORT = 15998 +GROUP_NAME = "test_nccl_group" +MASTER_PORT = DIST_PORT + 1 +HOST = network.gethostip() +RUN_SERVER_TIMEOUT = 180 + + +def check_server_health(base_url): + try: + response = requests.get(f"{base_url}/health", timeout=30) + return response.status_code == 200 + except requests.exceptions.RequestException: + return False + + +@pytest.fixture(scope="module") +def sglang_server_nccl(): + from realhf.base import seeding + + seeding.set_random_seed(1, EXPR_NAME) + cmd = SGLangConfig.build_cmd( + sglang_config=SGLangConfig( + mem_fraction_static=0.2, + model_path=MODEL_PATH, + skip_tokenizer_init=False, + log_level="info", + ), + tp_size=1, + base_gpu_id=1, + host=HOST, + port=PORT, + dist_init_addr=f"{HOST}:{DIST_PORT}", + ) + full_command = f"{cmd} --port {PORT}" + full_command = full_command.replace("\\\n", " ").replace("\\", " ") + os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" + + print(f"full_command to start sglang server: {full_command}", flush=True) + process = subprocess.Popen( + full_command.split(), + text=True, + stdout=sys.stdout, + stderr=sys.stdout, + ) + base_url = f"http://{HOST}:{PORT}" + tik = time.time() + while time.time() - tik < RUN_SERVER_TIMEOUT: + if check_server_health(base_url): + break + time.sleep(1) + if time.time() - tik > RUN_SERVER_TIMEOUT: + process.terminate() + raise RuntimeError("server launch failed") + yield + process.terminate() + + +def test_fsdpengine_nccl_weight_update_to_remote(tmp_path_factory, sglang_server_nccl): + # Set environment variables for torch distributed + os.environ["WORLD_SIZE"] = "1" + os.environ["RANK"] = "0" + os.environ["LOCAL_RANK"] = "0" + os.environ["MASTER_ADDR"] = HOST + os.environ["MASTER_PORT"] = str(MASTER_PORT) + + # Initialize FSDPEngine + engine_config = TrainEngineConfig( + experiment_name=EXPR_NAME, + trial_name=TRIAL_NAME, + path=MODEL_PATH, + optimizer=OptimizerConfig(), + ) + engine = FSDPEngine(engine_config) + ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=100, train_batch_size=2) + engine.initialize(None, ft_spec) + + # Initialize RemoteSGLangEngine + config = InferenceEngineConfig(experiment_name=EXPR_NAME, trial_name=TRIAL_NAME) + config.server_addrs = [f"{HOST}:{PORT}"] + remote_engine = RemoteSGLangEngine(config) + remote_engine.initialize(None, None) + + # Get WeightUpdateMeta + meta = WeightUpdateMeta.from_fsdp_nccl( + AllocationMode.from_str("sglang.d1p1t1+d1p1t1"), + engine, + nccl_group_name=GROUP_NAME, + ) + + # Broadcast weights + future = remote_engine.update_weights(meta) + print("got future", flush=True) + engine.upload_weights(meta) + print("uploaded wexights to remote engine", flush=True) + # Wait for remote engine to finish + future.result(timeout=120) + print("got result", flush=True) + remote_engine.destroy() + engine.destroy() diff --git a/areal/tests/test_sglang_engine.py b/areal/tests/test_sglang_engine.py new file mode 100644 index 0000000000..6fbc870a39 --- /dev/null +++ b/areal/tests/test_sglang_engine.py @@ -0,0 +1,211 @@ +import os +import subprocess +import sys +import time + +import pytest +import requests +import torch +from tensordict import TensorDict + +from areal.api.cli_args import ( + GenerationHyperparameters, + InferenceEngineConfig, + SGLangConfig, +) +from areal.api.io_struct import WeightUpdateMeta +from areal.utils import network +from realhf.api.core.data_api import load_hf_tokenizer + +EXPR_NAME = "test_sglang_engine" +TRIAL_NAME = "trial_0" +MODEL_PATH = "/storage/testing/models/Qwen__Qwen3-1.7B/" +if not os.path.exists(MODEL_PATH): + MODEL_PATH = "Qwen/Qwen2-0.5B" +PORT, DIST_PORT = network.find_free_ports(2) +HOST = network.gethostip() +# set a large timeout since we may need to download the model from hub +RUN_SERVER_TIMEOUT = 180 + + +def check_server_health(base_url): + try: + response = requests.get(f"{base_url}/health", timeout=30) + return response.status_code == 200 + except requests.exceptions.RequestException as e: + return False + + +@pytest.fixture(scope="module") +def sglang_server(): + from realhf.base import seeding + + seeding.set_random_seed(1, EXPR_NAME) + cmd = SGLangConfig.build_cmd( + sglang_config=SGLangConfig( + skip_tokenizer_init=True, + model_path=MODEL_PATH, + mem_fraction_static=0.3, + ), + host=HOST, + port=PORT, + tp_size=1, + base_gpu_id=0, + dist_init_addr=f"{HOST}:{DIST_PORT}", + ) + # Launch process + cmd = cmd.replace("\\\n", " ").replace("\\", " ") + process = subprocess.Popen( + cmd.split(), + text=True, + stdout=sys.stdout, + stderr=sys.stdout, + ) + base_url = f"http://{HOST}:{PORT}" + tik = time.time() + while time.time() - tik < RUN_SERVER_TIMEOUT: + if check_server_health(base_url): + break + time.sleep(1) + if time.time() - tik > RUN_SERVER_TIMEOUT: + raise RuntimeError("server launch failed") + yield + process.terminate() + + +@pytest.mark.parametrize("n_samples", [1, 2, 4]) +def test_remote_sglang_rollout(sglang_server, n_samples): + from areal.engine.sglang_remote import RemoteSGLangEngine + from areal.workflow.rlvr import RLVRWorkflow + + config = InferenceEngineConfig( + experiment_name=EXPR_NAME, + trial_name=TRIAL_NAME, + max_concurrent_rollouts=2, + consumer_batch_size=2, + ) + os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" + engine = RemoteSGLangEngine(config) + engine.initialize(None, None) + + gconfig = GenerationHyperparameters( + max_new_tokens=16, greedy=False, n_samples=n_samples + ) + tokenizer = load_hf_tokenizer(MODEL_PATH) + + workflow = RLVRWorkflow( + reward_fn=lambda **kwargs: 1.0, # Dummy reward function + gconfig=gconfig, + tokenizer=tokenizer, + enable_thinking=False, + ) + + data = { + "messages": [{"role": "user", "content": "Hello, how are you?"}], + } + result = engine.rollout_batch([data] * 2, workflow=workflow) + assert isinstance(result, TensorDict) + bs = result.batch_size + assert bs == torch.Size([2 * n_samples]) + engine.destroy() + + +@pytest.mark.parametrize("ofp", [1, 2, 4, 8, 16]) +@pytest.mark.parametrize("bs", [2, 4]) +@pytest.mark.parametrize("n_samples", [2, 1]) +def test_remote_sglang_staleness_control(sglang_server, bs, ofp, n_samples): + from areal.engine.sglang_remote import RemoteSGLangEngine + from areal.workflow.rlvr import RLVRWorkflow + + config = InferenceEngineConfig( + experiment_name=EXPR_NAME, + trial_name=TRIAL_NAME, + consumer_batch_size=bs, + max_head_offpolicyness=ofp, + ) + os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" + engine = RemoteSGLangEngine(config) + engine.initialize(None, None) + + gconfig = GenerationHyperparameters( + max_new_tokens=16, greedy=False, n_samples=n_samples + ) + tokenizer = load_hf_tokenizer(MODEL_PATH) + + workflow = RLVRWorkflow( + reward_fn=lambda **kwargs: 1.0, # Dummy reward function + gconfig=gconfig, + tokenizer=tokenizer, + enable_thinking=False, + ) + data = { + "messages": [{"role": "user", "content": "Hello, how are you?"}], + } + for _ in range(bs * 2): + engine.submit(data, workflow=workflow) + + # wait for some time + time.sleep(5) + assert engine.output_queue.qsize() == min(bs * 2, bs * (ofp + 1)) + + # Update model version + engine.set_version(1) + print("Updated model version", flush=True) + + # submit again + for _ in range(bs * 2): + engine.submit(data, workflow=workflow) + # wait for some time + time.sleep(5) + assert engine.output_queue.qsize() == min(bs * 4, bs * (ofp + 2)) + + # exit + engine.destroy() + + +def test_disk_update_weights_from_fsdp_engine(tmp_path_factory, sglang_server): + # setup FSDP engine + from areal.api.cli_args import OptimizerConfig, TrainEngineConfig + from areal.api.io_struct import FinetuneSpec + from areal.engine.fsdp_engine import FSDPEngine + + os.environ["WORLD_SIZE"] = "1" + os.environ["RANK"] = "0" + os.environ["LOCAL_RANK"] = "0" + os.environ["MASTER_ADDR"] = "localhost" + os.environ["MASTER_PORT"] = "7777" + + engine_config = TrainEngineConfig( + experiment_name=EXPR_NAME, + trial_name=TRIAL_NAME, + path=MODEL_PATH, + optimizer=OptimizerConfig(), + ) + engine = FSDPEngine(engine_config) + ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=100, train_batch_size=2) + engine.initialize(None, ft_spec) + engine.model_version = 100 + + # setup name resolve + import realhf.base.name_resolve as name_resolve + from realhf.api.cli_args import NameResolveConfig + + nfs_record_root = tmp_path_factory.mktemp("nfs_record_path") + name_resolve_config = NameResolveConfig(type="nfs", nfs_record_root=nfs_record_root) + name_resolve.reconfigure(name_resolve_config) + # initialize SGLang remote engine + from areal.api.cli_args import InferenceEngineConfig + from areal.engine.sglang_remote import RemoteSGLangEngine + + config = InferenceEngineConfig(experiment_name=EXPR_NAME, trial_name=TRIAL_NAME) + os.environ["AREAL_LLM_SERVER_ADDRS"] = f"{HOST}:{PORT}" + inf_engine = RemoteSGLangEngine(config) + inf_engine.initialize(None, None) + inf_engine.set_version(100) + # test update weights + path = tmp_path_factory.mktemp("upload_weights_from_disk") + update_weight_meta = WeightUpdateMeta(type="disk", path=str(path)) + future = inf_engine.update_weights(update_weight_meta) + engine.upload_weights(update_weight_meta) + future.result() + inf_engine.destroy() diff --git a/areal/tests/test_train_engine.py b/areal/tests/test_train_engine.py new file mode 100644 index 0000000000..5e05083c7f --- /dev/null +++ b/areal/tests/test_train_engine.py @@ -0,0 +1,163 @@ +# Copyright 2025 Ant Group Inc. +# Licensed under the Apache License, Version 2.0 + +"""Test script for Engine implementation.""" + +import os +from typing import Dict + +import pytest +import torch +from tensordict import TensorDict +from transformers import AutoTokenizer + +from areal.api.cli_args import MicroBatchSpec, OptimizerConfig, TrainEngineConfig +from areal.api.io_struct import FinetuneSpec, SaveLoadMeta + +VOCAB_SIZE = 100 +MODEL_PATH = "/storage/testing/models/Qwen__Qwen3-1.7B/" +if not os.path.exists(MODEL_PATH): + MODEL_PATH = "Qwen/Qwen2-0.5B" + + +@pytest.fixture(scope="module") +def mock_input( + batch_size=5, + min_seqlen=10, + max_seqlen=20, + device="cuda:0", +) -> Dict: + """Create mock padded input data (same format for huggingface) for testing. + Returns a dict with input_ids, attention_mask, and position_ids. + """ + pad_token_id = 0 + seqlens = torch.randint( + min_seqlen, max_seqlen, (batch_size,), dtype=torch.int, device=device + ) + max_seqlen = int(max(seqlens)) + input_ids = torch.randint( + 0, VOCAB_SIZE, (batch_size, max_seqlen), dtype=torch.long, device=device + ) + attn_mask = torch.zeros((batch_size, max_seqlen), dtype=torch.bool, device=device) + + attn_mask[ + torch.arange(0, max_seqlen, device=device).unsqueeze(0) < seqlens.unsqueeze(1) + ] = 1 + input_ids.masked_fill_(~attn_mask, pad_token_id) + + return TensorDict( + input_ids=input_ids, + attention_mask=attn_mask, + ) + + +def get_engine(engine_type: str, model_path: str): + from areal.engine.fsdp_engine import FSDPEngine + from areal.experimental.autotp_engine import DeepSpeedAutoTPEngine + + engine_cls = {"auto_tp": DeepSpeedAutoTPEngine, "fsdp": FSDPEngine}[engine_type] + + engine_config = TrainEngineConfig( + experiment_name=f"test-{engine_type}-engine", + trial_name="test0", + path=model_path, + optimizer=OptimizerConfig(), + ) + engine = engine_cls(engine_config) + ft_spec = FinetuneSpec(total_train_epochs=1, dataset_size=100, train_batch_size=2) + engine.initialize(None, ft_spec) + return engine + + +def mock_loss_fn(logits: torch.Tensor, input_data: Dict) -> torch.Tensor: + """Mock loss function for testing.""" + return torch.mean(logits) + + +@pytest.fixture(scope="module", params=["fsdp", "auto_tp"]) +def engine(request): + os.environ.update( + { + "WORLD_SIZE": "1", + "RANK": "0", + "LOCAL_RANK": "0", + "MASTER_ADDR": "localhost", + "MASTER_PORT": "7777", + } + ) + + engine = get_engine(request.param, MODEL_PATH) + print(f"✓ {request.param.upper()} Engine created successfully") + yield engine + + +@torch.no_grad() +def test_forward_microbatch(engine, mock_input): + engine.eval() + engine.config.mb_spec = MicroBatchSpec(n_mbs=2, max_tokens_per_mb=100) + x2 = engine.forward(input_=mock_input).squeeze(0).mean(-1) + engine.config.mb_spec = MicroBatchSpec(n_mbs=1, max_tokens_per_mb=100) + x1 = engine.forward(input_=mock_input).squeeze(0).mean(-1) + input_ids = mock_input["input_ids"] + assert x1.shape[:1] == input_ids.shape[:1] + assert x2.shape[:1] == input_ids.shape[:1] + assert torch.allclose(x1, x2, atol=1e-1, rtol=1e-2), (x1 - x2).abs().max().item() + + +@torch.no_grad() +def test_eval_batch(engine, mock_input): + engine.eval() + engine.config.mb_spec = MicroBatchSpec(n_mbs=2, max_tokens_per_mb=100) + eval_result = engine.eval_batch( + input_=mock_input, + loss_fn=mock_loss_fn, + loss_weight_fn=lambda x: x["cu_seqlens"][-1], + ) + assert isinstance(eval_result, torch.Tensor), "Evaluation should return a tensor" + assert eval_result.is_cuda, "Evaluation tensor should be on CUDA device" + assert eval_result is not None, "Evaluation should return a loss value" + print(f"✓ Evaluation successful, loss: {eval_result.item()}") + + +def test_train_batch(engine, mock_input): + engine.train() + engine.config.mb_spec = MicroBatchSpec(n_mbs=2, max_tokens_per_mb=100) + train_result = engine.train_batch( + input_=mock_input, + loss_fn=mock_loss_fn, + loss_weight_fn=lambda x: x["cu_seqlens"][-1], + ) + assert isinstance(train_result, dict), "Training should return a dictionary" + assert train_result["grad_norm"] is not None + assert train_result["lr"] is not None + print("✓ Training successful") + + +@torch.no_grad() +def test_hf_save_load_weights(tmp_path_factory, engine, mock_input): + from areal.experimental.autotp_engine import DeepSpeedAutoTPEngine + + if isinstance(engine, DeepSpeedAutoTPEngine): + print("AutoTP engine does not support HF save/load for now.") + return + + tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH) + path = tmp_path_factory.mktemp("hf_engine_test") + save_load_meta = SaveLoadMeta( + path=path, + weight_format="hf", + tokenizer=tokenizer, + with_optim=True, + base_model_path=None, + ) + + engine.config.mb_spec = MicroBatchSpec(n_mbs=1, max_tokens_per_mb=100) + old = engine.forward(input_=mock_input) + engine.save(save_load_meta) + + for name, param in engine.model.named_parameters(): + param.zero_() + + engine.load(save_load_meta) + new = engine.forward(input_=mock_input) + assert torch.allclose(old, new) diff --git a/areal/tests/test_utils.py b/areal/tests/test_utils.py new file mode 100644 index 0000000000..bb21ea04c9 --- /dev/null +++ b/areal/tests/test_utils.py @@ -0,0 +1,70 @@ +import pytest +import torch +from tensordict import TensorDict + +from areal.api.cli_args import MicroBatchSpec +from areal.utils.data import ( + pack_tensor_dict, + pad_and_stack_tensors_along_first_dim, + pad_sequences_to_tensors, + reorder_list, + split_padded_tensor_dict_into_mb_list, + unpack_sequence, +) + +BS = 16 +MAX_ANSWER_LEN = 16 +MAX_PROMPT_LEN = 8 +VOCAB_SIZE = 100 + + +@pytest.fixture +def mock_padded_data(): + prompt_lens = torch.randint(1, MAX_PROMPT_LEN, size=(BS,)) + answer_lens = torch.randint(1, MAX_ANSWER_LEN, size=(BS,)) + all_data = [] + for prompt_len, ans_len in zip(prompt_lens, answer_lens): + prompt_len = int(prompt_len) + ans_len = int(ans_len) + seq = dict( + input_ids=torch.randint(0, VOCAB_SIZE, size=(prompt_len + ans_len,)), + loss_mask=torch.tensor([0] * prompt_len + [1] * ans_len), + logprobs=torch.randn(prompt_len + ans_len), + position_ids=torch.arange(prompt_len + ans_len), + ) + all_data.append(TensorDict(seq)) + return pad_sequences_to_tensors(all_data) + + +@pytest.mark.parametrize("max_tokens_per_mb", [24, 36, 48, 100]) +@pytest.mark.parametrize("n_mbs", [1, 2, 4, 8]) +def test_micro_batch_split(mock_padded_data, n_mbs, max_tokens_per_mb): + mb_spec = MicroBatchSpec(n_mbs, max_tokens_per_mb) + + # Unpad and split to microbatches + packed_data = pack_tensor_dict(mock_padded_data) + original_lens = packed_data["cu_seqlens"][1:] - packed_data["cu_seqlens"][:-1] + assert torch.allclose(original_lens, mock_padded_data["attention_mask"].sum(1)) + split_result = split_padded_tensor_dict_into_mb_list(mock_padded_data, mb_spec) + split_result.mbs = [pack_tensor_dict(mb) for mb in split_result.mbs] + reordered_lens = [original_lens[i] for i in split_result.forward_indices] + + # assert microbatch split result does not violate requirements + assert len(split_result.mbs) >= n_mbs + + # test reorder back + for key in split_result.mbs[0].keys(): + if key in ["cu_seqlens", "max_seqlen"]: + continue + + # assert microbatch split result does not violate requirements + for mb in split_result.mbs: + assert mb[key].shape[0] <= max_tokens_per_mb + + x = torch.cat([mb[key] for mb in split_result.mbs]) + xs = unpack_sequence(x, lens=reordered_lens) + xs = reorder_list(xs, split_result.backward_indices) + x = torch.cat(xs) + assert torch.allclose(x, packed_data[key]) + y = pad_and_stack_tensors_along_first_dim(xs) + assert torch.allclose(mock_padded_data[key], y) diff --git a/areal/tests/test_wrapper.py b/areal/tests/test_wrapper.py new file mode 100644 index 0000000000..08649464e6 --- /dev/null +++ b/areal/tests/test_wrapper.py @@ -0,0 +1,48 @@ +from areal.utils.wrapper import ( + wrap, + wrap_get_method, + wrap_get_method_name, + wrap_remove_meta, + wrapable, +) + + +class Calculator: + def __init__(self): + self.remaining = 2 + + @wrapable() + def add(self, a, b): + return a + b + self.remaining + + @wrapable(name="multiply") + def mul(self, x, y): + return x * y * self.remaining + + +class EnhancedCalculator: + def __init__(self, obj: Calculator): + super().__init__() + self.prefix = "Result:" + self.calculator = obj + wrap(self, obj, self._wrap_call) + + def _wrap_call(self, *args, **kwargs): + method_name = wrap_get_method_name(kwargs) + method = wrap_get_method(kwargs) + kwargs = wrap_remove_meta(kwargs) + print("wrap method: ", method_name) + print("wrap method: ", method) + return method(*args, **kwargs) + + +def test_wrapper(): + calc = Calculator() + enhancer = EnhancedCalculator(calc) + + assert enhancer.add(2, 3) == 7 + assert enhancer.multiply(4, 5) == 40 + + +if __name__ == "__main__": + test_wrapper() diff --git a/areal/utils/__init__.py b/areal/utils/__init__.py new file mode 100644 index 0000000000..e69de29bb2 diff --git a/areal/utils/data.py b/areal/utils/data.py new file mode 100644 index 0000000000..13f3b16c1f --- /dev/null +++ b/areal/utils/data.py @@ -0,0 +1,542 @@ +# Copyright 2025 Ant Group Inc. +# Licensed under the Apache License, Version 2.0 + +# Pad/unpad operations are modified from flash-attention under BSD-3 license. +# Copyright (c) 2023, Tri Dao. + +from dataclasses import dataclass +from typing import Any, Callable, Dict, List, Optional, Tuple + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn.functional as F +from einops import rearrange +from tensordict import TensorDict + +from areal.api.cli_args import MicroBatchSpec +from realhf.base import datapack, logging + +logger = logging.getLogger("data utils") + + +def reorder_list(xs: List, indices: List[int]) -> List: + assert len(set(indices)) == len(xs) + return [xs[i] for i in indices] + + +def dict_map(x: Dict, fn: Callable) -> Dict: + return {k: fn(v) for k, v in x.items()} + + +def dict_of_list2list_of_dict( + dict_of_lists: Dict[str, List[Any]], +) -> List[Dict[str, Any]]: + if not dict_of_lists: + return [] + keys = list(dict_of_lists.keys()) + length = len(dict_of_lists[keys[0]]) + for key, value_list in dict_of_lists.items(): + if len(value_list) != length: + raise ValueError( + f"All lists must have the same length. Key '{key}' has length {len(value_list)}, expected {length}" + ) + return [{key: dict_of_lists[key][i] for key in keys} for i in range(length)] + + +def list_of_dict2dict_of_list( + list_of_dicts: List[Dict[str, Any]], +) -> Dict[str, List[Any]]: + if not list_of_dicts: + return {} + keys = list(list_of_dicts[0].keys()) + for i, dict_item in enumerate(list_of_dicts): + if set(dict_item.keys()) != set(keys): + raise ValueError( + f"All dictionaries must have the same keys. Dictionary at index {i} has keys {set(dict_item.keys())}, expected {set(keys)}" + ) + return {key: [dict_item[key] for dict_item in list_of_dicts] for key in keys} + + +def pad_sequences_to_tensors( + sequence_list: List[TensorDict], pad_value: float = 0.0 +) -> TensorDict: + if not sequence_list: + return TensorDict() + skip_keys = {"pixel_values", "image_grid_thw"} + max_length = max( + len(seq) + for item in sequence_list + for key, seq in item.items() + if key not in skip_keys + ) + result = {} + for key in sequence_list[0].keys(): + padded = [] + if key in skip_keys: + result[key] = [sequence_list[i][key] for i in range(len(sequence_list))] + continue + for item in sequence_list: + x = item[key] + if not torch.is_tensor(x): + x = torch.tensor(x) + padded_x = torch.nn.functional.pad( + x, (0, max_length - len(item[key])), value=pad_value + ) + padded.append(padded_x) + result[key] = torch.stack(padded) + attention_mask = [ + [1] * len(next(iter(item[key] for key in item.keys() if key not in skip_keys))) + + [0] + * ( + max_length + - len(next(iter(item[key] for key in item.keys() if key not in skip_keys))) + ) + for item in sequence_list + ] + result["attention_mask"] = torch.tensor(attention_mask, dtype=torch.bool) + return TensorDict(result, batch_size=[result["attention_mask"].shape[0]]) + + +def unpad_input( + hidden_states, attention_mask +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, int]: + seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32) + indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten() + max_seqlen_in_batch = seqlens_in_batch.max().item() + cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0)) + return ( + rearrange(hidden_states, "b s ... -> (b s) ...")[indices], + indices, + cu_seqlens, + max_seqlen_in_batch, + ) + + +def pad_input(hidden_states, indices, batch, seqlen): + output = hidden_states.new_zeros(batch * seqlen) + output[indices] = hidden_states + return rearrange(output, "(b s) ... -> b s ...", b=batch) + + +def concat_padded_tensors( + tensor_dicts: List[TensorDict], pad_value: float = 0.0 +) -> TensorDict: + """Concatenate and pad tensors from multiple padded tensor dictionaries.""" + if not tensor_dicts: + return TensorDict() + + batch_sizes = [tuple(d.batch_size) for d in tensor_dicts] + new_batch_size = [sum(x[0] for x in batch_sizes), *batch_sizes[0][1:]] + + # Find max sequence length across all dictionaries + assert all("attention_mask" in td for td in tensor_dicts) + max_length = max([x["attention_mask"].shape[1] for x in tensor_dicts]) + result = {} + + # Process each key + for key in tensor_dicts[0].keys(): + tensors_to_concat = [] + + for tensor_dict in tensor_dicts: + tensor = tensor_dict[key] + # Skip 1D tensors like rewards + if len(tensor.shape) == 1: + tensors_to_concat.append(tensor) + continue + current_length = tensor.shape[1] + if key == "pixel_values" or key == "image_grid_thw": + tensors_to_concat.append(tensor) + continue + if current_length < max_length: + # Pad tensor to max_length + pad_width = max_length - current_length + if key == "attention_mask": + # Pad attention mask with 0s + padding = torch.zeros( + (tensor.shape[0], pad_width), dtype=tensor.dtype + ) + + else: + # Pad feature tensors with pad_value + padding = torch.full( + (tensor.shape[0], pad_width), pad_value, dtype=tensor.dtype + ) + + tensor = torch.cat([tensor, padding], dim=1) + tensors_to_concat.append(tensor) + + result[key] = torch.cat(tensors_to_concat, dim=0) + return TensorDict(result, batch_size=new_batch_size) + + +def to_device(data: Dict[str, torch.Tensor | Any], device) -> Dict[str, torch.Tensor]: + """Move tensors in a dictionary to the specified device.""" + return { + key: value.to(device) if torch.is_tensor(value) else value + for key, value in data.items() + } + + +def unpack_sequence( + x: torch.Tensor, + cu_seqlens: Optional[torch.Tensor] = None, + lens: Optional[List[int]] = None, + dim: int = 0, +): + """Unpack a sequence tensor into a list of tensors based on cumulative sequence lengths.""" + if lens is not None: + return torch.split(x, lens, dim=dim) + if cu_seqlens is not None: + return torch.split( + x, (cu_seqlens[1:] - cu_seqlens[:-1]).cpu().numpy().tolist(), dim=dim + ) + raise ValueError("Either cu_seqlens or input_lens must be provided.") + + +def allocate_balanced_mbs(mb_spec: MicroBatchSpec, lens: List[int]) -> List[List[int]]: + group_indices = datapack.ffd_allocate( + lens, mb_spec.max_tokens_per_mb, min_groups=mb_spec.n_mbs + ) + group_indices = sorted([sorted(g) for g in group_indices]) + return group_indices + + +def allocate_balanced_mbs_synced( + mb_spec: MicroBatchSpec, + lens: List[int], + group: Optional[dist.ProcessGroup] = None, +) -> List[List[int]]: + group_indices = allocate_balanced_mbs(mb_spec, lens) + if not dist.is_initialized(): + return group_indices + + all_n_mbs = [None for _ in range(dist.get_world_size(group))] + dist.all_gather_object(all_n_mbs, len(group_indices), group=group) + if all(mbs == len(group_indices) for mbs in all_n_mbs): + return group_indices + return allocate_balanced_mbs_synced( + MicroBatchSpec.new(mb_spec, n_mbs=max(all_n_mbs)), lens + ) + + +def pack_tensor_dict(data: TensorDict): + """Pack a tensordict of shape [B, S, ...] into [total_length, ...], leaving other keys unchanged. + + Args: + data (Dict[str, Any]): Dictionary containing tensors to be packed. Should contain key "attention_mask" with shape [B, S]. + + Returns: + Dict[str, Any]: Dictionary with packed tensors. The "attention_mask" key will be replaced by "cu_seqlens" with shape [B+1]. + """ + + assert "attention_mask" in data, "Input data must contain 'attention_mask' key." + attention_mask = data["attention_mask"] + assert attention_mask.ndim == 2, "Attention mask must be a 2D tensor." + bs = attention_mask.shape[0] + seq_len = attention_mask.shape[1] + + # Calculate cumulative sequence lengths + lens = attention_mask.sum(dim=1, dtype=torch.int32) + max_seqlen = lens.max().item() + cu_seqlens = torch.cumsum(lens, dim=0) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + total_length = int(cu_seqlens[-1].item()) + # Pack tensors + packed_data = {} + packed_data["cu_seqlens"] = cu_seqlens + packed_data["max_seqlen"] = max_seqlen + for key, value in data.items(): + # if key == "attention_mask": + # packed_data["cu_seqlens"] = cu_seqlens + # packed_data["max_seqlen"] = max_seqlen + # # tensor and of shape [B, S, ...] + if ( + torch.is_tensor(value) + and value.ndim >= 2 + and value.shape[0] == bs + and value.shape[1] == seq_len + ): + packed_tensor = torch.empty( + (total_length, *value.shape[2:]), dtype=value.dtype, device=value.device + ) + # Fill the packed tensor with values from the original tensor + for i in range(bs): + start = cu_seqlens[i].item() + end = cu_seqlens[i + 1].item() + packed_tensor[start:end] = value[i][: end - start] + packed_data[key] = packed_tensor + else: + packed_data[key] = value + + return TensorDict(**packed_data) + + +def pad_and_stack_tensors_along_first_dim(tensor_list: List[torch.Tensor]): + max_length = max(tensor.shape[0] for tensor in tensor_list) + n_dim = tensor_list[0].ndim + assert all( + tensor.ndim == n_dim for tensor in tensor_list + ), "All tensors must have the same number of dimensions." + + padded_tensors = [] + for tensor in tensor_list: + pad_mode = (0,) * (2 * (n_dim - 1)) + (0, max_length - tensor.shape[0]) + padded_tensor = F.pad(tensor, pad_mode, value=0.0) + padded_tensors.append(padded_tensor) + return torch.stack(padded_tensors, dim=0) + + +@dataclass +class MicroBatchList: + data: TensorDict + mb_spec: MicroBatchSpec + mbs: List[TensorDict] + forward_indices: List[int] + backward_indices: List[int] + group_lens: List[int] + padded_mbs: Optional[List[TensorDict]] = None + padding_lengths: Optional[List[int]] = None + padded_to_lengths: Optional[List[int]] = None + + +DEFAULT_MAX_TOKENS_PER_MB = int(1e12) + + +def split_padded_tensor_dict_into_mb_list( + data: TensorDict, mb_spec: MicroBatchSpec, group: Optional[dist.ProcessGroup] = None +) -> MicroBatchList: + """Split a padded tensordict into micro-batches based on the attention mask. + + Args: + data (TensorDict): Dictionary containing padded tensors. + mb_spec (MicroBatchSpec): Specification for micro-batch splitting. + group (Optional[dist.ProcessGroup]): Process group for distributed synchronization. + + Returns: + MicroBatchList: A structure containing the split micro-batches and metadata. + """ + assert ( + "attention_mask" in data + ), "Input data must be padded and contain 'attention_mask' key." + if mb_spec.max_tokens_per_mb is None: + mb_spec = MicroBatchSpec.new( + mb_spec, max_tokens_per_mb=DEFAULT_MAX_TOKENS_PER_MB + ) + bs = data["attention_mask"].shape[0] + max_seqlen = data["attention_mask"].shape[1] + input_lens = data["attention_mask"].sum(1).long().cpu().numpy() + + # check tensor shape, split only 1d tensors with length "total_lens" + to_split = {} + not_to_split = {} + for key, value in data.items(): + if key == "image_grid_thw" or key == "pixel_values": + continue + if not torch.is_tensor(value) or value.numel() != bs * max_seqlen: + not_to_split[key] = value + else: + to_split[key] = value + + # split + group_indices = allocate_balanced_mbs_synced(mb_spec, input_lens, group=group) + splitted_lens = [ + [input_lens[i] for i in group_index] for group_index in group_indices + ] + group_n_seqs = [len(x) for x in splitted_lens] + group_lens = [sum(x) for x in splitted_lens] + + forward_indices = datapack.flat2d(group_indices) + backward_indices = np.zeros(bs, dtype=np.int64) + backward_indices[forward_indices] = np.arange(bs) + + def _split(tensor): + """Split and pad a tensor based on forward indices and lens.""" + # Unpack the sequence + unpacked = [tensor[i] for i in range(bs)] + # Reorder according to forward indices + reordered = reorder_list(unpacked, forward_indices) + reordered = torch.stack(reordered) + # Unpack again according to split lens + splitted = [] + offset = 0 + for _n_seqs in group_n_seqs: + splitted.append(reordered[offset : offset + _n_seqs]) + offset += _n_seqs + return splitted + + to_split = dict_map(to_split, lambda x: _split(x)) + if data.get("pixel_values", None) is not None: + pixel_values = data.get("pixel_values", []) + image_grid_thw = data.get("image_grid_thw", []) + + # Prepare the pixel_values and image_grid_thw for each group + pixel_values_split = [] + image_grid_thw_split = [] + + for group_index in group_indices: + group_pixel_values = [pixel_values[i] for i in group_index] + group_image_grid_thw = [image_grid_thw[i].squeeze() for i in group_index] + + # Stack pixel_values for each group (assuming pixel_values is a list of tensors) + pixel_values_split.append(torch.stack(group_pixel_values)) + image_grid_thw_split.append(torch.stack(group_image_grid_thw)) + + # Pack the split pixel_values and image_grid_thw back into the data + to_split["pixel_values"] = pixel_values_split + to_split["image_grid_thw"] = image_grid_thw_split + mbs = dict_of_list2list_of_dict(to_split) + + results = [] + # organize splitted micro batches + assert len(mbs) == len(splitted_lens), (len(mbs), len(splitted_lens)) + for i, (mb, lens) in enumerate(zip(mbs, splitted_lens)): + results.append(TensorDict(**mb, **not_to_split)) + return MicroBatchList( + data=data, + mbs=results, + mb_spec=mb_spec, + forward_indices=forward_indices, + backward_indices=backward_indices.tolist(), + group_lens=group_lens, + ) + + +def pad_packed_tensor_dict( + data: TensorDict, + pad_to_length: int, + pad_value: float = 0.0, +) -> Tuple[TensorDict, int]: + """Pad a packed tensor dict to a specified length. + This function assumes that the input data contains "cu_seqlens" and "max_seqlen" key, + and all other tensors of shape [total_length, ] will be padded to `pad_to_length`. + This function will pad a new sequence filled with `pad_value` to the end of each tensor, + and update the "cu_seqlens" and "max_seqlen" keys accordingly. + + Args: + data (TensorDict): Dictionary containing tensors to be packed. + pad_to_length (int): The length to pad the tensors to. All tensors + + Returns: + TensorDict: Dictionary with padded tensors and modified "cu_seqlens" and + "max_seqlen". + int: The pad length. + """ + assert "cu_seqlens" in data, "Input data must contain 'cu_seqlens' key." + assert "max_seqlen" in data, "Input data must contain 'max_seqlen' key." + total_length = data["cu_seqlens"][-1].item() + pad_length = pad_to_length - total_length + assert ( + pad_length >= 0 + ), f"pad_to_length {pad_to_length} must be greater than or equal to total length {total_length}." + cu_seqlens = data["cu_seqlens"] + max_seqlen = data["max_seqlen"] + new_cu_seqlens = F.pad(cu_seqlens, (0, 1), value=pad_to_length) + new_max_seqlen = max(max_seqlen, pad_length) + padded_data = {} + for key, value in data.items(): + if key == "cu_seqlens": + padded_data[key] = new_cu_seqlens + elif key == "max_seqlen": + padded_data[key] = new_max_seqlen + elif torch.is_tensor(value) and value.numel() == total_length: + # Pad the tensor to the new total length + if key == "position_ids": + # transformers will compute flash-attn arguments (e.g., cu_seqlens_q) + # according to this position ids. + pad = torch.arange(pad_length, dtype=torch.long, device=value.device) + padded_tensor = torch.cat([value, pad]) + else: + padded_tensor = torch.nn.functional.pad( + value, (0, pad_length), value=pad_value + ) + padded_data[key] = padded_tensor + else: + padded_data[key] = value + return TensorDict(padded_data, batch_size=data.batch_size), pad_length + + +def pad_mb_list( + mb_list: MicroBatchList, + pad_value: float = 0.0, + pad_to_maximum: bool = False, +) -> MicroBatchList: + padded_mb_inputs, pad_lengths = [], [] + pad_to_lengths = [] + if ( + pad_to_maximum + and mb_list.mb_spec.max_tokens_per_mb is not None + and mb_list.mb_spec.max_tokens_per_mb < DEFAULT_MAX_TOKENS_PER_MB + ): + pad_to_length = mb_list.mb_spec.max_tokens_per_mb + else: + if pad_to_maximum: + logger.warning( + f"Cannot pad to upper bound since mb_spec.max_tokens_per_mb is not set." + ) + # NOTE: GPU page size is 2MB + # Take hidden size 4096 with bf16 dtype as an example, + # the batch size of a page is 256 + pad_to_length = (int(max(mb_list.group_lens)) + 255) // 256 * 256 + for mb, l in zip(mb_list.mbs, mb_list.group_lens): + padded_mb, pad_len = pad_packed_tensor_dict( + mb, pad_to_length, pad_value=pad_value + ) + padded_mb_inputs.append(padded_mb) + pad_lengths.append(pad_len) + pad_to_lengths.append(pad_to_length) + mb_list.padded_mbs = padded_mb_inputs + mb_list.padding_lengths = pad_lengths + mb_list.padded_to_lengths = pad_to_lengths + return mb_list + + +def unsqueeze_packed_tensor_dict(data: TensorDict) -> TensorDict: + assert "cu_seqlens" in data, "Input data must contain 'cu_seqlens' key." + assert "max_seqlen" in data, "Input data must contain 'max_seqlen' key." + + total_length = data["cu_seqlens"][-1].item() + new_data = {} + for key, value in data.items(): + if ( + key + not in [ + "cu_seqlens", + "max_seqlen", + ] + and torch.is_tensor(value) + and value.numel() == total_length + ): + new_data[key] = value.unsqueeze(dim=0) + else: + new_data[key] = value + return TensorDict(new_data, batch_size=data.batch_size) + + +def unsqueeze_mb_list( + mb_list: MicroBatchList, +) -> MicroBatchList: + """Unsqueeze the packed tensordict in the micro-batch list.""" + new_mbs = [] + new_padded_mbs = [] + for i, mb in enumerate(mb_list.mbs): + new_mbs.append(unsqueeze_packed_tensor_dict(mb)) + if mb_list.padded_mbs is not None: + new_padded_mbs.append(unsqueeze_packed_tensor_dict(mb_list.padded_mbs[i])) + mb_list.padded_mbs = new_padded_mbs if mb_list.padded_mbs is not None else None + return mb_list + + +def amend_position_ids(data: TensorDict) -> TensorDict: + assert "attention_mask" in data, "Input data must contain 'attention_mask' key." + attn_mask = data["attention_mask"] + bs, seqlen = attn_mask.shape[:2] + position_ids = ( + torch.arange(0, seqlen, dtype=torch.long, device=attn_mask.device) + .unsqueeze(0) + .expand(bs, -1) + ) + position_ids.masked_fill(~attn_mask.bool(), 0) + data["position_ids"] = position_ids + return data diff --git a/areal/utils/device.py b/areal/utils/device.py new file mode 100644 index 0000000000..840ca33d99 --- /dev/null +++ b/areal/utils/device.py @@ -0,0 +1,31 @@ +from typing import Tuple + +import torch +import torch.distributed as dist + +from realhf.base import logging + +logger = logging.getLogger(__file__) + + +def _get_current_mem_info(unit: str = "GB", precision: int = 2) -> Tuple[str]: + """Get current memory usage.""" + assert unit in ["GB", "MB", "KB"] + divisor = 1024**3 if unit == "GB" else 1024**2 if unit == "MB" else 1024 + mem_allocated = torch.cuda.memory_allocated() + mem_reserved = torch.cuda.memory_reserved() + mem_free, mem_total = torch.cuda.mem_get_info() + mem_used = mem_total - mem_free + mem_allocated = f"{mem_allocated / divisor:.{precision}f}" + mem_reserved = f"{mem_reserved / divisor:.{precision}f}" + mem_used = f"{mem_used / divisor:.{precision}f}" + mem_total = f"{mem_total / divisor:.{precision}f}" + return mem_allocated, mem_reserved, mem_used, mem_total + + +# Adapted from verl +def log_gpu_stats(head: str, rank: int = 0): + if (not dist.is_initialized()) or (rank is None) or (dist.get_rank() == rank): + mem_allocated, mem_reserved, mem_used, mem_total = _get_current_mem_info() + message = f"{head}, memory allocated (GB): {mem_allocated}, memory reserved (GB): {mem_reserved}, device memory used/total (GB): {mem_used}/{mem_total}" + logger.info(msg=message) diff --git a/areal/utils/distributed.py b/areal/utils/distributed.py new file mode 100644 index 0000000000..e67e6f44dc --- /dev/null +++ b/areal/utils/distributed.py @@ -0,0 +1,73 @@ +import torch + + +# Copy from pytorch and OpenRLHF to allow creating multiple main groups. +# https://github.com/pytorch/pytorch/blob/main/torch/distributed/distributed_c10d.py +# https://github.com/OpenRLHF/OpenRLHF/blob/main/openrlhf/utils/distributed_util.py +def init_custom_process_group( + backend=None, + init_method=None, + timeout=None, + world_size=-1, + rank=-1, + store=None, + group_name=None, + pg_options=None, +): + from torch.distributed.distributed_c10d import ( + Backend, + PrefixStore, + _new_process_group_helper, + _world, + default_pg_timeout, + rendezvous, + ) + + assert (store is None) or ( + init_method is None + ), "Cannot specify both init_method and store." + + if store is not None: + assert world_size > 0, "world_size must be positive if using store" + assert rank >= 0, "rank must be non-negative if using store" + elif init_method is None: + init_method = "env://" + + if backend: + backend = Backend(backend) + else: + backend = Backend("undefined") + + if timeout is None: + timeout = default_pg_timeout + + # backward compatible API + if store is None: + rendezvous_iterator = rendezvous(init_method, rank, world_size, timeout=timeout) + store, rank, world_size = next(rendezvous_iterator) + store.set_timeout(timeout) + + # Use a PrefixStore to avoid accidental overrides of keys used by + # different systems (e.g. RPC) in case the store is multi-tenant. + store = PrefixStore(group_name, store) + + # NOTE: The pg_options parameter was renamed into backend_options in PyTorch 2.6.0 + # https://github.com/pytorch/pytorch/commit/a0c7029a75628cd5fa8df83c0de0ea98ee7fd844 + # We need to determine the appropriate parameter name based on PyTorch version + pg_options_param_name = ( + "backend_options" if str(torch.__version__) >= "2.6" else "pg_options" + ) + pg, _ = _new_process_group_helper( + world_size, + rank, + [], + backend, + store, + group_name=group_name, + **{pg_options_param_name: pg_options}, + timeout=timeout, + ) + + _world.pg_group_ranks[pg] = {i: i for i in range(world_size)} + + return pg diff --git a/areal/utils/evaluator.py b/areal/utils/evaluator.py new file mode 100644 index 0000000000..30e965a462 --- /dev/null +++ b/areal/utils/evaluator.py @@ -0,0 +1,30 @@ +from typing import Callable + +from areal.api.cli_args import EvaluatorConfig +from areal.api.io_struct import FinetuneSpec +from realhf.base import timeutil + + +class Evaluator: + + def __init__(self, config: EvaluatorConfig, ft_spec: FinetuneSpec): + self.config = config + self.ft_sepc = ft_spec + self.freq_ctl = timeutil.EpochStepTimeFreqCtl( + freq_epoch=config.freq_epochs, + freq_step=config.freq_steps, + freq_sec=config.freq_secs, + ) + + def evaluate( + self, + evaluate_fn: Callable, + epoch: int, + step: int, + global_step: int, + ): + if not self.freq_ctl.check( + epochs=int(step == self.ft_sepc.steps_per_epoch - 1), steps=1 + ): + return + evaluate_fn() diff --git a/areal/utils/fs.py b/areal/utils/fs.py new file mode 100644 index 0000000000..f790400358 --- /dev/null +++ b/areal/utils/fs.py @@ -0,0 +1,9 @@ +import getpass +import os + + +def get_user_tmp(): + user = getpass.getuser() + user_tmp = os.path.join("/home", user, ".cache", "realhf") + os.makedirs(user_tmp, exist_ok=True) + return user_tmp diff --git a/areal/utils/fsdp.py b/areal/utils/fsdp.py new file mode 100644 index 0000000000..34eaa23cbd --- /dev/null +++ b/areal/utils/fsdp.py @@ -0,0 +1,190 @@ +import math + +import torch +import torch.distributed as dist +import torch.nn as nn +from torch.distributed.device_mesh import init_device_mesh +from transformers import PreTrainedModel + +from realhf.base import logging, pkg_version + +logger = logging.getLogger("FSDPEngine") + +if pkg_version.is_version_greater_or_equal("torch", "2.6.0"): + from torch.distributed.fsdp import ( + CPUOffloadPolicy, + FSDPModule, + MixedPrecisionPolicy, + fully_shard, + ) +elif pkg_version.is_version_greater_or_equal("torch", "2.4.0"): + from torch.distributed._composable.fsdp import ( + CPUOffloadPolicy, + FSDPModule, + MixedPrecisionPolicy, + fully_shard, + ) +else: + CPUOffloadPolicy = None + FSDPModule = None + MixedPrecisionPolicy = None + fully_shard = None + logger.warning("Current PyTorch version < 2.4.0 is not supported for FSDPEngine.") + + +def fsdp2_clip_grad_norm_( + parameters, max_norm, norm_type=2.0, error_if_nonfinite=False, foreach=None +): + """torch.nn.utils.clip_grad_norm_ cann't run on cpu parameter DTensor""" + from torch.nn.utils.clip_grad import _clip_grads_with_norm_, _get_total_norm + + if isinstance(parameters, torch.Tensor): + parameters = [parameters] + else: + # prevent generators from being exhausted + parameters = list(parameters) + grads = [p.grad for p in parameters if p.grad is not None] + total_norm = _get_total_norm(grads, norm_type, error_if_nonfinite, foreach) + total_norm = total_norm.to(torch.cuda.current_device(), non_blocking=True) + + _clip_grads_with_norm_(parameters, max_norm, total_norm, foreach) + return total_norm + + +def create_fsdp_device_mesh(shard_size, world_size): + if shard_size < 0 or shard_size >= world_size: + device_mesh = init_device_mesh( + "cuda", mesh_shape=(world_size,), mesh_dim_names=("fsdp",) + ) + else: + device_mesh = init_device_mesh( + "cuda", + mesh_shape=(world_size // shard_size, shard_size), + mesh_dim_names=("ddp", "fsdp"), + ) + return device_mesh + + +def apply_fsdp2(model, fsdp_kwargs, wrap_policy): + """model: AutoModelForCausalLM""" + assert ( + CPUOffloadPolicy is not None + ), "PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)" + + default_transformer_cls_names_to_wrap = getattr(model, "_no_split_modules", list()) + fsdp_transformer_layer_cls_to_wrap = ( + wrap_policy.transformer_layer_cls_to_wrap if wrap_policy is not None else list() + ) + if not fsdp_transformer_layer_cls_to_wrap: + fsdp_transformer_layer_cls_to_wrap = default_transformer_cls_names_to_wrap + + if isinstance(fsdp_transformer_layer_cls_to_wrap, str): + fsdp_transformer_layer_cls_to_wrap = [fsdp_transformer_layer_cls_to_wrap] + + assert ( + len(fsdp_transformer_layer_cls_to_wrap) > 0 + and fsdp_transformer_layer_cls_to_wrap[0] is not None + ) + + modules = [] + for name, module in model.named_modules(): + if module.__class__.__name__ in fsdp_transformer_layer_cls_to_wrap or ( + isinstance(module, nn.Embedding) and not model.config.tie_word_embeddings + ): + modules.append(module) + + for idx, module in enumerate(modules): + fully_shard(module, **fsdp_kwargs) + fully_shard( + model, **fsdp_kwargs + ) # fsdp2 will not reshard_after_forward for root module + + +def fsdp2_load_full_state_dict( + model: PreTrainedModel, + full_state: dict, + cpu_offload=None, + tie_word_embeddings=False, +): + """ + Loads the full state dict (could be only on rank 0) into the sharded model. This is done by broadcasting the + parameters from rank 0 to all other ranks. This function modifies the model in-place. + + Args: + model (`torch.nn.Module`): The model to load the state dict into + full_state (`dict`): The full state dict to load, can only be on rank 0 + """ + from torch.distributed.checkpoint.state_dict import ( + StateDictOptions, + set_model_state_dict, + ) + + device = torch.cuda.current_device() + model = model.to(device=device, non_blocking=True) + cpu_offload = cpu_offload is not None + options = StateDictOptions( + full_state_dict=True, + cpu_offload=cpu_offload, + broadcast_from_rank0=True, + strict=not tie_word_embeddings, + ) + set_model_state_dict(model, full_state, options=options) + + if tie_word_embeddings: + model.tie_weights() + + # rotary_emb is not in state_dict, so we need to broadcast it manually + for name, buf in model.named_buffers(): + dist.broadcast(buf, src=0) + + if cpu_offload: + model.to("cpu", non_blocking=True) + for buf in model.buffers(): + buf.data = buf.data.to(device) + + +def get_cosine_schedule_with_warmup( + optimizer: torch.optim.Optimizer, + num_warmup_steps: int, + num_training_steps: int, + min_lr_ratio: float = 0.0, + num_cycles: float = 0.5, + last_epoch: int = -1, +): + """ + Create a schedule with a learning rate that decreases following the values of the cosine function between the + initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the + initial lr set in the optimizer. + Args: + optimizer (:class:`~torch.optim.Optimizer`): + The optimizer for which to schedule the learning rate. + num_warmup_steps (:obj:`int`): + The number of steps for the warmup phase. + num_training_steps (:obj:`int`): + The total number of training steps. + min_lr_ratio (:obj:`float`, `optional`, defaults to 0.0): + The minimum lr ratio w.r.t the maximum. + num_cycles (:obj:`float`, `optional`, defaults to 0.5): + The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0 + following a half-cosine). + last_epoch (:obj:`int`, `optional`, defaults to -1): + The index of the last epoch when resuming training. + Return: + :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule. + """ + assert min_lr_ratio >= 0 and min_lr_ratio <= 1.0 + coef = (1 - min_lr_ratio) * 0.5 + intercept = (1 + min_lr_ratio) * 0.5 + + def lr_lambda(current_step): + if current_step < num_warmup_steps: + return min_lr_ratio + (1.0 - min_lr_ratio) * ( + float(current_step) / float(max(1, num_warmup_steps)) + ) + progress = float(current_step - num_warmup_steps) / float( + max(1, num_training_steps - num_warmup_steps) + ) + x = math.cos(math.pi * float(num_cycles) * 2.0 * progress) + return max(min_lr_ratio, x * coef + intercept) + + return torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda, last_epoch) diff --git a/areal/utils/functional.py b/areal/utils/functional.py new file mode 100644 index 0000000000..3abafa8a6b --- /dev/null +++ b/areal/utils/functional.py @@ -0,0 +1,179 @@ +from typing import Dict, Optional, Tuple + +import numpy as np +import torch +import torch.distributed as dist + + +@torch.compile +def _gather_logprobs( + logits: torch.Tensor, labels: torch.Tensor, temperature: float = 1.0 +): + log_probs = torch.nn.functional.log_softmax(logits.float() / temperature, dim=-1) + log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) + return log_probs_labels + + +@torch.compile +def _gather_logprobs_entropy( + logits: torch.Tensor, labels: torch.Tensor, temperature: float = 1.0 +): + log_probs = torch.nn.functional.log_softmax(logits.float() / temperature, dim=-1) + entropy = -torch.sum(log_probs.exp() * log_probs, dim=-1) + log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)).squeeze(-1) + return log_probs_labels, entropy + + +def gather_logprobs( + logits: torch.Tensor, + labels: torch.Tensor, + temperature: float = 1.0, + chunk_size: int = 1024, +): + batch_size = logits.shape[0] + + if batch_size <= chunk_size: + return _gather_logprobs(logits, labels, temperature) + + log_probs_labels_list = [] + + for i in range(0, batch_size, chunk_size): + end_idx = min(i + chunk_size, batch_size) + chunk_logits = logits[i:end_idx] + chunk_labels = labels[i:end_idx] + + chunk_log_probs = _gather_logprobs(chunk_logits, chunk_labels, temperature) + + log_probs_labels_list.append(chunk_log_probs) + + return torch.cat(log_probs_labels_list) + + +def gather_logprobs_entropy( + logits: torch.Tensor, + labels: torch.Tensor, + temperature: float = 1.0, + chunk_size: int = 1024, +): + batch_size = logits.shape[0] + + if batch_size <= chunk_size: + return _gather_logprobs_entropy(logits, labels, temperature) + + log_probs_labels_list = [] + entropy_list = [] + + for i in range(0, batch_size, chunk_size): + end_idx = min(i + chunk_size, batch_size) + chunk_logits = logits[i:end_idx] + chunk_labels = labels[i:end_idx] + + chunk_log_probs, chunk_entropy = _gather_logprobs_entropy( + chunk_logits, chunk_labels, temperature + ) + + log_probs_labels_list.append(chunk_log_probs) + entropy_list.append(chunk_entropy) + + return torch.cat(log_probs_labels_list), torch.cat(entropy_list) + + +@torch.no_grad() +def masked_normalization( + x: torch.Tensor, + mask: Optional[torch.Tensor] = None, + dim=None, + unbiased=False, + eps=1e-5, + high_precision=True, + all_reduce=True, + reduce_group=None, +): + dtype = torch.float64 if high_precision else torch.float32 + x = x.to(dtype) + if dim is None: + dim = tuple(range(len(x.shape))) + if mask is None: + factor = torch.tensor( + np.prod([x.shape[d] for d in dim]), dtype=dtype, device=x.device + ) + else: + mask = mask.to(dtype) + x = x * mask + factor = mask.sum(dim, keepdim=True) + x_sum = x.sum(dim=dim, keepdim=True) + x_sum_sq = x.square().sum(dim=dim, keepdim=True) + if dist.is_initialized() and all_reduce: + dist.all_reduce(factor, op=dist.ReduceOp.SUM, group=reduce_group) + dist.all_reduce(x_sum, op=dist.ReduceOp.SUM, group=reduce_group) + dist.all_reduce( + x_sum_sq, + op=dist.ReduceOp.SUM, + group=reduce_group, + ) + mean = x_sum / factor + meansq = x_sum_sq / factor + var = meansq - mean**2 + if unbiased: + var *= factor / (factor - 1) + return ((x - mean) / (var.sqrt() + eps)).float() + + +def ppo_actor_loss_fn( + logprobs: torch.Tensor, + proximal_logprobs: torch.Tensor, + old_logprobs: torch.Tensor, + advantages: torch.Tensor, + eps_clip: float, + loss_mask: torch.Tensor, + c_clip: Optional[float] = None, + behav_imp_weight_cap: Optional[float] = None, +) -> Tuple[torch.Tensor, Dict]: + """ + When decoupled loss is disabled: + 1. if recompute logp, both old_logprobs and proximal_logprobs are recomputed logp; + 2. if no recomputation, both old_logp and proximal_logprobs are produced by the inference backend. + + When decoupled loss is enabled, proximal_logprobs is the recomputed logp, + old_logprobs is produced by the inference engine. + """ + loss_mask_count = loss_mask.count_nonzero() or 1 + ratio = torch.where(loss_mask, torch.exp(logprobs - proximal_logprobs), 0) + clipped_ratio = torch.clamp(ratio, 1.0 - eps_clip, 1.0 + eps_clip) + pg_loss1 = -advantages * ratio + pg_loss2 = -advantages * clipped_ratio + clip_mask = pg_loss1.detach() < pg_loss2.detach() + pg_loss = torch.max(pg_loss1, pg_loss2) + if c_clip is not None: + assert c_clip > 1.0, c_clip + pg_loss3 = torch.sign(advantages) * c_clip * advantages + dual_clip_mask = pg_loss3.detach() < pg_loss.detach() + pg_loss = torch.min(pg_loss, pg_loss3) + else: + dual_clip_mask = torch.zeros_like(clip_mask) + behav_kl = proximal_logprobs - old_logprobs + behav_imp_weight = behav_kl.exp() + behav_mask = ( + (behav_imp_weight <= behav_imp_weight_cap).logical_and(loss_mask) + if behav_imp_weight_cap is not None + else loss_mask + ) + behav_kl = torch.where(behav_mask, behav_kl, 0.0) + behav_imp_weight = torch.where(behav_mask, behav_imp_weight, 0.0) + pg_loss = pg_loss * behav_imp_weight + logging_loss = pg_loss.detach() + pg_loss = torch.where(loss_mask, pg_loss, 0).sum() / loss_mask_count + clip_mask.logical_and_(loss_mask) + dual_clip_mask.logical_and_(loss_mask) + stat = dict( + loss=logging_loss, + importance_weight=ratio.detach(), + approx_kl=(logprobs - proximal_logprobs).detach(), + clip_mask=clip_mask, + dual_clip_mask=dual_clip_mask, + ) + if proximal_logprobs is not None: + stat["behave_imp_weight"] = behav_imp_weight + stat["behave_approx_kl"] = behav_kl + stat["behave_mask"] = behav_mask + return pg_loss, stat diff --git a/areal/utils/http.py b/areal/utils/http.py new file mode 100644 index 0000000000..0ddca140da --- /dev/null +++ b/areal/utils/http.py @@ -0,0 +1,90 @@ +import asyncio +from typing import Any, Dict, Optional + +import aiohttp + +from realhf.base import logging + +DEFAULT_RETRIES = 1 +DEFAULT_REQUEST_TIMEOUT = 3600 + + +logger = logging.getLogger(__file__) + + +def get_default_connector(): + return aiohttp.TCPConnector(limit=0, use_dns_cache=False, force_close=True) + + +async def arequest_with_retry( + addr: str, + endpoint: str, + payload: Optional[Dict[str, Any]] = None, + session: aiohttp.ClientSession | None = None, + method: str = "POST", + max_retries: Optional[int] = None, + timeout: Optional[float] = None, + retry_delay: float = 1.0, + verbose=False, +) -> Dict: + timeout = timeout or DEFAULT_REQUEST_TIMEOUT + last_exception = None + max_retries = max_retries or DEFAULT_RETRIES + base_url = f"http://{addr}" + url = f"{base_url}{endpoint}" + + timeo = aiohttp.ClientTimeout( + total=timeout, + sock_connect=timeout, + connect=timeout, + ) + if session is None: + _session = aiohttp.ClientSession( + timeout=timeo, + read_bufsize=1024 * 1024 * 10, + connector=get_default_connector(), + ) + else: + _session = session + + for attempt in range(max_retries): + try: + if verbose: + logger.info("enter client session, start sending requests") + if method.upper() == "GET": + ctx = _session.get(url, timeout=timeo) + elif method.upper() == "POST": + ctx = _session.post(url, json=payload, timeout=timeo) + elif method.upper() == "PUT": + ctx = _session.put(url, json=payload, timeout=timeo) + elif method.upper() == "DELETE": + ctx = _session.delete(url, timeout=timeo) + else: + raise ValueError(f"Unsupported HTTP method: {method}") + + async with ctx as response: + if verbose: + logger.info("http requests return") + response.raise_for_status() + res = await response.json() + if verbose: + logger.info("get http result") + if session is None: + await _session.close() + return res + except ( + aiohttp.ClientError, + aiohttp.ClientResponseError, + asyncio.TimeoutError, + ) as e: + last_exception = e + if attempt < max_retries - 1: + await asyncio.sleep(retry_delay) + continue + if session is None: + await _session.close() + raise RuntimeError( + f"Failed after {max_retries} retries each. " + f"Payload: {payload}. Addr: {addr}. Endpoint: {endpoint}. " + f"Last error: {last_exception}" + ) diff --git a/areal/utils/image.py b/areal/utils/image.py new file mode 100644 index 0000000000..1d8b6f06e9 --- /dev/null +++ b/areal/utils/image.py @@ -0,0 +1,43 @@ +import base64 +from io import BytesIO +from typing import List + +from PIL import Image +from PIL.Image import Image as ImageObject + + +def image2base64(images: List[ImageObject] | ImageObject) -> List[str] | str: + + if isinstance(images, ImageObject): + images = [images] + + byte_images = [] + for image in images: + with BytesIO() as buffer: + image.save(buffer, format="PNG") + buffer.seek(0) + byte_image = base64.b64encode(buffer.read()).decode("utf-8") + byte_images.append(byte_image) + + return byte_images + + +def pad_images_batch_to_max_size(images): + max_width = max(image.size[0] for image in images) + max_height = max(image.size[1] for image in images) + + padded_images = [] + + for image in images: + + width, height = image.size + + padding_left = (max_width - width) // 2 + padding_top = (max_height - height) // 2 + + padded_image = Image.new("RGB", (max_width, max_height), (0, 0, 0)) + padded_image.paste(image, (padding_left, padding_top)) + + padded_images.append(padded_image) + + return padded_images diff --git a/areal/utils/launcher.py b/areal/utils/launcher.py new file mode 100644 index 0000000000..69e312fa98 --- /dev/null +++ b/areal/utils/launcher.py @@ -0,0 +1,102 @@ +import getpass +import os +import pathlib +import time +from typing import Dict, Optional + +from areal.api.io_struct import AllocationMode, AllocationType +from realhf.base import logging, name_resolve, names + +logger = logging.getLogger("Launcher Utils") + +LOCAL_CACHE_DIR = "/tmp/areal" +PYTORCH_KERNEL_CACHE_PATH = ( + f"{LOCAL_CACHE_DIR}/.cache/{getpass.getuser()}/torch/kernels/" +) +TRITON_CACHE_PATH = f"{LOCAL_CACHE_DIR}/.cache/{getpass.getuser()}/triton/" +os.makedirs(PYTORCH_KERNEL_CACHE_PATH, exist_ok=True) +os.makedirs(TRITON_CACHE_PATH, exist_ok=True) +BASE_ENVIRONS = { + "TOKENIZERS_PARALLELISM": "true", + "PYTORCH_KERNEL_CACHE_PATH": PYTORCH_KERNEL_CACHE_PATH, + "TRITON_CACHE_DIR": TRITON_CACHE_PATH, + "CUDA_DEVICE_MAX_CONNECTIONS": "1", + "PYTHONPATH": str(pathlib.Path(__file__).resolve().parent.parent.parent), +} +NA132_ENVIRONS = { + "NCCL_SOCKET_IFNAME": "bond0", + "NCCL_NET_PLUGIN": "", + "NCCL_IB_GID_INDEX": "3", + "NCCL_IB_TIMEOUT": "2", + "NCCL_IB_RETRY_CNT": "7", + "NCCL_IB_SL": "5", + "NCCL_IB_TC": "136", + "NCCL_IB_HCA": "mlx5_bond", + "NCCL_IB_QPS_PER_CONNECTION": "8", + "NCCL_SET_THREAD_NAME": "1", + "NCCL_DEBUG": "WARN", + "NCCL_DEBUG_SUBSYS": "INIT,TUNING,GRAPH", +} +SGLANG_SERVER_WAIT_TIMEOUT_SECONDS = 180 + + +def get_env_vars( + cluster_name: str, additional_env_vars: Optional[str] = None +) -> Dict[str, str]: + """Returns the environment variables for the cluster.""" + _additional_env_vars = ( + dict(item.split("=") for item in additional_env_vars.split(",")) + if additional_env_vars + else dict() + ) + if cluster_name == "na132": + return {**BASE_ENVIRONS, **NA132_ENVIRONS, **_additional_env_vars} + else: + return {**BASE_ENVIRONS, **_additional_env_vars} + + +def wait_sglang_server_addrs( + experiment_name: str, + trial_name: str, + n_sglang_servers: int, +): + # Get SGLang slurm nodes, find the hosts + name = names.gen_servers(experiment_name, trial_name) + start = time.perf_counter() + while True: + sglang_addrs = name_resolve.get_subtree(name) + if len(sglang_addrs) >= n_sglang_servers: + logger.info( + f"Found {len(sglang_addrs)} SGLang servers: {', '.join(sglang_addrs)}" + ) + break + + time.sleep(1) + if time.perf_counter() - start > SGLANG_SERVER_WAIT_TIMEOUT_SECONDS: + raise TimeoutError( + f"Timeout waiting for SGLang servers to be ready. " + f"Expected {n_sglang_servers} servers, found {len(sglang_addrs)}." + ) + return sglang_addrs + + +def validate_config_for_distributed_launcher(config): + n_nodes = config.cluster.n_nodes + n_gpus_per_node = config.cluster.n_gpus_per_node + allocation_mode = config.allocation_mode + allocation_mode = AllocationMode.from_str(allocation_mode) + assert ( + allocation_mode.gen_world_size + allocation_mode.train_world_size + == n_nodes * n_gpus_per_node + ), ( + f"#GPUs required for allocation mode {allocation_mode.gen_world_size + allocation_mode.train_world_size} " + f"is not equal to #GPUs in the config {n_nodes * n_gpus_per_node}." + ) + if allocation_mode.type_ == AllocationType.DECOUPLED_SGLANG: + # Launcher should launch SGLang servers according to allocation mode. + assert ( + allocation_mode.gen_pp_size == 1 + ), "Pipeline generation in SGLang is not supported for now." + assert ( + allocation_mode.gen_tp_size <= config.cluster.n_gpus_per_node + ), "Currently only support SGLang TP size less <= #GPUs per node." diff --git a/areal/utils/model.py b/areal/utils/model.py new file mode 100644 index 0000000000..5e60f9b54b --- /dev/null +++ b/areal/utils/model.py @@ -0,0 +1,16 @@ +import torch + +VALID_VISION_MODELS = [ + "qwen2_vl", + "qwen2_5_vl", +] +# This registry is used to check if a model is a vision model that we have checked it works with AReaL. +# As different vision models vary in their image processing, special tokens and keys, etc. We will add models to this registry as we test them. +# If you want to add a new vision model, please make sure it works with AReaL. + + +# Copied from trl +def disable_dropout_in_model(model: torch.nn.Module) -> None: + for module in model.modules(): + if isinstance(module, torch.nn.Dropout): + module.p = 0 diff --git a/areal/utils/network.py b/areal/utils/network.py new file mode 100644 index 0000000000..5e51826cbb --- /dev/null +++ b/areal/utils/network.py @@ -0,0 +1,100 @@ +import random +import socket +from typing import List, Set + + +def gethostname(): + return socket.gethostname() + + +def gethostip(): + return socket.gethostbyname(socket.gethostname()) + + +def find_free_ports( + count: int, port_range: tuple = (1024, 65535), exclude_ports: Set[int] | None = None +) -> List[int]: + """ + Find multiple free ports within a specified range. + + Args: + count: Number of free ports to find + port_range: Tuple of (min_port, max_port) to search within + exclude_ports: Set of ports to exclude from search + + Returns: + List of free port numbers + + Raises: + ValueError: If unable to find requested number of free ports + """ + if exclude_ports is None: + exclude_ports = set() + + min_port, max_port = port_range + free_ports = [] + attempted_ports = set() + + # Calculate available port range + available_range = max_port - min_port + 1 - len(exclude_ports) + + if count > available_range: + raise ValueError( + f"Cannot find {count} ports in range {port_range}. " + f"Only {available_range} ports available." + ) + + max_attempts = count * 10 # Reasonable limit to avoid infinite loops + attempts = 0 + + while len(free_ports) < count and attempts < max_attempts: + # Generate random port within range + port = random.randint(min_port, max_port) + + # Skip if port already attempted or excluded + if port in attempted_ports or port in exclude_ports: + attempts += 1 + continue + + attempted_ports.add(port) + + if is_port_free(port): + free_ports.append(port) + + attempts += 1 + + if len(free_ports) < count: + raise ValueError( + f"Could only find {len(free_ports)} free ports " + f"out of {count} requested after {max_attempts} attempts" + ) + + return sorted(free_ports) + + +def is_port_free(port: int) -> bool: + """ + Check if a port is free by attempting to bind to it. + + Args: + port: Port number to check + + Returns: + True if port is free, False otherwise + """ + # Check TCP + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + try: + sock.bind(("", port)) + sock.close() + except OSError: + return False + + # Check UDP + sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM) + try: + sock.bind(("", port)) + sock.close() + return True + except OSError: + return False diff --git a/areal/utils/ray.py b/areal/utils/ray.py new file mode 100644 index 0000000000..ef9834c636 --- /dev/null +++ b/areal/utils/ray.py @@ -0,0 +1,23 @@ +import ray +from ray.util.placement_group import PlacementGroup +from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy + +from areal.utils.network import find_free_ports, gethostip + + +def get_placement_group_master_ip_and_port(placement_group: PlacementGroup): + def _master_ip_and_port(): + host_ip = gethostip() + port = find_free_ports(1, (10000, 60000))[0] + return host_ip, port + + future = ray.remote( + num_cpus=1, + num_gpus=0, + memory=10 * 1024 * 1024, # Convert MB to bytes + scheduling_strategy=PlacementGroupSchedulingStrategy( + placement_group=placement_group, + placement_group_bundle_index=0, + ), + )(_master_ip_and_port).remote() + return ray.get(future) diff --git a/areal/utils/save_load.py b/areal/utils/save_load.py new file mode 100644 index 0000000000..19bc9850da --- /dev/null +++ b/areal/utils/save_load.py @@ -0,0 +1,96 @@ +import os +from pathlib import Path +from typing import Dict + +import torch + + +def get_state_dict_from_repo_id_or_path(repo_id_or_path: str) -> Dict: + """ + Obtain a state dictionary from either a Hugging Face repo ID or a local path. + + Args: + repo_id_or_path (str): Either a Hugging Face repo ID (e.g., 'username/model-name') + or a local path to a directory containing model weights. + + Returns: + Dict: The combined state dictionary from all .safetensors and .bin files. + """ + from safetensors.torch import load_file as safetensors_load + + state_dict = {} + + # Step 1: Identify if the input is a Hugging Face repo ID or local path + try: + from huggingface_hub.utils import HFValidationError, validate_repo_id + + try: + validate_repo_id(repo_id_or_path) + is_hf_repo = True + except HFValidationError: + is_hf_repo = False + except ImportError: + is_hf_repo = False + + if is_hf_repo: + from huggingface_hub import snapshot_download + + # Step 2: Download the repo if it's a Hugging Face repo ID + local_path = snapshot_download( + repo_id=repo_id_or_path, + ) + else: + # Assume it's a local path + local_path = repo_id_or_path + + # Step 3: Load all .safetensors and .bin files + file_paths_to_load = [] + if os.path.isdir(local_path): + for filename in os.listdir(local_path): + filepath = os.path.join(local_path, filename) + if filename.endswith(".safetensors") or filename.endswith(".bin"): + file_paths_to_load.append(filepath) + elif os.path.isfile(local_path): + file_paths_to_load.append(local_path) + else: + raise ValueError( + f"Local path {local_path} does not exist or is not a valid path, " + f"or {local_path} is a huggingface repo id but huggingface_hub is not installed." + ) + + def _load(filepath: str): + if filepath.endswith(".safetensors"): + state_dict = safetensors_load(filepath) + elif filepath.endswith(".bin"): + state_dict = torch.load(filepath, map_location="cpu") + else: + raise ValueError(f"{filepath} is not a torch bin or safetensor file.") + return state_dict + + state_dict = {} + + from concurrent.futures import ThreadPoolExecutor, as_completed + + with ThreadPoolExecutor( + max_workers=min(4, max(1, os.cpu_count() // 8)) + ) as executor: + future_to_checkpoint = { + executor.submit(_load, path): path for path in file_paths_to_load + } + + for future in as_completed(future_to_checkpoint): + path = future_to_checkpoint[future] + try: + sd = future.result() + state_dict.update(sd) + except Exception as e: + raise RuntimeError(f"Error loading checkpoint from {path}: {e}") + return state_dict + + +def is_existing_local_path(path: str) -> bool: + try: + path_obj = Path(path) + return path_obj.exists() and (path_obj.is_file() or path_obj.is_dir()) + except (ValueError, OSError): + return False diff --git a/areal/utils/saver.py b/areal/utils/saver.py new file mode 100644 index 0000000000..dd911a0369 --- /dev/null +++ b/areal/utils/saver.py @@ -0,0 +1,83 @@ +import getpass +import os + +from transformers import AutoProcessor, PreTrainedTokenizerFast + +from areal.api.cli_args import SaverConfig +from areal.api.engine_api import TrainEngine +from areal.api.io_struct import FinetuneSpec, SaveLoadMeta +from realhf.base import timeutil + + +class Saver: + + def __init__(self, config: SaverConfig, ft_spec: FinetuneSpec, for_recover: bool): + self.config = config + self.ft_sepc = ft_spec + self.for_recover = for_recover + self.freq_ctl = timeutil.EpochStepTimeFreqCtl( + freq_epoch=config.freq_epochs, + freq_step=config.freq_steps, + freq_sec=config.freq_secs, + ) + + @staticmethod + def get_save_checkpoint_root( + config: SaverConfig, + name: str = "default", + ): + path = os.path.join( + f"{config.fileroot}/checkpoints/{getpass.getuser()}/{config.experiment_name}/{config.trial_name}", + name, + ) + os.makedirs(path, exist_ok=True) + return path + + @staticmethod + def get_save_checkpoint_path( + config: SaverConfig, + epoch: int, + step: int, + globalstep: int, + name: str = "default", + ): + path = os.path.join( + Saver.get_save_checkpoint_root(config, name), + f"epoch{epoch}epochstep{step}globalstep{globalstep}", + ) + os.makedirs(path, exist_ok=True) + return path + + def save( + self, + engine: TrainEngine, + epoch: int, + step: int, + global_step: int, + name: str = "default", + tokenizer: PreTrainedTokenizerFast | None = None, + processor: AutoProcessor | None = None, + base_model_path: str | None = None, + ): + if not self.freq_ctl.check( + epochs=int(step == self.ft_sepc.steps_per_epoch - 1), steps=1 + ): + return + path = Saver.get_save_checkpoint_path( + self.config, epoch, step, global_step, name + ) + weight_format = "hf" + with_optim = False + if self.for_recover: + weight_format = "dcp" + with_optim = True + + meta = SaveLoadMeta( + path=path, + weight_format=weight_format, + with_optim=with_optim, + tokenizer=tokenizer, + processor=processor, + base_model_path=base_model_path, + ) + engine.save(meta) diff --git a/areal/utils/slurm.py b/areal/utils/slurm.py new file mode 100644 index 0000000000..f876c5e84a --- /dev/null +++ b/areal/utils/slurm.py @@ -0,0 +1,154 @@ +import subprocess +from typing import List, Literal, Optional + +from realhf.base import logging +from realhf.scheduler.client import JobInfo, JobState + +logger = logging.getLogger("Slurm Utils") + + +SQUEUE_FIELDS = [ + "JobID", + "State", + "SubmitTime", + "StartTime", + "Name", + "NodeList", + "UserName", + "MaxCPUs", + "cpus-per-task", + "NumTasks", + "tres-alloc", +] +STATUS_MAPPING = { + "RUNNING": JobState.RUNNING, + "COMPLETING": JobState.RUNNING, + "PENDING": JobState.PENDING, + "CANCELLED": JobState.CANCELLED, + "FAILED": JobState.FAILED, + "COMPLETED": JobState.COMPLETED, + "OUT_OF_MEMORY": JobState.FAILED, + "DEADLINE": JobState.COMPLETED, + "TIMEOUT": JobState.COMPLETED, +} + +SBATCH_SCRIPT_TEMPLATE = """#!/bin/bash +{sbatch_options} + +# Getting the node names +nodes=$(scontrol show hostnames "$SLURM_JOB_NODELIST") +echo nodes=$nodes + +nodes_array=($nodes) +echo node_array=$nodes_array + +head_node=${{nodes_array[0]}} +echo head_node=$head_node + +# Getting the head node IP address +head_node_ip=$(srun {srun_additional_args} --nodes=1 --ntasks=1 -n1 -c1 --mem=10M --nodelist="$head_node" hostname --ip-address) +echo head_node_ip=$head_node_ip + +# Find a free port on the head node +# Wonderful linux command to find a random free port (between 10000 and 60000) by deepseek +trainer_port=$(srun {srun_additional_args} --nodes=1 --ntasks=1 -n1 -c1 --mem=10M --nodelist="$head_node" bash -c "comm -23 <(seq 10000 60000 | sort) <(ss -tan | awk '{{print $4}}' | cut -d':' -f2 | grep '[0-9]\\{{1,5\\}}' | sort -u) | shuf | head -n 1") +echo trainer_port=$trainer_port + +# srun commands +{srun_cmds} + +wait +""" + +SRUN_CMD_TEMPLATE: str = """srun {additional_args} \\ + --nodelist=${{nodes_array[{node_id}]}} --nodes={nodes} --ntasks={ntasks} \\ + --gres=gpu:{n_gpus_per_node} --cpus-per-task={cpus_per_task} --mem-per-cpu={mem_per_cpu}M \\ + {cmd} & +""" + +APPTAINER_CMD_TEMPLATE: str = """singularity exec --no-home --writable-tmpfs --nv --pid \\ + --bind {container_mounts} \\ + {container_env_strings} \\ + {container_image} \\ + {cmd}""" + + +def cancel_jobs( + slurm_names: Optional[List[str]] = None, + slurm_ids: Optional[List[int]] = None, + signal: Literal["SIGINT", "SIGKILL"] = "SIGKILL", +): + assert ( + slurm_names is not None or slurm_ids is not None + ), "Must specify slurm_names or slurm_ids." + assert not ( + slurm_names and slurm_ids + ), "Cannot specify both slurm_names and slurm_ids." + cmd = ["scancel", "-s", signal] + if slurm_names is not None: + cmd += ["-n", ",".join(slurm_names)] + elif slurm_ids is not None: + cmd += [",".join(str(s) for s in slurm_ids)] + subprocess.check_call(cmd) + logger.info( + f"Cancelled Slurm job with signal {signal}: " + f"slurm identifiers {slurm_names if slurm_ids is None else slurm_ids}. CMD: {cmd}" + ) + + +def query_jobs( + slurm_names: Optional[List[str]] = None, + slurm_ids: Optional[List[int]] = None, + status: str = "all", + delimiter: str = "__PSI__", +) -> List[JobInfo]: + squeue_format = f":.{delimiter},".join(SQUEUE_FIELDS) + cmd = ["squeue", "-O", squeue_format, f"-t{status}"] + if slurm_names is not None: + cmd += ["-n", ",".join(slurm_names)] + if slurm_ids is not None: + cmd += ["-j", ",".join([str(s) for s in slurm_ids])] + + output = ( + subprocess.check_output(cmd, stderr=subprocess.DEVNULL).decode("ascii").strip() + ) + rs = [] + for line in output.split("\n")[1:]: + job_id, state, submit_time, start_time, slurm_name, nodelist, *_ = line.split( + delimiter + ) + rs.append( + JobInfo( + name=slurm_name, + state=STATUS_MAPPING[state], + host=nodelist, + submit_time=submit_time, + start_time=start_time, + slurm_id=int(job_id.strip()), + ) + ) + return rs + + +def parse_slurm_nodelist(nodelist: str) -> List[str]: + return ( + subprocess.check_output( + [ + "scontrol", + "show", + "hostnames", + nodelist, + ] + ) + .decode("utf-8") + .strip() + .split("\n") + ) + + +def get_slurm_host_ip(node: str, srun_addtional_args: str): + try: + cmd = f"srun {srun_addtional_args} --immediate=1 --nodes=1 --ntasks=1 -n1 -c1 --mem=10M --nodelist={node} hostname --ip-address" + return subprocess.check_output(cmd.split(" ")).decode("utf-8").strip() + except subprocess.CalledProcessError: + logger.warning(f"Get slurm host IP for node {node} failed.") diff --git a/areal/utils/stats_logger.py b/areal/utils/stats_logger.py new file mode 100644 index 0000000000..d801861aae --- /dev/null +++ b/areal/utils/stats_logger.py @@ -0,0 +1,94 @@ +import getpass +import os +import time +from typing import Dict, List + +import torch.distributed as dist +import wandb +from tensorboardX import SummaryWriter + +from areal.api.cli_args import StatsLoggerConfig +from areal.api.io_struct import FinetuneSpec +from realhf.api.core.data_api import tabulate_stats +from realhf.base import logging + +logger = logging.getLogger("StatsLogger", "system") + + +class StatsLogger: + + def __init__(self, config: StatsLoggerConfig, ft_spec: FinetuneSpec): + self.config = config + self.ft_spec = ft_spec + self.init() + + self._last_commit_step = 0 + + def init(self): + if dist.is_initialized() and dist.get_rank() != 0: + return + + self.start_time = time.perf_counter() + # wandb init, connect to remote wandb host + if self.config.wandb.mode != "disabled": + wandb.login() + wandb.init( + mode=self.config.wandb.mode, + entity=self.config.wandb.entity, + project=self.config.wandb.project or self.config.experiment_name, + name=self.config.wandb.name or self.config.trial_name, + job_type=self.config.wandb.job_type, + group=self.config.wandb.group + or f"{self.config.experiment_name}_{self.config.trial_name}", + notes=self.config.wandb.notes, + tags=self.config.wandb.tags, + config=self.config.wandb.config, + dir=self.get_log_path(self.config), + force=True, + id=f"{self.config.experiment_name}_{self.config.trial_name}_train", + resume="allow", + settings=wandb.Settings(start_method="fork"), + ) + # tensorboard logging + self.summary_writer = None + if self.config.tensorboard.path is not None: + self.summary_writer = SummaryWriter(log_dir=self.config.tensorboard.path) + + def close(self): + if dist.is_initialized() and dist.get_rank() != 0: + return + logger.info( + f"Training completes! Total time elapsed {time.monotonic() - self.start_time:.2f}." + ) + wandb.finish() + if self.summary_writer is not None: + self.summary_writer.close() + + def commit(self, epoch: int, step: int, global_step: int, data: Dict | List[Dict]): + if dist.is_initialized() and dist.get_rank() != 0: + return + logger.info( + f"Epoch {epoch+1}/{self.ft_spec.total_train_epochs} " + f"Step {step+1}/{self.ft_spec.steps_per_epoch} " + f"Train step {global_step + 1}/{self.ft_spec.total_train_steps} done." + ) + if isinstance(data, Dict): + data = [data] + log_step = max(global_step, self._last_commit_step + 1) + for i, item in enumerate(data): + logger.info(f"Stats ({i+1}/{len(data)}):") + self.print_stats(item) + wandb.log(item, step=log_step + i) + if self.summary_writer is not None: + for key, val in item.items(): + self.summary_writer.add_scalar(f"{key}", val, log_step + i) + self._last_commit_step = log_step + len(data) - 1 + + def print_stats(self, stats: Dict[str, float]): + logger.info("\n" + tabulate_stats(stats)) + + @staticmethod + def get_log_path(config: StatsLoggerConfig): + path = f"{config.fileroot}/logs/{getpass.getuser()}/{config.experiment_name}/{config.trial_name}" + os.makedirs(path, exist_ok=True) + return path diff --git a/areal/utils/wrapper.py b/areal/utils/wrapper.py new file mode 100644 index 0000000000..3ffdfdd4c8 --- /dev/null +++ b/areal/utils/wrapper.py @@ -0,0 +1,47 @@ +import inspect +from functools import partial +from typing import Callable, Dict, Optional + + +def wrapable(name: Optional[str] = None): + def decorator(func): + func.__wrap_meta__ = { + "name": name or func.__name__, + } + return func + + return decorator + + +def wrap(target: object, source: object, transform: Optional[Callable] = None): + for name, member in inspect.getmembers(source): + if callable(member) and hasattr(member, "__wrap_meta__"): + meta = member.__wrap_meta__ + method_name = meta["name"] + + if transform: + setattr( + target, + method_name, + partial( + transform, + wrap_method_name=method_name, + wrap_original_method=member, + ), + ) + else: + setattr(target, method_name, member) + + +def wrap_get_method_name(kwargs) -> str: + return kwargs["wrap_method_name"] + + +def wrap_get_method(kwargs) -> Callable: + return kwargs["wrap_original_method"] + + +def wrap_remove_meta(kwargs) -> Dict: + del kwargs["wrap_method_name"] + del kwargs["wrap_original_method"] + return kwargs diff --git a/areal/workflow/multi_turn.py b/areal/workflow/multi_turn.py new file mode 100644 index 0000000000..1681d44c66 --- /dev/null +++ b/areal/workflow/multi_turn.py @@ -0,0 +1,87 @@ +import uuid + +import torch +from transformers import PreTrainedTokenizerFast + +from areal.api.cli_args import GenerationHyperparameters +from areal.api.engine_api import InferenceEngine +from areal.api.io_struct import LLMRequest +from areal.api.workflow_api import RolloutWorkflow +from areal.utils.data import concat_padded_tensors + + +class MultiTurnWorkflow(RolloutWorkflow): + def __init__( + self, + reward_fn, + gconfig: GenerationHyperparameters, + tokenizer: PreTrainedTokenizerFast, + max_turns: int, + turn_discount: float, + ): + self.reward_fn = reward_fn + self.gconfig = gconfig + self.tokenizer = tokenizer + self.max_turns = max_turns + self.turn_discount = turn_discount + + async def arun_episode(self, engine: InferenceEngine, data): + # Placeholders for the results + seq, logprobs, loss_mask, versions = [], [], [], [] + messages = data["messages"] + # Run multi-turn rollout until correct + t = reward = 0 + discount = 1 + rid = uuid.uuid4().hex + while reward == 0 and t < self.max_turns: + # Amend a prompt if the previous answer is incorrect + if t > 0: + messages += [ + {"role": "asistant", "content": completions_str}, + { + "role": "user", + "content": "Your answer is not correct. Please try to answer it again.", + }, + ] + # Convert the prompt into input_ids + input_ids = self.tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + ) + # Send generate request to get the response. + req = LLMRequest( + rid=rid, + input_ids=input_ids, + gconfig=self.gconfig.new(n_samples=1), + ) + resp = await engine.agenerate(req) + # compute reward: 1 for correct and 0 otherwise + prompt_str = self.tokenizer.decode(input_ids) + completions_str = self.tokenizer.decode(resp.output_tokens) + reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) + # Amend results + input_len = len(resp.input_tokens) - len(seq) + seq += resp.input_tokens[-input_len:] + resp.output_tokens + logprobs += [0.0] * input_len + resp.output_logprobs + loss_mask += [0] * input_len + [1] * resp.output_len + versions += [-1] * input_len + resp.output_versions + # Increase counter + t += 1 + discount *= self.turn_discount + res = dict( + seq=torch.tensor(seq), + logprobs=torch.tensor(logprobs), + loss_mask=torch.tensor(loss_mask), + versions=torch.tensor(versions), + rewards=torch.tensor([float(reward * discount)]), + attetion_mask=torch.ones(len(seq), dtype=torch.bool), + ) + res = {k: v.unsqueeze(0) for k, v in res.items()} + return concat_padded_tensors([res]) diff --git a/areal/workflow/rlvr.py b/areal/workflow/rlvr.py new file mode 100644 index 0000000000..d0676412b9 --- /dev/null +++ b/areal/workflow/rlvr.py @@ -0,0 +1,137 @@ +import asyncio +import functools +import os +import uuid +from concurrent.futures import ProcessPoolExecutor + +import colorama +import torch +from tensordict import TensorDict +from transformers import PreTrainedTokenizerFast + +from areal.api.cli_args import GenerationHyperparameters +from areal.api.engine_api import InferenceEngine +from areal.api.io_struct import LLMRequest +from areal.api.workflow_api import RolloutWorkflow +from areal.utils.data import concat_padded_tensors +from realhf.base import logging + +logger = logging.getLogger("RLVR workflow") + +REWARD_TIMEOUT_SECONDS = 15 + + +class RLVRWorkflow(RolloutWorkflow): + def __init__( + self, + reward_fn, + gconfig: GenerationHyperparameters, + tokenizer: PreTrainedTokenizerFast, + enable_thinking: bool, + dump_dir: str | None = None, + ): + self.reward_fn = reward_fn + self.gconfig = gconfig + self.tokenizer = tokenizer + self.enable_thinking = enable_thinking + self.dump_dir = dump_dir + self.rw_executor = ProcessPoolExecutor(max_workers=4) + if self.dump_dir is not None and not os.path.exists(self.dump_dir): + os.makedirs(self.dump_dir, exist_ok=True) + + async def arun_episode(self, engine: InferenceEngine, data): + input_ids = self.tokenizer.apply_chat_template( + data["messages"], + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + + n_samples = self.gconfig.n_samples + req = LLMRequest( + rid=uuid.uuid4().hex, + input_ids=input_ids, + gconfig=self.gconfig.new(n_samples=1), + ) + resps = await asyncio.gather(*[engine.agenerate(req) for _ in range(n_samples)]) + + version = engine.get_version() + prompt_strs = [] + completions_strs = [] + rewards = [] + seqlens = [] + + results = [] + loop = asyncio.get_event_loop() + for resp in resps: + seq = resp.input_tokens + resp.output_tokens + logprobs = [0.0] * resp.input_len + resp.output_logprobs + loss_mask = [0] * resp.input_len + [1] * resp.output_len + versions = [-1] * resp.input_len + resp.output_versions + + prompt_str = self.tokenizer.decode(input_ids) + completions_str = self.tokenizer.decode(resp.output_tokens) + prompt_strs.append(prompt_str) + completions_strs.append(completions_str) + seqlens.append(len(seq)) + try: + reward = await asyncio.wait_for( + loop.run_in_executor( + self.rw_executor, + functools.partial( + self.reward_fn, + prompt_str, + completions_str, + resp.input_tokens, + resp.output_tokens, + **data, + ), + ), + timeout=REWARD_TIMEOUT_SECONDS, + ) + except asyncio.TimeoutError: + logger.warning( + f"Computing reward timeout after {REWARD_TIMEOUT_SECONDS}s. Set reward to 0." + ) + reward = 0 + rewards.append(reward) + res = dict( + # unsqueeze to add an additional batch dimension + input_ids=torch.tensor(seq).unsqueeze(0), + loss_mask=torch.tensor(loss_mask).unsqueeze(0), + logprobs=torch.tensor(logprobs).unsqueeze(0), + versions=torch.tensor(versions).unsqueeze(0), + attention_mask=torch.ones(len(seq), dtype=torch.bool).unsqueeze(0), + # reward + rewards=torch.tensor([float(reward)]), + ) + results.append(TensorDict(res, batch_size=[1])) + + if self.dump_dir is not None: + os.makedirs(os.path.join(self.dump_dir, str(version)), exist_ok=True) + # Get the unique identifier for this prompt + qid = None + for key in ["query_id", "id", "qid"]: + qid = data.get(key, None) + if qid is not None: + break + qid = qid or uuid.uuid4().hex + + # Dump rollout to file + with open( + os.path.join(self.dump_dir, str(version), f"{qid}.txt"), "a" + ) as f: + n_samples = self.gconfig.n_samples + for i, (p, c, r, sl) in enumerate( + zip(prompt_strs, completions_strs, rewards, seqlens) + ): + info = "\n".join( + [ + f"idx: {i + 1} / {n_samples}, seqlen: {sl}, reward is {r}.", + f"prompt is \n{colorama.Fore.YELLOW + colorama.Style.DIM}{p}{colorama.Style.RESET_ALL}", + f"sequence is: \n{colorama.Fore.YELLOW + colorama.Style.DIM}{c}{colorama.Style.RESET_ALL}", + ] + ) + f.write(info + "\n") + + return concat_padded_tensors(results) diff --git a/areal/workflow/vision_rlvr.py b/areal/workflow/vision_rlvr.py new file mode 100644 index 0000000000..ed739e68cb --- /dev/null +++ b/areal/workflow/vision_rlvr.py @@ -0,0 +1,121 @@ +import asyncio +import os +import uuid + +import colorama +import torch +from tensordict import TensorDict +from transformers import AutoProcessor, PreTrainedTokenizerFast + +from areal.api.cli_args import GenerationHyperparameters +from areal.api.io_struct import VLMRequest +from areal.utils.data import concat_padded_tensors +from areal.utils.image import image2base64, pad_images_batch_to_max_size +from areal.workflow.rlvr import RLVRWorkflow + + +class VisionRLVRWorkflow(RLVRWorkflow): + def __init__( + self, + reward_fn, + gconfig: GenerationHyperparameters, + tokenizer: PreTrainedTokenizerFast, + processor: AutoProcessor, + enable_thinking: bool, + dump_dir: str | None = None, + ): + super().__init__(reward_fn, gconfig, tokenizer, enable_thinking, dump_dir) + self.processor = processor + + async def arun_episode(self, engine, data): + + padded_images = pad_images_batch_to_max_size(data["images"]) + + processed_input = self.processor( + images=padded_images, + text=data["messages"], + padding=False, + return_tensors="pt", + ) + + input_ids = processed_input["input_ids"].tolist()[0] + + n_samples = self.gconfig.n_samples + + byte_images = image2base64(padded_images) + + req = VLMRequest( + rid=uuid.uuid4().hex, + input_ids=input_ids, + image_data=byte_images, + gconfig=self.gconfig.new(n_samples=1), + ) + resps = await asyncio.gather(*[engine.agenerate(req) for _ in range(n_samples)]) + + version = engine.get_version() + prompt_strs = [] + completions_strs = [] + rewards = [] + seqlens = [] + + results = [] + for resp in resps: + seq = resp.input_tokens + resp.output_tokens + logprobs = [0.0] * resp.input_len + resp.output_logprobs + loss_mask = [0] * resp.input_len + [1] * resp.output_len + versions = [-1] * resp.input_len + resp.output_versions + + prompt_str = self.tokenizer.decode(input_ids) + completions_str = self.tokenizer.decode(resp.output_tokens) + prompt_strs.append(prompt_str) + completions_strs.append(completions_str) + seqlens.append(len(seq)) + reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) + rewards.append(reward) + res = dict( + # unsqueeze to add an additional batch dimension + input_ids=torch.tensor(seq).unsqueeze(0), + loss_mask=torch.tensor(loss_mask).unsqueeze(0), + pixel_values=processed_input["pixel_values"].unsqueeze(0), + image_grid_thw=processed_input["image_grid_thw"].unsqueeze(0), + logprobs=torch.tensor(logprobs).unsqueeze(0), + versions=torch.tensor(versions).unsqueeze(0), + attention_mask=torch.ones(len(seq), dtype=torch.bool).unsqueeze(0), + # reward + rewards=torch.tensor([reward]), + ) + results.append(TensorDict(res, batch_size=[1])) + if self.dump_dir is not None: + os.makedirs(os.path.join(self.dump_dir, str(version)), exist_ok=True) + # Get the unique identifier for this prompt + qid = None + for key in ["query_id", "id", "qid"]: + qid = data.get(key, None) + if qid is not None: + break + qid = qid or uuid.uuid4().hex + + # Dump rollout to file + with open( + os.path.join(self.dump_dir, str(version), f"{qid}.txt"), "a" + ) as f: + n_samples = self.gconfig.n_samples + for i, (p, c, r, sl) in enumerate( + zip(prompt_strs, completions_strs, rewards, seqlens) + ): + info = "\n".join( + [ + f"idx: {i + 1} / {n_samples}, seqlen: {sl}, reward is {r}.", + f"prompt is \n{colorama.Fore.YELLOW + colorama.Style.DIM}{p}{colorama.Style.RESET_ALL}", + f"sequence is: \n{colorama.Fore.YELLOW + colorama.Style.DIM}{c}{colorama.Style.RESET_ALL}", + ] + ) + f.write(info + "\n") + + return concat_padded_tensors(results) diff --git a/assets/areal_lite_layers.png b/assets/areal_lite_layers.png new file mode 100644 index 0000000000..f57fdab3b5 Binary files /dev/null and b/assets/areal_lite_layers.png differ diff --git a/assets/logo_lite.png b/assets/logo_lite.png new file mode 100644 index 0000000000..fd9c61bebe Binary files /dev/null and b/assets/logo_lite.png differ diff --git a/blog/AReaL_v0_3.md b/blog/AReaL_v0_3.md index 0222343c6a..9ae8334f43 100644 --- a/blog/AReaL_v0_3.md +++ b/blog/AReaL_v0_3.md @@ -6,59 +6,81 @@ We now release AReaL v0.3, featuring three major milestones: -- **A fully asynchronous RL training pipeline with system and RL algorithm co-design**, achieving over 2.77x speedup without any performance drop -- **SOTA coding models**, i.e., a 14B model with a **69.1 score on LCB-v5**. Reproducible results with fully open-sourced datasets are also provided +- **A fully asynchronous RL training pipeline with system and RL algorithm co-design**, + achieving over 2.77x speedup without any performance drop +- **SOTA coding models**, i.e., a 14B model with a **69.1 score on LCB-v5**. + Reproducible results with fully open-sourced datasets are also provided - Experimental support for **multi-turn agentic RL training** ### Performance Results -| **Model (8B)** | **LiveCodeBench v5** (2024.10-2025.2) | **Codeforce** | **CodeContests** | -|:---:|:---:|:---:|:---:| -| Qwen3-8B | 58.8 | 1879/96.7% | 31.4 | -| DeepSeek-R1-0528-Qwen3-8B | 58.4 | 1945/97.3% | 31.0 | -| [🤗 AReaL-boba²-8B-Open](https://huggingface.co/inclusionAI/AReaL-boba-2-8B-subset) | 62.0 | 1933/97.2% | **41.4** | -| [🤗 AReaL-boba²-8B](https://huggingface.co/inclusionAI/AReaL-boba-2-8B) | **63.0** | **1962/97.5%** | 40.8 | - -| **Model (14B)** | **LiveCodeBench v5** (2024.10-2025.2) | **Codeforce** | **CodeContests** | -|:---:|:---:|:---:|:---:| -| Qwen3-14B | 65.4 | 1978/97.7% | 38.3 | -| DeepCoder-14B-Preview | 60.6 | 1936/95.3% | 40.1 | -| [🤗 AReaL-boba²-14B-Open](https://huggingface.co/inclusionAI/AReaL-boba-2-14B-subset) | 67.3 | 1990/97.8% | **46.2** | -| [🤗 AReal-boba²-14B](https://huggingface.co/inclusionAI/AReaL-boba-2-14B) | **69.1** | **2044/98.2%** | 46.1 | - -| **Larger Models** | **LiveCodeBench v5** (2024.10-2025.2) | **Codeforce** | **CodeContests** | -|:---:|:---:|:---:|:---:| -| Qwen3-235B | 70.7 | 2056 | - | -| DeepSeek-R1 | 64.3 | 2029 | - | -| OpenAI-o3-mini (Medium) | 66.3 | 2036 | - | - -*Table 1: Coding Task Performance Comparison. AReaL-boba²-8B/14B-Open denotes training results on open-sourced data. AReaL-boba²-8B/14B models are trained with an additional small amount of internal data and achieve SOTA performance on LiveCodeBench, Codeforce & CodeContests.* - -To access our latest models and training data, please visit this [Huggingface Link](https://huggingface.co/collections/inclusionAI/areal-boba-2-683f0e819ccb7bb2e1b2f2d5). +| **Model (8B)** | **LiveCodeBench v5** (2024.10-2025.2) | **Codeforce** | **CodeContests** | +| :---------------------------------------------------------------------------------: | :-----------------------------------: | :------------: | :--------------: | +| Qwen3-8B | 58.8 | 1879/96.7% | 31.4 | +| DeepSeek-R1-0528-Qwen3-8B | 58.4 | 1945/97.3% | 31.0 | +| [🤗 AReaL-boba²-8B-Open](https://huggingface.co/inclusionAI/AReaL-boba-2-8B-subset) | 62.0 | 1933/97.2% | **41.4** | +| [🤗 AReaL-boba²-8B](https://huggingface.co/inclusionAI/AReaL-boba-2-8B) | **63.0** | **1962/97.5%** | 40.8 | + +| **Model (14B)** | **LiveCodeBench v5** (2024.10-2025.2) | **Codeforce** | **CodeContests** | +| :-----------------------------------------------------------------------------------: | :-----------------------------------: | :------------: | :--------------: | +| Qwen3-14B | 65.4 | 1978/97.7% | 38.3 | +| DeepCoder-14B-Preview | 60.6 | 1936/95.3% | 40.1 | +| [🤗 AReaL-boba²-14B-Open](https://huggingface.co/inclusionAI/AReaL-boba-2-14B-subset) | 67.3 | 1990/97.8% | **46.2** | +| [🤗 AReal-boba²-14B](https://huggingface.co/inclusionAI/AReaL-boba-2-14B) | **69.1** | **2044/98.2%** | 46.1 | + +| **Larger Models** | **LiveCodeBench v5** (2024.10-2025.2) | **Codeforce** | **CodeContests** | +| :---------------------: | :-----------------------------------: | :-----------: | :--------------: | +| Qwen3-235B | 70.7 | 2056 | - | +| DeepSeek-R1 | 64.3 | 2029 | - | +| OpenAI-o3-mini (Medium) | 66.3 | 2036 | - | + +*Table 1: Coding Task Performance Comparison. AReaL-boba²-8B/14B-Open denotes training +results on open-sourced data. AReaL-boba²-8B/14B models are trained with an additional +small amount of internal data and achieve SOTA performance on LiveCodeBench, Codeforce & +CodeContests.* + +To access our latest models and training data, please visit this +[Huggingface Link](https://huggingface.co/collections/inclusionAI/areal-boba-2-683f0e819ccb7bb2e1b2f2d5). ## Motivation for Asynchronous RL System ### Inference devices are underutilized -During the synchronous RL training process, the generation step must wait until the finish of the longest output within a batch. Due to the varying output lengths for LRM, a synchronous RL system suffers from severe training inefficiency. +During the synchronous RL training process, the generation step must wait until the +finish of the longest output within a batch. Due to the varying output lengths for LRM, +a synchronous RL system suffers from severe training inefficiency. -Some works ([DeepCoder](https://pretty-radio-b75.notion.site/DeepCoder-A-Fully-Open-Source-14B-Coder-at-O3-mini-Level-1cf81902c14680b3bee5eb349a512a51), [Intellect](https://www.primeintellect.ai/blog/intellect-2)) use outputs generated from a previous model version to update the current model. For the best performances, the model version used for rollout generation is limited to only one or two steps older. However, all these systems still follow a batched generation setting, the issue of system inefficiency during the generation phase still remains unaddressed. +Some works +([DeepCoder](https://pretty-radio-b75.notion.site/DeepCoder-A-Fully-Open-Source-14B-Coder-at-O3-mini-Level-1cf81902c14680b3bee5eb349a512a51), +[Intellect](https://www.primeintellect.ai/blog/intellect-2)) use outputs generated from +a previous model version to update the current model. For the best performances, the +model version used for rollout generation is limited to only one or two steps older. +However, all these systems still follow a batched generation setting, the issue of +system inefficiency during the generation phase still remains unaddressed. ![](/assets/sync_one_step_gen.png) -*Fig.1. Left: Execution timeline of a synchronous RL training. Right: Execution timeline of one-step overlap RL system.* +*Fig.1. Left: Execution timeline of a synchronous RL training. Right: Execution timeline +of one-step overlap RL system.* ### Scalability is poor in synchronous RL systems -Synchronous systems distribute generation across all devices, reducing the per-GPU decoding batch size. This pushes the decoding process into a memory-IO-bound regime where additional devices fail to improve throughput. +Synchronous systems distribute generation across all devices, reducing the per-GPU +decoding batch size. This pushes the decoding process into a memory-IO-bound regime +where additional devices fail to improve throughput. ![](/assets/gen_scaling_trend.png) -*Fig2. Left: Strong scaling of batched generation throughput for a 1.5B LRM. Right: Generation becomes memory-IO bound as GPU count increases.* +*Fig2. Left: Strong scaling of batched generation throughput for a 1.5B LRM. Right: +Generation becomes memory-IO bound as GPU count increases.* ## Asynchronous AReaL Overview -We design a system that fully decouples generation and training across separate GPU clusters. This system should be hardware-efficient, scalable, and equipped with the flexibility for a customized RL workflow. We implement these principles in AReaL. Fig.3 presents the architecture and data flow of AREAL. The system comprises 4 core components: +We design a system that fully decouples generation and training across separate GPU +clusters. This system should be hardware-efficient, scalable, and equipped with the +flexibility for a customized RL workflow. We implement these principles in AReaL. Fig.3 +presents the architecture and data flow of AREAL. The system comprises 4 core +components: ![](/assets/arch.png) @@ -66,13 +88,34 @@ We design a system that fully decouples generation and training across separate ### Core Components -- **Interruptible Rollout Worker** handles two types of requests: (1) The *generate* request generates responses given prompts. (2) The *update_weights* request interrupts all ongoing generations and loads parameters of new versions. Upon the interruption, the rollout workers discard KV caches computed by old weights, and re-compute them using the new weights. Afterwards, the rollout workers continue to decode the unfinished sequences until the next interruption or termination. We emphasize that such interruptions and in-flight weight updates would result in trajectories composed of segments produced by different model versions. This introduces a novel algorithmic challenge, which will be addressed in the next section. - -- **Reward Service** evaluates the accuracy of the responses generated by the model. For example, in the coding task, this service extracts the code and executes unit tests to verify its accuracy. - -- **Trainer Workers** continuously sample from the replay buffer, accumulating data until reaching the configured training batch size. They then perform PPO updates and store the resulting parameters in distributed storage. To ensure data freshness, data from the replay buffer is used only once. - -- **Rollout Controller** serves as a critical bridge between the rollout workers, reward service, and the model workers. During the training process, it reads data from the dataset and invokes the rollout worker's *generate* request. The received response is then sent to the reward service to obtain the reward. The trajectory, along with the reward, is stored in the replay buffer, waiting to be trained by the model worker. After the model worker updates the parameters, the controller calls the rollout worker's *update_weight*. We illustrate the generation and training management in Fig.4. This asynchronous pipeline ensures continuous full utilization of both generation and training resources. +- **Interruptible Rollout Worker** handles two types of requests: (1) The *generate* + request generates responses given prompts. (2) The *update_weights* request interrupts + all ongoing generations and loads parameters of new versions. Upon the interruption, + the rollout workers discard KV caches computed by old weights, and re-compute them + using the new weights. Afterwards, the rollout workers continue to decode the + unfinished sequences until the next interruption or termination. We emphasize that + such interruptions and in-flight weight updates would result in trajectories composed + of segments produced by different model versions. This introduces a novel algorithmic + challenge, which will be addressed in the next section. + +- **Reward Service** evaluates the accuracy of the responses generated by the model. For + example, in the coding task, this service extracts the code and executes unit tests to + verify its accuracy. + +- **Trainer Workers** continuously sample from the replay buffer, accumulating data + until reaching the configured training batch size. They then perform PPO updates and + store the resulting parameters in distributed storage. To ensure data freshness, data + from the replay buffer is used only once. + +- **Rollout Controller** serves as a critical bridge between the rollout workers, reward + service, and the model workers. During the training process, it reads data from the + dataset and invokes the rollout worker's *generate* request. The received response is + then sent to the reward service to obtain the reward. The trajectory, along with the + reward, is stored in the replay buffer, waiting to be trained by the model worker. + After the model worker updates the parameters, the controller calls the rollout + worker's *update_weight*. We illustrate the generation and training management in + Fig.4. This asynchronous pipeline ensures continuous full utilization of both + generation and training resources. ![](/assets/async_timeline.png) @@ -82,49 +125,80 @@ We design a system that fully decouples generation and training across separate ### Challenges -While the asynchronous system design offers significant acceleration through improved device utilization, it introduces several technical challenges that require algorithmic considerations. +While the asynchronous system design offers significant acceleration through improved +device utilization, it introduces several technical challenges that require algorithmic +considerations. -- **Challenge 1: Data Staleness** - Due to the asynchronous nature of AREAL, each training batch contains data from multiple prior policy versions. Data staleness would lead to a distribution gap between the training data and the latest model. In asynchronous RL training for LRMs, this issue could be even more severe for long trajectories due to extended decoding time. +- **Challenge 1: Data Staleness** - Due to the asynchronous nature of AREAL, each + training batch contains data from multiple prior policy versions. Data staleness would + lead to a distribution gap between the training data and the latest model. In + asynchronous RL training for LRMs, this issue could be even more severe for long + trajectories due to extended decoding time. -- **Challenge 2: Inconsistent Policy Versions** - As shown in Fig.4, a single generated trajectory may involve segments produced by different policy versions. This inconsistency fundamentally violates the formulation of standard PPO that assumes all actions being generated by a single policy. +- **Challenge 2: Inconsistent Policy Versions** - As shown in Fig.4, a single generated + trajectory may involve segments produced by different policy versions. This + inconsistency fundamentally violates the formulation of standard PPO that assumes all + actions being generated by a single policy. ### Solutions To overcome these two challenges, we propose two solutions: -- **Solution 1: Staleness-Aware Training** - To avoid negative impact of the data staleness issue, we constrain the version discrepancy between the policy version of generated trajectories and the training policy. We introduce a hyperparameter η representing the maximum permitted staleness. When η = 0, the system degenerates to the synchronous RL setting where generation exactly matches training batches. When η = 1, the system recovers to the previous one-step overlap methods. +- **Solution 1: Staleness-Aware Training** - To avoid negative impact of the data + staleness issue, we constrain the version discrepancy between the policy version of + generated trajectories and the training policy. We introduce a hyperparameter η + representing the maximum permitted staleness. When η = 0, the system degenerates to + the synchronous RL setting where generation exactly matches training batches. When η = + 1, the system recovers to the previous one-step overlap methods. -- **Solution 2: Decoupled PPO Objective** - We apply a decoupled PPO objective that disentangles the behavior policy and the proximal policy. The behavior policy represents the policy used for sampling trajectories and the proxy policy is a proximal policy serving as a recent target to regularize the update of online policy. +- **Solution 2: Decoupled PPO Objective** - We apply a decoupled PPO objective that + disentangles the behavior policy and the proximal policy. The behavior policy + represents the policy used for sampling trajectories and the proxy policy is a + proximal policy serving as a recent target to regularize the update of online policy. ![](/assets/decoupled_ppo_obj.png) ## Validating Asynchronous AReaL -We first evaluate Asynchronous AReaL on a math task to validate our system and algorithmic choices. Building on our previous release, AReaL-boba, we adopt R1-Distill-Qwen as the base model and [AReaL-boba-106k](https://huggingface.co/datasets/inclusionAI/AReaL-boba-Data/blob/main/AReaL-boba-106k.jsonl) as the training dataset. +We first evaluate Asynchronous AReaL on a math task to validate our system and +algorithmic choices. Building on our previous release, AReaL-boba, we adopt +R1-Distill-Qwen as the base model and +[AReaL-boba-106k](https://huggingface.co/datasets/inclusionAI/AReaL-boba-Data/blob/main/AReaL-boba-106k.jsonl) +as the training dataset. ### End-to-End Performance Comparison -We compared synchronous and asynchronous training on 1.5B and 7B models. Under equal resource constraints and training steps, the asynchronous system was over twice as fast as the synchronous one. Evaluations on AIME24 confirmed that this acceleration does not compromise performance. - -| **Model (1.5B)** | **AIME24** | **# Nodes** | **PPO Steps** | **Training Hours** | -|:---:|:---:|:---:|:---:|:---:| -| basemodel | 29.3 | - | - | - | -| Synchronous AReaL | 42.0 | 16 | 1000 | 41.0 | -| Asynchronous AReaL | 42.2 | 16 | 1000 | **14.8** | - -| **Model (7B)** | **AIME24** | **# Nodes** | **PPO Steps** | **Training Hours** | -|:---:|:---:|:---:|:---:|:---:| -| basemodel | 54.3 | - | - | - | -| Synchronous AReaL | 63.0 | 24 | 1000 | 57.7 | -| Asynchronous AReaL | 63.1 | 24 | 1000 | **25.4** | - -*Table 2: End-to-end comparison between asynchronous AReaL and synchronous AReaL. Following last release AReaL-boba, we adopt R1-Distill-Qwen as the basemodel and [AReaL-boba-106k](https://huggingface.co/datasets/inclusionAI/AReaL-boba-Data/blob/main/AReaL-boba-106k.jsonl) as the training data. Each performance number represents the average results of 32 sampling responses.* +We compared synchronous and asynchronous training on 1.5B and 7B models. Under equal +resource constraints and training steps, the asynchronous system was over twice as fast +as the synchronous one. Evaluations on AIME24 confirmed that this acceleration does not +compromise performance. + +| **Model (1.5B)** | **AIME24** | **# Nodes** | **PPO Steps** | **Training Hours** | +| :----------------: | :--------: | :---------: | :-----------: | :----------------: | +| basemodel | 29.3 | - | - | - | +| Synchronous AReaL | 42.0 | 16 | 1000 | 41.0 | +| Asynchronous AReaL | 42.2 | 16 | 1000 | **14.8** | + +| **Model (7B)** | **AIME24** | **# Nodes** | **PPO Steps** | **Training Hours** | +| :----------------: | :--------: | :---------: | :-----------: | :----------------: | +| basemodel | 54.3 | - | - | - | +| Synchronous AReaL | 63.0 | 24 | 1000 | 57.7 | +| Asynchronous AReaL | 63.1 | 24 | 1000 | **25.4** | + +*Table 2: End-to-end comparison between asynchronous AReaL and synchronous AReaL. +Following last release AReaL-boba, we adopt R1-Distill-Qwen as the basemodel and +[AReaL-boba-106k](https://huggingface.co/datasets/inclusionAI/AReaL-boba-Data/blob/main/AReaL-boba-106k.jsonl) +as the training data. Each performance number represents the average results of 32 +sampling responses.* ### Ablation Study #### System Ablations -We ablate interruptible generation and present the resulting generation throughput in Fig.5. Without interruptible generation, the controller must wait for the longest response. In particular, interruptible generation leads to a 12% and 17% throughput increase for 1.5B and 7B models respectively on 4 nodes. +We ablate interruptible generation and present the resulting generation throughput in +Fig.5. Without interruptible generation, the controller must wait for the longest +response. In particular, interruptible generation leads to a 12% and 17% throughput +increase for 1.5B and 7B models respectively on 4 nodes. ![](/assets/interrupt_gen_ablation.png) @@ -132,58 +206,92 @@ We ablate interruptible generation and present the resulting generation throughp #### Algorithm Ablations -We vary the maximum allowed staleness η and compare configurations with and without the decoupled PPO objective. Fig.6 shows the learning curves after 1000 PPO updates, demonstrating that naive PPO fails to match the performance of the synchronous RL oracle (i.e., the performance when η = 0). Even slight staleness can significantly degrade final performance due to the improper clipping center and policy changes during interruptible generation. Furthermore, increasing data staleness consistently degrades learning performance. +We vary the maximum allowed staleness η and compare configurations with and without the +decoupled PPO objective. Fig.6 shows the learning curves after 1000 PPO updates, +demonstrating that naive PPO fails to match the performance of the synchronous RL oracle +(i.e., the performance when η = 0). Even slight staleness can significantly degrade +final performance due to the improper clipping center and policy changes during +interruptible generation. Furthermore, increasing data staleness consistently degrades +learning performance. ![](/assets/algo_ablation.png) -*Fig.6 Ablation Study on Decoupled PPO Objective with DeepSeek-R1-Distill-Qwen-1.5B. Left: Learning curves with naive PPO. Right: Learning curves with decoupled PPO objective.* - -| η | **AIME 24** | | **AIME 25** | | -|:---:|:---:|:---:|:---:|:---:| -| | naive PPO | + Decoupled Objective | naive PPO | + Decoupled Objective | -| 0 | 42.0 | | 32.9 | | -| 1 | 41.8 | 42.1 | 30.7 | 31.9 | -| 2 | 40.0 | 41.8 | 32.1 | 32.5 | -| 4 | 23.3 | 42.2 | 23.1 | 32.0 | -| 8 | 35.7 | 41.0 | 27.8 | 31.1 | -| 16 | 35.8 | 38.7 | 26.2 | 32.5 | -| ∞ | 34.0 | 36.9 | 26.9 | 29.9 | +*Fig.6 Ablation Study on Decoupled PPO Objective with DeepSeek-R1-Distill-Qwen-1.5B. +Left: Learning curves with naive PPO. Right: Learning curves with decoupled PPO +objective.* -*Table 3: Evaluation scores when varying data staleness, comparing performance with and without the decoupled objective.* +| η | **AIME 24** | | **AIME 25** | | +| :-: | :---------: | :-------------------: | :---------: | :-------------------: | +| | naive PPO | + Decoupled Objective | naive PPO | + Decoupled Objective | +| 0 | 42.0 | | 32.9 | | +| 1 | 41.8 | 42.1 | 30.7 | 31.9 | +| 2 | 40.0 | 41.8 | 32.1 | 32.5 | +| 4 | 23.3 | 42.2 | 23.1 | 32.0 | +| 8 | 35.7 | 41.0 | 27.8 | 31.1 | +| 16 | 35.8 | 38.7 | 26.2 | 32.5 | +| ∞ | 34.0 | 36.9 | 26.9 | 29.9 | +*Table 3: Evaluation scores when varying data staleness, comparing performance with and +without the decoupled objective.* ```{note} In practice, the effect of staleness could depend on various factors, such as learning rate, batch size, the number of mini batches, generation length, and also the RL algorithm. Here we aim to illustrate the impact of staleness on the training performance of a synchronous RL system, and show that stalenss control and decoupled PPO could well mitigate the negative impact of staleness. Here we follow a large batch size setting with `batch_size=512`, `n_rollouts=16`, `n_mini_batch=4`, `max_new_tokens=8192`, and `lr=2e-5`. The staleness is defined as the gap of model versions, where the model version increases by 1 when a batch of `batch_size` X `n_rollouts` samples are trained on. ``` -However, the decoupled PPO objective substantially improves training stability when handling stale data. Notably, even with the decoupled objective, unbounded staleness (maximum staleness → ∞) still results in inferior performance compared to the zero-staleness oracle. When properly constrained, moderate staleness (e.g., η ≤ 4) has minimal impact on final performance while significantly accelerating training through the asynchronous pipeline, as demonstrated in Tab.3 and Fig.7. +However, the decoupled PPO objective substantially improves training stability when +handling stale data. Notably, even with the decoupled objective, unbounded staleness +(maximum staleness → ∞) still results in inferior performance compared to the +zero-staleness oracle. When properly constrained, moderate staleness (e.g., η ≤ 4) has +minimal impact on final performance while significantly accelerating training through +the asynchronous pipeline, as demonstrated in Tab.3 and Fig.7. ![](/assets/staleness_throughput.png) -*Fig.7 The relationship between η and training throughput. Larger η leads to higher throughput.* +*Fig.7 The relationship between η and training throughput. Larger η leads to higher +throughput.* ## Training Recipe with Asynchronous AReaL ### Dataset Curation -We train the 7B and 14B model using open-source data from [DeepCoder](https://pretty-radio-b75.notion.site/DeepCoder-A-Fully-Open-Source-14B-Coder-at-O3-mini-Level-1cf81902c14680b3bee5eb349a512a51). After deduplication and quality filtering, approximately 7,600 coding problems remained. Additionally, we have around 10,000 internal data entries that can further improve the model's performance. We are actively working on open-sourcing this portion of the data. +We train the 7B and 14B model using open-source data from +[DeepCoder](https://pretty-radio-b75.notion.site/DeepCoder-A-Fully-Open-Source-14B-Coder-at-O3-mini-Level-1cf81902c14680b3bee5eb349a512a51). +After deduplication and quality filtering, approximately 7,600 coding problems remained. +Additionally, we have around 10,000 internal data entries that can further improve the +model's performance. We are actively working on open-sourcing this portion of the data. ### Rollout Strategy -During the rollout phase, the LLM generates 16 responses per question. To minimize output truncation, we set the maximum generation length to 27K tokens. +During the rollout phase, the LLM generates 16 responses per question. To minimize +output truncation, we set the maximum generation length to 27K tokens. ### Filtering Strategy -In the design of our asynchronous system, we can naturally incorporate strategies to filter some rollouts. Specifically, we exclude the questions where all 16 answers are either entirely correct or entirely incorrect. +In the design of our asynchronous system, we can naturally incorporate strategies to +filter some rollouts. Specifically, we exclude the questions where all 16 answers are +either entirely correct or entirely incorrect. ### Key Hyperparameters -| **Parameter** | **Value** | -|:---:|:---:| -| PPO Minibatches | 1 | -| Learning Rate | 2e-5 | -| Adam ε | 1e-5 | -| Batch Size | 2,048 | -| max_head_offpolicyness η | 16 | -| success_rate_ub | 0.95 | -| success_rate_lb | 0.05 | \ No newline at end of file +| **Parameter** | **Value** | +| :----------------------: | :-------: | +| PPO Minibatches | 1 | +| Learning Rate | 2e-5 | +| Adam ε | 1e-5 | +| Batch Size | 2,048 | +| max_head_offpolicyness η | 16 | +| success_rate_ub | 0.95 | +| success_rate_lb | 0.05 | + +## Resources + +- [boba² Code dataset](https://huggingface.co/datasets/inclusionAI/AReaL-boba-2-RL-Code) +- **Reproduce boba² Code Models** + - 🤗 **Model weights**: [8B-code](https://huggingface.co/inclusionAI/AReaL-boba-2-8B), + [14B-code](https://huggingface.co/inclusionAI/AReaL-boba-2-14B), + [8B-code-open](https://huggingface.co/inclusionAI/AReaL-boba-2-8B-subset), + [14B-code-open](https://huggingface.co/inclusionAI/AReaL-boba-2-14B-subset) + - [Evaluation Guide](https://inclusionai.github.io/AReaL/tutorial/eval.html) + - [Training configs](https://github.com/inclusionAI/AReaL/tree/main/examples/configs/v0.3-qwen3-code) + and [instructions](https://inclusionai.github.io/AReaL/references/reproduce.html) +- [Scripts for Benchmark Training Throughput](https://github.com/inclusionAI/AReaL/tree/main/benchmark/verl_v0_3_0_post1_76084d3) diff --git a/ci/test_areal.sh b/ci/test_areal.sh new file mode 100644 index 0000000000..88c98cf326 --- /dev/null +++ b/ci/test_areal.sh @@ -0,0 +1,47 @@ +#!/usr/bin/env bash + +set -e + +RUN_ID="test-areal-$(openssl rand -hex 6)" + +# Calculate environment hash from pyproject.toml +ENV_SHA=$(sha256sum pyproject.toml | awk '{print $1}') + +# Build environment image if it doesn't exist +if ! docker images --format '{{.Repository}}:{{.Tag}}' | grep -q "areal-env:$ENV_SHA"; then + echo "Building image areal-env:$ENV_SHA..." + + # Build the environment image + docker run \ + --name $RUN_ID \ + -v $(pwd):/workspace \ + -w /workspace \ + nvcr.io/nvidia/pytorch:25.01-py3 \ + bash -c " + pip config set global.index-url https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple + pip config unset global.extra-index-url + pip uninstall -y --ignore-installed transformer-engine + pip uninstall -y --ignore-installed torch-tensorrt + pip uninstall -y --ignore-installed nvidia-dali-cuda120 + bash examples/env/scripts/setup-pip-deps.sh + " || { docker rm -f $RUN_ID; exit 1; } + + # Commit the container as the environment image + docker commit $RUN_ID areal-env:$ENV_SHA + docker rm -f $RUN_ID +else + echo "Image areal-env:$ENV_SHA already exists, skipping build." +fi + +# Run tests using the environment image +echo "Running tests on image areal-env:$ENV_SHA..." +docker run \ + -e HF_ENDPOINT=https://hf-mirror.com \ + --name $RUN_ID \ + --gpus all \ + --shm-size=8g \ + --rm \ + -v $(pwd):/workspace \ + -w /workspace \ + areal-env:$ENV_SHA \ + python -m pytest -s areal/tests/ diff --git a/docs/_toc.yml b/docs/_toc.yml index fecabda369..eaa524219c 100644 --- a/docs/_toc.yml +++ b/docs/_toc.yml @@ -8,8 +8,12 @@ parts: chapters: - file: tutorial/installation - file: tutorial/quickstart + - file: tutorial/quickstart_legacy - file: tutorial/eval - file: tutorial/troubleshooting + - caption: Getting Started with AReaL-lite + chapters: + - file: lite/gsm8k_grpo - caption: References chapters: - file: references/benchmark @@ -19,7 +23,7 @@ parts: - file: customization/dataset - file: customization/agent - file: customization/algorithm - - caption: Code Walkthrough + - caption: Code Walkthrough (Legacy) chapters: - file: developer/overview - file: developer/trainer @@ -35,3 +39,8 @@ parts: - caption: Contributing chapters: - file: contrib + - caption: Customization (Legacy) + chapters: + - file: legacy/customization/dataset + - file: legacy/customization/agent + - file: legacy/customization/algorithm diff --git a/docs/contrib.md b/docs/contrib.md index 70b0877355..6a150d6963 100644 --- a/docs/contrib.md +++ b/docs/contrib.md @@ -1,15 +1,18 @@ # Contribution Guide -Thank you for your interest in contributing to AReaL! We welcome contributions from everyone, whether you're fixing bugs, improving documentation, or adding new system and algorithmic features. +Thank you for your interest in contributing to AReaL! We welcome contributions from +everyone, whether you're fixing bugs, improving documentation, or adding new system and +algorithmic features. ## Setting Up Your Development Environment -New contributors do not have write permissions to the official repository. Please fork the repository and clone your fork locally. AReaL is fully Python-based, making installation straightforward. +New contributors do not have write permissions to the official repository. Please fork +the repository and clone your fork locally. AReaL is fully Python-based, making +installation straightforward. ```bash git clone https://github.com/${your-username}/AReaL cd AReaL -pip3 install -r requirements.txt pip3 install -e . ``` @@ -17,16 +20,22 @@ pip3 install -e . ### Issue Templates -Please follow the [issue template on GitHub](https://github.com/inclusionAI/AReaL/tree/main/.github/ISSUE_TEMPLATE). Issues can be: +Please follow the +[issue template on GitHub](https://github.com/inclusionAI/AReaL/tree/main/.github/ISSUE_TEMPLATE). +Issues can be: + - Bug reports -- Feature requests +- Feature requests - Refactor requests -The required fields in the template help reduce communication overhead when resolving issues. **Issues with arbitrary formatting may be ignored.** +The required fields in the template help reduce communication overhead when resolving +issues. **Issues with arbitrary formatting may be ignored.** ## Pull Request Guidelines -There are no specific PR templates, but **pull requests should be related to a well-templated issue**. Your PR should: +There are no specific PR templates, but **pull requests should be related to a +well-templated issue**. Your PR should: + - Explain how the issue is resolved - Describe the benefits this change will provide - Reference the related issue number @@ -46,24 +55,23 @@ isort . && black . AReaL's unit tests are based on the `pytest` framework: ```bash -# Run all tests (excluding GPU tests) -pytest -m "not gpu" - -# Run a specific test case -pytest tests/test_something.py +# Run all tests +pytest -s -v areal/tests/ ``` **Note**: Running all tests may take several hours to complete. ## Documentation -Writing documentation is an excellent starting point for new contributors. The documentation is located in the `docs` folder and built using [Jupyter Book](https://jupyterbook.org/en/stable/intro.html). +Writing documentation is an excellent starting point for new contributors. The +documentation is located in the `docs` folder and built using +[Jupyter Book](https://jupyterbook.org/en/stable/intro.html). ### Adding New Documentation 1. Create your documentation files in the `docs` folder -2. Add the file path to `docs/_toc.yaml` -3. Build the documentation: +1. Add the file path to `docs/_toc.yaml` +1. Build the documentation: ```bash jb build docs @@ -71,4 +79,5 @@ jb build docs 4. Preview your changes by opening the HTML files in `docs/_build/html` -This process allows you to see how your documentation will appear before submitting your contribution. \ No newline at end of file +This process allows you to see how your documentation will appear before submitting your +contribution. diff --git a/docs/customization/agent.md b/docs/customization/agent.md index 8b13b440e5..8a96bbe955 100644 --- a/docs/customization/agent.md +++ b/docs/customization/agent.md @@ -1,135 +1,265 @@ # Rollout and Agentic RL -This guide provides an example of modifying the rollout behavior for PPO training. +This guide shows you how to create custom rollout behaviors for RL training by building +a multi-turn math agent with **AReaL-lite**. This agent keeps trying to solve math +problems until it finds the correct answer. -In particular, we implement a multi-turn math agent using end-to-end RL. The math agent will continuously attempt to think through and solve math problems until it reaches the correct answer. +You can find the complete implementation in `areal/workflow/multi_turn.py`. -## Define Your Agent +## Step 1: Define Your Workflow -Create a new file under `realhf/impl/agent/`, for example, `math_multi_turn_agent.py`. Your `Agent` must implement the interface defined in `realhf/api/core/agent.py`, which requires implementing a single method: `collect_trajectory`. +AReaL-lite gives you flexibility in how you design your agents to run **an episode**. +**An episode** defines how your agent rollouts a complete training sample from an input +prompt, using tools, reward functions, and (multi-turn) generation. Instead of rigid +`Agent` classes that might constrain your agent's capabilities, AReaL-lite captures all +rollout behavior in a `RolloutWorkflow` class. This approach allows you to customize +your agent's behavior however you need. ```python -class MathMultiTurnAgent(Agent): - - async def collect_trajectory( - self, - prompt: SequenceSample, - env: EnvironmentService, - obs_queue: asyncio.Queue, - act_queue: asyncio.Queue, - ): - ... +# areal/api/workflow_api.py +class RolloutWorkflow: + async def arun_episode( + self, engine: InferenceEngine, data: Dict[str, Any] + ) -> TensorDict: + """Run a single episode of the workflow. + + See concrete example implementations under the `areal/workflow` directory. + """ + raise NotImplementedError() ``` -## Implement the `collect_trajectory` Logic +The workflow exposes an `arun_episode` method that runs and collects data from a single +episode. This method takes two key arguments: + +1. **InferenceEngine**: Provides the `agenerate` method for generating responses to user + inputs +1. **data**: The prompt data loaded from your RL dataset + +Within this method, you have complete control over how your agent and environment +interact. -The `collect_trajectory` function takes a task prompt, an environment, and two queues as input, then produces several trajectories for the RL trainer. Within this function, you can create arbitrary data processing logic to produce the input for the inference engine (i.e., via `obs_queue`) and extract the action (i.e., via `act_queue`) from the generated tokens. +> **Note**: Each `arun_episode` call takes a single prompt and outputs the trajectories +> generated from that prompt—it's not batched. However, you can generate multiple +> trajectories from a single prompt (for example, with GRPO or tree search). -In this example, the initial observation is the math problem itself. We put the token IDs and generation config into `obs_queue` and wait for the action produced by the inference engine from `act_queue`. After the inference engine returns, we extract the generated answers and send them to the environment. +### Setting Up the Multi-turn Math Workflow + +Let's build a multi-turn rollout workflow for solving math problems. First, we'll define +the `__init__` method to set up what we need during rollout: + +> **Note**: You have complete flexibility in defining the `__init__` method. Pass +> whatever arguments you need to construct your workflow. If you want to use tools, pass +> the corresponding environment here so your agent can call it in the `arun_episode` +> method. ```python -for turn in range(self.num_turns): - await obs_queue.put((qid, token_ids, self.gconfig)) - act: BundledGenerationOutputs = await act_queue.get() - _, success, *_ = await env.step((qid, answers)) - ... +class MultiTurnWorkflow(RolloutWorkflow): + def __init__( + self, + reward_fn, + gconfig: GenerationHyperparameters, # aka sampling_params + tokenizer: PreTrainedTokenizerFast, + max_turns: int, + turn_discount: float, + ): + self.reward_fn = reward_fn + self.gconfig = gconfig + self.tokenizer = tokenizer + self.max_turns = max_turns + # Discount rewards if the agent takes longer to find the correct answer + self.turn_discount = turn_discount ``` -The environment is similar to a [gym environment](https://github.com/Farama-Foundation/Gymnasium), which defines two methods: `reset` and `step`. However, to maintain efficiency, we use an asynchronous implementation to avoid mutual blocking across different environment instances. +### Implementing the Episode Logic -The math environment is stateless and essentially serves as a wrapper around the reward function: +Now let's implement the `arun_episode` method. We'll start by tokenizing the prompt data +and converting it into an `LLMRequest` object for the inference engine: ```python -class MathCodeSingleStepEnv(EnvironmentService): - - async def step(self, action: Tuple[str, List[str]]): - qid, answers = action - ... - # Make `math_verify_call` async - format_rewards = await asyncio.to_thread( - math_verify_call, - answers, - ... - ) - return None, format_rewards, True, False, {} +class MultiTurnWorkflow(RolloutWorkflow): + # ... __init__ method above ... + + async def arun_episode(self, engine: InferenceEngine, data) -> TensorDict: + # Initialize result containers + seq, logprobs, loss_mask, versions = [], [], [], [] + messages = data["messages"] + # Run multi-turn rollout until we get the correct answer + turn_index = 0 + reward = 0 + discount = 1.0 + rid = uuid.uuid4().hex + while reward == 0 and turn_index < self.max_turns: + # Convert the conversation into input tokens + input_ids = self.tokenizer.apply_chat_template( + messages, + tokenize=True, + add_generation_prompt=True, + ) + # Generate response from the model + req = LLMRequest( + rid=rid, + input_ids=input_ids, + gconfig=self.gconfig.new(n_samples=1), + ) + resp = await engine.agenerate(req) + # ... continue processing ... ``` -After `env.step` returns the reward for the current step, we can check whether the answer is correct. If not, we can append a user prompt and send it to `obs_queue` again to enter the next round. +> **Note**: This example uses the "messages" key from the prompt data to get +> OpenAI-compatible messages. This isn't required—the key and prompt format depend +> entirely on your implementation. For instance, if your dataset stores prompt strings +> in a "prompt" column, you could get input token IDs with +> `self.tokenizer.encode(data["prompt"])`. + +> **Note**: The `rid` field in `LLMRequest` is the request ID. Requests with the same ID +> will reuse the LLM inference server's KV caches for better efficiency. + +### Handling Multi-turn Conversations + +Next, we'll check if the current answer is correct using our `reward_fn`. This function +should return 1 for correct answers and 0 otherwise. When the answer is wrong, we'll +apply a discount, add feedback to the conversation, and let the model try again: ```python -for turn in range(self.num_turns): - ... - feedback = None - if success[0]: - feedback = "Congratulations! You are correct!" - else: - feedback = "Unfortunately your answer is wrong. Let's try again." - - feedback = "\n" + self.tokenizer.apply_chat_template( - [dict(content=feedback, role="user")], - add_generation_prompt=True, - tokenize=False, - ) - feedback = self.tokenizer(feedback)["input_ids"] - token_ids.extend(feedback) +class MultiTurnWorkflow(RolloutWorkflow): + # ... previous methods ... + + async def arun_episode(self, engine: InferenceEngine, data) -> TensorDict: + # ... initialization code ... + while reward == 0 and t < self.max_turns: + # Add feedback if the previous answer was incorrect + if t > 0: + messages += [ + {"role": "assistant", "content": completions_str}, + { + "role": "user", + "content": "Your answer is not correct. Please try to answer it again." + }, + ] + # Generate response (code from above) + # ... + # Evaluate the response + prompt_str = self.tokenizer.decode(input_ids) + completions_str = self.tokenizer.decode(resp.output_tokens) + reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) + # Update counters + t += 1 + discount *= self.turn_discount ``` -## Modify the Configuration +### Reward Function Signature -You're now close to running the end-to-end RL loop. The final step is to register and import your implementation, then modify the experiment configuration. +To make it easier to switch between different reward functions, we recommend following +this signature: ```python -# in realhf/impl/agent/math_multi_turn_agent.py -register_agent("math-multi-turn", MathMultiTurnAgent) +def reward_fn( + prompt: str, + completions: str, + prompt_ids: List[int], + completion_ids: List[int], + **kwargs, +): + """Reward function for evaluating agent performance. + + This signature is recommended for compatibility with predefined workflows, + but you can modify it freely in custom implementations. + + Args: + prompt: The task description string + completions: The agent's response string + prompt_ids: Tokenized prompt + completion_ids: Tokenized response + **kwargs: Additional dataset attributes (solutions, input_outputs, etc.) + + Returns: + float: Reward value (typically 1.0 for correct, 0.0 for incorrect) + """ ``` +While this signature is convenient, you're not restricted to it in custom +workflows—modify as needed for your specific use case. + +### Collecting Training Data + +Finally, let's complete the implementation by collecting trajectories in the +`TensorDict` format: + ```python -# in realhf/impl/agent/__init__.py -import realhf.impl.agent.math_multi_turn_agent -``` +class MultiTurnWorkflow(RolloutWorkflow): + # ... previous methods ... -In `realhf/experiments/async_exp/async_math_ppo.py`: - -```diff -@dataclasses.dataclass -class AsyncPPOMATHConfig(AsyncRLExperimentConfig, PPOMATHConfig): -+ # New CLI arguments are defined here -+ my_param: float = 1.0 - - # in realhf/experiments/async_exp/async_ppo_math_exp.py - @property - def agent(self) -> AgentAbstraction: - return AgentAbstraction( -- "math-single-step", -+ "math-multi-turn", # Your registered name - args=dict( -- ... -+ # Any configurations for your __init__ method -+ my_param=my_param, - ), - ) + async def arun_episode(self, engine: InferenceEngine, data) -> TensorDict: + # ... episode logic above ... + + while reward == 0 and t < self.max_turns: + # ... generation and evaluation ... - @property - def env(self) -> EnvServiceAbstraction: -- return EnvServiceAbstraction( -- "math-code-single-step", args=dict(dataset_path=self.dataset.path) -- ) -+ # Change to your customized environment if necessary -+ return EnvServiceAbstraction( -+ "my-env", args=dict(...) -+ ) + # Collect trajectory data + input_len = len(resp.input_tokens) - len(seq) + seq += resp.input_tokens[-input_len:] + resp.output_tokens + logprobs += [0.0] * input_len + resp.output_logprobs + loss_mask += [0] * input_len + [1] * resp.output_len + versions += [-1] * input_len + resp.output_versions + + # Package results + res = dict( + input_ids=torch.tensor(seq), + logprobs=torch.tensor(logprobs), + loss_mask=torch.tensor(loss_mask), + versions=torch.tensor(versions), + rewards=torch.tensor(float(reward * discount)), + attention_mask=torch.ones(len(seq), dtype=torch.bool), + ) + res = {k: v.unsqueeze(0) for k, v in res.items()} + return concat_padded_tensors([res]) ``` -## Run Training +> **Important**: The returned `TensorDict` must follow HuggingFace's padded data format, +> where each tensor has shape `[batch_size, sequence_length, *]`. This allows AReaL-lite +> to automatically batch multiple trajectories for training. Since this example returns +> a single trajectory, we use `unsqueeze(0)` to create a batch of size 1. -Please follow the guide in [quickstart](../tutorial/quickstart.md). Generally, start your experiments by running: +> **Note**: You're not restricted to specific keys in your `TensorDict`—different +> algorithms need different keys. This example targets the GRPO algorithm, so we include +> `input_ids`, `loss_mask`, `attention_mask`, and `logprobs` (needed for computing +> importance ratios). -```bash -python3 training/main_async_ppo.py my_param=5.0 # and any additional CLI arguments -``` +## Step 2: Training with Your Custom Workflow + +Using your custom workflow is straightforward—just create it in your training script and +pass it to the `rollout_batch` or `prepare_batch` method: -The training reward of our trial is shown below: +```python +def main(args): + # ... setup code ... + + # Create your custom workflow + workflow = MultiTurnWorkflow( + reward_fn=gsm8k_reward_fn, + gconfig=config.gconfig, + tokenizer=tokenizer, + turn_discount=0.9, + max_turns=5, + ) -![](multiturn_reward.png) + # Run training—no other changes needed! + data_generator = itertools.cycle(train_dataloader) + for global_step in range(max_steps): + with stats_tracker.record_timing("rollout"): + if config.async_training: + batch = rollout.prepare_batch(train_dataloader, workflow=workflow) + else: + batch = rollout.rollout_batch(next(data_generator), workflow=workflow) + # ... continue with training loop ... +``` -Happy coding! \ No newline at end of file +That's it! Your custom multi-turn math agent is now ready for reinforcement learning +training. The workflow will automatically handle the multi-turn conversations, reward +computation, and data collection needed for effective RL training. diff --git a/docs/customization/algorithm.md b/docs/customization/algorithm.md index bf5030b767..b15a3fd168 100644 --- a/docs/customization/algorithm.md +++ b/docs/customization/algorithm.md @@ -1,305 +1,198 @@ # Training Algorithm -An algorithm is encapsulated in a `ModelInterface`, which primarily defines three methods: +> **Note**: We recommend the user to first read the +> [agent customization guide](agent.md). -```python -# in realhf/api/core/model_api.py -class ModelInterface(abc.ABC): - """An interface for model training, inference, and generation. - - This interface is designed to follow the dependency injection pattern. - We pass the model to the interface and call its methods, ensuring that model APIs - and algorithms are fully decoupled. For example, REINFORCE and PPO can exhibit - different behaviors during training. Separate interfaces can be written for these - algorithms while using the same model that provides basic forward-backward-update - functionality (i.e., :class:`PipelinableEngine`). - - During runtime, the master worker requests model workers to execute a specific - interface type (e.g., generate) on a specific model. The model worker locates - the corresponding model, passes it into the requested interface, performs the - computation, and returns the result. - """ - - def inference( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> SequenceSample | None: - raise NotImplementedError() - - def generate( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> SequenceSample | None: - raise NotImplementedError() - - def train_step( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> Dict | List[Dict]: - raise NotImplementedError() -``` - -When the dataflow is fixed, it's usually sufficient to modify or add the file that defines the algorithm interface. +**AReaL-lite** structures RL algorithms around two core components: -We provide two examples: (1) changing PPO's global advantage normalization to grouped normalization in GRPO, and (2) changing the original PPO loss to the decoupled PPO loss in AReaL's paper. +- **RolloutWorkflow**: Defines what data to generate during rollouts +- **TrainEngine**: Defines how to process the generated data for training -```{note} -We recommend using asynchronous RL, so that you can customize the generation behavior by [modifying your RL agent](agent.md) and don't need to modify the `generate` method of model interfaces. -``` +We'll demonstrate this by implementing an RL algorithm similar to ReMax. -## Grouped Advantage Normalization +## Step 1: Implementing the RolloutWorkflow -The PPO algorithm is written in a single file `ppo_interface.py`. The method we are going to modify is the `train_step` method in `PPOActorInterface`. PPO's global advantage normalization looks like: +The rollout workflow generates both greedy and sampled completions, then uses the reward +difference as the final training signal: ```python -@dataclass -class PPOActorInterface(ModelInterface): - def train_step( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> Dict | List[Dict]: - ... - if self.adv_norm: - advantages = masked_normalization(advantages, loss_mask) - ... -``` +class ReMaxRLVRWorkflow(RolloutWorkflow): + async def arun_episode(self, engine: InferenceEngine, data): + # Prepare input tokens from chat messages + input_ids = self.tokenizer.apply_chat_template( + data["messages"], + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) -### An Additional Note on Data Management - -We need to explain how data in each batch is organized. - -Usually, each data batch (i.e., the `data` variable) includes multiple prompts. The number of prompts is called "batch size". Additionally, each prompt may have multiple corresponding answers. The number of answers is called "group_size". Therefore, there are batch_size × group_size sequences in each batch. - -These sequences have different lengths, but they are concatenated (or packed) together as a 1D tensor. The inner dimension is the "group" with the same prompt, and the outer dimension consists of answers from different prompts. Similar to flash-attention, we use `cu_seqlens` to mark the boundary of each sequence. `cu_seqlens` is the cumulative sum of sequence lengths across the batch. - -Each token in the sequence has a corresponding reward and advantage, so `advantages` is also a packed 1D tensor just like the tokens (i.e., `packed_input_ids`). However, the "sequences" of advantages are all one step shorter than tokens due to the auto-regressive nature of LLMs. We can only compute the loss on tokens except for the first one in each sequence. - -### Implementation - -For grouped advantage normalization, we need to partition the advantages into groups and run normalization within the tensor chunk of each group: - -```diff -@dataclass -class PPOActorInterface(ModelInterface): -+ group_adv_norm: bool = False - - def train_step( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> Dict | List[Dict]: - ... - if self.adv_norm: -- advantages = masked_normalization(advantages, loss_mask) -+ if not self.group_adv_norm: -+ advantages = masked_normalization(advantages, loss_mask) -+ else: -+ n_samples = data.bs -+ adv_list = [] -+ for i in range(0, n_samples, self.group_size): -+ # Start and end of the chunk -+ s = short1cu_seqlens[i] -+ e = short1cu_seqlens[i + self.group_size] -+ # Get advantages within each group of the same prompt -+ adv = advantages[s: e] -+ mask = loss_mask[s: e] -+ # Run normalization -+ advn = masked_normalization(adv, mask, all_reduce=False) -+ adv_list.append(advn) -+ advantages = torch.cat(adv_list, 0) - ... -``` + n_samples = self.gconfig.n_samples + rid = uuid.uuid4().hex -### Modify Your Experiment Configuration - -To make our new argument `group_adv_norm` effective in CLI args, we should make the following changes to the `PPOMathConfig` under `realhf/experiments/common/ppo_math_exp.py`: - -```diff -@dataclasses.dataclass -class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): -+ group_adv_norm: bool = False - - @property - def rpcs(self): - ... - # interfaces - actor_interface = ModelInterfaceAbstraction( - "ppo_actor", - args={ - **copy.deepcopy(self.ppo_kwargs), -+ "group_adv_norm": self.group_adv_norm, - ... - }, + # Create requests for both sampled and greedy generation + sample_req = LLMRequest( + rid=rid, + input_ids=input_ids, + gconfig=self.gconfig, + ) + greedy_req = LLMRequest( + rid=rid, + input_ids=input_ids, + gconfig=self.gconfig.new(greedy=True), ) -``` - -## The Decoupled PPO Loss -![decoupled loss](decoupled_loss.png) + # Generate both responses concurrently + resp, greedy_resp = await asyncio.gather( + engine.agenerate(sample_req), + engine.agenerate(greedy_req), + ) -As mentioned in AReaL's paper, we implement this loss by recomputing the probabilities before mini-batched updates, and use this value as π_prox to compute the above loss. + # Calculate rewards for both completions + prompt_str = self.tokenizer.decode(input_ids) + completions_str = self.tokenizer.decode(resp.output_tokens) -### Probability Recomputation + sample_reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) -Recomputation involves a single forward pass, which has already been implemented by `PPOActorInterface.inference`. We need to call this method in the `train_step` method: + greedy_completions = self.tokenizer.decode(greedy_resp.output_tokens) + greedy_reward = self.reward_fn( + prompt=prompt_str, + completions=greedy_completions, + prompt_ids=greedy_resp.input_tokens, + completion_ids=greedy_resp.output_tokens, + **data, + ) -```diff -@dataclass -class PPOActorInterface(ModelInterface): -+ use_decoupled_loss: bool = False + # Package results for training + res = dict( + # Add batch dimension + input_ids=torch.tensor(resp.input_tokens + resp.output_tokens).unsqueeze(0), + loss_mask=torch.tensor([0] * resp.input_len + [1] * resp.output_len).unsqueeze(0), + versions=torch.tensor([-1] * resp.input_len + resp.output_versions).unsqueeze(0), + attention_mask=torch.ones(resp.input_len + resp.output_len, dtype=torch.bool).unsqueeze(0), + # Use reward difference across all tokens + rewards=torch.tensor([float(sample_reward - greedy_reward)] * (resp.input_len + resp.output_len)), + ) - def train_step( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> Dict | List[Dict]: -+ if self.use_decoupled_loss: -+ s: SequenceSample = self.inference(model, data, mb_spec) -+ prox_logp = s.data["logprobs"] - ... + return TensorDict(res, batch_size=[1]) ``` -Next, we need to pass `prox_logp` to loss computation: - -```diff -@dataclass -class PPOActorInterface(ModelInterface): - ... - - def train_step( - self, - model: Model, - data: SequenceSample, - mb_spec: MicroBatchSpec, - ) -> Dict | List[Dict]: - # Prepare data to be split into mini-batches. - flat_data = dict( - advantages=advantages, - old_logp=old_logp, - ppo_loss_mask=loss_mask, - packed_input_ids=input_.data["packed_input_ids"], - kl_rewards=kl_rewards, - ) -+ if self.use_decoupled_loss: -+ flat_data["prox_logp"] = prox_logp.float() +> **Note**: For detailed guidance on customizing rollout workflows, see the +> [agent customization guide](agent.md). - flat_input = SequenceSample.from_default( - ids=list(range(input_.bs * self.group_size)), - data=flat_data, - seqlens=[int(x) for x in input_lens.cpu().numpy().tolist()], - ) - ... - datas = flat_input.split_with_spec(spec) - ... - for mb_i, data in enumerate(datas): - train_stat = module.train_batch( - input_=data, - mb_spec=mb_spec, - version_steps=model.version.global_step, - loss_fn=_loss_fn, - loss_weight_fn=lambda x: x.data[ - "ppo_loss_mask" - ].count_nonzero(), - token_normalize_scope=self.token_normalize_scope, - ) -``` +## Step 2: Implementing the REINFORCE Training Algorithm + +Training algorithms are implemented by subclassing `TrainEngine` and using its atomic +operations like `forward`, `train_batch`, and `eval_batch`. -The `flat_input` variable will be divided into mini-batches. Each mini-batch of data will be passed into the `train_batch` method to run distributed training. The data included in this `SequenceSample` object will all be passed into the `_loss_fn`. In this case, `_loss_fn` is a wrapper over `_ppo_actor_loss_from_model_outputs`: +First, let's define the REINFORCE loss function: ```python -def _ppo_actor_loss_from_model_outputs( - logits: torch.FloatTensor, # [tot_seqlen, vocab_size] - input_: SequenceSample, - ... -) -> torch.Tensor: - ... -``` +def reinforce_loss_fn(logits, data): + input_ids = data["input_ids"] + loss_mask = data["loss_mask"].bool() + rewards = data["rewards"] -`logits` is the output of model forward, and `input_` is exactly the `input_` we passed into `train_batch`. So now we can retrieve the `prox_logp` via: - -```diff -def _ppo_actor_loss_from_model_outputs( - logits: torch.FloatTensor, # [tot_seqlen, vocab_size] - input_: SequenceSample, - ... -) -> torch.Tensor: - ... -+ prox_logp = input_.data["prox_logp"] - loss, ppo_stat = ppo_functional.actor_loss_fn( - logprobs=logprobs, - old_logprobs=old_logp, - advantages=advantages, - eps_clip=eps_clip, - loss_mask=ppo_loss_mask, - c_clip=c_clip, -+ proximal_logprobs=prox_logp, - behav_imp_weight_cap=behav_imp_weight_cap, + logprobs = gather_logprobs( + logits, torch.roll(input_ids, shifts=-1, dims=-1) ) + loss = -logprobs * rewards + loss = torch.where(loss_mask, loss, 0.0) + + return loss.sum() / loss_mask.count_nonzero() ``` -We have successfully recomputed the probability and passed it into the loss function. Next we should revise the loss computation code. - -### Modifying the PPO Loss - -```diff -def actor_loss_fn( - logprobs: torch.FloatTensor, - old_logprobs: torch.FloatTensor, - advantages: torch.FloatTensor, - eps_clip: float, - loss_mask: Optional[torch.BoolTensor] = None, - c_clip: Optional[float] = None, -+ proximal_logprobs: Optional[torch.FloatTensor] = None, - behav_imp_weight_cap: Optional[torch.FloatTensor] = None, -) -> Tuple[torch.Tensor, Dict]: - ... -+ if proximal_logprobs is not None: -+ denorm_logprobs = proximal_logprobs -+ else: -+ denorm_logprobs = old_logprobs - ... - loss_mask_count = loss_mask.count_nonzero() or 1 - # For numerical stability. -- ratio = torch.where(loss_mask, torch.exp(logprobs - old_logprobs), 0) -+ ratio = torch.where(loss_mask, torch.exp(logprobs - denorm_logprobs), 0) - ... -+ if proximal_logprobs is not None: -+ behav_kl = proximal_logprobs - old_logprobs -+ behav_imp_weight = behav_kl.exp() -+ behav_kl = torch.where(loss_mask, behav_kl, 0.0) -+ behav_imp_weight = torch.where(loss_mask, behav_imp_weight, 0.0) -+ pg_loss = pg_loss * behav_imp_weight - ... - return pg_loss, stat +```{note} +To decrease memory usage, AReaL-lite automatically packs multiple sequences in an 1D tensor before forward passes. Hence, the loss function should assume handling 1D *packed* tensors instead of *padded* tensors. ``` -### Modify the Experiment Configuration - -```diff -@dataclasses.dataclass -class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): -+ use_decoupled_loss: bool = False - - @property - def rpcs(self): - ... - # interfaces - actor_interface = ModelInterfaceAbstraction( - "ppo_actor", - args={ - **copy.deepcopy(self.ppo_kwargs), -+ "use_decoupled_loss": self.use_decoupled_loss, - ... - }, +Next, we implement the training engine. We use a two-class design to maintain backend +compatibility: + +```python +class ReinforceActor: + def __init__(self, engine: TrainEngine): + self.engine = engine + + def train_reinforce(self, data: TensorDict): + # Enable gradient checkpointing + self.engine.train() + return self.engine.train_batch( + data, + loss_fn=reinforce_loss_fn, + loss_weight_fn=lambda x: x["loss_mask"].count_nonzero(), ) -``` \ No newline at end of file + +class FSDPReinforceActor(FSDPEngine): + def __init__(self): + self.actor = ReinforceActor(self) + + def train_reinforce(self, *args, **kwargs): + return self.actor.train_reinforce(*args, **kwargs) +``` + +**Why two classes?** This design separates concerns: + +1. **Backend Agnostic Logic**: `ReinforceActor` contains the core REINFORCE algorithm + that works with any backend (FSDP, DeepSpeed, Megatron) since they share the same + `train_batch` API. + +1. **Backend-Specific Features**: `FSDPReinforceActor` inherits from `FSDPEngine` to + provide backend-specific utilities like `save`, `load`, and `upload_weights`. For + other backends, you'd create `MegatronReinforceActor`, etc. + +> **Note**: This pattern is similar to interfaces in Go or traits in Rust, adapted for +> Python's object model. + +## Step 3: Composing the Complete Training Loop + +The main training loop brings everything together: + +```python +def main(args): + # Initialize inference engine for rollouts + rollout = RemoteSGLangEngine(config.rollout) + rollout.initialize(None, ft_spec) + + # Initialize training engine + actor = FSDPReinforceActor(config=config.actor) + actor.initialize(None, ft_spec) + + # Create rollout workflow + workflow = ReMaxRLVRWorkflow( + reward_fn=gsm8k_reward_fn, + gconfig=config.gconfig, + tokenizer=tokenizer, + ) + + # Main training loop + data_generator = itertools.cycle(dataloader) + for global_step in range(max_steps): + # Generate training data + with stats_tracker.record_timing("rollout"): + batch = rollout.rollout_batch(next(data_generator), workflow=workflow) + + batch = batch.to(actor.device) + + # Synchronize all processes + dist.barrier() + torch.cuda.synchronize() + + # Training step + with ( + stats_tracker.record_timing("train_step"), + stats_tracker.scope("actor"), + ): + stats = actor.train_reinforce(batch) + actor.step_lr_scheduler() + + # Update model weights + with stats_tracker.record_timing("update_weights"): + # Weight update logic here + ... +``` diff --git a/docs/customization/dataset.md b/docs/customization/dataset.md index 9d4b4e0221..5e0d066d84 100644 --- a/docs/customization/dataset.md +++ b/docs/customization/dataset.md @@ -1,134 +1,144 @@ # Dataset -This guide provides detailed examples of how to create custom datasets in AReaL for model training. +**AReaL-lite** directly integrates with the `Dataset` class from the HuggingFace +`datasets` package. This gives you full flexibility to load, process, and filter your +data before training. -## Define Your Dataset +The required columns in your dataset depend on the specific implementation of the +`RolloutWorkflow` (for online reinforcement learning) or the training engines (for +offline training, such as `LMEngine` for Supervised Fine-Tuning (SFT)). -Create a new file under `realhf/impl/dataset/`, for example, `my_custom_dataset.py`. Your `Dataset` must implement the `torch.utils.data.Dataset` interface and follow the framework's conventions. +Here are two concrete examples from the existing implementation: -```python -class MyCustomDataset(torch.utils.data.Dataset): - - def __init__( - self, - util: data_api.DatasetUtility, - max_length: Optional[int] = None, - dataset_path: Optional[str] = None, - dataset_builder: Optional[Callable[[], List[Dict]]] = None, - # Your custom parameters - custom_param: float = 1.0, - ): - """Custom dataset initialization - - Args: - util: Dataset utility class containing tokenizer, seed, distributed info, etc. - max_length: Maximum sequence length - dataset_path: Path to dataset file (optional) - dataset_builder: Data construction function (optional, alternative to dataset_path) - custom_param: Your custom parameter - """ - self._util = util - self.max_length = max_length - - # Load and split dataset - data = data_api.load_shuffle_split_dataset(util, dataset_path, dataset_builder) - - # Your custom data processing logic - ... -``` - -## Implement Core Methods +## SFT (Offline Training) -Every dataset class must implement the following two core methods: +In the SFT example, we see that the loaded data is directly passed to the `train_lm` +method: -### 1. `__len__` Method +```python +# examples/lite/gsm8k_sft.py +def main(args): + ... + # Create dataset and dataloaders + train_dataloader = StatefulDataLoader( + get_gsm8k_dataset("train", tokenizer, rank, world_size), + collate_fn=pad_sequences_to_tensors, + ) + ... + # Run training loop + for epoch in range(total_epochs): + for step, data in enumerate(train_dataloader): + stats = engine.train_lm(data) +``` -Returns the size of the dataset: +In this case, the `train_lm` method requires the keys "input_ids", "attention_mask", and +"loss_mask" to function. We first tokenize the dataset to extract the "input_ids" and +"loss_mask". Then, the `pad_sequences_to_tensors` method is used to batch multiple +sequences and append the "attention_mask": ```python -def __len__(self): - return len(self.data_samples) +def process_gsm8k_sft_dataset(dataset: Dataset, tokenizer): + def process(sample): + seq_token = tokenizer.encode( + sample["question"] + sample["answer"] + tokenizer.eos_token + ) + prompt_token = tokenizer.encode(sample["question"]) + loss_mask = [0] * len(prompt_token) + [1] * (len(seq_token) - len(prompt_token)) + return {"input_ids": seq_token, "loss_mask": loss_mask} + + # Remove unnecessary columns to avoid errors during collation + dataset = dataset.map(process).remove_columns(["question", "answer"]) + return dataset + +def get_gsm8k_dataset(split, tokenizer, rank, world_size): + dataset = load_dataset(path="openai/gsm8k", name="main", split=split) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + return process_gsm8k_sft_dataset(dataset, tokenizer) ``` -### 2. `__getitem__` Method +## GRPO (Online Training) -Returns the sample at the specified index, must return a `SequenceSample` object: +In the GRPO example, the loaded data is passed to the `InferenceEngine`, rather than the +`TrainEngine`: ```python -def __getitem__(self, idx): - # Get raw data - sample = self.data_samples[idx] - - # Process data +# examples/lite/gsm8k_ppo.py +def main(args): ... - - # Return SequenceSample object - return data_api.SequenceSample.from_default( - ids=[sample["id"]], - seqlens=[len(processed_data["input_ids"])], - data=dict( - packed_prompts=torch.tensor(processed_data["input_ids"], dtype=torch.long), - # Other necessary data fields - ), + # Create dataset and dataloaders + train_dataloader = StatefulDataLoader( + get_gsm8k_dataset("train", rank, world_size), + collate_fn=lambda x: x, + ) + # Initialize inference engine + rollout = RemoteSGLangEngine(config.rollout) + workflow = RLVRWorkflow( + reward_fn=gsm8k_reward_fn, + ... ) + # Run training loop + ... + for global_step in range(max_steps): + batch = rollout.rollout_batch(data, workflow=workflow) + ... ``` -### Dataset Examples - -We provide some examples of dataset under `realhf/impl/dataset/`: -- For SFT, please refer `prompt_answer_dataset.py`. -- For Reward model training, please refer `rw_paired_dataset.py` -- For RL training, please refer `math_code_dataset.py` +Note that the `collate_fn` here is an identity function, meaning it simply returns the +list of individual data items as a batch. In `rollout_batch`, the data is then +dispatched to multiple concurrent executions of `workflow.arun_episode`, where each +dispatched data corresponds to a single episode. -## Data Format Requirements +The `RLVRWorkflow` implementation extracts the "messages" field from the data dictionary +as the prompt for generating a response. Additionally, this data is passed to the +`reward_fn` as keyword arguments, which allows the reward function to make use of other +dataset fields, like "answers". Here’s an example: -### JSONL File Format - -Your data file should be in JSONL format, with one JSON object per line. -If you are using our PromptDataset implementation, your data should be like: -- Math Data -```json -{"qid": "sample_1", "prompt": "Solve this math problem: 2+2=", "solutions": ["\\boxed{4}"]} -``` -- Code Data -```json -{"qid": "sample_2", "prompt": "Code problem", "input_output": "{\"inputs\": [\"5\\n2 3 5 10 12\\n\"], \"outputs\": [\"17\\n\"]}"} +```python +class RLVRWorkflow(RolloutWorkflow): + + async def arun_episode(self, engine: InferenceEngine, data): + input_ids = self.tokenizer.apply_chat_template( + data["messages"], + tokenize=True, + add_generation_prompt=True, + enable_thinking=self.enable_thinking, + ) + req = LLMRequest( + input_ids=input_ids, + ... + ) + ... + reward = self.reward_fn( + prompt=prompt_str, + completions=completions_str, + prompt_ids=resp.input_tokens, + completion_ids=resp.output_tokens, + **data, + ) ``` -- `qid`: Unique identifier for the sample -- `prompt`: Input prompt text -- `task`: Task type, used to distinguish how to calculate the reward. ("math" and "code" are supported now.) - -Note: There is no format restriction for a customized dataset as long as it can be loaded by your custom code. - -## Registration and Configuration - -### Register Dataset - -Register your dataset at the end of your dataset file: +Thus, the "messages" field must be constructed when loading the dataset, and the reward +function should be defined to handle the dataset's specific fields. Here’s how you can +process the dataset for this example: ```python -# in realhf/impl/dataset/my_custom_dataset.py -data_api.register_dataset("my-custom", MyCustomDataset) -``` +def process_gsm8k_rl_dataset(dataset: Dataset): + def process(sample): + messages = [{"role": "user", "content": sample["question"]}] + return {"messages": messages} -### Modify Experiment Configuration + # The dataset has two fields "messages" and "answer" + dataset = dataset.map(process).remove_columns(["question"]) + return dataset -Use your new dataset in the experiment configuration (refer to `realhf/experiments/common/*_exp.py`): +def get_gsm8k_dataset(split, rank, world_size): + dataset = load_dataset(path="openai/gsm8k", name="main", split=split) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + return process_gsm8k_rl_dataset(dataset) -```python -# in your experiment config file -@property -def datasets(self) -> List[DatasetAbstraction]: - return [ - DatasetAbstraction( - "my-custom", # Your registered name - args=dict( - dataset_path=self.dataset_path, - max_length=self.max_length, - custom_param=self.custom_param, - # Other initialization parameters - ), - ) - ] +def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs): + # "answer" is passed in through "**data" + from realhf.impl.dataset.math_parser import process_results + + return int(process_results(completions, answer)[0]) ``` diff --git a/docs/intro.md b/docs/intro.md index 2bbd9f023e..cd7337fb97 100644 --- a/docs/intro.md +++ b/docs/intro.md @@ -1,6 +1,6 @@ # Overview -## Welcome to AReaL’s documentation! +## Welcome to AReaL's documentation! ```{tableofcontents} ``` diff --git a/docs/legacy/customization/agent.md b/docs/legacy/customization/agent.md new file mode 100644 index 0000000000..1e46228042 --- /dev/null +++ b/docs/legacy/customization/agent.md @@ -0,0 +1,162 @@ +# Rollout and Agentic RL (Legacy) + +> **Note**: While this legacy approach works, we strongly recommend using AReaL-lite for +> new projects. It provides better flexibility, cleaner abstractions, and easier +> maintenance. + +## Step 1: Define Your Agent Class + +Create a new file under `realhf/impl/agent/`, such as `math_multi_turn_agent.py`. Your +`Agent` must implement the interface defined in `realhf/api/core/agent.py`, which +requires a single method: `collect_trajectory`. + +```python +class MathMultiTurnAgent(Agent): + async def collect_trajectory( + self, + prompt: SequenceSample, + env: EnvironmentService, + obs_queue: asyncio.Queue, + act_queue: asyncio.Queue, + ): + # Implementation goes here + ... +``` + +## Step 2: Implement the Trajectory Collection Logic + +The `collect_trajectory` method takes a task prompt, an environment, and two +communication queues. Within this method, you control the data flow between your agent +and the inference engine using these queues: + +- **obs_queue**: Send observations (token IDs and generation config) to the inference + engine +- **act_queue**: Receive actions (generated responses) from the inference engine + +Here's how the multi-turn conversation works: + +```python +for turn in range(self.num_turns): + # Send the current state to the inference engine + await obs_queue.put((qid, token_ids, self.gconfig)) + + # Get the generated response + act: BundledGenerationOutputs = await act_queue.get() + + # Evaluate the response through the environment + success, rewards = await env.step((qid, answers)) + # ... process results ... +``` + +### Environment Integration + +The environment follows a +[Gym-like interface](https://github.com/Farama-Foundation/Gymnasium) with `reset` and +`step` methods, but uses asynchronous implementations to prevent blocking across +different environment instances. + +For math problems, the environment is typically stateless and acts as a wrapper around +your reward function: + +```python +class MathCodeSingleStepEnv(EnvironmentService): + async def step(self, action: Tuple[str, List[str]]): + qid, answers = action + # ... setup code ... + + # Run reward computation asynchronously + format_rewards = await asyncio.to_thread( + math_verify_call, + answers, + # ... other parameters ... + ) + return None, format_rewards, True, False, {} +``` + +### Handling Multi-Turn Feedback + +After receiving the reward from `env.step`, check if the answer is correct. If not, +provide feedback and continue to the next turn: + +```python +for turn in range(self.num_turns): + # ... generation and evaluation code ... + + # Provide feedback based on the result + if success[0]: + feedback = "Congratulations! You are correct!" + else: + feedback = "Unfortunately your answer is wrong. Let's try again." + + # Format feedback as a user message + feedback = "\n" + self.tokenizer.apply_chat_template( + [{"content": feedback, "role": "user"}], + add_generation_prompt=True, + tokenize=False, + ) + + # Add feedback tokens to the conversation + feedback_tokens = self.tokenizer(feedback)["input_ids"] + token_ids.extend(feedback_tokens) +``` + +## Step 3: Register and Configure Your Agent + +First, register your agent implementation: + +```python +# in realhf/impl/agent/math_multi_turn_agent.py +register_agent("math-multi-turn", MathMultiTurnAgent) +``` + +```python +# in realhf/impl/agent/__init__.py +import realhf.impl.agent.math_multi_turn_agent +``` + +Then update your experiment configuration in +`realhf/experiments/async_exp/async_math_ppo.py`: + +```python +@dataclasses.dataclass +class AsyncPPOMATHConfig(AsyncRLExperimentConfig, PPOMATHConfig): + # Add any new CLI arguments your agent needs + my_param: float = 1.0 + + @property + def agent(self) -> AgentAbstraction: + return AgentAbstraction( + "math-multi-turn", # Your registered agent name + args=dict( + # Pass any arguments needed for your __init__ method + my_param=self.my_param, + # ... other configuration ... + ), + ) + + @property + def env(self) -> EnvServiceAbstraction: + # Update to use your custom environment if needed + return EnvServiceAbstraction( + "math-code-single-step", + args=dict(dataset_path=self.dataset.path) + ) +``` + +## Step 4: Run Training + +Follow the standard training procedure outlined in the +[quickstart guide](../../tutorial/quickstart_legacy.md). Launch your experiment with: + +```bash +python3 training/main_async_ppo.py my_param=5.0 # plus any additional CLI arguments +``` + +## Training Results + +Here's an example of the training reward curve from our multi-turn math agent: + +![Multi-turn Training Rewards](multiturn_reward.png) + +The agent successfully learns to solve math problems with improved accuracy over time, +demonstrating the effectiveness of the multi-turn approach. diff --git a/docs/legacy/customization/algorithm.md b/docs/legacy/customization/algorithm.md new file mode 100644 index 0000000000..1bf0768772 --- /dev/null +++ b/docs/legacy/customization/algorithm.md @@ -0,0 +1,259 @@ +# Training Algorithm (Legacy) + +> **Note**: The AReaL-lite approach is more recommended for new implementations due to +> its cleaner separation of concerns and better maintainability. + +The legacy approach encapsulates algorithms in a `ModelInterface` with three core +methods: + +```python +# From realhf/api/core/model_api.py +class ModelInterface(abc.ABC): + """Interface for model training, inference, and generation. + + This interface follows the dependency injection pattern, allowing + algorithms like REINFORCE and PPO to use the same underlying model + while exhibiting different training behaviors. + """ + + def inference( + self, + model: Model, + data: SequenceSample, + mb_spec: MicroBatchSpec, + ) -> SequenceSample | None: + raise NotImplementedError() + + def generate( + self, + model: Model, + data: SequenceSample, + mb_spec: MicroBatchSpec, + ) -> SequenceSample | None: + raise NotImplementedError() + + def train_step( + self, + model: Model, + data: SequenceSample, + mb_spec: MicroBatchSpec, + ) -> Dict | List[Dict]: + raise NotImplementedError() +``` + +When the dataflow is fixed, you typically only need to modify the algorithm interface +file. + +> **Note**: We recommend using asynchronous RL so you can customize generation behavior +> by [modifying your RL agent](agent.md) instead of the `generate` method. + +## Example 1: Grouped Advantage Normalization + +Let's modify PPO's global advantage normalization to use grouped normalization (GRPO +approach). + +### Understanding Data Organization + +Each batch contains multiple prompts (batch size) and each prompt may have multiple +responses (group size). So total sequences = batch_size × group_size. + +Sequences have different lengths but are packed into a 1D tensor. We use `cu_seqlens` +(cumulative sequence lengths) to mark boundaries, similar to flash-attention. + +### Implementation + +The standard PPO normalization looks like: + +```python +@dataclass +class PPOActorInterface(ModelInterface): + def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: + # ... + if self.adv_norm: + advantages = masked_normalization(advantages, loss_mask) + # ... +``` + +For grouped normalization, we partition advantages by group: + +```python +@dataclass +class PPOActorInterface(ModelInterface): + group_adv_norm: bool = False + + def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: + # ... + if self.adv_norm: + if not self.group_adv_norm: + advantages = masked_normalization(advantages, loss_mask) + else: + n_samples = data.bs + adv_list = [] + for i in range(0, n_samples, self.group_size): + # Define chunk boundaries + s = short1cu_seqlens[i] + e = short1cu_seqlens[i + self.group_size] + + # Extract advantages for this group + adv = advantages[s:e] + mask = loss_mask[s:e] + + # Normalize within group + advn = masked_normalization(adv, mask, all_reduce=False) + adv_list.append(advn) + + advantages = torch.cat(adv_list, 0) + # ... +``` + +### Configuration Changes + +Update the experiment configuration to expose the new parameter: + +```python +@dataclasses.dataclass +class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): + group_adv_norm: bool = False + + @property + def rpcs(self): + # ... + actor_interface = ModelInterfaceAbstraction( + "ppo_actor", + args={ + **copy.deepcopy(self.ppo_kwargs), + "group_adv_norm": self.group_adv_norm, + # ... + }, + ) +``` + +## Example 2: Decoupled PPO Loss + +The decoupled PPO loss (from AReaL's paper) recomputes probabilities before mini-batch +updates and uses this as π_prox: + +![decoupled loss](decoupled_loss.png) + +### Probability Recomputation + +We recompute probabilities using the existing `inference` method: + +```python +@dataclass +class PPOActorInterface(ModelInterface): + use_decoupled_loss: bool = False + + def train_step(self, model: Model, data: SequenceSample, mb_spec: MicroBatchSpec) -> Dict | List[Dict]: + if self.use_decoupled_loss: + s: SequenceSample = self.inference(model, data, mb_spec) + prox_logp = s.data["logprobs"] + + # Prepare mini-batch data + flat_data = dict( + advantages=advantages, + old_logp=old_logp, + ppo_loss_mask=loss_mask, + packed_input_ids=input_.data["packed_input_ids"], + kl_rewards=kl_rewards, + ) + + if self.use_decoupled_loss: + flat_data["prox_logp"] = prox_logp.float() + + flat_input = SequenceSample.from_default( + ids=list(range(input_.bs * self.group_size)), + data=flat_data, + seqlens=[int(x) for x in input_lens.cpu().numpy().tolist()], + ) + + # Split into mini-batches and train + datas = flat_input.split_with_spec(spec) + for mb_i, data in enumerate(datas): + train_stat = module.train_batch( + input_=data, + mb_spec=mb_spec, + version_steps=model.version.global_step, + loss_fn=_loss_fn, + loss_weight_fn=lambda x: x.data["ppo_loss_mask"].count_nonzero(), + token_normalize_scope=self.token_normalize_scope, + ) +``` + +### Modifying the Loss Function + +Update the loss computation to use the recomputed probabilities: + +```python +def _ppo_actor_loss_from_model_outputs( + logits: torch.FloatTensor, # [tot_seqlen, vocab_size] + input_: SequenceSample, + ... +) -> torch.Tensor: + # ... + prox_logp = input_.data.get("prox_logp") + + loss, ppo_stat = ppo_functional.actor_loss_fn( + logprobs=logprobs, + old_logprobs=old_logp, + advantages=advantages, + eps_clip=eps_clip, + loss_mask=ppo_loss_mask, + c_clip=c_clip, + proximal_logprobs=prox_logp, + behav_imp_weight_cap=behav_imp_weight_cap, + ) +``` + +And in the core loss function: + +```python +def actor_loss_fn( + logprobs: torch.FloatTensor, + old_logprobs: torch.FloatTensor, + advantages: torch.FloatTensor, + eps_clip: float, + loss_mask: Optional[torch.BoolTensor] = None, + c_clip: Optional[float] = None, + proximal_logprobs: Optional[torch.FloatTensor] = None, + behav_imp_weight_cap: Optional[torch.FloatTensor] = None, +) -> Tuple[torch.Tensor, Dict]: + # Use proximal probabilities if available, otherwise use old probabilities + denorm_logprobs = proximal_logprobs if proximal_logprobs is not None else old_logprobs + + loss_mask_count = loss_mask.count_nonzero() or 1 + + # Compute importance weights + ratio = torch.where(loss_mask, torch.exp(logprobs - denorm_logprobs), 0) + + # Apply behavioral importance weighting for decoupled loss + if proximal_logprobs is not None: + behav_kl = proximal_logprobs - old_logprobs + behav_imp_weight = behav_kl.exp() + behav_kl = torch.where(loss_mask, behav_kl, 0.0) + behav_imp_weight = torch.where(loss_mask, behav_imp_weight, 0.0) + pg_loss = pg_loss * behav_imp_weight + + # ... + return pg_loss, stat +``` + +### Configuration Update + +```python +@dataclasses.dataclass +class PPOMATHConfig(CommonExperimentConfig, PPOMATHExperimentOptions): + use_decoupled_loss: bool = False + + @property + def rpcs(self): + # ... + actor_interface = ModelInterfaceAbstraction( + "ppo_actor", + args={ + **copy.deepcopy(self.ppo_kwargs), + "use_decoupled_loss": self.use_decoupled_loss, + # ... + }, + ) +``` diff --git a/docs/legacy/customization/dataset.md b/docs/legacy/customization/dataset.md new file mode 100644 index 0000000000..36d20f35b1 --- /dev/null +++ b/docs/legacy/customization/dataset.md @@ -0,0 +1,146 @@ +# Dataset (Legacy) + +> **Note**: While this legacy approach works, we strongly recommend using AReaL-lite for +> new projects. It provides better flexibility, cleaner abstractions, and easier +> maintenance. + +## Define Your Dataset + +Create a new file under `realhf/impl/dataset/`, for example, `my_custom_dataset.py`. +Your `Dataset` must implement the `torch.utils.data.Dataset` interface and follow the +framework's conventions. + +```python +class MyCustomDataset(torch.utils.data.Dataset): + + def __init__( + self, + util: data_api.DatasetUtility, + max_length: Optional[int] = None, + dataset_path: Optional[str] = None, + dataset_builder: Optional[Callable[[], List[Dict]]] = None, + # Your custom parameters + custom_param: float = 1.0, + ): + """Custom dataset initialization + + Args: + util: Dataset utility class containing tokenizer, seed, distributed info, etc. + max_length: Maximum sequence length + dataset_path: Path to dataset file (optional) + dataset_builder: Data construction function (optional, alternative to dataset_path) + custom_param: Your custom parameter + """ + self._util = util + self.max_length = max_length + + # Load and split dataset + data = data_api.load_shuffle_split_dataset(util, dataset_path, dataset_builder) + + # Your custom data processing logic + ... +``` + +## Implement Core Methods + +Every dataset class must implement the following two core methods: + +### 1. `__len__` Method + +Returns the size of the dataset: + +```python +def __len__(self): + return len(self.data_samples) +``` + +### 2. `__getitem__` Method + +Returns the sample at the specified index, must return a `SequenceSample` object: + +```python +def __getitem__(self, idx): + # Get raw data + sample = self.data_samples[idx] + + # Process data + ... + + # Return SequenceSample object + return data_api.SequenceSample.from_default( + ids=[sample["id"]], + seqlens=[len(processed_data["input_ids"])], + data=dict( + packed_prompts=torch.tensor(processed_data["input_ids"], dtype=torch.long), + # Other necessary data fields + ), + ) +``` + +### Dataset Examples + +We provide some examples of dataset under `realhf/impl/dataset/`: + +- For SFT, please refer `prompt_answer_dataset.py`. +- For Reward model training, please refer `rw_paired_dataset.py` +- For RL training, please refer `math_code_dataset.py` + +## Data Format Requirements + +### JSONL File Format + +Your data file should be in JSONL format, with one JSON object per line. If you are +using our PromptDataset implementation, your data should be like: + +- Math Data + +```json +{"qid": "sample_1", "prompt": "Solve this math problem: 2+2=", "solutions": ["\\boxed{4}"]} +``` + +- Code Data + +```json +{"qid": "sample_2", "prompt": "Code problem", "input_output": "{\"inputs\": [\"5\\n2 3 5 10 12\\n\"], \"outputs\": [\"17\\n\"]}"} +``` + +- `qid`: Unique identifier for the sample +- `prompt`: Input prompt text +- `task`: Task type, used to distinguish how to calculate the reward. ("math" and "code" + are supported now.) + +Note: There is no format restriction for a customized dataset as long as it can be +loaded by your custom code. + +## Registration and Configuration + +### Register Dataset + +Register your dataset at the end of your dataset file: + +```python +# in realhf/impl/dataset/my_custom_dataset.py +data_api.register_dataset("my-custom", MyCustomDataset) +``` + +### Modify Experiment Configuration + +Use your new dataset in the experiment configuration (refer to +`realhf/experiments/common/*_exp.py`): + +```python +# in your experiment config file +@property +def datasets(self) -> List[DatasetAbstraction]: + return [ + DatasetAbstraction( + "my-custom", # Your registered name + args=dict( + dataset_path=self.dataset_path, + max_length=self.max_length, + custom_param=self.custom_param, + # Other initialization parameters + ), + ) + ] +``` diff --git a/docs/customization/decoupled_loss.png b/docs/legacy/customization/decoupled_loss.png similarity index 100% rename from docs/customization/decoupled_loss.png rename to docs/legacy/customization/decoupled_loss.png diff --git a/docs/customization/multiturn_reward.png b/docs/legacy/customization/multiturn_reward.png similarity index 100% rename from docs/customization/multiturn_reward.png rename to docs/legacy/customization/multiturn_reward.png diff --git a/docs/lite/gsm8k_grpo.md b/docs/lite/gsm8k_grpo.md new file mode 100644 index 0000000000..e41db2d1ce --- /dev/null +++ b/docs/lite/gsm8k_grpo.md @@ -0,0 +1,589 @@ +# Running GRPO on GSM8K Dataset + +This guide introduces how AReaL-lite runs the GRPO algorithm on the GSM8K dataset, using +the training script +[examples/lite/gsm8k_grpo.py](https://github.com/inclusionAI/AReaL/blob/main/examples/lite/gsm8k_grpo.py) +and configuration file +[examples/lite/configs/gsm8k_grpo.yaml](https://github.com/inclusionAI/AReaL/blob/main/examples/lite/configs/gsm8k_grpo.yaml). + +## How AReaL-lite Works + +The following figure illustrates the launching and one asynchronous training step of the +GRPO algorithm on the GSM8K dataset on AReaL-lite. Compared with the old AReaL +implementation, AReaL-lite runs inference servers and a SPMD training script instead of +a bunch of various workers. In a training step, AReaL-lite: + +1. Submits prompts from the dataset to `RemoteSGLangEngine`, who runs `RLVRWorkflow` in + a streaming manner. +1. Completes `RLVRWorkflow` by interacting with remote `SGLangServer` instances to + generate sequences, and computing rewards with the reward function. +1. Once there are enough outputs from `RLVRWorkflow`, aggregates them into a data batch + for algorithm-specific training engine `FSDPPPOActor`. +1. Computes losses and update weights in `FSDPPPOActor`. +1. Transfers the updated weights to remote `SGLangServer` instances. + +![AReaL-lite-gsm8k-example](gsm8k_grpo.png) + +In the following sections, we will walk you through the code to explain concepts and +show you how these steps are done in details. + +## Launching the Experiment + +As shown in the [quickstart guide](../tutorial/quickstart.md), experiments in AReaL-lite +are launched using standalone launchers with the following commands: + +``` +# Local Launcher +python -m areal.launcher.local --config +# Ray Launcher +python -m areal.launcher.ray --config +# Slurm Launcher +python -m areal.launcher.slurm --config +``` + +In AReaL-lite: + +- The **training script** is an SPMD python script that serves as the experiment entry + point. +- The launcher runs the training script with its distributed backend (`subprocess` for + `LocalLauncher`, `ray.remote` for `RayLauncher`, `srun` for `SlurmLauncher`). +- The launcher also manages inference servers (currently only supporting + `SGLangServer`). The number and parallelization strategies (e.g. tensor parallel) are + determined by the option + [allocation_mode](https://github.com/inclusionAI/AReaL/blob/main/areal/api/cli_args.py#L797). +- For distributed launchers (`RayLauncher` and `SlurmLauncher`), inference servers run + with a wrapper + [areal/launcher/sglang_server.py](https://github.com/inclusionAI/AReaL/blob/main/areal/launcher/sglang_server.py) + to handle addresses and ports in distributed settings. +- After `SGLangServer` instances are started, launchers collect their addresses and + ports to set the `AREAL_LLM_SERVER_ADDRS` environment variable for training scripts to + access these inference servers. + +The **configuration file** is a YAML file that sets the options provided in +[areal/api/cli_args.py](https://github.com/inclusionAI/AReaL/blob/main/areal/api/cli_args.py). +It could be modified via CLI arguments such as `actor.path=Qwen/Qwen3-1.7B` and +`+sglang.attention_backend=triton`. The training scripts parse the config with CLI +arguments into the config class defined in +[areal/api/cli_args.py](https://github.com/inclusionAI/AReaL/blob/main/areal/api/cli_args.py). + +``` +config, _ = load_expr_config(args, GRPOConfig) +config: GRPOConfig +``` + +## Loading and Preprocessing Dataset + +We use the `datasets` and `torchdata` packages to load and preprocess the dataset into +our dataloader. First, we download `openai/gsm8k` from Hugging Face and split it by data +parallel ranks, then map it to our desired format: + +```python +def process_gsm8k_rl_dataset(dataset: Dataset): + def process(sample): + messages = [{"role": "user", "content": sample["question"]}] + return {"messages": messages} + dataset = dataset.map(process).remove_columns(["question"]) + return dataset + +def get_gsm8k_dataset(split, rank, world_size): + dataset = load_dataset(path="openai/gsm8k", name="main", split=split) + dataset = split_dataset_by_node(dataset, rank=rank, world_size=world_size) + return process_gsm8k_rl_dataset(dataset) +``` + +We then prepare training and evaluation dataloaders with `torchdata.StatefulDataLoader`: + +```python +train_dataloader = torchdata.StatefulDataLoader( + get_gsm8k_dataset("train", rank, world_size), + batch_size=config.train_dataset.batch_size // world_size, + shuffle=config.train_dataset.shuffle, + num_workers=config.train_dataset.num_workers, + collate_fn=lambda x: x, + drop_last=config.train_dataset.drop_last, +) +valid_dataloader = ... +``` + +If you wish to use your own huggingface datasets or datasets on your local storage, +please refers to [Customization: Dataset](../customization/dataset.md) for further +details. + +## Rollout + +### Inference Engine: `RemoteSGLangEngine` + +In AReaL-lite, generation tasks are offloaded to remote inference servers, which operate +on separate GPUs from those used for training. The `RemoteSGLangEngine` acts as a client +that interacts with the servers. `RemoteSGLangEngine` runs in a SPMD manner on every +training process, without occupying any GPUs. + +`RemoteSGLangEngine` provides two core APIs that access the remote servers, `agenerate` +and `update_weights_async`. It is worth mentioning that, in asynchronous RL experiment +in AReaL-lite, inference-side weight update could happen **in the middle of** generation +of one prompt. With that being said, one output sequence could be generated by multiple +versions of models. Let us glimpse into code of `agenerate` and `update_weights_async` +for a better understanding. + +In `update_weights_async`, the engine first send `pause_generation` requests to all +inference servers, notifying them a weight update is about to happen. Upon receiveing +`pause_generation`, inference servers will immediately stop generating and respond with +already generated tokens. Then, the engine sends `update_weights_from_distributed` (for +NCCL update) or `update_weights_from_disk` (for disk update). After the update is +finished, the engine sends `continue_generation` to inference server telling them to +start working again. + +```python +class RemoteSGLangEngine: + ... + def update_weights_async(self, meta: WeightUpdateMeta): + # `update_weights_async` is completely async. + # It submits task to a ProcessPoolExecutor and returns a future + for addr in self.addresses: + res = requests.post(f"http://{addr}/pause_generation") + if meta.type == "nccl": + future = self.executor.submit( + # a function that send `update_weights_from_distributed` request + update_weights_from_distributed, + ) + elif meta.type == "disk": + ... + + def callback(future): + for addr in self.addresses + requests.post(f"http://{addr}/continue_generation") + + future.add_done_callback(callback) + return future +``` + +`agenerate` takes an `LLMRequest` with `input_ids` of **a single prompt** and generation +hyperparameters, and returns the final generation result, an `LLMResponse` with +`output_tokens` and other outputs. Since the generation could be interrupted, +`agenerate` iteratively prepares payload, sends requests and receives responses until +the generation finishes. + +```python +class RemoteSGLangEngine: + ... + async def agenerate(self, req: LLMRequest): + payload = ... # prepare payload for request + # If request is from the same workflow, choose old server + # to allow KVCache reuse. Otherwise choose server in a round + # robin manner. + server_addr = self.choose_server(req) + stop_reason = None + # other outputs are omitted for simplicity + output_tokens = [] + while (stop_reason != "stop" and len(output_tokens) < max_new_tokens): + # Request is interrupted, wait to avoid contention + if stop_reason is not None: + await asyncio.sleep(0.5) + # send request to remote sever + result = await arequest_with_retry( + addr=server_addr, + endpoint="/generate", + payload=payload, + method="POST" + ) + output_tokens.extend(result["output_ids"]) + # prepare payload for the next request + payload["input_ids"] += results["output_ids"] + payload["sample_params"]["max_new_tokens"] -= len(results["output_ids"]) + return LLMResponse( + input_tokens=req.input_ids, + output_tokens=output_tokens, + ... + ) + +``` + +The `InferenceEngine` class is designed to be extensible, supporting not just SGLang but +also other backends like vLLM. While different inference engines may be used, the +rollout management logic remains consistent. This common functionality is abstracted +into the `WorkflowExecutor`, which will be introduced in the following section. + +### `RLVRWorkflow` and `WorkflowExecutor` + +The rollout data lifecycle is controlled by an `RLVRWorkflow`, which defines how data +progresses from prompts to complete rollout data containing all fields required for +training. Our example shows a single-turn RLVR workflow with a math reward function. The +core logic of the workflow is implemented in an async method `arun_episode`, which takes +a prompt, generate answers with `RemoteSGLangEngine`, computes rewards, and populates +additional fields to produce finalized training data. + +```python +class RLVRWorkflow(RolloutWorkflow): + def __init__( + self, reward_fn, gconfig, tokenizer, ... + ): + self.reward_fn = reward_fn + self.gconfig = gconfig + self.tokenizer = tokenizer + + async def arun_episode(self, engine, data): + # rollout data with inference engine + input_ids = self.tokenizer.apply_chat_template(data["message"], ...) + req = LLMRequest(rid=..., input_ids=input_ids, gconfig=self.gconfig.new(n_samples=1)) + resps = await asyncio.gather( + *[engine.agenerate(req) for _ in range(self.gconfig.n_samples)] + ) + # post process rollout responses + results = [] + for resp in resps: + reward = self.reward_fn(...) + ... # other required fields for training + res = dict( + input_ids=..., + rewards=..., + ... # other required fields for training + ) + results.append(res) + # return padded `self.gconfig.n_samples` samples with prompt `data["message"]` + return concat_padded_tensors(results) + +def gsm8k_reward_fn(completions, answer): + ... + +tokenizer = load_hf_tokenizer(config.tokenizer_path) +workflow = RLVRWorkflow( + reward_fn=gsm8k_reward_fn, + gconfig=config.gconfig, + tokenizer=tokenizer, + ... +) +``` + +`WorkflowExecutor` is responsible for managing the data streaming through rollout +workflows, and collates completed rollout data into batched training samples. When +initializing, it launches a rollout thread that runs rollout workflows as `asyncio` +tasks. The following code shows the simplified version of rollout thread implementation, +which iteratively: + +- Checks available capacity. The capacity controls current number of rollout workflows + to limit concurrency and **data off-policyness** (The difference between the model + version used by generation and the model version updated by the trainer). +- If there is capacity left and rollout is not paused for weight update, continuously + obtains data from `input_queue` and creates `asyncio` tasks to run the workflows. +- Waits for rollout workflows to finish. +- Gathers data from finished workflows and puts them into `output_queue` + +```python +class WorkflowExecutor: + ... + async def _rollout_thread_async(self): + rid = 0 + try: + while not self.exiting.is_set(): + # Check capacity + capacity = self.get_capacity() + # Create rollout tasks with data obtained from input_queue + while ( + capacity > 0 + and not self.paused.is_set() + and self.input_queue.qsize() > 0 + ): + data, workflow = self.input_queue.get_nowait() + task = asyncio.create_task( + workflow.arun_episode(self, data), name=str(rid) + ) + rollout_tasks[str(rid)] = task + self.rollout_stat.submitted += 1 + self.rollout_stat.running += 1 + capacity -= 1 + rid += 1 + # Wait for rollout completion + tasks = list(rollout_tasks.values()) + completed_tasks = [] + if tasks: + completed_tasks, _ = await asyncio.wait( + tasks, + timeout=ROLLOUT_POLL_WAIT_TIME, + return_when=asyncio.FIRST_COMPLETED, + ) + # Collect done results, put the results into output queue + for task in completed_tasks: + traj = await task + task_rid = task.get_name() + rollout_tasks.pop(task_rid) + self.rollout_stat.accepted += 1 + self.output_queue.put_nowait(traj) + self.rollout_stat.running -= 1 + await asyncio.sleep(1) + ... +``` + +With this rollout thread running, the training script (the main thread) submits prompts +into `input_queue` and collates rollout data from `output_queue` into training batches +with `prepare_batch` (for asynchronous RL) or `rollout_batch` (for synchronous RL). The +following code shows the implementation of `prepare_batch`: + +```python +def prepare_batch( + self, + dataloader: StatefulDataLoader, + workflow: "RolloutWorkflow", +): + if not hasattr(self, "data_generator"): + self.data_generator = itertools.cycle(dataloader) + assert dataloader.batch_size is not None + while True: + # Submit at least two batches to allow maximum overlap + if ( + self.get_capacity() + dataloader.batch_size > 0 + and self.input_queue.qsize() + dataloader.batch_size + < self.input_queue.maxsize + ): + data = next(self.data_generator) + for item in data: + # submit data into input_queue + self.submit(item, workflow=workflow) + try: + # wait for dataloader.batch_size data from output_queue + return self.wait(dataloader.batch_size, timeout=1) + except TimeoutError: + pass +``` + +The `RemoteSGLangEngine` exposes `rollout_batch` and `prepare_batch` by calling them in +the workflow executor: + +```python +class RemoteSGLangEngine(InferenceEngine): + ... + def prepare_batch(self, *args, **kwargs): + return self.workflow_executor.prepare_batch(*args, **kwargs) +``` + +The usage of `RemoteSGLangEngine` in the training script is simple: + +```python +rollout = RemoteSGLangEngine(config.inf_engine) +rollout.initialize() +eval_rollout = ... + +data_generator = itertools.cycle(train_dataloader) +for global_step in range(max_steps): + # rollout batched training data for current step + if config.async_training: + batch = rollout.prepare_batch(train_dataloader, workflow=workflow) + else: + batch = rollout.rollout_batch(next(data_generator), workflow=workflow) +``` + +If you want to use rollout workflows with custom reward functions or agentic tool +calling, see [Customization: Rollout Workflows](../customization/agent.md) for more +details. + +## Training + +After obtaining the training batch, we use `FSDPPPOActor` to calculate losses and update +weights. Each train engine corresponds to one model, therefore we need an additional +engine for the reference model. Note that `torch.distributed` process groups will be +lazily initialized using `init_process_group` when the first train engine is +initialized. The initialization of train engine will also load model weights from paths +specified by the configuration. + +```python +actor = FSDPPPOActor(config=config.actor) +actor.initialize(None, ft_spec) +ref = None +if config.actor.kl_ctl > 0 and config.ref is not None: + ref = FSDPPPOActor(config=config.ref) + ref.initialize(None, ft_spec) +``` + +`FSDPPPOActor` is a high-level engine with algorithm-specific APIs, such as +`compute_logp`,`compute_advantages` and `ppo_update`. `FSDPPPOActor` is powered by the +lower-level train engine `FSDPEngine`, which use **pytorch FSDP2** to provide basic APIs +for the model such as `train_batch` and `forward`. The following code shows a GRPO +training step: + +```python +logp = actor.compute_logp(batch) +batch["prox_logp"] = logp +if ref is not None: + batch["ref_logp"] = ref.compute_logp(batch) + log_gpu_stats("ref logp") +actor.compute_advantages(batch) +stats = actor.ppo_update(batch) +actor.step_lr_scheduler() +``` + +If you want to customize your own training algorithm, see +[Customize algorithms](https://inclusionai.github.io/AReaL/customization/algorithm.html) +for more details. + +## Transferring Weights to Inference Servers + +After training, we transfer updated model weights to remote inference servers through +cooperation between `FSDPPPOActor` and `RemoteSGLangEngine`. We provide options to +transfer model weights from shared storage or NCCL. In our example training script, we +first prepare `WeightUpdateMeta` for NCCL backend on all training processes. + +```python +# NOTE: Weight update meta only requires address and free port of rank 0, +# but `WeightUpdateMeta.from_fsdp_nccl` has to be executed on all ranks +# due to `engine.get_param_specs()`. +# Therefore, we create weight update meta on all ranks, then broadcast the one on rank 0. +weight_update_meta = [ + WeightUpdateMeta.from_fsdp_nccl( + AllocationMode.from_str(config.allocation_mode), actor + ) +] +dist.broadcast_object_list(weight_update_meta, src=0) +weight_update_meta = weight_update_meta[0] +``` + +If you wish to transfer model weights from shared storage, you can use: + +```python +weight_update_meta = WeightUpdateMeta.from_disk(config.saver) +``` + +After a training step is finished, we transfer new weights from actor engine to remote +inference servers: + +1. The rollout engine needs to stop sending generation requests to remote servers + (`rollout.pause()`) before weight update to avoid server-side congestion. +1. Since we need to invoke weight update on the trainer engine and remote inference + servers at the same time, in the training script, we asynchronously send requests to + remote inference servers, and then immediately upload weights on the trainer engine. + +```python +rollout.pause() +if dist.get_rank() == 0: + future = rollout.update_weights_async(weight_update_meta) +actor.upload_weights(weight_update_meta) +if dist.get_rank() == 0: + future.result() +dist.barrier(device_ids=[actor.device.index]) +torch.cuda.synchronize() +rollout.resume() +actor.set_version(global_step + 1) +rollout.set_version(global_step + 1) +``` + +Now a complete GRPO training step in AReaL-lite is done! The core logic of our example +training script can be summarized as: + +```python +data_generator = itertools.cycle(train_dataloader) +for global_step in range(max_steps): + if config.async_training: + batch = rollout.prepare_batch(train_dataloader, workflow=workflow) + else: + batch = rollout.rollout_batch(next(data_generator), workflow=workflow) + + logp = actor.compute_logp(batch) + batch["prox_logp"] = logp + if ref is not None: + batch["ref_logp"] = ref.compute_logp(batch) + log_gpu_stats("ref logp") + actor.compute_advantages(batch) + stats = actor.ppo_update(batch) + actor.step_lr_scheduler() + + rollout.pause() + if dist.get_rank() == 0: + future = rollout.update_weights_async(weight_update_meta) + actor.upload_weights(weight_update_meta) + if dist.get_rank() == 0: + future.result() + rollout.resume() + actor.set_version(global_step + 1) + rollout.set_version(global_step + 1) +``` + +## Utilities + +In AReaL-lite, we provide a wide range of utilities for basic functionalities required +for observing and tuning your experiments. + +### `Saver` and `Evaluator` + +`Saver` +([areal/utils/saver.py](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/saver.py)) +and `Evaluator` +([areal/utils/evaluator.py](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/evaluator.py)) +manage the frequency to save and evaluate the model with the train engine. + +In our example, we call `saver.save` and `evaluator.evaluate` after every training step. +these two methods will automatically check if it is time to save or evaluate the model, +according to the experiment configuration. + +### `stats_tracker` + +`stats_tracker` +([realhf/base/stats_tracker.py](https://github.com/inclusionAI/AReaL/blob/main/realhf/base/stats_tracker.py)) +gathers training statistics across parallel ranks and reduce them. + +1. **Scalar-type statistics** are recorded by `stats_tracker.scalar(key=value)` and will + be averaged by the number of scalars with the same key when reduced. +1. **Tensor-type statistics** require `denominator` and `reduce_type` to decide how to + reduce statistics under the same key. + +- `denominator` is a bool tensor that masks the elements in the tensor that we do not + want to record. +- `reduce_type` includes average, sum, min and max. By default, the average, min and max + are all calculated. + +For example, if we want to record the length of sequences with correct and incorrect +answers in a training batch: + +```python +seqlens = ... # tensor of shape [#seqs,] +reward_score = ... # tensor of shape [#seqs,] + +result_denominators = { + "correct_n_seqs": (reward_score > 0).bool(), + "incorrect_n_seqs": (reward_score <= 0).bool(), +} +# register the denominator +stats_tracker.denominator(**result_denominators) +# record the correct and incorrect sequence length +stats_tracker.stat( + correct_seq_len=seqlens.float(), denominator="correct_n_seqs" +) +stats_tracker.stat( + incorrect_seq_len=seqlens.float(), denominator="incorrect_n_seqs" +) +``` + +`stats_tracker` offers timer context to record time cost of a code block as a scalar. +And there is also a scope context to manage keys of statistics. + +```python +with stats_tracker.record_timing("train_step"): + # training step + ... + +with stats_tracker.scope("A"): + stats_tracker.scalar(c=123) # key="A/c", value=123 + with stats_tracker.scope("B"): + stats_tracker.scalar(c=234) # key="A/B/c", value=234 +``` + +After recording sufficient data, e.g. after a `train_batch` is finished, +`stats_tracker.export` is called to aggregate all statistics and dump them into a +dictionary. + +```python +stats = stats_tracker.export() +``` + +### `StatsLogger` + +`StatsLogger` +([areal/utils/stats_logger.py](https://github.com/inclusionAI/AReaL/blob/main/areal/utils/stats_logger.py)) +logs gathered training data to recorders like `wandb` and `tensorboard` on rank 0. In +our example script, after finishing a training step, +`logger.commit(epoch, step, global_step, stats)` is called to record all statistics from +`stats_tracker` to print them as well as log them into the recorders set by the +configuration. + +## Next Steps + +- [Customize dataset](../customization/dataset.md) +- [Customize Agentic/RVLR rollout workflows](../customization/agent.md) +- [Customize algorithms](../customization/algorithm.md) diff --git a/docs/lite/gsm8k_grpo.png b/docs/lite/gsm8k_grpo.png new file mode 100644 index 0000000000..51e5bb7244 Binary files /dev/null and b/docs/lite/gsm8k_grpo.png differ diff --git a/docs/tutorial/installation.md b/docs/tutorial/installation.md index ce563b5858..16e968584e 100644 --- a/docs/tutorial/installation.md +++ b/docs/tutorial/installation.md @@ -10,31 +10,36 @@ The following hardware configuration has been extensively tested: - **CPU**: 64 cores per node - **Memory**: 1TB per node - **Network**: NVSwitch + RoCE 3.2 Tbps -- **Storage**: +- **Storage**: - 1TB local storage for single-node experiments - 10TB shared storage (NAS) for distributed experiments ### Software Requirements -| Component | Version | -|---|:---:| -| Operating System | CentOS 7 / Ubuntu 22.04 or any system meeting the requirements below | -| NVIDIA Driver | 550.127.08 | -| CUDA | 12.8 | -| Git LFS | Required for downloading models, datasets, and AReaL code. See [installation guide](https://docs.github.com/en/repositories/working-with-files/managing-large-files/installing-git-large-file-storage) | -| Docker | 27.5.1 | -| NVIDIA Container Toolkit | See [installation guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) | -| AReaL Image | `ghcr.io/inclusionai/areal-runtime:v0.3.0.post1` (includes runtime dependencies and Ray components) | +| Component | Version | +| ------------------------ | :----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------: | +| Operating System | CentOS 7 / Ubuntu 22.04 or any system meeting the requirements below | +| NVIDIA Driver | 550.127.08 | +| CUDA | 12.8 | +| Git LFS | Required for downloading models, datasets, and AReaL code. See [installation guide](https://docs.github.com/en/repositories/working-with-files/managing-large-files/installing-git-large-file-storage) | +| Docker | 27.5.1 | +| NVIDIA Container Toolkit | See [installation guide](https://docs.nvidia.com/datacenter/cloud-native/container-toolkit/latest/install-guide.html) | +| AReaL Image | `ghcr.io/inclusionai/areal-runtime:v0.3.0.post2` (includes runtime dependencies and Ray components) | -**Note**: This tutorial does not cover the installation of NVIDIA Drivers, CUDA, or shared storage mounting, as these depend on your specific node configuration and system version. Please complete these installations independently. +**Note**: This tutorial does not cover the installation of NVIDIA Drivers, CUDA, or +shared storage mounting, as these depend on your specific node configuration and system +version. Please complete these installations independently. ## Runtime Environment -**For multi-node training**: Ensure a shared storage path is mounted on every node (and mounted to the container if you are using Docker). This path will be used to save checkpoints and logs. +**For multi-node training**: Ensure a shared storage path is mounted on every node (and +mounted to the container if you are using Docker). This path will be used to save +checkpoints and logs. ### Option 1: Docker (Recommended) -We recommend using Docker with our provided image. The Dockerfile is available in the top-level directory of the AReaL repository. +We recommend using Docker with our provided image. The Dockerfile is available in the +top-level directory of the AReaL repository. ```bash docker pull ghcr.io/inclusionai/areal-runtime:v0.3.0.post1 @@ -50,9 +55,10 @@ bash examples/env/scripts/setup-container-deps.sh ### Option 2: Custom Environment Installation -1. Install [Miniconda](https://www.anaconda.com/docs/getting-started/miniconda/install) or [Anaconda](https://www.anaconda.com/docs/getting-started/anaconda/install). +1. Install [Miniconda](https://www.anaconda.com/docs/getting-started/miniconda/install) + or [Anaconda](https://www.anaconda.com/docs/getting-started/anaconda/install). -2. Create a conda virtual environment: +1. Create a conda virtual environment: ```bash conda create -n areal python=3.12 @@ -67,9 +73,11 @@ cd AReaL bash examples/env/scripts/setup-pip-deps.sh ``` + ## (Optional) Launch Ray Cluster for Distributed Training @@ -89,8 +97,9 @@ ray start --address $RAY_HEAD_IP You should see the Ray resource status displayed when running `ray status`. -Properly set the `n_nodes` argument in AReaL's training command, then AReaL's training script will automatically detect the resources and allocate workers to the cluster. +Properly set the `n_nodes` argument in AReaL's training command, then AReaL's training +script will automatically detect the resources and allocate workers to the cluster. ## Next Steps -Check the [quickstart section](quickstart.md) to launch your first AReaL job. \ No newline at end of file +Check the [quickstart section](quickstart.md) to launch your first AReaL job. diff --git a/docs/tutorial/quickstart.md b/docs/tutorial/quickstart.md index 899b377828..776b31aaf1 100644 --- a/docs/tutorial/quickstart.md +++ b/docs/tutorial/quickstart.md @@ -1,121 +1,134 @@ # Quickstart -This guide walks you through a simple example of training an LLM to solve math problems. Please ensure you have properly [installed dependencies and set up the runtime environment](installation.md) before proceeding. +Welcome to the **AReaL-lite** Quickstart Guide! This guide demonstrates how to run an +AReaL-lite experiment training an LLM on the GSM8K dataset using the GRPO algorithm with +function-based rewards. Ensure you've completed +[the installation and environment setup](installation.md) before proceeding. -## Dataset +## Running the Experiment (on a single node) -Use `huggingface-cli` to download our open-source dataset: +To run the experiment, you will need: -```bash -huggingface-cli download --repo-type=dataset inclusionAI/AReaL-RL-Data -``` - -> **Note**: The command above will display the path of the downloaded dataset. You'll need to pass this path to the training command. +- Training script: + [examples/lite/gsm8k_grpo.py](https://github.com/inclusionAI/AReaL/blob/main/examples/lite/gsm8k_grpo.py) +- Config YAML: + [examples/lite/configs/gsm8k_grpo.yaml](https://github.com/inclusionAI/AReaL/blob/main/examples/lite/configs/gsm8k_grpo.yaml) -## Model +Our training scripts will automatically download the dataset (openai/gsm8k) and model +(Qwen/Qwen2-1.5B-Instruct). To run the example with default configuration, execute from +the repository directory: -We train using open-source models available on Hugging Face Hub. You can either download the model in advance or use the model identifier when running the experiment. - -```bash -# If you want to download it in advance -huggingface-cli download Qwen/Qwen3-1.7B ``` - -Refer to the [official documentation](https://huggingface.co/docs/huggingface_hub/guides/cli) for more information on using `huggingface-cli`. - -## Training - -From the repository directory, run: - -```bash -# examples/run_async_ppo.sh -python3 training/main_async_ppo.py \ - n_nodes=1 n_gpus_per_node=8 \ - allocation_mode=sglang.d4p1m1+d2p2m1 \ - cluster.fileroot=/path/to/save/logs/checkpoints/ \ - actor.type._class=qwen3 \ - actor.path=Qwen/Qwen3-1.7B \ - ref.type._class=qwen3 \ - ref.path=Qwen/Qwen3-1.7B \ - dataset.path=/path/to/boba_106k_0319.jsonl \ - dataset.train_bs_n_seqs=32 \ - group_size=8 \ - ppo.gen.max_new_tokens=4096 \ - ppo.ppo_n_minibatches=4 \ - actor_train.mb_spec.max_tokens_per_mb=32768 \ - actor_inf.mb_spec.max_tokens_per_mb=32768 \ - max_concurrent_rollouts=16 \ - max_head_offpolicyness=4 +python3 -m areal.launcher.local examples/lite/gsm8k_grpo.py --config examples/lite/configs/gsm8k_grpo.yaml experiment_name= trial_name= ``` -::::{important} -Running `main_async_ppo.py` with `ppo.recompute_logprob=False`, `ppo.use_decoupled_loss=False`, and `max_head_offpolicyness=0` will essentially replicate the behavior of synchronous PPO. Therefore, it's usually not recommended to run synchronous PPO directly (i.e., `main_sync_ppo.py`). The workflow of asynchronous RL is more stable and easier to customize. -:::: +> **Note**: The command above uses `LocalLauncher`, which only works for a single node +> (`cluster.n_nodes == 1`). For distributed experiments, see +> [Distributed Experiments with Ray or Slurm](#distributed-experiments-with-ray-or-slurm). -## Command Line Options +## Modifying configuration -To view all available options: +All available configuration options are listed in +[areal/api/cli_args.py](https://github.com/inclusionAI/AReaL/blob/main/areal/api/cli_args.py). +To customize the experiment (models, resources, algorithm options), you can: -```bash -python3 training/main_sync_ppo.py --help -``` +1. Edit the YAML file directly at + [examples/lite/configs/gsm8k_grpo.yaml](https://github.com/inclusionAI/AReaL/blob/main/examples/lite/configs/gsm8k_grpo.yaml). +1. Add command-line options: + - For existing options in the YAML file, directly add the option: + `actor.path=Qwen/Qwen3-1.7B`. + - For other options in `cli_args.py`, but not in the YAML file, add with a prefix + "+": `+sglang.attention_backend=triton`. -### Configuration Parameters +For example, here is the command to launch a customized configuration, based on our +GSM8K GRPO example: -- **`experiment_name`**: The name of your project. -- **`trial_name`**: The name of this trial in your project. -- **`{actor|ref}.path`**: The path to the model files. -- **`dataset.path`**: The path to the dataset JSONL file. -- **`cluster.fileroot`**: The root path for saving training outputs (logs and checkpoints). -- **`n_nodes`**: The number of nodes in the cluster. -- **`n_gpus_per_node`**: The number of GPUs per node. -- **`allocation_mode`**: The GPU allocation strategy and 3D parallelism configuration for the experiment. Format: - - `sglang.d${DP1}m${TP1}p${PP1}+d${DP2}m${TP2}p${PP2}`: Configures parallel strategies for SGLang generation and training respectively. Generation and training use separate GPU sets, and the total GPU count must equal: DP1×TP1×PP1 + DP2×TP2×PP2 = #GPUs. +``` +python3 -m areal.launcher.local examples/lite/gsm8k_grpo.py \ + --config examples/lite/configs/gsm8k_grpo.yaml \ + experiment_name= \ + trial_name= \ + allocation_mode=sglang.d2p1t1+d2p1t1 \ + cluster.n_nodes=1 \ + cluster.n_gpus_per_node=4 \ + gconfig.max_new_tokens=2048 \ + train_dataset.batch_size=1024 \ + +sglang.attention_backend=triton +``` -### Training Control +::::{important} We're currently refactoring from legacy AReaL to AReaL-lite, which +introduces some configuration differences. We provide a **config converter** to transfer +old AReaL config into AReaL-lite YAML file for users' convenience. +[Click here](#switching-from-legacy-areal-to-areal-lite) to learn how to use the +**config converter**. :::: -- **`exp_ctrl.total_train_epochs`**: Number of training epochs (complete dataset iterations). -- **`exp_ctrl.save_freq_{epochs|steps|secs}`**: Frequency for saving model parameters to persistent storage. Set to null to disable saving. -- **`exp_ctrl.ckpt_freq_{epochs|steps|secs}`**: Frequency for saving temporary parameters for restart capability. -- **`dataset.train_bs_n_seqs`**: Training batch size (number of prompts sampled per training iteration). -- **`group_size`**: Number of responses sampled per prompt. +## Distributed Experiments with Ray or Slurm -### Memory and Performance +AReaL-lite provides standalone launchers for distributed experiments. After setting up +your Ray or Slurm cluster, launch experiments similarly to `LocalLauncher`: -- **`{actor_train|ref_inf|actor_inf}.mb_spec.max_tokens_per_mb`**: Maximum tokens per mini-batch for forward/backward passes during reference model inference and actor model training. Reduce this value to avoid OOM errors. -- **`max_concurrent_rollouts`**: The maximum number of concurrent rollouts. SGLang will run out of memory if this value is too large. Defaults to `dataset.train_bs_n_seqs`. +``` +# Launch with Ray launcher. 4 nodes (4 GPUs each), 3 nodes for generation, 1 node for training. +python3 -m areal.launcher.ray examples/lite/gsm8k_grpo.py \ + --config examples/lite/configs/gsm8k_grpo.yaml \ + experiment_name= \ + trial_name= \ + allocation_mode=sglang.d12p1t1+d4p1t1 \ + cluster.n_nodes=4 \ + cluster.n_gpus_per_node=4 \ + +# Launch with Slurm launcher. 16 nodes (8 GPUs each), 12 nodes for generation, 4 nodes for training +python3 -m areal.launcher.slurm examples/lite/gsm8k_grpo.py \ + --config examples/lite/configs/gsm8k_grpo.yaml \ + experiment_name= \ + trial_name= \ + allocation_mode=sglang.d96p1t1+d32p1t1 \ + cluster.n_nodes=16 \ + cluster.n_gpus_per_node=8 \ +``` -### Algorithm Configuration +Additional references: -- **`max_head_offpolicyness`**: The allowed maximum data staleness. 0 recovers synchronous training. A large value will increase generation throughput but degrade final performance. We recommend keeping this value at 8 or below. -- **`ppo.recompute_logprob`**: Whether to compute proximal log probabilities for training. Defaults to True for asynchronous experiments and False for synchronous baselines. -- **`ppo.use_decoupled_loss`**: Use decoupled loss to stabilize asynchronous training. Defaults to True. -- **`ppo.gen.max_new_tokens`**: Maximum tokens to generate per prompt. -- **`ppo.ppo_n_minibatches`**: Number of mini-batches for dividing data during each PPO update. -- **`success_rate_ub`**: Upper bound of success rate. Prompts with a higher success rate will be filtered out. -- **`success_rate_lb`**: Lower bound of success rate. Prompts with a lower success rate will be filtered out. +- For more options for launchers, check `LauncherConfig` in + [areal/api/cli_args.py](https://github.com/inclusionAI/AReaL/blob/main/areal/api/cli_args.py). +- [Ray cluster setup guide](./installation.md#optional-launch-ray-cluster-for-distributed-training) + for a guide on how to set up a ray cluster. -## Monitoring the Training Process +> **Important Notes**: +> +> 1. Ensure `allocation_mode` matches your cluster configuration +> (`#GPUs == cluster.n_nodes * cluster.n_gpus_per_node`) +> 1. Ray or Slurm launchers only works for more than 1 node (`cluster.n_nodes > 1`). For +> single node scenario, please use `LocalLauncher`. +> 1. In Ray or Slurm launchers, GPUs are allocated at node granularity, which means +> #GPUs for generation or training must be integer multiples of +> `cluster.n_gpus_per_node`. -+ We recommend using [Weights & Biases (wandb)](https://github.com/wandb/wandb) or [SwanLab](https://github.com/SwanHubX/SwanLab) for monitoring—run `wandb login` or `swanlab login`, or set the corresponding environment variable API key (`WANDB_API_KEY` or `SWANLAB_API_KEY`). Set `wandb.mode="online"` or `swanlab.mode="cloud"` in your configuration to upload training statistics. If you cannot connect to the server, you can also use `wandb.mode="offline"` or `swanlab.mode="local"` to save data locally without uploading. + +## Switching from legacy AReaL to AReaL-lite -You can also use TensorBoard by setting the `tensorboard.path` parameter. +We also provide a convenient script to convert your AReaL YAML config into AReaL-lite +config in one command line. First you need to locate your AReaL config either modified +from files from `examples` folder, or generated when you run your experiments in +`//` folder. Runs: -The main log will be saved to `${fileroot}/logs/${USER}/${experiment_name}/${trial_name}/main.log` and contains the statistics uploaded to wandb. +``` +python arealite/utils/convert.py --convert_src AReaL --src_config_path --template_path examples/arealite/configs/gsm8k_grpo.yaml --output_path +``` -If SwanLab is enabled, logs will be saved to the directory specified by `swanlab.logdir`. +Then you should be able to run experiments with your old settings on AReaL-lite! -### Key Training Statistics +## Next Steps -- **`Epoch 1/5`**: Indicates the total epochs required and the current epoch being trained. -- **`step 6/19`**: Shows that the current epoch has 19 steps, with the 6th step just completed. -- **`global step 6`**: Step count across all epochs. -- **`ppo_actor/task_reward/avg`**: Average reward value of all sampled responses in this step. This should steadily increase during training and eventually stabilize. -- **`ppo_actor/importance_weight/avg`**: Average importance sampling ratio across all tokens in the PPO loss. This is typically close to 1.0. -- **`ppo_actor/actor_clip_ratio/avg`**: Ratio of clipped tokens in PPO loss to total tokens. This is usually less than 0.1. -- **`ppo_actor/actor_loss/avg`**: PPO loss value. **This does not show clear trends during training** and should not be used as a performance indicator. +Check [Getting Started with AReaL-lite](../lite/gsm8k_grpo.md) for a complete code +walkthrough on the GRPO GSM8K Example. -## Next Steps +Customization guides: -[Evaluate your model](eval.md) or check the [troubleshooting section](troubleshooting.md) if you encounter any issues. \ No newline at end of file +- [Custom dataset](../customization/dataset.md) +- [Custom agentic/RVLR rollout workflows](../customization/agent.md) +- [Custom algorithms](../customization/algorithm.md) diff --git a/docs/tutorial/quickstart_legacy.md b/docs/tutorial/quickstart_legacy.md new file mode 100644 index 0000000000..f4e792330c --- /dev/null +++ b/docs/tutorial/quickstart_legacy.md @@ -0,0 +1,169 @@ +# Quickstart (Legacy) + +> **Note**: This is a quickstart guide for launching AReaL experiment with legacy code +> in `realhf/`. We strongly recommend users to try AReaL-lite for better experiences. +> [Click here](quickstart.md) for AReaL-lite quickstart guide! + +This guide walks you through a simple example of training an LLM to solve math problems. +Please ensure you have properly +[installed dependencies and set up the runtime environment](installation.md) before +proceeding. + +## Dataset + +Use `huggingface-cli` to download our open-source dataset: + +```bash +huggingface-cli download --repo-type=dataset inclusionAI/AReaL-RL-Data +``` + +> **Note**: The command above will display the path of the downloaded dataset. You'll +> need to pass this path to the training command. + +## Model + +We train using open-source models available on Hugging Face Hub. You can either download +the model in advance or use the model identifier when running the experiment. + +```bash +# If you want to download it in advance +huggingface-cli download Qwen/Qwen3-1.7B +``` + +Refer to the +[official documentation](https://huggingface.co/docs/huggingface_hub/guides/cli) for +more information on using `huggingface-cli`. + +## Training + +From the repository directory, run: + +```bash +# examples/run_async_ppo.sh +python3 training/main_async_ppo.py \ + n_nodes=1 n_gpus_per_node=8 \ + allocation_mode=sglang.d4p1m1+d2p2m1 \ + cluster.fileroot=/path/to/save/logs/checkpoints/ \ + actor.type._class=qwen3 \ + actor.path=Qwen/Qwen3-1.7B \ + ref.type._class=qwen3 \ + ref.path=Qwen/Qwen3-1.7B \ + dataset.path=/path/to/boba_106k_0319.jsonl \ + dataset.train_bs_n_seqs=32 \ + group_size=8 \ + ppo.gen.max_new_tokens=4096 \ + ppo.ppo_n_minibatches=4 \ + actor_train.mb_spec.max_tokens_per_mb=32768 \ + actor_inf.mb_spec.max_tokens_per_mb=32768 \ + max_concurrent_rollouts=16 \ + max_head_offpolicyness=4 +``` + +::::{important} Running `main_async_ppo.py` with `ppo.recompute_logprob=False`, +`ppo.use_decoupled_loss=False`, and `max_head_offpolicyness=0` will essentially +replicate the behavior of synchronous PPO. Therefore, it's usually not recommended to +run synchronous PPO directly (i.e., `main_sync_ppo.py`). The workflow of asynchronous RL +is more stable and easier to customize. :::: + +## Command Line Options + +To view all available options: + +```bash +python3 training/main_sync_ppo.py --help +``` + +### Configuration Parameters + +- **`experiment_name`**: The name of your project. +- **`trial_name`**: The name of this trial in your project. +- **`{actor|ref}.path`**: The path to the model files. +- **`dataset.path`**: The path to the dataset JSONL file. +- **`cluster.fileroot`**: The root path for saving training outputs (logs and + checkpoints). +- **`n_nodes`**: The number of nodes in the cluster. +- **`n_gpus_per_node`**: The number of GPUs per node. +- **`allocation_mode`**: The GPU allocation strategy and 3D parallelism configuration + for the experiment. Format: + - `sglang.d${DP1}m${TP1}p${PP1}+d${DP2}m${TP2}p${PP2}`: Configures parallel strategies + for SGLang generation and training respectively. Generation and training use + separate GPU sets, and the total GPU count must equal: DP1×TP1×PP1 + DP2×TP2×PP2 = + #GPUs. + +### Training Control + +- **`exp_ctrl.total_train_epochs`**: Number of training epochs (complete dataset + iterations). +- **`exp_ctrl.save_freq_{epochs|steps|secs}`**: Frequency for saving model parameters to + persistent storage. Set to null to disable saving. +- **`exp_ctrl.ckpt_freq_{epochs|steps|secs}`**: Frequency for saving temporary + parameters for restart capability. +- **`dataset.train_bs_n_seqs`**: Training batch size (number of prompts sampled per + training iteration). +- **`group_size`**: Number of responses sampled per prompt. + +### Memory and Performance + +- **`{actor_train|ref_inf|actor_inf}.mb_spec.max_tokens_per_mb`**: Maximum tokens per + mini-batch for forward/backward passes during reference model inference and actor + model training. Reduce this value to avoid OOM errors. +- **`max_concurrent_rollouts`**: The maximum number of concurrent rollouts. SGLang will + run out of memory if this value is too large. Defaults to `dataset.train_bs_n_seqs`. + +### Algorithm Configuration + +- **`max_head_offpolicyness`**: The allowed maximum data staleness. 0 recovers + synchronous training. A large value will increase generation throughput but degrade + final performance. We recommend keeping this value at 8 or below. +- **`ppo.recompute_logprob`**: Whether to compute proximal log probabilities for + training. Defaults to True for asynchronous experiments and False for synchronous + baselines. +- **`ppo.use_decoupled_loss`**: Use decoupled loss to stabilize asynchronous training. + Defaults to True. +- **`ppo.gen.max_new_tokens`**: Maximum tokens to generate per prompt. +- **`ppo.ppo_n_minibatches`**: Number of mini-batches for dividing data during each PPO + update. +- **`success_rate_ub`**: Upper bound of success rate. Prompts with a higher success rate + will be filtered out. +- **`success_rate_lb`**: Lower bound of success rate. Prompts with a lower success rate + will be filtered out. + +## Monitoring the Training Process + +- We recommend using [Weights & Biases (wandb)](https://github.com/wandb/wandb) or + [SwanLab](https://github.com/SwanHubX/SwanLab) for monitoring—run `wandb login` or + `swanlab login`, or set the corresponding environment variable API key + (`WANDB_API_KEY` or `SWANLAB_API_KEY`). Set `wandb.mode="online"` or + `swanlab.mode="cloud"` in your configuration to upload training statistics. If you + cannot connect to the server, you can also use `wandb.mode="offline"` or + `swanlab.mode="local"` to save data locally without uploading. + +You can also use TensorBoard by setting the `tensorboard.path` parameter. + +The main log will be saved to +`${fileroot}/logs/${USER}/${experiment_name}/${trial_name}/main.log` and contains the +statistics uploaded to wandb. + +If SwanLab is enabled, logs will be saved to the directory specified by +`swanlab.logdir`. + +### Key Training Statistics + +- **`Epoch 1/5`**: Indicates the total epochs required and the current epoch being + trained. +- **`step 6/19`**: Shows that the current epoch has 19 steps, with the 6th step just + completed. +- **`global step 6`**: Step count across all epochs. +- **`ppo_actor/task_reward/avg`**: Average reward value of all sampled responses in this + step. This should steadily increase during training and eventually stabilize. +- **`ppo_actor/importance_weight/avg`**: Average importance sampling ratio across all + tokens in the PPO loss. This is typically close to 1.0. +- **`ppo_actor/actor_clip_ratio/avg`**: Ratio of clipped tokens in PPO loss to total + tokens. This is usually less than 0.1. +- **`ppo_actor/actor_loss/avg`**: PPO loss value. **This does not show clear trends + during training** and should not be used as a performance indicator. + +## Next Steps + +[Evaluate your model](eval.md) or check the +[troubleshooting section](troubleshooting.md) if you encounter any issues. diff --git a/examples/env/scripts/setup-container-deps.sh b/examples/env/scripts/setup-container-deps.sh deleted file mode 100644 index 2b07e31601..0000000000 --- a/examples/env/scripts/setup-container-deps.sh +++ /dev/null @@ -1,11 +0,0 @@ -#!/bin/sh -AREAL_PATH=$PWD -cd /sglang -git apply $AREAL_PATH/patch/sglang/v0.4.6.post4.patch -cd $AREAL_PATH - -# Package used for calculating math reward -pip install -e evaluation/latex2sympy - -# Install AReaL -pip install -e . \ No newline at end of file diff --git a/examples/env/scripts/setup-eval-pip-deps.sh b/examples/env/scripts/setup-eval-pip-deps.sh index fbdf4ee4f5..659ed73460 100644 --- a/examples/env/scripts/setup-eval-pip-deps.sh +++ b/examples/env/scripts/setup-eval-pip-deps.sh @@ -1,9 +1,9 @@ #/bin/bash # basic dependencies pip install -U pip -pip uninstall deepspeed flash-attn pynvml cugraph-dgl dask-cuda cugraph-service-server raft-dask cugraph cuml cugraph-pyg -y -pip install nvidia-ml-py +pip uninstall pynvml cugraph-dgl dask-cuda cugraph-service-server raft-dask cugraph cuml cugraph-pyg -y +pip install pynvml nvidia-ml-py pip install -e evaluation/latex2sympy pip install vllm==0.8.5 --no-build-isolation -pip install flash_attn --no-build-isolation +pip install "flash-attn<=2.7.3" --no-build-isolation pip install -r evaluation/requirements.txt \ No newline at end of file diff --git a/examples/env/scripts/setup-pip-deps.sh b/examples/env/scripts/setup-pip-deps.sh index 8fa203bf97..e907109e7b 100644 --- a/examples/env/scripts/setup-pip-deps.sh +++ b/examples/env/scripts/setup-pip-deps.sh @@ -1,24 +1,14 @@ #!/bin/bash # basic dependencies pip install -U pip -pip uninstall torch deepspeed flash-attn pynvml cugraph-dgl dask-cuda cugraph-service-server raft-dask cugraph cuml cugraph-pyg -y -pip install torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 -pip install "sglang[all]==0.4.6.post4" +pip uninstall pynvml cugraph-dgl dask-cuda cugraph-service-server raft-dask cugraph cuml cugraph-pyg -y +pip install torch==2.7.1 torchaudio==2.7.1 torchvision==0.22.1 "deepspeed>=0.17.2" pynvml +pip install "sglang[all]==0.4.9.post2" pip install megatron-core==0.11.0 nvidia-ml-py pip install git+https://github.com/garrett4wade/cugae --no-build-isolation --verbose -pip install "flash-attn<=2.7.3" --no-build-isolation +pip install "flash-attn<=2.8.2" --no-build-isolation # Package used for calculating math reward pip install -e evaluation/latex2sympy - -# Install an editable sglang -rm -rf ./sglang -git clone -b v0.4.6.post4 https://github.com/sgl-project/sglang -AREAL_PATH=$PWD -cd sglang -git apply ../patch/sglang/v0.4.6.post4.patch -pip install -e "python[all]" --no-deps -cd $AREAL_PATH - # Install AReaL pip install -e . diff --git a/examples/env/validate_installation.py b/examples/env/validate_installation.py index 61ef6f7c5d..67ae740cdf 100644 --- a/examples/env/validate_installation.py +++ b/examples/env/validate_installation.py @@ -67,6 +67,7 @@ def test_torch_cuda(self, torch_module): def test_flash_attn_functionality(self, flash_attn_module): """Test flash attention functionality.""" # Try to import key functions + import flash_attn_2_cuda from flash_attn import flash_attn_func, flash_attn_varlen_func print(" - Flash attention functions imported successfully") @@ -79,12 +80,12 @@ def test_sglang_functionality(self, sglang_module): """Test SGLang basic functionality.""" # Basic import test is sufficient for CI import sgl_kernel - from sglang import launch_server - assert Version(get_version("sglang")) == Version("0.4.6.post4") + from sglang import Engine, launch_server + assert Version(get_version("sglang")) == Version("0.4.9.post2"), "SGLang version should be v0.4.9.post2" print(" - SGLang imported successfully") def test_transformers(self, transformers_module): - assert Version(get_version("transformers")) == Version("4.51.1") + assert Version(get_version("transformers")) == Version("4.53.1"), "transformers version should be 4.53.1" print(" - transformers imported successfully") def validate_critical_dependencies(self): @@ -140,7 +141,7 @@ def validate_optional_dependencies(self): self.test_import("flashattn_hopper", required=False) # Optional utilities - self.test_import("tensorboardx", required=False) + self.test_import("tensorboardX", required=False) self.test_import("swanlab", required=False) self.test_import("matplotlib", required=False) self.test_import("seaborn", required=False) diff --git a/examples/lite/clevr_count_70k_grpo.py b/examples/lite/clevr_count_70k_grpo.py new file mode 100644 index 0000000000..1b78fa1038 --- /dev/null +++ b/examples/lite/clevr_count_70k_grpo.py @@ -0,0 +1,263 @@ +import os +import re +import sys + +import torch +import torch.distributed as dist +import wandb +from torch.utils.data import Subset +from torchdata.stateful_dataloader import StatefulDataLoader + +from areal.api.cli_args import GRPOConfig, load_expr_config +from areal.api.io_struct import FinetuneSpec, WeightUpdateMeta +from areal.dataset.__init__ import get_custom_dataset +from areal.engine.ppo.actor import FSDPPPOActor +from areal.engine.sglang_remote import RemoteSGLangEngine +from areal.utils.device import log_gpu_stats +from areal.utils.evaluator import Evaluator +from areal.utils.saver import Saver +from areal.utils.stats_logger import StatsLogger +from areal.workflow.vision_rlvr import VisionRLVRWorkflow +from realhf.api.core.data_api import load_hf_processor_and_tokenizer +from realhf.base import seeding, stats_tracker + + +def extract_answer(pred_str, data_name, use_last_number=True): + match = re.findall(r"\[([0-9\.]+)\]", pred_str) + if match: + return match[-1] + + return "" + + +def clevr_count_70k_reward_fn( + prompt, completions, prompt_ids, completion_ids, answer, **kwargs +): + sol = extract_answer(completions, data_name="") # str number + ans = answer + + if sol is None: + return 0 + if ans is None: + return 0 + + if sol.strip() == ans.strip(): + print(f"completions: {completions}, answer: {answer}") + return 1 + + return 0 + + +def main(args): + + wandb.init(project="clevr_70k") + + config, _ = load_expr_config(args, GRPOConfig) + config: GRPOConfig + + rank = int(os.getenv("RANK")) + world_size = int(os.getenv("WORLD_SIZE")) + + seeding.set_random_seed(config.seed, f"trainer{rank}") + + processor, tokenizer = load_hf_processor_and_tokenizer(config.tokenizer_path) + train_dataset = get_custom_dataset( + path=config.train_dataset.path, + rank=rank, + world_size=world_size, + split="train", + type=config.train_dataset.type, + processor=processor, + ) + + train_size = len(train_dataset) + subset_size = int(1.0 * train_size) + + random_indices = torch.randperm(train_size).tolist()[:subset_size] + + subset_train_dataset = Subset(train_dataset, random_indices) + + valid_dataset = get_custom_dataset( + path=config.valid_dataset.path, + rank=rank, + world_size=world_size, + split="test", + type=config.valid_dataset.type, + processor=processor, + ) + # Create dataset and dataloaders + train_dataloader = StatefulDataLoader( + subset_train_dataset, + batch_size=config.train_dataset.batch_size // world_size, + shuffle=config.train_dataset.shuffle, + num_workers=config.train_dataset.num_workers, + collate_fn=lambda x: x, + drop_last=config.train_dataset.drop_last, + ) + valid_dataloader = StatefulDataLoader( + valid_dataset, + batch_size=config.valid_dataset.batch_size // world_size, + shuffle=config.valid_dataset.shuffle, + num_workers=config.valid_dataset.num_workers, + collate_fn=lambda x: x, + drop_last=config.valid_dataset.drop_last, + ) + ft_spec = FinetuneSpec( + total_train_epochs=config.total_train_epochs, + dataset_size=len(train_dataloader) * config.train_dataset.batch_size, + train_batch_size=config.train_dataset.batch_size, + ) + + # Initialize inference engine + rollout = RemoteSGLangEngine(config.rollout) + rollout.initialize(None, ft_spec) + eval_rollout = RemoteSGLangEngine(config.rollout) + eval_rollout.initialize(None, ft_spec) + # NOTE: set a large version such that eval does not have any offpolicyness control + eval_rollout.set_version(int(1e12)) + + # Initialize train engine + actor = FSDPPPOActor(config=config.actor) + actor.initialize(None, ft_spec) + ref = None + if config.actor.kl_ctl > 0 and config.ref is not None: + ref = FSDPPPOActor(config=config.ref) + ref.initialize(None, ft_spec) + + # NOTE: Weight update meta only requires address and free port of rank 0, + # but `WeightUpdateMeta.from_fsdp_nccl` has to be executed on all ranks + # due to `engine.get_param_specs()`. + # Therefore, we create weight update meta on all ranks, then broadcast the one on rank 0. + weight_update_meta = [WeightUpdateMeta.from_disk(config.saver)] + dist.broadcast_object_list(weight_update_meta, src=0) + weight_update_meta = weight_update_meta[0] + + # Create rollout workflow + if tokenizer.pad_token_id not in config.gconfig.stop_token_ids: + config.gconfig.stop_token_ids.append(tokenizer.pad_token_id) + if tokenizer.eos_token_id not in config.gconfig.stop_token_ids: + config.gconfig.stop_token_ids.append(tokenizer.eos_token_id) + + workflow = VisionRLVRWorkflow( + reward_fn=clevr_count_70k_reward_fn, + gconfig=config.gconfig, + tokenizer=tokenizer, + processor=processor, + enable_thinking=False, + ) + + # Run training. + saver = Saver(config.saver, ft_spec, for_recover=False) + stats_logger = StatsLogger(config.stats_logger, ft_spec) + evaluator = Evaluator(config.evaluator, ft_spec) + + total_epochs = config.total_train_epochs + steps_per_epoch = len(train_dataloader) + max_steps = total_epochs * steps_per_epoch + + data_generator = iter(train_dataloader) + for global_step in range(max_steps): + epoch = global_step // steps_per_epoch + step = global_step % steps_per_epoch + + with stats_tracker.record_timing("rollout"): + if config.async_training: + batch = rollout.prepare_batch(train_dataloader, workflow=workflow) + else: + try: + data = next(data_generator) + except StopIteration: + data_generator = iter(train_dataloader) + data = next(data_generator) + batch = rollout.rollout_batch(data, workflow=workflow) + + batch = batch.to(actor.device) + # Create barrier to synchronize all rollout processes. + dist.barrier(device_ids=[actor.device.index]) + torch.cuda.synchronize() + + if config.actor.recompute_logprob or config.actor.use_decoupled_loss: + with stats_tracker.record_timing("recompute_logp"): + logp = actor.compute_logp(batch) + batch["prox_logp"] = logp + log_gpu_stats("recompute logp") + + if ref is not None: + with stats_tracker.record_timing("ref_logp"): + batch["ref_logp"] = ref.compute_logp(batch) + + log_gpu_stats("ref logp") + + with stats_tracker.record_timing("compute_advantage"): + actor.compute_advantages(batch) + log_gpu_stats("compute advantages") + + with ( + stats_tracker.record_timing("train_step"), + stats_tracker.scope("grpo_actor"), + ): + stats = actor.ppo_update(batch) + wandb.log({"final_reward": stats[0]["grpo_actor/final_reward/avg"]}) + wandb.log({"task_reward": stats[0]["grpo_actor/task_reward/avg"]}) + actor.step_lr_scheduler() + log_gpu_stats("ppo update") + + with stats_tracker.record_timing("update_weights"): + rollout.pause() + if dist.get_rank() == 0: + future = rollout.update_weights(weight_update_meta) + actor.upload_weights(weight_update_meta) + if dist.get_rank() == 0: + future.result() + dist.barrier(device_ids=[actor.device.index]) + torch.cuda.synchronize() + rollout.resume() + actor.set_version(global_step + 1) + rollout.set_version(global_step + 1) + + with stats_tracker.record_timing("save"): + saver.save(actor, epoch, step, global_step) + + with stats_tracker.record_timing("eval"): + + def evaluate_fn(): + rollout.pause() + cnt = 0 + for data in valid_dataloader: + for item in data: + eval_rollout.submit(item, workflow) + cnt += 1 + batch = eval_rollout.wait(cnt, timeout=None) + rewards = batch["rewards"].float().to(actor.device) + wandb.log({"eval_reward": rewards.mean().item()}) + with stats_tracker.scope("grpo-eval"): + stats_tracker.denominator( + n_seqs=torch.ones( + rewards.shape[0], + device=rewards.device, + dtype=torch.bool, + ) + ) + stats_tracker.stat(task_reward=rewards, denominator="n_seqs") + rollout.resume() + + evaluator.evaluate( + evaluate_fn, + epoch, + step, + global_step, + ) + + stats_logger.commit(epoch, step, global_step, stats) + + stats_logger.close() + eval_rollout.destroy() + rollout.destroy() + if ref is not None: + ref.destroy() + actor.destroy() + wandb.finish() + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/examples/lite/clevr_count_70k_sft.py b/examples/lite/clevr_count_70k_sft.py new file mode 100644 index 0000000000..21c0255789 --- /dev/null +++ b/examples/lite/clevr_count_70k_sft.py @@ -0,0 +1,126 @@ +import os +import sys + +from torchdata.stateful_dataloader import StatefulDataLoader + +from areal.api.cli_args import SFTConfig, load_expr_config +from areal.api.io_struct import FinetuneSpec +from areal.dataset.__init__ import get_custom_dataset +from areal.engine.sft.lm_engine import FSDPLMEngine +from areal.utils.data import pad_sequences_to_tensors +from areal.utils.evaluator import Evaluator +from areal.utils.saver import Saver +from areal.utils.stats_logger import StatsLogger +from realhf.api.core.data_api import load_hf_processor_and_tokenizer +from realhf.base import seeding, stats_tracker + + +def main_sft(): + config, _ = load_expr_config(sys.argv[1:], SFTConfig) + config: SFTConfig + + rank = int(os.getenv("RANK")) + world_size = int(os.getenv("WORLD_SIZE")) + + seeding.set_random_seed(config.seed, f"trainer{rank}") + + processor, tokenizer = load_hf_processor_and_tokenizer(config.tokenizer_path) + train_dataset = get_custom_dataset( + path=config.train_dataset.path, + rank=rank, + world_size=world_size, + split="train", + type=config.train_dataset.type, + tokenizer=tokenizer, + processor=processor, + ) + valid_dataset = get_custom_dataset( + path=config.valid_dataset.path, + rank=rank, + world_size=world_size, + split="test", + type=config.valid_dataset.type, + tokenizer=tokenizer, + processor=processor, + ) + + # Create dataset and dataloaders + train_dataloader = StatefulDataLoader( + train_dataset, + batch_size=config.train_dataset.batch_size // world_size, + shuffle=config.train_dataset.shuffle, + num_workers=config.train_dataset.num_workers, + collate_fn=pad_sequences_to_tensors, + drop_last=config.train_dataset.drop_last, + ) + + valid_dataloader = StatefulDataLoader( + valid_dataset, + batch_size=config.valid_dataset.batch_size // world_size, + shuffle=config.valid_dataset.shuffle, + num_workers=config.valid_dataset.num_workers, + collate_fn=pad_sequences_to_tensors, + drop_last=config.valid_dataset.drop_last, + ) + + # Initialize engine + ft_spec = FinetuneSpec( + total_train_epochs=config.total_train_epochs, + dataset_size=len(train_dataloader) * config.train_dataset.batch_size, + train_batch_size=config.train_dataset.batch_size, + ) + engine = FSDPLMEngine(config=config.model) + engine.initialize(None, ft_spec) + + # Run training. + saver = Saver(config.saver, ft_spec, for_recover=False) + stats_logger = StatsLogger(config.stats_logger, ft_spec) + evaluator = Evaluator(config.evaluator, ft_spec) + + total_epochs = config.total_train_epochs + len(train_dataloader) + + global_step = 0 + for epoch in range(total_epochs): + for step, data in enumerate(train_dataloader): + + with ( + stats_tracker.record_timing("train_step"), + stats_tracker.scope("sft"), + ): + stats = engine.train_lm(data) + engine.step_lr_scheduler() + stats_tracker.scalar(**stats) + + with stats_tracker.record_timing("save"): + saver.save(engine, epoch, step, global_step) + + with stats_tracker.record_timing("eval"): + # No need to log anything. Logging will be handled outside + # via stats_tracker.export(). + def evaluate_fn(): + with stats_tracker.scope("sft-eval"): + for data in valid_dataloader: + engine.evaluate_lm(data) + + evaluator.evaluate( + evaluate_fn, + epoch, + step, + global_step, + ) + + stats_logger.commit( + epoch, + step, + global_step, + stats_tracker.export(reduce_group=engine.parallelism_group), + ) + global_step += 1 + + stats_logger.close() + engine.destroy() + + +if __name__ == "__main__": + main_sft() diff --git a/examples/lite/config_converter.py b/examples/lite/config_converter.py new file mode 100644 index 0000000000..13749886c0 --- /dev/null +++ b/examples/lite/config_converter.py @@ -0,0 +1,906 @@ +import abc +import argparse +import os +import shlex +from dataclasses import dataclass, field +from typing import Any, Callable, Optional, Type + +import yaml +from deepmerge import always_merger +from termcolor import colored + +CVRT_WARNING = "This parameter has no exact mapping in areal" +FEAT_WARNING = ( + "This parameter is not supported in areal, but should be supported in the future" +) + + +@dataclass +class ArgSpec: + + arg: str = "" + arg_type: Type[Any] = field(default=str) + description: str = "" + map_fn: Optional[Callable] = field(default=None) + + +class Converter(abc.ABC): + """ + Base class for converters. + """ + + def parse(self) -> dict: + """ + Parse the arguments from other framework format to a areal config. + """ + + def convert(self) -> dict: + """ + Convert the arguments to a new format. + """ + + def get_lite_template(self, template_path: str): + """ + Load the areal template from the specified file. + """ + with open(template_path, "r", encoding="utf-8") as f: + return yaml.safe_load(f) + + def flatten_dict(self, d, parent_key="", sep="."): + def _flatten_dict(d, parent_key="", sep="."): + """ + flatten a nested dictionary into a single-level dictionary + """ + items = [] + for k, v in d.items(): + new_key = parent_key + sep + k if parent_key else k + if isinstance(v, dict): + items.extend(_flatten_dict(v, new_key, sep=sep).items()) + else: + items.append((new_key, v)) + return dict(items) + + return _flatten_dict(d, parent_key, sep) + + def convert_to_nested_args(self, args, ARG_MAP: dict) -> dict: + """ + Convert flat arguments to nested dict based on ARG_MAP. + alert warning: This function will print warnings for unmapped arguments. + """ + cfg = {} + unmapped = {} + for k, v in args.items(): + if k in ARG_MAP: + argspec = ARG_MAP[k] + if not argspec.arg: + print( + colored(f"## Warning: For ", "yellow") + + colored(f"{k:>40}", "yellow", attrs=["bold"]) + + colored(f", # {argspec.description}!", "yellow") + ) + continue + + keys = argspec.arg.split(".") + arg_type = argspec.arg_type + map_fn = argspec.map_fn + + if v is not None: + try: + # type conversion + if arg_type == bool: + v = ( + bool(v) + if isinstance(v, bool) + else v.lower() in ("1", "true", "yes", "on") + ) + elif arg_type == int: + v = int(v) + elif arg_type == float: + v = float(v) + elif arg_type == str: + v = str(v) + else: + raise ValueError(f"Unsupported type: {arg_type}") + except Exception as e: + print( + colored(f"## Error: For ", "red") + + colored(f"{k:>40} {v}", "red", attrs=["bold"]) + + colored(f", # {e}!", "red") + ) + exit(-1) + + # apply map function if exists + if map_fn: + v = map_fn(v) + + self.set_nested(cfg, keys, v) + else: + unmapped[k] = v + print( + colored(f"## Warning: For ", "yellow") + + colored(f"{k:>50}", "yellow", attrs=["bold"]) + + colored(f", # {CVRT_WARNING}!", "yellow") + ) + # print("Unmapped arguments:", list(unmapped.keys())) + + return cfg + + def set_nested(self, d: dict, keys, value): + """ + Set a value in a nested dictionary using a list of keys. + """ + for k in keys[:-1]: + d = d.setdefault(k, {}) + d[keys[-1]] = value + + +class OpenRLHFConverter(Converter): + + ARG_MAP = { + # Ray and vLLM + "ref_num_nodes": ArgSpec("", int, CVRT_WARNING), + "ref_num_gpus_per_node": ArgSpec("", int, CVRT_WARNING), + "reward_num_nodes": ArgSpec("", int, CVRT_WARNING), + "reward_num_gpus_per_node": ArgSpec("", int, CVRT_WARNING), + "colocate_actor_ref": ArgSpec("", bool, CVRT_WARNING), + "actor_num_nodes": ArgSpec("cluster.n_nodes", int, ""), + "actor_num_gpus_per_node": ArgSpec("cluster.n_gpus_per_node", int, ""), + "critic_num_nodes": ArgSpec("", int, CVRT_WARNING), + "critic_num_gpus_per_node": ArgSpec("", int, CVRT_WARNING), + "colocate_critic_reward": ArgSpec("", bool, CVRT_WARNING), + "colocate_all_models": ArgSpec("", bool, CVRT_WARNING), + "vllm_num_engines": ArgSpec("allocation_mode.sglang.d", int, ""), + "vllm_tensor_parallel_size": ArgSpec("allocation_mode.sglang.t", int, ""), + "vllm_sync_backend": ArgSpec("", str, FEAT_WARNING), + "vllm_sync_with_ray": ArgSpec("", bool, CVRT_WARNING), + "enable_prefix_caching": ArgSpec("", bool, CVRT_WARNING), + "enforce_eager": ArgSpec("", bool, CVRT_WARNING), + "vllm_enable_sleep": ArgSpec("", bool, CVRT_WARNING), + "vllm_gpu_memory_utilization": ArgSpec( + "sglang.mem_fraction_static", float, "Not equivalent to SGLang" + ), + # Async training + "async_train": ArgSpec( + "async_training", bool + ), # TODO: convert areal offpolicyness > 0 + # Checkpoints + "eval_steps": ArgSpec("", int, CVRT_WARNING), + "save_steps": ArgSpec("", int, FEAT_WARNING), + "logging_steps": ArgSpec("", int, CVRT_WARNING), + "ckpt_path": ArgSpec( + "cluster.fileroot", str, "", lambda x: os.path.dirname(x) + ), # TODO: convert to areal's cluster.fileroot + "save_hf_ckpt": ArgSpec("", bool, CVRT_WARNING), + "disable_ds_ckpt": ArgSpec("", bool, CVRT_WARNING), + "max_ckpt_num": ArgSpec("", int, CVRT_WARNING), + "max_ckpt_mem": ArgSpec("", int, CVRT_WARNING), + "load_checkpoint": ArgSpec("", bool, CVRT_WARNING), + "use_ds_universal_ckpt": ArgSpec("", bool, CVRT_WARNING), + # DeepSpeed + "local_rank": ArgSpec("", int, CVRT_WARNING), + "zero_stage": ArgSpec("", int, CVRT_WARNING), + "gradient_checkpointing": ArgSpec("actor.gradient_checkpointing", bool, ""), + "deepcompile": ArgSpec("", bool, CVRT_WARNING), + "bf16": ArgSpec( + "actor.dtype", str, "", lambda x: "bfloat16" if x else "float16" + ), + "enable_ema": ArgSpec("", bool, CVRT_WARNING), + "ema_beta": ArgSpec("", float, CVRT_WARNING), + "zpg": ArgSpec("", int, CVRT_WARNING), + "adam_offload": ArgSpec("actor.optimizer.offload", bool, ""), + "actor_init_on_gpu": ArgSpec("", bool, CVRT_WARNING), + "flash_attn": ArgSpec( + "actor.attn_impl", + str, + "default use flash_attention_2", + lambda: "flash_attention_2", + ), + "use_liger_kernel": ArgSpec("", bool, CVRT_WARNING), + "grad_accum_dtype": ArgSpec("", str, CVRT_WARNING), + "overlap_comm": ArgSpec("", bool, CVRT_WARNING), + "gradient_checkpointing_use_reentrant": ArgSpec( + "", bool, CVRT_WARNING + ), # default not use reentrant + "disable_fast_tokenizer": ArgSpec("", bool, CVRT_WARNING), + "deepspeed_enable_sleep": ArgSpec("", bool, CVRT_WARNING), + "ds_tensor_parallel_size": ArgSpec("allocation_mode.engine.t", int, ""), + # packing samples + "packing_samples": ArgSpec("", bool, "default packing"), + # LoRA + "load_in_4bit": ArgSpec("", bool, CVRT_WARNING), + "lora_rank": ArgSpec("", int, CVRT_WARNING), + "lora_alpha": ArgSpec("", int, CVRT_WARNING), + "target_modules": ArgSpec("", list, CVRT_WARNING), + "lora_dropout": ArgSpec("", float, CVRT_WARNING), + # PPO + "save_path": ArgSpec("saver.fileroot", str, "", lambda x: os.path.dirname(x)), + "num_episodes": ArgSpec("", int, ""), + "rollout_batch_size": ArgSpec("rollout.consumer_batch_size", int, ""), + "vllm_generate_batch_size": ArgSpec("", int, CVRT_WARNING), + "micro_rollout_batch_size": ArgSpec("", int, CVRT_WARNING), + "max_epochs": ArgSpec("total_train_epochs", int, ""), + "prompt_max_len": ArgSpec("", int, FEAT_WARNING), # train_dataset.max_seqlen + "generate_max_len": ArgSpec("gconfig.max_new_tokens", int, ""), # TODO: check + "max_len": ArgSpec("", int, CVRT_WARNING), + "max_samples": ArgSpec("", int, CVRT_WARNING), + "max_norm": ArgSpec("actor.optimizer.gradient_clipping", float, ""), + "l2": ArgSpec("actor.optimizer.weight_decay", float, ""), + "ptx_coef": ArgSpec("", float, CVRT_WARNING), + "eps_clip": ArgSpec("actor.eps_clip", float, ""), + "eps_clip_low_high": ArgSpec("", list, CVRT_WARNING), # check + "value_clip": ArgSpec("", float, CVRT_WARNING), + "lambd": ArgSpec("actor.gae_lambda", float, ""), + "gamma": ArgSpec("actor.discount", float, ""), + "micro_train_batch_size": ArgSpec( + "actor.mb_spec.n_mbs", + int, + "default to 1, try actor.mb_spec.max_tokens_per_mb", + lambda x: 1, + ), + "train_batch_size": ArgSpec("train_dataset.batch_size", int, ""), + "normalize_reward": ArgSpec("", bool, "Default normalize reward"), + "top_p": ArgSpec("gconfig.top_p", float, ""), + "temperature": ArgSpec("gconfig.temperature", float, ""), + "seed": ArgSpec("seed", int, ""), + "freezing_actor_steps": ArgSpec("", int, CVRT_WARNING), + "n_samples_per_prompt": ArgSpec("gconfig.n_samples", int, ""), + "save_value_network": ArgSpec("", bool, CVRT_WARNING), + "actor_learning_rate": ArgSpec("actor.optimizer.lr", float, ""), + "critic_learning_rate": ArgSpec("", float, "no critic"), + "lr_warmup_ratio": ArgSpec( + "actor.optimizer.warmup_steps_proportion", float, "" + ), + "lr_scheduler": ArgSpec( + "actor.optimizer.lr_scheduler_type", str, "" + ), # should be mapped to areal's scheduler + "kl_target": ArgSpec("", float, CVRT_WARNING), + "kl_horizon": ArgSpec("", int, CVRT_WARNING), + "init_kl_coef": ArgSpec("actor.kl_ctl", float, ""), + "kl_estimator": ArgSpec("", str, CVRT_WARNING), + "aux_loss_coef": ArgSpec("", float, CVRT_WARNING), # moe + "entropy_loss_coef": ArgSpec("", float, CVRT_WARNING), # entropy loss + "adam_betas": ArgSpec("", list, CVRT_WARNING), # should decouple + "reward_clip_range": ArgSpec("actor.reward_clip", list, ""), + "advantage_estimator": ArgSpec("", str, "only support grpo"), + "use_kl_loss": ArgSpec("", bool, "actor.kl_ctl > 0 means use kl loss"), + "no_advantage_std_norm": ArgSpec("", bool, CVRT_WARNING), + # Context Parallel + "ring_attn_size": ArgSpec("", int, CVRT_WARNING), + "ring_head_stride": ArgSpec("", int, CVRT_WARNING), + # Models + "pretrain": ArgSpec("actor.path", str, ""), + "reward_pretrain": ArgSpec("", str, CVRT_WARNING), # no reward model + "remote_rm_url": ArgSpec("", str, CVRT_WARNING), + "critic_pretrain": ArgSpec("critic.path", str, CVRT_WARNING), # no critic model + "value_head_prefix": ArgSpec("", str, CVRT_WARNING), + "ref_reward_offload": ArgSpec("", bool, CVRT_WARNING), + "agent_func_path": ArgSpec("", str, CVRT_WARNING), + # Custom dataset + "prompt_data": ArgSpec("", str, CVRT_WARNING), + "prompt_data_probs": ArgSpec("", str, CVRT_WARNING), + "prompt_split": ArgSpec("", str, CVRT_WARNING), + "eval_dataset": ArgSpec("", str, CVRT_WARNING), + "eval_split": ArgSpec("", str, CVRT_WARNING), + "eval_temperature": ArgSpec("", float, CVRT_WARNING), + "eval_n_samples_per_prompt": ArgSpec("", int, CVRT_WARNING), + "input_key": ArgSpec("", str, CVRT_WARNING), + "label_key": ArgSpec("", str, CVRT_WARNING), + "input_template": ArgSpec("", str, CVRT_WARNING), + "apply_chat_template": ArgSpec("", bool, CVRT_WARNING), + # wandb + "use_wandb": ArgSpec( + "stats_logger.wandb.mode", str, "", lambda x: "online" + ), # set WANDB_API_KEY, WANDB_BASE_URL + "wandb_org": ArgSpec("stats_logger.wandb.entity", str, ""), + "wandb_group": ArgSpec("stats_logger.wandb.group", str, ""), + "wandb_project": ArgSpec("stats_logger.wandb.project", str, ""), + "wandb_run_name": ArgSpec("stats_logger.wandb.name", str, ""), + # Dynamic filtering + "dynamic_filtering": ArgSpec("", bool, CVRT_WARNING), + "dynamic_filtering_reward_range": ArgSpec("", list, CVRT_WARNING), + # TensorBoard + "use_tensorboard": ArgSpec( + "stats_logger.tensorboard.path", str, "" + ), # set default path + # performance ArgsSpectuning + "perf": ArgSpec("", bool, CVRT_WARNING), + # ModelScope + "use_ms": ArgSpec("", bool, CVRT_WARNING), + } + + def __init__( + self, src_script_path: str, template_path: str, command_start: str = "python" + ): + """ + Args: + src_script_path (str): The path to the source OpenRLHF script file. + template_path (str): The path to the areal template file. + command_start (str): The beginning of the command to look for, + e.g., "python", "python3 my_script.py". + """ + self.src_script_path = src_script_path + self.template_path = template_path + self.command_start = command_start + + def parse(self) -> dict: + """ + Parse the arguments from OpenRLHF script to a dict. + """ + args = self._parse_args_from_script(self.src_script_path, self.command_start) + return args + + def convert(self) -> dict: + """ + Convert the parsed arguments to + """ + args = self.parse() + args = self.convert_to_nested_args(args, self.ARG_MAP) + args = post_process_args(args) + template = self.get_lite_template(self.template_path) + lite_args = always_merger.merge(template, args) + return lite_args + + def _parse_args_from_script( + self, script_path: str, command_start: str = "python" + ) -> dict: + """ + Parses arguments for a command in a shell script. + + It finds a command block starting with `command_start`, handles line + continuations, and then extracts all key-value arguments. + + Args: + script_path (str): The path to the .sh script file. + command_start (str): The beginning of the command to look for, + e.g., "python", "python3 my_script.py". + + Returns: + dict: A dictionary containing all the parsed arguments and their values. + """ + full_command = "" + in_command_block = False + + try: + with open(script_path, "r", encoding="utf-8") as f: + for line in f: + stripped_line = line.strip() + + # Find the start of the command block + if not in_command_block and stripped_line.startswith(command_start): + in_command_block = True + + if in_command_block: + if stripped_line.startswith("#") or not stripped_line: + continue + + if stripped_line.endswith("\\"): + full_command += stripped_line[:-1].strip() + " " + else: + full_command += stripped_line + in_command_block = False + break + except FileNotFoundError: + print(f"Error: File not found at '{script_path}'") + return {} + + if not full_command: + print(f"Error: Could not find a command starting with '{command_start}'.") + return {} + + # Tokenize the entire command string using shlex + tokens = shlex.split(full_command) + + # Find the index of the first argument (starts with '-' or '--') + args_start_index = -1 + for i, token in enumerate(tokens): + if token.startswith("-"): + args_start_index = i + break + + if args_start_index == -1: + # No arguments found + return {} + + # The arguments are the rest of the list + args_list = tokens[args_start_index:] + + # --- The rest of the parsing logic is the same --- + params = {} + i = 0 + while i < len(args_list): + arg = args_list[i] + if arg.startswith("--"): + key = arg.lstrip("-") + if (i + 1 < len(args_list)) and not args_list[i + 1].startswith("-"): + params[key] = args_list[i + 1] + i += 2 + else: + params[key] = True + i += 1 + else: + # Also handle single-dash arguments if needed, or just skip + i += 1 + + return params + + +def post_process_args(args: dict): + + if "allocation_mode" in args: + # convert allocation_mode to sglang.dX.tY.pZ + dp = args["cluster"]["n_nodes"] * args["cluster"]["n_gpus_per_node"] + allocation_mode = "" + if "sglang" not in args["allocation_mode"]: + allocation_mode = f"sglang.d{dp}t1p1" + else: + if "d" not in args["allocation_mode"]["sglang"]: + allocation_mode = f"sglang.d{dp}" + else: + allocation_mode += f"sglang.d{args['allocation_mode']['sglang']['d']}" + + if "t" not in args["allocation_mode"]["sglang"]: + allocation_mode += "t1" + else: + allocation_mode += f"t{args['allocation_mode']['sglang']['t']}" + allocation_mode += f"p1" + allocation_mode += "+" + if "engine" not in args["allocation_mode"]: + allocation_mode += f"d{dp}t1p1" + else: + allocation_mode += f"d{args['allocation_mode']['engine']['d']}" + if "t" not in args["allocation_mode"]["engine"]: + allocation_mode += "t1" + else: + allocation_mode += f".t{args['allocation_mode']['engine']['t']}" + allocation_mode += f"p1" + + args["allocation_mode"] = allocation_mode + args["cluster"]["n_nodes"] = args["cluster"]["n_nodes"] * 2 + return args + + +class AReaLConverter(Converter): + """ + Convert realhf arguments to AReaL-lite arguments. + """ + + ARG_MAP = { + # Top-level experiment info + "experiment_name": ArgSpec("experiment_name", str, ""), + "trial_name": ArgSpec("trial_name", str, ""), + "mode": ArgSpec("", str, CVRT_WARNING), + "debug": ArgSpec("", bool, CVRT_WARNING), + "metric_discovery_port": ArgSpec("", int, CVRT_WARNING), + "partition": ArgSpec("", str, CVRT_WARNING), + "schedule_strategy": ArgSpec("", str, CVRT_WARNING), + "recover_mode": ArgSpec("", bool, CVRT_WARNING), + "recover_retries": ArgSpec("", int, CVRT_WARNING), + "recover_after": ArgSpec("", int, CVRT_WARNING), + "ignore_worker_error": ArgSpec("", bool, CVRT_WARNING), + "allocation_mode": ArgSpec( + "allocation_mode", str, "", lambda x: x.replace("m", "t") + ), + "n_nodes": ArgSpec("cluster.n_nodes", int, ""), # TODO: check cluster.n_nodes + "n_gpus_per_node": ArgSpec( + "cluster.n_gpus_per_node", int, "" + ), # TODO: check cluster + "node_list": ArgSpec("", str, CVRT_WARNING), + "exclude": ArgSpec("", str, CVRT_WARNING), + "seed": ArgSpec("seed", int, ""), + # Cluster configuration + "cluster.config_path": ArgSpec("", str, CVRT_WARNING), + "cluster.name_resolve.type": ArgSpec("cluster.name_resolve.type", str, ""), + "cluster.name_resolve.nfs_record_root": ArgSpec( + "cluster.name_resolve.nfs_record_root", str, "" + ), + "cluster.name_resolve.etcd3_addr": ArgSpec( + "cluster.name_resolve.etcd3_addr", str, "" + ), + "cluster.name_resolve.ray_actor_name": ArgSpec( + "cluster.name_resolve.ray_actor_name", str, "" + ), + "cluster.cluster_name": ArgSpec("cluster.cluster_name", str, ""), + "cluster.fileroot": ArgSpec("cluster.fileroot", str, ""), + "cluster.gpu_type": ArgSpec("", str, CVRT_WARNING), + "cluster.mount": ArgSpec("launcher.slurm.mount", str, ""), + "cluster.gpu_image": ArgSpec("", str, "Do not convert image"), + "cluster.cpu_image": ArgSpec("", str, "Do not convert image"), + "cluster.gpu_infer_image": ArgSpec("", str, "Do not convert image"), + "cluster.node_name_prefix": ArgSpec("", str, CVRT_WARNING), + "cluster.n_nodes": ArgSpec("", int, CVRT_WARNING), + "cluster.n_gpus_per_node": ArgSpec("cluster.n_gpus_per_node", int, ""), + # CPU resource allocation + "cpus_per_master_worker": ArgSpec("", int, CVRT_WARNING), + "mem_per_master_worker": ArgSpec("", int, CVRT_WARNING), + "cpus_per_model_worker": ArgSpec("launcher.trainer_cpus_per_gpu", int, ""), + "mem_per_model_worker": ArgSpec("launcher.trainer_mem_per_gpu", int, ""), + # Actor model + "actor.type._class": ArgSpec("", str, CVRT_WARNING), + "actor.type.is_critic": ArgSpec("", bool, CVRT_WARNING), + "actor.path": ArgSpec("actor.path", str, ""), + "actor.init_from_scratch": ArgSpec("actor.init_from_scratch", bool, ""), + "actor.init_critic_from_actor": ArgSpec( + "actor.init_critic_from_actor", bool, "" + ), + "actor.backend": ArgSpec( + "actor.backend", str, "Only support fsdp now", map_fn=lambda x: "fsdp" + ), + "actor.gradient_checkpointing": ArgSpec( + "actor.gradient_checkpointing", bool, "" + ), + "actor.bf16": ArgSpec( + "actor.dtype", bool, "", lambda x: "bfloat16" if x else "float16" + ), + # "actor.dtype": ArgSpec("actor.dtype", str, ""), + "actor.optimizer.type": ArgSpec("actor.optimizer.type", str, ""), + "actor.optimizer.lr": ArgSpec("actor.optimizer.lr", float, ""), + "actor.optimizer.weight_decay": ArgSpec( + "actor.optimizer.weight_decay", float, "" + ), + "actor.optimizer.beta1": ArgSpec("actor.optimizer.beta1", float, ""), + "actor.optimizer.beta2": ArgSpec("actor.optimizer.beta2", float, ""), + "actor.optimizer.eps": ArgSpec("actor.optimizer.eps", float, ""), + "actor.optimizer.min_lr_ratio": ArgSpec( + "actor.optimizer.min_lr_ratio", float, "" + ), + "actor.optimizer.lr_scheduler_type": ArgSpec( + "actor.optimizer.lr_scheduler_type", str, "" + ), + "actor.optimizer.warmup_steps_proportion": ArgSpec( + "actor.optimizer.warmup_steps_proportion", float, "" + ), + "actor.optimizer.offload": ArgSpec("actor.optimizer.offload", bool, ""), + "actor.optimizer.initial_loss_scale": ArgSpec( + "actor.optimizer.initial_loss_scale", float, "" + ), + "actor.optimizer.min_loss_scale": ArgSpec( + "actor.optimizer.min_loss_scale", float, "" + ), + "actor.optimizer.loss_scale_window": ArgSpec( + "actor.optimizer.loss_scale_window", int, "" + ), + "actor.optimizer.hysteresis": ArgSpec("actor.optimizer.hysteresis", int, ""), + "actor.optimizer.gradient_clipping": ArgSpec( + "actor.optimizer.gradient_clipping", float, "" + ), + # Megatron not implemented in areal + "actor.megatron.ddp.grad_reduce_in_fp32": ArgSpec("", bool, CVRT_WARNING), + "actor.megatron.ddp.overlap_grad_reduce": ArgSpec("", bool, CVRT_WARNING), + "actor.megatron.ddp.overlap_param_gather": ArgSpec("", bool, CVRT_WARNING), + "actor.megatron.ddp.align_param_gather": ArgSpec("", bool, CVRT_WARNING), + "actor.megatron.ddp.use_distributed_optimizer": ArgSpec("", bool, CVRT_WARNING), + "actor.megatron.ddp.check_for_nan_in_grad": ArgSpec("", bool, CVRT_WARNING), + "actor.megatron.ddp.bucket_size": ArgSpec("", int, CVRT_WARNING), + "actor.megatron.ddp.average_in_collective": ArgSpec("", bool, CVRT_WARNING), + "actor.megatron.ddp.fp8_param_gather": ArgSpec("", bool, CVRT_WARNING), + "actor.megatron.overlap_param_gather_with_optimizer_step": ArgSpec( + "", bool, CVRT_WARNING + ), + "actor.megatron.use_precision_aware_optimizer": ArgSpec("", bool, CVRT_WARNING), + "actor.megatron.main_grads_dtype": ArgSpec("", str, CVRT_WARNING), + "actor.megatron.main_params_dtype": ArgSpec("", str, CVRT_WARNING), + "actor.megatron.exp_avg_dtype": ArgSpec("", str, CVRT_WARNING), + "actor.megatron.exp_avg_sq_dtype": ArgSpec("", str, CVRT_WARNING), + # SGLang + "actor.sglang.disable_cuda_graph": ArgSpec( + "sglang.disable_cuda_graph", bool, "" + ), + "actor.sglang.disable_radix_cache": ArgSpec( + "sglang.disable_radix_cache", bool, "" + ), + "actor.sglang.disable_cuda_graph_padding": ArgSpec( + "sglang.disable_cuda_graph_padding", bool, "" + ), + "actor.sglang.enable_nccl_nvls": ArgSpec("sglang.enable_nccl_nvls", bool, ""), + "actor.sglang.disable_outlines_disk_cache": ArgSpec( + "sglang.disable_outlines_disk_cache", bool, "" + ), + "actor.sglang.disable_custom_all_reduce": ArgSpec( + "sglang.disable_custom_all_reduce", bool, "" + ), + "actor.sglang.disable_overlap_schedule": ArgSpec( + "sglang.disable_overlap_schedule", bool, "" + ), + "actor.sglang.enable_mixed_chunk": ArgSpec( + "sglang.enable_mixed_chunk", bool, "" + ), + "actor.sglang.enable_dp_attention": ArgSpec( + "sglang.enable_dp_attention", bool, "" + ), + "actor.sglang.enable_ep_moe": ArgSpec("sglang.enable_ep_moe", bool, ""), + "actor.sglang.enable_torch_compile": ArgSpec( + "sglang.enable_torch_compile", bool, "" + ), + "actor.sglang.torch_compile_max_bs": ArgSpec( + "sglang.torch_compile_max_bs", int, "" + ), + "actor.sglang.cuda_graph_max_bs": ArgSpec("sglang.cuda_graph_max_bs", int, ""), + "actor.sglang.cuda_graph_bs": ArgSpec("sglang.cuda_graph_bs", list, ""), + "actor.sglang.torchao_config": ArgSpec("sglang.torchao_config", str, ""), + "actor.sglang.enable_nan_detection": ArgSpec( + "sglang.enable_nan_detection", bool, "" + ), + "actor.sglang.enable_p2p_check": ArgSpec("sglang.enable_p2p_check", bool, ""), + "actor.sglang.triton_attention_reduce_in_fp32": ArgSpec( + "sglang.triton_attention_reduce_in_fp32", bool, "" + ), + "actor.sglang.triton_attention_num_kv_splits": ArgSpec( + "sglang.triton_attention_num_kv_splits", int, "" + ), + "actor.sglang.num_continuous_decode_steps": ArgSpec( + "sglang.num_continuous_decode_steps", int, "" + ), + "actor.sglang.enable_memory_saver": ArgSpec( + "sglang.enable_memory_saver", bool, "" + ), + "actor.sglang.allow_auto_truncate": ArgSpec( + "sglang.allow_auto_truncate", bool, "" + ), + "actor.sglang.attention_backend": ArgSpec("sglang.attention_backend", str, ""), + "actor.sglang.sampling_backend": ArgSpec("sglang.sampling_backend", str, ""), + "actor.sglang.context_length": ArgSpec("sglang.context_length", int, ""), + "actor.sglang.mem_fraction_static": ArgSpec( + "sglang.mem_fraction_static", float, "" + ), + "actor.sglang.max_running_requests": ArgSpec( + "sglang.max_running_requests", int, "" + ), + "actor.sglang.chunked_prefill_size": ArgSpec( + "sglang.chunked_prefill_size", int, "" + ), + "actor.sglang.max_prefill_tokens": ArgSpec( + "sglang.max_prefill_tokens", int, "" + ), + "actor.sglang.schedule_policy": ArgSpec("sglang.schedule_policy", str, ""), + "actor.sglang.schedule_conservativeness": ArgSpec( + "sglang.schedule_conservativeness", float, "" + ), + "actor.sglang.cpu_offload_gb": ArgSpec("sglang.cpu_offload_gb", int, ""), + "actor.sglang.hybrid_train": ArgSpec("", bool, CVRT_WARNING), + "actor.sglang.dtype": ArgSpec("sglang.dtype", str, ""), + "actor.sglang.kv_cache_dtype": ArgSpec("sglang.kv_cache_dtype", str, ""), + "actor.sglang.log_level": ArgSpec("sglang.log_level", str, ""), + "actor.sglang.log_level_http": ArgSpec("sglang.log_level_http", str, ""), + "actor.sglang.log_requests": ArgSpec("sglang.log_requests", bool, ""), + "actor.sglang.log_requests_level": ArgSpec( + "sglang.log_requests_level", int, "" + ), + "actor.sglang.show_time_cost": ArgSpec("sglang.show_time_cost", bool, ""), + "actor.sglang.enable_metrics": ArgSpec("sglang.enable_metrics", bool, ""), + "actor.sglang.decode_log_interval": ArgSpec( + "sglang.decode_log_interval", int, "" + ), + "group_size": ArgSpec("gconfig.n_samples", int, ""), + "generation_size": ArgSpec("", int, CVRT_WARNING), + "group_adv_norm": ArgSpec("actor.group_adv_norm", bool, ""), + "mask_no_eos_with_zero": ArgSpec("actor.mask_no_eos_with_zero", bool, ""), + "mask_too_long": ArgSpec("", bool, CVRT_WARNING), + "check_verifier_status": ArgSpec("", bool, CVRT_WARNING), + "ref_ema_eta": ArgSpec("", float, CVRT_WARNING), + "rw_type": ArgSpec("", str, CVRT_WARNING), + "check_xml_format": ArgSpec("", bool, CVRT_WARNING), + "dataset_filter_threshold": ArgSpec("", float, CVRT_WARNING), + "dataset_max_filter_percentage": ArgSpec("", float, CVRT_WARNING), + "success_rate_ub": ArgSpec("", float, CVRT_WARNING), + "success_rate_lb": ArgSpec("", float, CVRT_WARNING), + "no_training": ArgSpec("", bool, CVRT_WARNING), + # Ref model + "ref.type._class": ArgSpec("", str, CVRT_WARNING), + "ref.type.is_critic": ArgSpec("", bool, CVRT_WARNING), + "ref.path": ArgSpec("ref.path", str, ""), + "ref.init_from_scratch": ArgSpec("ref.init_from_scratch", bool, ""), + "ref.backend": ArgSpec("ref.backend", str, ""), + "ref.bf16": ArgSpec( + "ref.dtype", bool, "", lambda x: "bfloat16" if x else "float16" + ), + # TODO: remaining ref args + # MFC + "actor_train.mb_spec.n_mbs": ArgSpec("actor.mb_spec.n_mbs", int, ""), + "actor_train.mb_spec.max_tokens_per_mb": ArgSpec( + "actor.mb_spec.max_tokens_per_mb", int, "" + ), + "actor_train.parallel": ArgSpec("", int, CVRT_WARNING), + "actor_train.device_mesh": ArgSpec("", str, CVRT_WARNING), + "actor_gen.mb_spec.n_mbs": ArgSpec("", int, CVRT_WARNING), + "actor_gen.mb_spec.max_tokens_per_mb": ArgSpec("", int, CVRT_WARNING), + "actor_gen.parallel": ArgSpec("", int, CVRT_WARNING), + "actor_gen.device_mesh": ArgSpec("", str, CVRT_WARNING), + "actor_inf.mb_spec.n_mbs": ArgSpec("", int, CVRT_WARNING), + "actor_inf.mb_spec.max_tokens_per_mb": ArgSpec("", int, CVRT_WARNING), + "actor_inf.parallel": ArgSpec("", int, CVRT_WARNING), + "actor_inf.device_mesh": ArgSpec("", str, CVRT_WARNING), + "ref_inf.mb_spec.n_mbs": ArgSpec("ref.mb_spec.n_mbs", int, ""), + "ref_inf.mb_spec.max_tokens_per_mb": ArgSpec( + "ref.mb_spec.max_tokens_per_mb", int, "" + ), + "ref_inf.parallel": ArgSpec("", int, CVRT_WARNING), + "ref_inf.device_mesh": ArgSpec("", str, CVRT_WARNING), + # Dataset + "dataset.path": ArgSpec("", str, CVRT_WARNING), + "dataset.max_prompt_len": ArgSpec("", int, "Should be implemented"), + "dataset.train_bs_n_seqs": ArgSpec("train_dataset.batch_size", int, ""), + "dataset.fill_to_max_length": ArgSpec("", bool, CVRT_WARNING), + "shuffle_dataset": ArgSpec("train_dataset.shuffle", bool, ""), + # PPO/gen + "ppo.gen.n": ArgSpec("", int, "Use group_size to control"), + "ppo.gen.max_new_tokens": ArgSpec("gconfig.max_new_tokens", int, ""), + "ppo.gen.min_new_tokens": ArgSpec("gconfig.min_new_tokens", int, ""), + "ppo.gen.greedy": ArgSpec("gconfig.greedy", bool, ""), + "ppo.gen.top_p": ArgSpec("gconfig.top_p", float, ""), + "ppo.gen.top_k": ArgSpec("gconfig.top_k", int, ""), + "ppo.gen.temperature": ArgSpec("gconfig.temperature", float, ""), + # PPO core + "ppo.ppo_n_minibatches": ArgSpec("actor.ppo_n_minibatches", int, ""), + "ppo.eps_clip": ArgSpec("actor.eps_clip", float, ""), + "ppo.c_clip": ArgSpec("actor.c_clip", float, ""), + "ppo.value_eps_clip": ArgSpec("", float, CVRT_WARNING), + "ppo.early_stop_imp_ratio": ArgSpec("", float, CVRT_WARNING), + "ppo.actor_sample_reuse": ArgSpec("", bool, CVRT_WARNING), + "ppo.critic_sample_reuse": ArgSpec("", bool, CVRT_WARNING), + "ppo.max_reward_clip": ArgSpec("actor.reward_clip", float, ""), + "ppo.reward_output_scaling": ArgSpec("actor.reward_scaling", float, ""), + "ppo.reward_output_bias": ArgSpec("actor.reward_bias", float, ""), + "ppo.fuse_rew_ref": ArgSpec("", bool, CVRT_WARNING), + "ppo.discount": ArgSpec("actor.discount", float, ""), + "ppo.gae_lambda": ArgSpec("actor.gae_lambda", float, ""), + "ppo.adv_norm": ArgSpec("actor.adv_norm", bool, ""), + "ppo.kl_ctl": ArgSpec("actor.kl_ctl", float, ""), + "ppo.use_adaptive_kl_ctl": ArgSpec("", bool, CVRT_WARNING), + "ppo.disable_value": ArgSpec("", bool, CVRT_WARNING), + "ppo.value_norm": ArgSpec("", bool, CVRT_WARNING), + "ppo.value_norm_type": ArgSpec("", str, CVRT_WARNING), + "ppo.value_norm_beta": ArgSpec("", float, CVRT_WARNING), + "ppo.value_norm_eps": ArgSpec("", float, CVRT_WARNING), + "ppo.recompute_logprob": ArgSpec("actor.recompute_logprob", bool, ""), + "ppo.use_decoupled_loss": ArgSpec("actor.use_decoupled_loss", bool, ""), + "ppo.behav_imp_weight_cap": ArgSpec("actor.behav_imp_weight_cap", float, ""), + # Async + "schedule_policy": ArgSpec("rollout.schedule_policy", str, ""), + "new_tokens_per_chunk": ArgSpec("", int, CVRT_WARNING), + "max_head_offpolicyness": ArgSpec("rollout.max_head_offpolicyness", int, ""), + "n_rollout_workers": ArgSpec("", int, CVRT_WARNING), + "max_concurrent_rollouts": ArgSpec("rollout.max_concurrent_rollouts", int, ""), + "flush_request_timeout": ArgSpec("rollout.request_timeout", float, ""), + "cpus_per_generation_server": ArgSpec( + "launcher.inference_server_cpus_per_gpu", int, "" + ), + "mem_per_generation_server": ArgSpec( + "launcher.inference_server_mem_per_gpu", int, "" + ), + "cpus_per_gserver_manager": ArgSpec("", int, CVRT_WARNING), + "mem_per_gserver_manager": ArgSpec("", int, CVRT_WARNING), + "cpus_per_rollout_worker": ArgSpec("", int, CVRT_WARNING), + "mem_per_rollout_worker": ArgSpec("", int, CVRT_WARNING), + # Saver/checkpointer/evaluator + "saver.freq_epochs": ArgSpec("saver.freq_epochs", int, ""), + "saver.freq_steps": ArgSpec("saver.freq_steps", int, ""), + "saver.freq_secs": ArgSpec("saver.freq_secs", int, ""), + "checkpointer.freq_epochs": ArgSpec("checkpointer.freq_epochs", int, ""), + "checkpointer.freq_steps": ArgSpec("checkpointer.freq_steps", int, ""), + "checkpointer.freq_secs": ArgSpec("checkpointer.freq_secs", int, ""), + "evaluator.freq_epochs": ArgSpec("evaluator.freq_epochs", int, ""), + "evaluator.freq_steps": ArgSpec("evaluator.freq_steps", int, ""), + "evaluator.freq_secs": ArgSpec("evaluator.freq_secs", int, ""), + # Logging + "wandb.mode": ArgSpec("stats_logger.wandb.mode", str, ""), + "wandb.entity": ArgSpec("stats_logger.wandb.entity", str, ""), + "wandb.project": ArgSpec("stats_logger.wandb.project", str, ""), + "wandb.name": ArgSpec("stats_logger.wandb.name", str, ""), + "wandb.job_type": ArgSpec("stats_logger.wandb.job_type", str, ""), + "wandb.group": ArgSpec("stats_logger.wandb.group", str, ""), + "wandb.notes": ArgSpec("stats_logger.wandb.notes", str, ""), + "wandb.tags": ArgSpec("stats_logger.wandb.tags", list, ""), + "wandb.config": ArgSpec("stats_logger.wandb.config", dict, ""), + "swanlab.project": ArgSpec("stats_logger.swanlab.project", str, ""), + "swanlab.name": ArgSpec("stats_logger.swanlab.name", str, ""), + "swanlab.config": ArgSpec("stats_logger.swanlab.config", dict, ""), + "swanlab.log_dir": ArgSpec("stats_logger.swanlab.log_dir", str, ""), + "swanlab.mode": ArgSpec("stats_logger.swanlab.mode", str, ""), + "swanlab.api_key": ArgSpec("stats_logger.swanlab.api_key", str, ""), + "tensorboard.path": ArgSpec("stats_logger.tensorboard.path", str, ""), + # Exp Control + "exp_ctrl.total_train_epochs": ArgSpec("total_train_epochs", int, ""), + "exp_ctrl.save_freq_epochs": ArgSpec("saver.freq_epochs", int, ""), + "exp_ctrl.save_freq_steps": ArgSpec("saver.freq_steps", int, ""), + "exp_ctrl.save_freq_secs": ArgSpec("saver.freq_secs", int, ""), + "exp_ctrl.ckpt_freq_epochs": ArgSpec("checkpointer.freq_epochs", int, ""), + "exp_ctrl.ckpt_freq_steps": ArgSpec("checkpointer.freq_steps", int, ""), + "exp_ctrl.ckpt_freq_secs": ArgSpec("checkpointer.freq_secs", int, ""), + "exp_ctrl.eval_freq_epochs": ArgSpec("evaluator.freq_epochs", int, ""), + "exp_ctrl.eval_freq_steps": ArgSpec("evaluator.freq_steps", int, ""), + "exp_ctrl.eval_freq_secs": ArgSpec("evaluator.freq_secs", int, ""), + "exp_ctrl.benchmark_steps": ArgSpec("", int, CVRT_WARNING), + "exp_ctrl.benchmark_n_seqs": ArgSpec("", int, CVRT_WARNING), + # Auto Evaluation + "auto_eval": ArgSpec("", bool, CVRT_WARNING), + "auto_eval_config.data_names": ArgSpec("", list, CVRT_WARNING), + "auto_eval_config.max_gen_tokens": ArgSpec("", int, CVRT_WARNING), + "auto_eval_config.max_concurrent_jobs": ArgSpec("", int, CVRT_WARNING), + "auto_eval_config.eval_job_image": ArgSpec("", str, CVRT_WARNING), + "auto_eval_config.initial_checkpoint_path": ArgSpec("", str, CVRT_WARNING), + "auto_eval_config.prompt_type": ArgSpec("", str, CVRT_WARNING), + # Miscellaneous + "cache_clear_freq": ArgSpec("", int, CVRT_WARNING), + "torch_cache_mysophobia": ArgSpec("", bool, CVRT_WARNING), + "ray_temp_path": ArgSpec("", str, CVRT_WARNING), + } + + def __init__(self, src_config_path: str, template_path: str): + self.src_config_path = src_config_path + self.template_path = template_path + + def parse(self) -> dict: + with open(self.src_config_path, "r", encoding="utf-8") as f: + cfg = yaml.safe_load(f) + return cfg + + def convert(self) -> dict: + args = self.parse() + args = self.flatten_dict(args) + args = self.convert_to_nested_args(args, self.ARG_MAP) # realite args format + template = self.get_lite_template(self.template_path) + lite_args = always_merger.merge(template, args) + return lite_args + + +def parse_args(): + parser = argparse.ArgumentParser(description="Convert to areal YAML config") + parser.add_argument( + "--convert_src", + type=str, + choices=["OpenRLHF", "AReaL"], + default="OpenRLHF", + help="source config type to convert from", + ) + parser.add_argument( + "--src_script_path", type=str, help="path to the source OpenRLHF script file" + ) + parser.add_argument( + "--src_config_path", type=str, help="path to the src AReaL config file" + ) + parser.add_argument( + "--template_path", + type=str, + required=True, + help="path to the areal template file", + ) + parser.add_argument( + "--output_path", + type=str, + default="output.yaml", + help="path to save the converted YAML config", + ) + return parser.parse_args() + + +def main(): + """ + Usage: + - OpenRLHF: python areal/utils/convert.py --convert_src OpenRLHF --src_script_path --template_path --output_path + - AReaL: python areal/utils/convert.py --convert_srv AReaL --src_config_path --template_path --output_path + """ + args = parse_args() + converters = {"OpenRLHF": OpenRLHFConverter, "AReaL": AReaLConverter} + converter_args = { + "OpenRLHF": { + "src_script_path": args.src_script_path, + "template_path": args.template_path, + }, + "AReaL": { + "src_config_path": args.src_config_path, + "template_path": args.template_path, + }, + } + converter: Converter = converters[args.convert_src]( + **converter_args[args.convert_src] + ) + lite_args = converter.convert() + yaml_str = yaml.dump(lite_args, sort_keys=False, allow_unicode=True) + with open(args.output_path, "w", encoding="utf-8") as f: + yaml.dump(lite_args, f, sort_keys=False, allow_unicode=True) + print(f"Converted areal config saved to {args.output_path}") + # print(yaml_str) + + +if __name__ == "__main__": + main() diff --git a/examples/lite/configs/clevr_count_70k_grpo.yaml b/examples/lite/configs/clevr_count_70k_grpo.yaml new file mode 100644 index 0000000000..91cbfc23e1 --- /dev/null +++ b/examples/lite/configs/clevr_count_70k_grpo.yaml @@ -0,0 +1,139 @@ +experiment_name: clevr_count_70k-grpo +trial_name: trial1 + + +seed: 1 +total_train_epochs: 3 +tokenizer_path: ${actor.path} +async_training: true + +cluster: + n_nodes: 1 + n_gpus_per_node: 8 + cluster_name: na132 + fileroot: /storage/openpsi/experiments + name_resolve: + type: nfs + nfs_record_root: /storage/openpsi/experiments/name_resolve/clevr_count_70k-grpo + etcd3_addr: etcd-client.openpsi-etcd.svc.sigma-na130-lingbo.na130.wl-robby.local:2379 + +allocation_mode: sglang.d1p1t1+d7p1t1 + +rollout: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 256 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 4 + enable_rollout_tracing: false + +gconfig: + n_samples: 4 + min_new_tokens: 0 + max_new_tokens: 512 + greedy: false + temperature: 1.0 + +actor: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: /storage/openpsi/models/Qwen2.5-VL-3B-Instruct + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: false + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 10240 + optimizer: + type: adam + lr: 2e-6 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + backend: fsdp + + group_size: ${gconfig.n_samples} + group_adv_norm: false + eps_clip: 0.4 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: true + use_decoupled_loss: true + behav_imp_weight_cap: 5.0 + +ref: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 10240 + optimizer: null + backend: fsdp + +# SGLang +server_only: false +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + max_running_requests: null + context_length: 32768 + mem_fraction_static: 0.8 + +# datasets +train_dataset: + batch_size: 32 + shuffle: true + pin_memory: true + num_workers: 4 + path: /storage/openpsi/data/clevr_count_70k/ + type: rl + +valid_dataset: + batch_size: 32 + shuffle: true + pin_memory: true + num_workers: 4 + path: /storage/openpsi/data/clevr_count_70k/ + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +checkpointer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} diff --git a/examples/lite/configs/clevr_count_70k_sft.yaml b/examples/lite/configs/clevr_count_70k_sft.yaml new file mode 100644 index 0000000000..9f864d9ce5 --- /dev/null +++ b/examples/lite/configs/clevr_count_70k_sft.yaml @@ -0,0 +1,81 @@ +experiment_name: clevr_count_70k-sft +trial_name: trial0 + +cluster: + n_nodes: 1 + n_gpus_per_node: 1 + cluster_name: na132 + fileroot: /storage/openpsi/experiments + name_resolve: + type: nfs + nfs_record_root: /storage/openpsi/experiments/name_resolve/clevr_count_70k-sft + etcd3_addr: etcd-client.openpsi-etcd.svc.sigma-na130-lingbo.na130.wl-robby.local:2379 +seed: 1 +total_train_epochs: 3 +tokenizer_path: ${model.path} + +model: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: /storage/openpsi/models/Qwen2-VL-7B + dtype: bfloat16 + init_from_scratch: false + gradient_checkpointing: false + mb_spec: + max_tokens_per_mb: 4096 + optimizer: + type: adam + lr: 2e-5 + weight_decay: 0.05 + beta1: 0.9 + beta2: 0.95 + eps: 1e-5 + lr_scheduler_type: cosine + gradient_clipping: 1.0 + backend: fsdp + +train_dataset: + batch_size: 128 + shuffle: true + pin_memory: true + num_workers: 4 + path: /storage/openpsi/data/clevr_count_70k/ + type: sft + +valid_dataset: + batch_size: 128 + shuffle: true + pin_memory: true + num_workers: 4 + path: /storage/openpsi/data/clevr_count_70k/ + type: sft + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +checkpointer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: null + freq_steps: 1 + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} diff --git a/examples/lite/configs/gsm8k_grpo.yaml b/examples/lite/configs/gsm8k_grpo.yaml new file mode 100644 index 0000000000..855da02306 --- /dev/null +++ b/examples/lite/configs/gsm8k_grpo.yaml @@ -0,0 +1,136 @@ +experiment_name: gsm8k-grpo +trial_name: trial0 +allocation_mode: sglang.d4p1t1+d4p1t1 +cluster: + fileroot: /tmp/areal/experiments + n_nodes: 1 + n_gpus_per_node: 8 + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve +seed: 1 +total_train_epochs: 10 +tokenizer_path: ${actor.path} +async_training: true + +rollout: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + max_concurrent_rollouts: 256 + queue_size: null + consumer_batch_size: ${train_dataset.batch_size} + max_head_offpolicyness: 4 + enable_rollout_tracing: false + +gconfig: + n_samples: 4 + min_new_tokens: 0 + max_new_tokens: 512 + greedy: false + temperature: 1.0 + +actor: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen2-1.5B-Instruct + init_from_scratch: false + disable_dropout: true + gradient_checkpointing: false + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 10240 + optimizer: + type: adam + lr: 1e-5 + weight_decay: 0.01 + beta1: 0.9 + beta2: 0.999 + eps: 1e-8 + lr_scheduler_type: constant + gradient_clipping: 1.0 + warmup_steps_proportion: 0.001 + backend: fsdp + + group_size: ${gconfig.n_samples} + group_adv_norm: false + eps_clip: 0.4 + temperature: ${gconfig.temperature} + reward_scaling: 10.0 + reward_bias: -0.5 + kl_ctl: 0.0 + ppo_n_minibatches: 1 + recompute_logprob: true + use_decoupled_loss: true + behav_imp_weight_cap: 5.0 + +ref: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: ${actor.path} + init_from_scratch: false + disable_dropout: true + dtype: ${actor.dtype} + mb_spec: + max_tokens_per_mb: 10240 + optimizer: null + backend: fsdp + +# SGLang +server_only: false +sglang: + model_path: ${actor.path} + random_seed: ${seed} + skip_tokenizer_init: true + dtype: ${actor.dtype} + max_running_requests: null + context_length: 32768 + mem_fraction_static: 0.8 + +# datasets +train_dataset: + batch_size: 256 + shuffle: true + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + +valid_dataset: + batch_size: 256 + shuffle: true + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: rl + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +checkpointer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: disabled diff --git a/examples/lite/configs/gsm8k_sft.yaml b/examples/lite/configs/gsm8k_sft.yaml new file mode 100644 index 0000000000..db4dabfc7a --- /dev/null +++ b/examples/lite/configs/gsm8k_sft.yaml @@ -0,0 +1,82 @@ +experiment_name: gsm8k-sft +trial_name: trial0 + +allocation_mode: d8p1t1 +cluster: + fileroot: /tmp/areal/experiments + n_nodes: 1 + n_gpus_per_node: 8 + name_resolve: + type: nfs + nfs_record_root: /tmp/areal/name_resolve +seed: 1 +total_train_epochs: 1 +tokenizer_path: ${model.path} + +model: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + path: Qwen/Qwen3-1.7B + init_from_scratch: false + gradient_checkpointing: false + dtype: bfloat16 + mb_spec: + max_tokens_per_mb: 4096 + optimizer: + type: adam + lr: 2e-5 + weight_decay: 0.05 + beta1: 0.9 + beta2: 0.95 + eps: 1e-5 + lr_scheduler_type: cosine + gradient_clipping: 1.0 + backend: fsdp + +train_dataset: + batch_size: 128 + shuffle: true + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: sft + +valid_dataset: + batch_size: 128 + shuffle: true + pin_memory: true + num_workers: 4 + path: openai/gsm8k + type: sft + +# Utilities +saver: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: null + +checkpointer: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: 1 + freq_steps: null + freq_secs: 3600 + +evaluator: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + freq_epochs: null + freq_steps: 1 + freq_secs: null + +stats_logger: + experiment_name: ${experiment_name} + trial_name: ${trial_name} + fileroot: ${cluster.fileroot} + wandb: + mode: online diff --git a/examples/lite/gsm8k_grpo.py b/examples/lite/gsm8k_grpo.py new file mode 100644 index 0000000000..86154615df --- /dev/null +++ b/examples/lite/gsm8k_grpo.py @@ -0,0 +1,226 @@ +import itertools +import os +import sys + +import torch +import torch.distributed as dist +from torchdata.stateful_dataloader import StatefulDataLoader + +from areal.api.cli_args import GRPOConfig, load_expr_config +from areal.api.io_struct import AllocationMode, FinetuneSpec, WeightUpdateMeta +from areal.dataset.__init__ import get_custom_dataset +from areal.engine.ppo.actor import FSDPPPOActor +from areal.engine.sglang_remote import RemoteSGLangEngine +from areal.utils.device import log_gpu_stats +from areal.utils.evaluator import Evaluator +from areal.utils.saver import Saver +from areal.utils.stats_logger import StatsLogger +from areal.workflow.rlvr import RLVRWorkflow +from realhf.api.core.data_api import load_hf_tokenizer +from realhf.base import seeding, stats_tracker + + +def gsm8k_reward_fn(prompt, completions, prompt_ids, completion_ids, answer, **kwargs): + from realhf.impl.dataset.math_parser import process_results + + return int(process_results(completions, answer)[0]) + + +def main(args): + config, _ = load_expr_config(args, GRPOConfig) + config: GRPOConfig + + rank = int(os.getenv("RANK")) + world_size = int(os.getenv("WORLD_SIZE")) + tokenizer = load_hf_tokenizer(config.tokenizer_path) + + seeding.set_random_seed(config.seed, key=f"trainer{rank}") + + train_dataset = get_custom_dataset( + path=config.train_dataset.path, + rank=rank, + world_size=world_size, + split="train", + type=config.train_dataset.type, + tokenizer=tokenizer, + ) + valid_dataset = get_custom_dataset( + path=config.valid_dataset.path, + rank=rank, + world_size=world_size, + split="test", + type=config.valid_dataset.type, + tokenizer=tokenizer, + ) + + # Create dataset and dataloaders + train_dataloader = StatefulDataLoader( + train_dataset, + batch_size=config.train_dataset.batch_size // world_size, + shuffle=config.train_dataset.shuffle, + num_workers=config.train_dataset.num_workers, + collate_fn=lambda x: x, + drop_last=config.train_dataset.drop_last, + ) + valid_dataloader = StatefulDataLoader( + valid_dataset, + batch_size=config.valid_dataset.batch_size // world_size, + shuffle=config.valid_dataset.shuffle, + num_workers=config.valid_dataset.num_workers, + collate_fn=lambda x: x, + drop_last=config.valid_dataset.drop_last, + ) + ft_spec = FinetuneSpec( + total_train_epochs=config.total_train_epochs, + dataset_size=len(train_dataloader) * config.train_dataset.batch_size, + train_batch_size=config.train_dataset.batch_size, + ) + + # Initialize inference engine + rollout = RemoteSGLangEngine(config.rollout) + rollout.initialize(None, ft_spec) + eval_rollout = RemoteSGLangEngine(config.rollout) + eval_rollout.initialize(None, ft_spec) + # NOTE: set a large version such that eval does not have any offpolicyness control + eval_rollout.set_version(int(1e12)) + + # Initialize train engine + actor = FSDPPPOActor(config=config.actor) + actor.initialize(None, ft_spec) + ref = None + if config.actor.kl_ctl > 0 and config.ref is not None: + ref = FSDPPPOActor(config=config.ref) + ref.initialize(None, ft_spec) + + # NOTE: Weight update meta only requires address and free port of rank 0, + # but `WeightUpdateMeta.from_fsdp_nccl` has to be executed on all ranks + # due to `engine.get_param_specs()`. + # Therefore, we create weight update meta on all ranks, then broadcast the one on rank 0. + weight_update_meta = [ + WeightUpdateMeta.from_fsdp_nccl( + AllocationMode.from_str(config.allocation_mode), actor + ) + ] + dist.broadcast_object_list(weight_update_meta, src=0) + weight_update_meta = weight_update_meta[0] + + # Create rollout workflow + if tokenizer.pad_token_id not in config.gconfig.stop_token_ids: + config.gconfig.stop_token_ids.append(tokenizer.pad_token_id) + if tokenizer.eos_token_id not in config.gconfig.stop_token_ids: + config.gconfig.stop_token_ids.append(tokenizer.eos_token_id) + workflow = RLVRWorkflow( + reward_fn=gsm8k_reward_fn, + gconfig=config.gconfig, + tokenizer=tokenizer, + enable_thinking=False, + dump_dir=os.path.join( + StatsLogger.get_log_path(config.stats_logger), "generated" + ), + ) + + # Run training. + saver = Saver(config.saver, ft_spec, for_recover=False) + stats_logger = StatsLogger(config.stats_logger, ft_spec) + evaluator = Evaluator(config.evaluator, ft_spec) + + total_epochs = config.total_train_epochs + steps_per_epoch = len(train_dataloader) + max_steps = total_epochs * steps_per_epoch + + data_generator = itertools.cycle(train_dataloader) + for global_step in range(max_steps): + epoch = global_step // steps_per_epoch + step = global_step % steps_per_epoch + + with stats_tracker.record_timing("rollout"): + if config.async_training: + batch = rollout.prepare_batch(train_dataloader, workflow=workflow) + else: + batch = rollout.rollout_batch(next(data_generator), workflow=workflow) + + batch = batch.to(actor.device) + # Create barrier to synchronize all rollout processes. + dist.barrier(device_ids=[actor.device.index]) + torch.cuda.synchronize() + + if config.actor.recompute_logprob or config.actor.use_decoupled_loss: + with stats_tracker.record_timing("recompute_logp"): + logp = actor.compute_logp(batch) + batch["prox_logp"] = logp + log_gpu_stats("recompute logp") + + if ref is not None: + with stats_tracker.record_timing("ref_logp"): + batch["ref_logp"] = ref.compute_logp(batch) + log_gpu_stats("ref logp") + + with stats_tracker.record_timing("compute_advantage"): + actor.compute_advantages(batch) + log_gpu_stats("compute advantages") + + with ( + stats_tracker.record_timing("train_step"), + stats_tracker.scope("grpo_actor"), + ): + stats = actor.ppo_update(batch) + actor.step_lr_scheduler() + log_gpu_stats("ppo update") + + with stats_tracker.record_timing("update_weights"): + rollout.pause() + if dist.get_rank() == 0: + future = rollout.update_weights(weight_update_meta) + actor.upload_weights(weight_update_meta) + if dist.get_rank() == 0: + future.result() + dist.barrier(device_ids=[actor.device.index]) + torch.cuda.synchronize() + rollout.resume() + actor.set_version(global_step + 1) + rollout.set_version(global_step + 1) + + with stats_tracker.record_timing("save"): + saver.save(actor, epoch, step, global_step) + + with stats_tracker.record_timing("eval"): + + def evaluate_fn(): + rollout.pause() + cnt = 0 + for data in valid_dataloader: + for item in data: + eval_rollout.submit(item, workflow) + cnt += 1 + batch = eval_rollout.wait(cnt, timeout=None) + rewards = batch["rewards"].float().to(actor.device) + with stats_tracker.scope("grpo-eval"): + stats_tracker.denominator( + n_seqs=torch.ones( + rewards.shape[0], + device=rewards.device, + dtype=torch.bool, + ) + ) + stats_tracker.stat(task_reward=rewards, denominator="n_seqs") + rollout.resume() + + evaluator.evaluate( + evaluate_fn, + epoch, + step, + global_step, + ) + + stats_logger.commit(epoch, step, global_step, stats) + + stats_logger.close() + eval_rollout.destroy() + rollout.destroy() + if ref is not None: + ref.destroy() + actor.destroy() + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/examples/lite/gsm8k_sft.py b/examples/lite/gsm8k_sft.py new file mode 100644 index 0000000000..b8a725968b --- /dev/null +++ b/examples/lite/gsm8k_sft.py @@ -0,0 +1,122 @@ +import os +import sys + +from torchdata.stateful_dataloader import StatefulDataLoader + +from areal.api.cli_args import SFTConfig, load_expr_config +from areal.api.io_struct import FinetuneSpec +from areal.dataset.__init__ import get_custom_dataset +from areal.engine.sft.lm_engine import FSDPLMEngine +from areal.utils.data import pad_sequences_to_tensors +from areal.utils.evaluator import Evaluator +from areal.utils.saver import Saver +from areal.utils.stats_logger import StatsLogger +from realhf.api.core.data_api import load_hf_tokenizer +from realhf.base import seeding, stats_tracker + + +def main(args): + config, _ = load_expr_config(args, SFTConfig) + config: SFTConfig + + rank = int(os.getenv("RANK")) + world_size = int(os.getenv("WORLD_SIZE")) + tokenizer = load_hf_tokenizer(config.tokenizer_path) + + seeding.set_random_seed(config.seed, f"trainer{rank}") + + train_dataset = get_custom_dataset( + path=config.train_dataset.path, + rank=rank, + world_size=world_size, + split="train", + type=config.train_dataset.type, + tokenizer=tokenizer, + ) + valid_dataset = get_custom_dataset( + path=config.valid_dataset.path, + rank=rank, + world_size=world_size, + split="test", + type=config.valid_dataset.type, + tokenizer=tokenizer, + ) + + # Create dataset and dataloaders + train_dataloader = StatefulDataLoader( + train_dataset, + batch_size=config.train_dataset.batch_size // world_size, + shuffle=config.train_dataset.shuffle, + num_workers=config.train_dataset.num_workers, + collate_fn=pad_sequences_to_tensors, + drop_last=config.train_dataset.drop_last, + ) + valid_dataloader = StatefulDataLoader( + valid_dataset, + batch_size=config.valid_dataset.batch_size // world_size, + shuffle=config.valid_dataset.shuffle, + num_workers=config.valid_dataset.num_workers, + collate_fn=pad_sequences_to_tensors, + drop_last=config.valid_dataset.drop_last, + ) + + # Initialize engine + ft_spec = FinetuneSpec( + total_train_epochs=config.total_train_epochs, + dataset_size=len(train_dataloader) * config.train_dataset.batch_size, + train_batch_size=config.train_dataset.batch_size, + ) + engine = FSDPLMEngine(config=config.model) + engine.initialize(None, ft_spec) + + # Run training. + saver = Saver(config.saver, ft_spec, for_recover=False) + stats_logger = StatsLogger(config.stats_logger, ft_spec) + evaluator = Evaluator(config.evaluator, ft_spec) + + total_epochs = config.total_train_epochs + len(train_dataloader) + + global_step = 0 + for epoch in range(total_epochs): + for step, data in enumerate(train_dataloader): + with ( + stats_tracker.record_timing("train_step"), + stats_tracker.scope("sft"), + ): + stats = engine.train_lm(data) + engine.step_lr_scheduler() + stats_tracker.scalar(**stats) + + with stats_tracker.record_timing("save"): + saver.save(engine, epoch, step, global_step) + + with stats_tracker.record_timing("eval"): + # No need to log anything. Logging will be handled outside + # via stats_tracker.export(). + def evaluate_fn(): + with stats_tracker.scope("sft-eval"): + for data in valid_dataloader: + engine.evaluate_lm(data) + + evaluator.evaluate( + evaluate_fn, + epoch, + step, + global_step, + ) + + stats_logger.commit( + epoch, + step, + global_step, + stats_tracker.export(reduce_group=engine.parallelism_group), + ) + global_step += 1 + + stats_logger.close() + engine.destroy() + + +if __name__ == "__main__": + main(sys.argv[1:]) diff --git a/functioncall/math/verify.py b/functioncall/math/verify.py index 2635e0423d..e0fd146eaf 100644 --- a/functioncall/math/verify.py +++ b/functioncall/math/verify.py @@ -10,7 +10,12 @@ def math_verify( - id2info, generateds: List, query_ids: List, batch_size=10, timeout=1000 + id2info, + generateds: List, + query_ids: List, + batch_size=10, + timeout=1000, + max_workers=None, ) -> List: assert len(generateds) == len(query_ids), ( len(generateds), diff --git a/patch/sglang/v0.4.6.post2.patch b/patch/sglang/v0.4.6.post2.patch deleted file mode 100644 index 6bf47bf7b4..0000000000 --- a/patch/sglang/v0.4.6.post2.patch +++ /dev/null @@ -1,144 +0,0 @@ -diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index 174656b2..33fe0a5f 100644 ---- a/python/sglang/srt/managers/io_struct.py -+++ b/python/sglang/srt/managers/io_struct.py -@@ -687,10 +687,21 @@ class FlushCacheReqOutput: - success: bool - - -+@dataclass -+class InterruptAllReqInput: -+ pass -+ -+ -+@dataclass -+class InterruptAllReqOutput: -+ num_interrupted_requests: int -+ -+ - @dataclass - class UpdateWeightFromDiskReqInput: - # The model path with the new weights - model_path: str -+ allow_interrupt: bool = False - # The format to load the weights - load_format: Optional[str] = None - -diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index 8891115c..843a8a82 100644 ---- a/python/sglang/srt/managers/scheduler.py -+++ b/python/sglang/srt/managers/scheduler.py -@@ -70,6 +70,8 @@ from sglang.srt.managers.io_struct import ( - HealthCheckOutput, - InitWeightsUpdateGroupReqInput, - InitWeightsUpdateGroupReqOutput, -+ InterruptAllReqInput, -+ InterruptAllReqOutput, - OpenSessionReqInput, - OpenSessionReqOutput, - ProfileReq, -@@ -419,6 +421,7 @@ class Scheduler( - # Init request dispatcher - self._request_dispatcher = TypeBasedDispatcher( - [ -+ (InterruptAllReqInput, self.interrupt_all_requests), - (TokenizedGenerateReqInput, self.handle_generate_request), - (TokenizedEmbeddingReqInput, self.handle_embedding_request), - (FlushCacheReqInput, self.flush_cache_wrapped), -@@ -1938,6 +1941,15 @@ class Scheduler( - def _pause_engine(self) -> Tuple[List[Req], int]: - raise NotImplementedError() - -+ def interrupt_all_requests(self, recv_req: InterruptAllReqInput): -+ num = len(self.waiting_queue) + len(self.running_batch.reqs) -+ for req in self.waiting_queue: -+ req.sampling_params.max_new_tokens = 0 -+ for req in self.running_batch.reqs: -+ req.sampling_params.max_new_tokens = len(req.output_ids) -+ logger.info(f"Interrupt {num} requests.") -+ return InterruptAllReqOutput(num) -+ - def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): - """In-place update of the weights from disk.""" - success, message = self.tp_worker.update_weights_from_disk(recv_req) -diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index 82709b09..bfab3ce7 100644 ---- a/python/sglang/srt/managers/tokenizer_manager.py -+++ b/python/sglang/srt/managers/tokenizer_manager.py -@@ -76,6 +76,8 @@ from sglang.srt.managers.io_struct import ( - HealthCheckOutput, - InitWeightsUpdateGroupReqInput, - InitWeightsUpdateGroupReqOutput, -+ InterruptAllReqInput, -+ InterruptAllReqOutput, - OpenSessionReqInput, - OpenSessionReqOutput, - ProfileReq, -@@ -265,6 +267,9 @@ class TokenizerManager: - self.resume_memory_occupation_communicator = _Communicator( - self.send_to_scheduler, server_args.dp_size - ) -+ self.interrupt_requests_communicator = _Communicator( -+ self.send_to_scheduler, server_args.dp_size -+ ) - self.flush_cache_communicator = _Communicator( - self.send_to_scheduler, server_args.dp_size - ) -@@ -294,6 +299,10 @@ class TokenizerManager: - UpdateWeightFromDiskReqOutput, - self._handle_update_weights_from_disk_req_output, - ), -+ ( -+ InterruptAllReqOutput, -+ self.interrupt_requests_communicator.handle_recv, -+ ), - ( - InitWeightsUpdateGroupReqOutput, - self.init_weights_update_group_communicator.handle_recv, -@@ -767,6 +776,13 @@ class TokenizerManager: - ) -> Tuple[bool, str]: - self.auto_create_handle_loop() - -+ if obj.allow_interrupt: -+ num_interrupted_requests = await self.interrupt_all_requests( -+ InterruptAllReqInput() -+ ) -+ # Set a break point to wait for the interrupt to finish -+ await asyncio.sleep(0.1) -+ - # default the load format to the server_args - if obj.load_format is None: - obj.load_format = self.server_args.load_format -@@ -776,7 +792,12 @@ class TokenizerManager: - # Hold the lock if it is not async. This means that weight sync - # cannot run while requests are in progress. - async with self.model_update_lock.writer_lock: -- return await self._wait_for_model_update_from_disk(obj) -+ success, message, n_paused = ( -+ await self._wait_for_model_update_from_disk(obj) -+ ) -+ if obj.allow_interrupt: -+ return success, message, num_interrupted_requests -+ return success, message, n_paused - - async def _wait_for_model_update_from_disk( - self, obj: UpdateWeightFromDiskReqInput -@@ -849,6 +870,18 @@ class TokenizerManager: - result = (await self.update_weights_from_tensor_communicator(obj))[0] - return result.success, result.message - -+ async def interrupt_all_requests( -+ self, -+ obj: InterruptAllReqInput, -+ request: Optional[fastapi.Request] = None, -+ ) -> Tuple[bool, str]: -+ self.auto_create_handle_loop() -+ result = await self.interrupt_requests_communicator(obj) -+ if self.server_args.dp_size == 1: -+ return result[0].num_interrupted_requests -+ else: -+ return [r.num_interrupted_requests for r in result] -+ - async def get_weights_by_name( - self, obj: GetWeightsByNameReqInput, request: Optional[fastapi.Request] = None - ): diff --git a/patch/sglang/v0.4.6.post4.patch b/patch/sglang/v0.4.6.post4.patch deleted file mode 100644 index b7dbd09cf7..0000000000 --- a/patch/sglang/v0.4.6.post4.patch +++ /dev/null @@ -1,144 +0,0 @@ -diff --git a/python/sglang/srt/managers/io_struct.py b/python/sglang/srt/managers/io_struct.py -index 5390668c..db370d19 100644 ---- a/python/sglang/srt/managers/io_struct.py -+++ b/python/sglang/srt/managers/io_struct.py -@@ -687,10 +687,21 @@ class FlushCacheReqOutput: - success: bool - - -+@dataclass -+class InterruptAllReqInput: -+ pass -+ -+ -+@dataclass -+class InterruptAllReqOutput: -+ num_interrupted_requests: int -+ -+ - @dataclass - class UpdateWeightFromDiskReqInput: - # The model path with the new weights - model_path: str -+ allow_interrupt: bool = False - # The format to load the weights - load_format: Optional[str] = None - -diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py -index 1178eec5..318dee33 100644 ---- a/python/sglang/srt/managers/scheduler.py -+++ b/python/sglang/srt/managers/scheduler.py -@@ -73,6 +73,8 @@ from sglang.srt.managers.io_struct import ( - HealthCheckOutput, - InitWeightsUpdateGroupReqInput, - InitWeightsUpdateGroupReqOutput, -+ InterruptAllReqInput, -+ InterruptAllReqOutput, - OpenSessionReqInput, - OpenSessionReqOutput, - ProfileReq, -@@ -427,6 +429,7 @@ class Scheduler( - # Init request dispatcher - self._request_dispatcher = TypeBasedDispatcher( - [ -+ (InterruptAllReqInput, self.interrupt_all_requests), - (TokenizedGenerateReqInput, self.handle_generate_request), - (TokenizedEmbeddingReqInput, self.handle_embedding_request), - (FlushCacheReqInput, self.flush_cache_wrapped), -@@ -1971,6 +1974,15 @@ class Scheduler( - def _pause_engine(self) -> Tuple[List[Req], int]: - raise NotImplementedError() - -+ def interrupt_all_requests(self, recv_req: InterruptAllReqInput): -+ num = len(self.waiting_queue) + len(self.running_batch.reqs) -+ for req in self.waiting_queue: -+ req.sampling_params.max_new_tokens = 0 -+ for req in self.running_batch.reqs: -+ req.sampling_params.max_new_tokens = len(req.output_ids) -+ logger.info(f"Interrupt {num} requests.") -+ return InterruptAllReqOutput(num) -+ - def update_weights_from_disk(self, recv_req: UpdateWeightFromDiskReqInput): - """In-place update of the weights from disk.""" - success, message = self.tp_worker.update_weights_from_disk(recv_req) -diff --git a/python/sglang/srt/managers/tokenizer_manager.py b/python/sglang/srt/managers/tokenizer_manager.py -index b646fae1..c668728b 100644 ---- a/python/sglang/srt/managers/tokenizer_manager.py -+++ b/python/sglang/srt/managers/tokenizer_manager.py -@@ -80,6 +80,8 @@ from sglang.srt.managers.io_struct import ( - HealthCheckOutput, - InitWeightsUpdateGroupReqInput, - InitWeightsUpdateGroupReqOutput, -+ InterruptAllReqInput, -+ InterruptAllReqOutput, - OpenSessionReqInput, - OpenSessionReqOutput, - ProfileReq, -@@ -279,6 +281,9 @@ class TokenizerManager: - self.slow_down_communicator = _Communicator( - self.send_to_scheduler, server_args.dp_size - ) -+ self.interrupt_requests_communicator = _Communicator( -+ self.send_to_scheduler, server_args.dp_size -+ ) - self.flush_cache_communicator = _Communicator( - self.send_to_scheduler, server_args.dp_size - ) -@@ -309,6 +314,10 @@ class TokenizerManager: - UpdateWeightFromDiskReqOutput, - self._handle_update_weights_from_disk_req_output, - ), -+ ( -+ InterruptAllReqOutput, -+ self.interrupt_requests_communicator.handle_recv, -+ ), - ( - InitWeightsUpdateGroupReqOutput, - self.init_weights_update_group_communicator.handle_recv, -@@ -799,6 +808,13 @@ class TokenizerManager: - ) -> Tuple[bool, str]: - self.auto_create_handle_loop() - -+ if obj.allow_interrupt: -+ num_interrupted_requests = await self.interrupt_all_requests( -+ InterruptAllReqInput() -+ ) -+ # Set a break point to wait for the interrupt to finish -+ await asyncio.sleep(0.1) -+ - # default the load format to the server_args - if obj.load_format is None: - obj.load_format = self.server_args.load_format -@@ -808,7 +824,12 @@ class TokenizerManager: - # Hold the lock if it is not async. This means that weight sync - # cannot run while requests are in progress. - async with self.model_update_lock.writer_lock: -- return await self._wait_for_model_update_from_disk(obj) -+ success, message, n_paused = ( -+ await self._wait_for_model_update_from_disk(obj) -+ ) -+ if obj.allow_interrupt: -+ return success, message, num_interrupted_requests -+ return success, message, n_paused - - async def _wait_for_model_update_from_disk( - self, obj: UpdateWeightFromDiskReqInput -@@ -881,6 +902,18 @@ class TokenizerManager: - result = (await self.update_weights_from_tensor_communicator(obj))[0] - return result.success, result.message - -+ async def interrupt_all_requests( -+ self, -+ obj: InterruptAllReqInput, -+ request: Optional[fastapi.Request] = None, -+ ) -> Tuple[bool, str]: -+ self.auto_create_handle_loop() -+ result = await self.interrupt_requests_communicator(obj) -+ if self.server_args.dp_size == 1: -+ return result[0].num_interrupted_requests -+ else: -+ return [r.num_interrupted_requests for r in result] -+ - async def get_weights_by_name( - self, obj: GetWeightsByNameReqInput, request: Optional[fastapi.Request] = None - ): diff --git a/pyproject.toml b/pyproject.toml index 7597e137a0..4e9e558ef7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,10 +31,9 @@ dependencies = [ "huggingface_hub", "datasets", "accelerate", - "transformers==4.51.1", + "transformers==4.53.1", # Scientific computing - "numpy<2.0.0", "scipy", "pandas", "matplotlib", @@ -53,6 +52,12 @@ dependencies = [ "hydra-core==1.4.0.dev1", "packaging", "tabulate", + "gymnasium>=1.1.1", + "torchdata", + "autoflake", + "tensordict", + "pybase64", + "msgspec", # Monitoring and logging "wandb", @@ -82,6 +87,7 @@ dependencies = [ # Distributed computing "ray", "redis", + "deepspeed>=0.17.2", # Web frameworks "fastapi>=0.115.12", diff --git a/realhf/api/core/data_api.py b/realhf/api/core/data_api.py index ce6d9bf955..5222353a5c 100644 --- a/realhf/api/core/data_api.py +++ b/realhf/api/core/data_api.py @@ -8,6 +8,7 @@ import random import time from contextlib import contextmanager +from functools import lru_cache # NOTE: We don't sue wildcard importing here because the type # `Sequence` has a very similar name to `SequenceSample`. @@ -47,6 +48,7 @@ RL_TASKS = ["math", "code", "rlhf", "stem"] +@lru_cache(maxsize=8) def load_hf_tokenizer( model_name_or_path: str, fast_tokenizer=True, @@ -67,6 +69,28 @@ def load_hf_tokenizer( return tokenizer +@lru_cache(maxsize=8) +def load_hf_processor_and_tokenizer( + model_name_or_path: str, + fast_tokenizer=True, + padding_side: Optional[str] = None, +) -> Tuple["transformers.ProcessorMixin", transformers.PreTrainedTokenizerFast]: + """Load a tokenizer and processor from Hugging Face.""" + # NOTE: use the raw type annoation will trigger cuda initialization + tokenizer = load_hf_tokenizer(model_name_or_path, fast_tokenizer, padding_side) + try: + processor = transformers.AutoProcessor.from_pretrained( + model_name_or_path, trust_remote_code=True, force_download=True + ) + except Exception: + processor = None + logger.warning( + f"Failed to load processor for {model_name_or_path}. " + "Using tokenizer only. This may cause issues with some models." + ) + return processor, tokenizer + + @pdclasses.dataclass class SequenceSplitSpec: partitions: Optional[List[Tuple[int, int]]] = None diff --git a/realhf/base/name_resolve.py b/realhf/base/name_resolve.py index 9c6ec25ba5..ef9879e7aa 100644 --- a/realhf/base/name_resolve.py +++ b/realhf/base/name_resolve.py @@ -1360,7 +1360,9 @@ def _keepalive_thread_run(self): def make_repository(args: "NameResolveConfig"): if args.type == "nfs": - return NfsNameRecordRepository(args.nfs_record_root) + repo = NfsNameRecordRepository(args.nfs_record_root) + os.makedirs(repo.record_root, exist_ok=True) + return repo elif args.type == "etcd3": host, port = args.etcd3_addr.split(":") return Etcd3NameRecordRepository(host=host, port=int(port)) diff --git a/realhf/base/names.py b/realhf/base/names.py index f66db5da0a..7584bafbf5 100644 --- a/realhf/base/names.py +++ b/realhf/base/names.py @@ -107,3 +107,7 @@ def training_samples(experiment_name, trial_name): def experiment_status(experiment_name, trial_name): return f"{USER_NAMESPACE}/{experiment_name}/{trial_name}/experiment_status" + + +def update_weights_from_disk(experiment_name, trial_name, model_version): + return f"{USER_NAMESPACE}/{experiment_name}/{trial_name}/update_weights_from_disk/{model_version}" diff --git a/realhf/base/stats_tracker.py b/realhf/base/stats_tracker.py index 0ecc7af8ba..7fea2935d5 100644 --- a/realhf/base/stats_tracker.py +++ b/realhf/base/stats_tracker.py @@ -1,4 +1,6 @@ +import time from collections import defaultdict +from contextlib import contextmanager from enum import Enum, auto from typing import Dict @@ -49,6 +51,17 @@ def _get_full_key(self, key): return key return "/".join(self.scope_stack + [key]) + @contextmanager + def record_timing(self, key): + start_time = time.perf_counter() + try: + yield + finally: + # NOTE: timing records are fixed under the "timeperf" scope + full_key = f"timeperf/{key}" + self._set_reduce_type(full_key, ReduceType.SCALAR) + self.stats[full_key].append(time.perf_counter() - start_time) + def denominator(self, **kwargs): for key, value in kwargs.items(): if not isinstance(value, torch.Tensor) or value.dtype != torch.bool: @@ -252,3 +265,4 @@ def _max_of(self, key, reduce_group): export = DEFAULT_TRACKER.export scope = DEFAULT_TRACKER.scope scalar = DEFAULT_TRACKER.scalar +record_timing = DEFAULT_TRACKER.record_timing diff --git a/realhf/impl/model/backend/sglang.py b/realhf/impl/model/backend/sglang.py index c0ef278eec..6fcb1783e8 100644 --- a/realhf/impl/model/backend/sglang.py +++ b/realhf/impl/model/backend/sglang.py @@ -117,20 +117,28 @@ async def _do_generate( ) most_recent_timestamps[output_idx] = timestamp - output.output_ids = [data[SGLANG_TOKEN_OUTPUT_IDENTIFIER]] - finish_reason = data["meta_info"]["finish_reason"] - if req.return_logprob: - output.output_logprobs = [ - [ - x[0] - for x in data["meta_info"]["output_token_logprobs"] - ] - ] + meta_info = data["meta_info"] + finish_reason = meta_info["finish_reason"] assert finish_reason["type"] in [ "length", "stop", + "abort", ], finish_reason - output.no_eos = [finish_reason["type"] == "length"] + + if meta_info.get("output_token_logprobs"): + output.output_ids = [ + [x[1] for x in meta_info["output_token_logprobs"]] + ] + if req.return_logprob: + output.output_logprobs = [ + [x[0] for x in meta_info["output_token_logprobs"]] + ] + else: + output.output_ids = [[]] + if req.return_logprob: + output.output_logprobs = [[]] + + output.no_eos = [finish_reason["type"] in ["length", "abort"]] output.latency = latency output_idx += 1 diff --git a/realhf/scheduler/client.py b/realhf/scheduler/client.py index 71cfdf8f21..3662b7c244 100644 --- a/realhf/scheduler/client.py +++ b/realhf/scheduler/client.py @@ -41,12 +41,12 @@ def __init__(self, run_name, worker_type, host, reason: JobState): class JobInfo: name: str state: JobState - host: str = ( + host: Optional[str] = ( None # The host on which the job is/was running. None if the job had not run. ) - submit_time: str = None - start_time: str = None - slurm_id: str = None # Slurm only. The Slurm id of the job. + submit_time: Optional[str] = None + start_time: Optional[str] = None + slurm_id: Optional[int] = None # Slurm only. The Slurm id of the job. class SchedulerClient: diff --git a/realhf/system/generation_server.py b/realhf/system/generation_server.py index 227554dad6..5d09707b36 100644 --- a/realhf/system/generation_server.py +++ b/realhf/system/generation_server.py @@ -40,42 +40,11 @@ def execute_shell_command(command: str) -> subprocess.Popen: ) -def apply_sglang_path(): - p = Path(os.path.dirname(__file__)) - patch_path = str( - p.parent.parent - / "patch" - / "sglang" - / f"v{pkg_version.get_version('sglang')}.patch" - ) - - target_path = "" - sglang_meta = subprocess.check_output( - "python3 -m pip show sglang", shell=True - ).decode("ascii") - for line in sglang_meta.split("\n"): - line = line.strip() - if line.startswith("Editable project location: "): - target_path = str(Path(line.split(": ")[1]).parent) - - if target_path: - proc = subprocess.Popen( - ["git", "apply", patch_path], - cwd=target_path, - stderr=sys.stdout, - stdout=sys.stdout, - ) - proc.wait() - logger.info(f"Applied SGLang patch at {target_path}") - - def launch_server_cmd(command: str, port: int = 30000): """ Launch the server using the given command. If no port is specified, a free port is reserved. """ - if not ray.is_initialized(): - apply_sglang_path() assert port is not None full_command = f"{command} --port {port}" process = execute_shell_command(full_command) diff --git a/realhf/system/gserver_manager.py b/realhf/system/gserver_manager.py index e746fe976b..a789ae7fe0 100644 --- a/realhf/system/gserver_manager.py +++ b/realhf/system/gserver_manager.py @@ -155,39 +155,22 @@ def check_new_params(self) -> str | None: return None - async def flush_requests_and_update_weights( - self, server_url, new_param_path, update_weights_retries=5 - ): - server_index = self.server_urls.index(server_url) - success = False - for _ in range(update_weights_retries): - async with aiohttp.ClientSession( - server_url, - timeout=aiohttp.ClientTimeout( - total=self.config.flush_request_timeout, - sock_connect=self.config.flush_request_timeout, - ), - ) as session: - async with session.post( - f"/update_weights_from_disk", - json=dict(model_path=new_param_path, allow_interrupt=True), - ) as resp: - if resp.status == 200: - res = await resp.json() - success = res["success"] - if success: - if "num_paused_requests" in res: - logger.info( - f"{res['num_paused_requests']} requests are interrupted " - f"during updating weights for server {server_index}: {server_url}" - ) - return - logger.warning( - f"Update weights failed: {res['message']}. Retrying." - ) - logger.warning(f"Update weights failed: {resp.reason}. Retrying.") - time.sleep(0.1) - raise RuntimeError("Update weights failed.") + async def flush_requests_and_update_weights(self, server_url, new_param_path): + async with aiohttp.ClientSession( + server_url, + timeout=aiohttp.ClientTimeout( + total=self.config.flush_request_timeout, + sock_connect=self.config.flush_request_timeout, + ), + ) as session: + (await session.post("/pause_generation")).raise_for_status() + async with session.post( + f"/update_weights_from_disk", + json=dict(model_path=new_param_path), + ) as resp: + resp.raise_for_status() + assert (await resp.json())["success"] + (await session.post("/continue_generation")).raise_for_status() def _round_robin_schedule(self, req_meta: GenReqMeta) -> int: if not hasattr(self, "round_robin_idx"): diff --git a/realhf/utils.py b/realhf/utils.py index 47009aee19..7a2aa41c4d 100644 --- a/realhf/utils.py +++ b/realhf/utils.py @@ -35,6 +35,7 @@ def load_hf_or_local_file(path: str) -> str: => /root/.cache/huggingface/hub/models--inclusionAI--AReaL-RL-Data/data/boba_106k_0319.jsonl """ + path = str(path) if path.startswith("hf://") or path.startswith("hf-dataset://"): repo_type = "dataset" if path.startswith("hf-dataset://") else "model" hf_path = path.strip().split("://")[1] diff --git a/requirements.txt b/requirements.txt index 4ffab7e39b..7cde02b42e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -25,7 +25,6 @@ numba packaging pandas pybind11>=2.10.0 -numpy<2.0.0 psutil pynvml pytest @@ -69,4 +68,12 @@ word2number Pebble timeout-decorator prettytable -swanlab[dashboard] \ No newline at end of file +gymnasium>=1.1.1 +swanlab[dashboard] +torchdata +autoflake +tensordict +deepspeed>=0.17.2 +pybase64 +msgspec +transformers==4.53.1 \ No newline at end of file diff --git a/setup.py b/setup.py index 2de55bcfd0..6b5745a27c 100644 --- a/setup.py +++ b/setup.py @@ -5,7 +5,7 @@ import setuptools setuptools.setup( - name="realhf", + name="areal", packages=setuptools.find_packages(), include_package_data=True, )