mirror of https://github.com/hpcaitech/ColossalAI
ver217
2 years ago
committed by
GitHub
64 changed files with 3756 additions and 0 deletions
@ -0,0 +1,146 @@
|
||||
# Byte-compiled / optimized / DLL files |
||||
__pycache__/ |
||||
*.py[cod] |
||||
*$py.class |
||||
|
||||
# C extensions |
||||
*.so |
||||
|
||||
# Distribution / packaging |
||||
.Python |
||||
build/ |
||||
develop-eggs/ |
||||
dist/ |
||||
downloads/ |
||||
eggs/ |
||||
.eggs/ |
||||
lib/ |
||||
lib64/ |
||||
parts/ |
||||
sdist/ |
||||
var/ |
||||
wheels/ |
||||
pip-wheel-metadata/ |
||||
share/python-wheels/ |
||||
*.egg-info/ |
||||
.installed.cfg |
||||
*.egg |
||||
MANIFEST |
||||
|
||||
# PyInstaller |
||||
# Usually these files are written by a python script from a template |
||||
# before PyInstaller builds the exe, so as to inject date/other infos into it. |
||||
*.manifest |
||||
*.spec |
||||
|
||||
# Installer logs |
||||
pip-log.txt |
||||
pip-delete-this-directory.txt |
||||
|
||||
# Unit test / coverage reports |
||||
htmlcov/ |
||||
.tox/ |
||||
.nox/ |
||||
.coverage |
||||
.coverage.* |
||||
.cache |
||||
nosetests.xml |
||||
coverage.xml |
||||
*.cover |
||||
*.py,cover |
||||
.hypothesis/ |
||||
.pytest_cache/ |
||||
|
||||
# Translations |
||||
*.mo |
||||
*.pot |
||||
|
||||
# Django stuff: |
||||
*.log |
||||
local_settings.py |
||||
db.sqlite3 |
||||
db.sqlite3-journal |
||||
|
||||
# Flask stuff: |
||||
instance/ |
||||
.webassets-cache |
||||
|
||||
# Scrapy stuff: |
||||
.scrapy |
||||
|
||||
# Sphinx documentation |
||||
docs/_build/ |
||||
docs/.build/ |
||||
|
||||
# PyBuilder |
||||
target/ |
||||
|
||||
# Jupyter Notebook |
||||
.ipynb_checkpoints |
||||
|
||||
# IPython |
||||
profile_default/ |
||||
ipython_config.py |
||||
|
||||
# pyenv |
||||
.python-version |
||||
|
||||
# pipenv |
||||
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control. |
||||
# However, in case of collaboration, if having platform-specific dependencies or dependencies |
||||
# having no cross-platform support, pipenv may install dependencies that don't work, or not |
||||
# install all needed dependencies. |
||||
#Pipfile.lock |
||||
|
||||
# PEP 582; used by e.g. github.com/David-OConnor/pyflow |
||||
__pypackages__/ |
||||
|
||||
# Celery stuff |
||||
celerybeat-schedule |
||||
celerybeat.pid |
||||
|
||||
# SageMath parsed files |
||||
*.sage.py |
||||
|
||||
# Environments |
||||
.env |
||||
.venv |
||||
env/ |
||||
venv/ |
||||
ENV/ |
||||
env.bak/ |
||||
venv.bak/ |
||||
|
||||
# Spyder project settings |
||||
.spyderproject |
||||
.spyproject |
||||
|
||||
# Rope project settings |
||||
.ropeproject |
||||
|
||||
# mkdocs documentation |
||||
/site |
||||
|
||||
# mypy |
||||
.mypy_cache/ |
||||
.dmypy.json |
||||
dmypy.json |
||||
|
||||
# Pyre type checker |
||||
.pyre/ |
||||
|
||||
# IDE |
||||
.idea/ |
||||
.vscode/ |
||||
|
||||
# macos |
||||
*.DS_Store |
||||
#data/ |
||||
|
||||
docs/.build |
||||
|
||||
# pytorch checkpoint |
||||
*.pt |
||||
|
||||
# ignore version.py generated by setup.py |
||||
colossalai/version.py |
||||
@ -0,0 +1,202 @@
|
||||
Copyright 2021- HPC-AI Technology Inc. All rights reserved. |
||||
Apache License |
||||
Version 2.0, January 2004 |
||||
http://www.apache.org/licenses/ |
||||
|
||||
TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION |
||||
|
||||
1. Definitions. |
||||
|
||||
"License" shall mean the terms and conditions for use, reproduction, |
||||
and distribution as defined by Sections 1 through 9 of this document. |
||||
|
||||
"Licensor" shall mean the copyright owner or entity authorized by |
||||
the copyright owner that is granting the License. |
||||
|
||||
"Legal Entity" shall mean the union of the acting entity and all |
||||
other entities that control, are controlled by, or are under common |
||||
control with that entity. For the purposes of this definition, |
||||
"control" means (i) the power, direct or indirect, to cause the |
||||
direction or management of such entity, whether by contract or |
||||
otherwise, or (ii) ownership of fifty percent (50%) or more of the |
||||
outstanding shares, or (iii) beneficial ownership of such entity. |
||||
|
||||
"You" (or "Your") shall mean an individual or Legal Entity |
||||
exercising permissions granted by this License. |
||||
|
||||
"Source" form shall mean the preferred form for making modifications, |
||||
including but not limited to software source code, documentation |
||||
source, and configuration files. |
||||
|
||||
"Object" form shall mean any form resulting from mechanical |
||||
transformation or translation of a Source form, including but |
||||
not limited to compiled object code, generated documentation, |
||||
and conversions to other media types. |
||||
|
||||
"Work" shall mean the work of authorship, whether in Source or |
||||
Object form, made available under the License, as indicated by a |
||||
copyright notice that is included in or attached to the work |
||||
(an example is provided in the Appendix below). |
||||
|
||||
"Derivative Works" shall mean any work, whether in Source or Object |
||||
form, that is based on (or derived from) the Work and for which the |
||||
editorial revisions, annotations, elaborations, or other modifications |
||||
represent, as a whole, an original work of authorship. For the purposes |
||||
of this License, Derivative Works shall not include works that remain |
||||
separable from, or merely link (or bind by name) to the interfaces of, |
||||
the Work and Derivative Works thereof. |
||||
|
||||
"Contribution" shall mean any work of authorship, including |
||||
the original version of the Work and any modifications or additions |
||||
to that Work or Derivative Works thereof, that is intentionally |
||||
submitted to Licensor for inclusion in the Work by the copyright owner |
||||
or by an individual or Legal Entity authorized to submit on behalf of |
||||
the copyright owner. For the purposes of this definition, "submitted" |
||||
means any form of electronic, verbal, or written communication sent |
||||
to the Licensor or its representatives, including but not limited to |
||||
communication on electronic mailing lists, source code control systems, |
||||
and issue tracking systems that are managed by, or on behalf of, the |
||||
Licensor for the purpose of discussing and improving the Work, but |
||||
excluding communication that is conspicuously marked or otherwise |
||||
designated in writing by the copyright owner as "Not a Contribution." |
||||
|
||||
"Contributor" shall mean Licensor and any individual or Legal Entity |
||||
on behalf of whom a Contribution has been received by Licensor and |
||||
subsequently incorporated within the Work. |
||||
|
||||
2. Grant of Copyright License. Subject to the terms and conditions of |
||||
this License, each Contributor hereby grants to You a perpetual, |
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable |
||||
copyright license to reproduce, prepare Derivative Works of, |
||||
publicly display, publicly perform, sublicense, and distribute the |
||||
Work and such Derivative Works in Source or Object form. |
||||
|
||||
3. Grant of Patent License. Subject to the terms and conditions of |
||||
this License, each Contributor hereby grants to You a perpetual, |
||||
worldwide, non-exclusive, no-charge, royalty-free, irrevocable |
||||
(except as stated in this section) patent license to make, have made, |
||||
use, offer to sell, sell, import, and otherwise transfer the Work, |
||||
where such license applies only to those patent claims licensable |
||||
by such Contributor that are necessarily infringed by their |
||||
Contribution(s) alone or by combination of their Contribution(s) |
||||
with the Work to which such Contribution(s) was submitted. If You |
||||
institute patent litigation against any entity (including a |
||||
cross-claim or counterclaim in a lawsuit) alleging that the Work |
||||
or a Contribution incorporated within the Work constitutes direct |
||||
or contributory patent infringement, then any patent licenses |
||||
granted to You under this License for that Work shall terminate |
||||
as of the date such litigation is filed. |
||||
|
||||
4. Redistribution. You may reproduce and distribute copies of the |
||||
Work or Derivative Works thereof in any medium, with or without |
||||
modifications, and in Source or Object form, provided that You |
||||
meet the following conditions: |
||||
|
||||
(a) You must give any other recipients of the Work or |
||||
Derivative Works a copy of this License; and |
||||
|
||||
(b) You must cause any modified files to carry prominent notices |
||||
stating that You changed the files; and |
||||
|
||||
(c) You must retain, in the Source form of any Derivative Works |
||||
that You distribute, all copyright, patent, trademark, and |
||||
attribution notices from the Source form of the Work, |
||||
excluding those notices that do not pertain to any part of |
||||
the Derivative Works; and |
||||
|
||||
(d) If the Work includes a "NOTICE" text file as part of its |
||||
distribution, then any Derivative Works that You distribute must |
||||
include a readable copy of the attribution notices contained |
||||
within such NOTICE file, excluding those notices that do not |
||||
pertain to any part of the Derivative Works, in at least one |
||||
of the following places: within a NOTICE text file distributed |
||||
as part of the Derivative Works; within the Source form or |
||||
documentation, if provided along with the Derivative Works; or, |
||||
within a display generated by the Derivative Works, if and |
||||
wherever such third-party notices normally appear. The contents |
||||
of the NOTICE file are for informational purposes only and |
||||
do not modify the License. You may add Your own attribution |
||||
notices within Derivative Works that You distribute, alongside |
||||
or as an addendum to the NOTICE text from the Work, provided |
||||
that such additional attribution notices cannot be construed |
||||
as modifying the License. |
||||
|
||||
You may add Your own copyright statement to Your modifications and |
||||
may provide additional or different license terms and conditions |
||||
for use, reproduction, or distribution of Your modifications, or |
||||
for any such Derivative Works as a whole, provided Your use, |
||||
reproduction, and distribution of the Work otherwise complies with |
||||
the conditions stated in this License. |
||||
|
||||
5. Submission of Contributions. Unless You explicitly state otherwise, |
||||
any Contribution intentionally submitted for inclusion in the Work |
||||
by You to the Licensor shall be under the terms and conditions of |
||||
this License, without any additional terms or conditions. |
||||
Notwithstanding the above, nothing herein shall supersede or modify |
||||
the terms of any separate license agreement you may have executed |
||||
with Licensor regarding such Contributions. |
||||
|
||||
6. Trademarks. This License does not grant permission to use the trade |
||||
names, trademarks, service marks, or product names of the Licensor, |
||||
except as required for reasonable and customary use in describing the |
||||
origin of the Work and reproducing the content of the NOTICE file. |
||||
|
||||
7. Disclaimer of Warranty. Unless required by applicable law or |
||||
agreed to in writing, Licensor provides the Work (and each |
||||
Contributor provides its Contributions) on an "AS IS" BASIS, |
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or |
||||
implied, including, without limitation, any warranties or conditions |
||||
of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A |
||||
PARTICULAR PURPOSE. You are solely responsible for determining the |
||||
appropriateness of using or redistributing the Work and assume any |
||||
risks associated with Your exercise of permissions under this License. |
||||
|
||||
8. Limitation of Liability. In no event and under no legal theory, |
||||
whether in tort (including negligence), contract, or otherwise, |
||||
unless required by applicable law (such as deliberate and grossly |
||||
negligent acts) or agreed to in writing, shall any Contributor be |
||||
liable to You for damages, including any direct, indirect, special, |
||||
incidental, or consequential damages of any character arising as a |
||||
result of this License or out of the use or inability to use the |
||||
Work (including but not limited to damages for loss of goodwill, |
||||
work stoppage, computer failure or malfunction, or any and all |
||||
other commercial damages or losses), even if such Contributor |
||||
has been advised of the possibility of such damages. |
||||
|
||||
9. Accepting Warranty or Additional Liability. While redistributing |
||||
the Work or Derivative Works thereof, You may choose to offer, |
||||
and charge a fee for, acceptance of support, warranty, indemnity, |
||||
or other liability obligations and/or rights consistent with this |
||||
License. However, in accepting such obligations, You may act only |
||||
on Your own behalf and on Your sole responsibility, not on behalf |
||||
of any other Contributor, and only if You agree to indemnify, |
||||
defend, and hold each Contributor harmless for any liability |
||||
incurred by, or claims asserted against, such Contributor by reason |
||||
of your accepting any such warranty or additional liability. |
||||
|
||||
END OF TERMS AND CONDITIONS |
||||
|
||||
APPENDIX: How to apply the Apache License to your work. |
||||
|
||||
To apply the Apache License to your work, attach the following |
||||
boilerplate notice, with the fields enclosed by brackets "[]" |
||||
replaced with your own identifying information. (Don't include |
||||
the brackets!) The text should be enclosed in the appropriate |
||||
comment syntax for the file format. We also recommend that a |
||||
file or class name and description of purpose be included on the |
||||
same "printed page" as the copyright notice for easier |
||||
identification within third-party archives. |
||||
|
||||
Copyright 2021- HPC-AI Technology Inc. |
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License"); |
||||
you may not use this file except in compliance with the License. |
||||
You may obtain a copy of the License at |
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0 |
||||
|
||||
Unless required by applicable law or agreed to in writing, software |
||||
distributed under the License is distributed on an "AS IS" BASIS, |
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
||||
See the License for the specific language governing permissions and |
||||
limitations under the License. |
||||
@ -0,0 +1,80 @@
|
||||
# RLHF - ColossalAI |
||||
|
||||
Implementation of RLHF (Reinforcement Learning with Human Feedback) powered by ColossalAI. It supports distributed training and offloading, which can fit extremly large models. |
||||
|
||||
<p align="center"> |
||||
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chatgpt/chatgpt.png" width=700/> |
||||
</p> |
||||
|
||||
## Training process (step 3) |
||||
<p align="center"> |
||||
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chatgpt/experience.jpg" width=500/> |
||||
</p> |
||||
<p align="center"> |
||||
<img src="https://raw.githubusercontent.com/hpcaitech/public_assets/main/applications/chatgpt/train.jpg" width=500/> |
||||
</p> |
||||
|
||||
|
||||
## Install |
||||
```shell |
||||
pip install . |
||||
``` |
||||
|
||||
|
||||
## Usage |
||||
|
||||
The main entrypoint is `Trainer`. We only support PPO trainer now. We support many training strategies: |
||||
|
||||
- NaiveStrategy: simplest strategy. Train on single GPU. |
||||
- DDPStrategy: use `torch.nn.parallel.DistributedDataParallel`. Train on multi GPUs. |
||||
- ColossalAIStrategy: use Gemini and Zero of ColossalAI. It eliminates model duplication on each GPU and supports offload. It's very useful when training large models on multi GPUs. |
||||
|
||||
Simplest usage: |
||||
|
||||
```python |
||||
from chatgpt.trainer import PPOTrainer |
||||
from chatgpt.trainer.strategies import ColossalAIStrategy |
||||
|
||||
strategy = ColossalAIStrategy() |
||||
|
||||
with strategy.model_init_context(): |
||||
# init your model here |
||||
actor = Actor() |
||||
critic = Critic() |
||||
|
||||
trainer = PPOTrainer(actor = actor, critic= critic, strategy, ...) |
||||
|
||||
trainer.fit(dataset, ...) |
||||
``` |
||||
|
||||
For more details, see `examples/`. |
||||
|
||||
We also support training reward model with true-world data. See `examples/train_reward_model.py`. |
||||
|
||||
## Todo |
||||
|
||||
- [x] implement PPO training |
||||
- [x] implement training reward model |
||||
- [x] support LoRA |
||||
- [ ] implement PPO-ptx fine-tuning |
||||
- [ ] integrate with Ray |
||||
- [ ] support more RL paradigms, like Implicit Language Q-Learning (ILQL) |
||||
|
||||
## Citations |
||||
|
||||
```bibtex |
||||
@article{Hu2021LoRALA, |
||||
title = {LoRA: Low-Rank Adaptation of Large Language Models}, |
||||
author = {Edward J. Hu and Yelong Shen and Phillip Wallis and Zeyuan Allen-Zhu and Yuanzhi Li and Shean Wang and Weizhu Chen}, |
||||
journal = {ArXiv}, |
||||
year = {2021}, |
||||
volume = {abs/2106.09685} |
||||
} |
||||
|
||||
@article{ouyang2022training, |
||||
title={Training language models to follow instructions with human feedback}, |
||||
author={Ouyang, Long and Wu, Jeff and Jiang, Xu and Almeida, Diogo and Wainwright, Carroll L and Mishkin, Pamela and Zhang, Chong and Agarwal, Sandhini and Slama, Katarina and Ray, Alex and others}, |
||||
journal={arXiv preprint arXiv:2203.02155}, |
||||
year={2022} |
||||
} |
||||
``` |
||||
@ -0,0 +1,94 @@
|
||||
# Benchmarks |
||||
|
||||
## Benchmark GPT on dummy prompt data |
||||
|
||||
We provide various GPT models (string in parentheses is the corresponding model name used in this script): |
||||
|
||||
- GPT2-S (s) |
||||
- GPT2-M (m) |
||||
- GPT2-L (l) |
||||
- GPT2-XL (xl) |
||||
- GPT2-4B (4b) |
||||
- GPT2-6B (6b) |
||||
- GPT2-8B (8b) |
||||
- GPT2-10B (10b) |
||||
- GPT2-12B (12b) |
||||
- GPT2-15B (15b) |
||||
- GPT2-18B (18b) |
||||
- GPT2-20B (20b) |
||||
- GPT2-24B (24b) |
||||
- GPT2-28B (28b) |
||||
- GPT2-32B (32b) |
||||
- GPT2-36B (36b) |
||||
- GPT2-40B (40b) |
||||
- GPT3 (175b) |
||||
|
||||
We also provide various training strategies: |
||||
|
||||
- ddp: torch DDP |
||||
- colossalai_gemini: ColossalAI GeminiDDP with `placement_policy="cuda"`, like zero3 |
||||
- colossalai_gemini_cpu: ColossalAI GeminiDDP with `placement_policy="cpu"`, like zero3-offload |
||||
- colossalai_zero2: ColossalAI zero2 |
||||
- colossalai_zero2_cpu: ColossalAI zero2-offload |
||||
- colossalai_zero1: ColossalAI zero1 |
||||
- colossalai_zero1_cpu: ColossalAI zero1-offload |
||||
|
||||
We only support `torchrun` to launch now. E.g. |
||||
|
||||
```shell |
||||
# run GPT2-S on single-node single-GPU with min batch size |
||||
torchrun --standalone --nproc_pero_node 1 benchmark_gpt_dummy.py --model s --strategy ddp --experience_batch_size 1 --train_batch_size 1 |
||||
# run GPT2-XL on single-node 4-GPU |
||||
torchrun --standalone --nproc_per_node 4 benchmark_gpt_dummy.py --model xl --strategy colossalai_zero2 |
||||
# run GPT3 on 8-node 8-GPU |
||||
torchrun --nnodes 8 --nproc_per_node 8 \ |
||||
--rdzv_id=$JOB_ID --rdzv_backend=c10d --rdzv_endpoint=$HOST_NODE_ADDR \ |
||||
benchmark_gpt_dummy.py --model 175b --strategy colossalai_gemini |
||||
``` |
||||
|
||||
> ⚠ Batch sizes in CLI args and outputed throughput/TFLOPS are all values of per GPU. |
||||
|
||||
In this benchmark, we assume the model architectures/sizes of actor and critic are the same for simplicity. But in practice, to reduce training cost, we may use a smaller critic. |
||||
|
||||
We also provide a simple shell script to run a set of benchmarks. But it only supports benchmark on single node. However, it's easy to run on multi-nodes by modifying launch command in this script. |
||||
|
||||
Usage: |
||||
|
||||
```shell |
||||
# run for GPUS=(1 2 4 8) x strategy=("ddp" "colossalai_zero2" "colossalai_gemini" "colossalai_zero2_cpu" "colossalai_gemini_cpu") x model=("s" "m" "l" "xl" "2b" "4b" "6b" "8b" "10b") x batch_size=(1 2 4 8 16 32 64 128 256) |
||||
./benchmark_gpt_dummy.sh |
||||
# run for GPUS=2 x strategy=("ddp" "colossalai_zero2" "colossalai_gemini" "colossalai_zero2_cpu" "colossalai_gemini_cpu") x model=("s" "m" "l" "xl" "2b" "4b" "6b" "8b" "10b") x batch_size=(1 2 4 8 16 32 64 128 256) |
||||
./benchmark_gpt_dummy.sh 2 |
||||
# run for GPUS=2 x strategy=ddp x model=("s" "m" "l" "xl" "2b" "4b" "6b" "8b" "10b") x batch_size=(1 2 4 8 16 32 64 128 256) |
||||
./benchmark_gpt_dummy.sh 2 ddp |
||||
# run for GPUS=2 x strategy=ddp x model=l x batch_size=(1 2 4 8 16 32 64 128 256) |
||||
./benchmark_gpt_dummy.sh 2 ddp l |
||||
``` |
||||
|
||||
## Benchmark OPT with LoRA on dummy prompt data |
||||
|
||||
We provide various OPT models (string in parentheses is the corresponding model name used in this script): |
||||
|
||||
- OPT-125M (125m) |
||||
- OPT-350M (350m) |
||||
- OPT-700M (700m) |
||||
- OPT-1.3B (1.3b) |
||||
- OPT-2.7B (2.7b) |
||||
- OPT-3.5B (3.5b) |
||||
- OPT-5.5B (5.5b) |
||||
- OPT-6.7B (6.7b) |
||||
- OPT-10B (10b) |
||||
- OPT-13B (13b) |
||||
|
||||
We only support `torchrun` to launch now. E.g. |
||||
|
||||
```shell |
||||
# run OPT-125M with no lora (lora_rank=0) on single-node single-GPU with min batch size |
||||
torchrun --standalone --nproc_pero_node 1 benchmark_opt_lora_dummy.py --model 125m --strategy ddp --experience_batch_size 1 --train_batch_size 1 --lora_rank 0 |
||||
# run OPT-350M with lora_rank=4 on single-node 4-GPU |
||||
torchrun --standalone --nproc_per_node 4 benchmark_opt_lora_dummy.py --model 350m --strategy colossalai_zero2 --lora_rank 4 |
||||
``` |
||||
|
||||
> ⚠ Batch sizes in CLI args and outputed throughput/TFLOPS are all values of per GPU. |
||||
|
||||
In this benchmark, we assume the model architectures/sizes of actor and critic are the same for simplicity. But in practice, to reduce training cost, we may use a smaller critic. |
||||
@ -0,0 +1,183 @@
|
||||
import argparse |
||||
from copy import deepcopy |
||||
|
||||
import torch |
||||
import torch.distributed as dist |
||||
import torch.nn as nn |
||||
from chatgpt.nn import GPTActor, GPTCritic, RewardModel |
||||
from chatgpt.nn.generation_utils import gpt_prepare_inputs_fn, update_model_kwargs_fn |
||||
from chatgpt.trainer import PPOTrainer |
||||
from chatgpt.trainer.callbacks import PerformanceEvaluator |
||||
from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, Strategy |
||||
from torch.optim import Adam |
||||
from transformers.models.gpt2.configuration_gpt2 import GPT2Config |
||||
from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer |
||||
|
||||
from colossalai.nn.optimizer import HybridAdam |
||||
|
||||
|
||||
def get_model_numel(model: nn.Module, strategy: Strategy) -> int: |
||||
numel = sum(p.numel() for p in model.parameters()) |
||||
if isinstance(strategy, ColossalAIStrategy) and strategy.stage == 3 and strategy.shard_init: |
||||
numel *= dist.get_world_size() |
||||
return numel |
||||
|
||||
|
||||
def preprocess_batch(samples) -> dict: |
||||
input_ids = torch.stack(samples) |
||||
attention_mask = torch.ones_like(input_ids, dtype=torch.long) |
||||
return {'input_ids': input_ids, 'attention_mask': attention_mask} |
||||
|
||||
|
||||
def print_rank_0(*args, **kwargs) -> None: |
||||
if dist.get_rank() == 0: |
||||
print(*args, **kwargs) |
||||
|
||||
|
||||
def print_model_numel(model_dict: dict) -> None: |
||||
B = 1024**3 |
||||
M = 1024**2 |
||||
K = 1024 |
||||
outputs = '' |
||||
for name, numel in model_dict.items(): |
||||
outputs += f'{name}: ' |
||||
if numel >= B: |
||||
outputs += f'{numel / B:.2f} B\n' |
||||
elif numel >= M: |
||||
outputs += f'{numel / M:.2f} M\n' |
||||
elif numel >= K: |
||||
outputs += f'{numel / K:.2f} K\n' |
||||
else: |
||||
outputs += f'{numel}\n' |
||||
print_rank_0(outputs) |
||||
|
||||
|
||||
def get_gpt_config(model_name: str) -> GPT2Config: |
||||
model_map = { |
||||
's': GPT2Config(), |
||||
'm': GPT2Config(n_embd=1024, n_layer=24, n_head=16), |
||||
'l': GPT2Config(n_embd=1280, n_layer=36, n_head=20), |
||||
'xl': GPT2Config(n_embd=1600, n_layer=48, n_head=25), |
||||
'2b': GPT2Config(n_embd=2048, n_layer=40, n_head=16), |
||||
'4b': GPT2Config(n_embd=2304, n_layer=64, n_head=16), |
||||
'6b': GPT2Config(n_embd=4096, n_layer=30, n_head=16), |
||||
'8b': GPT2Config(n_embd=4096, n_layer=40, n_head=16), |
||||
'10b': GPT2Config(n_embd=4096, n_layer=50, n_head=16), |
||||
'12b': GPT2Config(n_embd=4096, n_layer=60, n_head=16), |
||||
'15b': GPT2Config(n_embd=4096, n_layer=78, n_head=16), |
||||
'18b': GPT2Config(n_embd=4096, n_layer=90, n_head=16), |
||||
'20b': GPT2Config(n_embd=8192, n_layer=25, n_head=16), |
||||
'24b': GPT2Config(n_embd=8192, n_layer=30, n_head=16), |
||||
'28b': GPT2Config(n_embd=8192, n_layer=35, n_head=16), |
||||
'32b': GPT2Config(n_embd=8192, n_layer=40, n_head=16), |
||||
'36b': GPT2Config(n_embd=8192, n_layer=45, n_head=16), |
||||
'40b': GPT2Config(n_embd=8192, n_layer=50, n_head=16), |
||||
'175b': GPT2Config(n_positions=2048, n_embd=12288, n_layer=96, n_head=96), |
||||
} |
||||
try: |
||||
return model_map[model_name] |
||||
except KeyError: |
||||
raise ValueError(f'Unknown model "{model_name}"') |
||||
|
||||
|
||||
def main(args): |
||||
if args.strategy == 'ddp': |
||||
strategy = DDPStrategy() |
||||
elif args.strategy == 'colossalai_gemini': |
||||
strategy = ColossalAIStrategy(stage=3, placement_policy='cuda', initial_scale=2**5) |
||||
elif args.strategy == 'colossalai_gemini_cpu': |
||||
strategy = ColossalAIStrategy(stage=3, placement_policy='cpu', initial_scale=2**5) |
||||
elif args.strategy == 'colossalai_zero2': |
||||
strategy = ColossalAIStrategy(stage=2, placement_policy='cuda') |
||||
elif args.strategy == 'colossalai_zero2_cpu': |
||||
strategy = ColossalAIStrategy(stage=2, placement_policy='cpu') |
||||
elif args.strategy == 'colossalai_zero1': |
||||
strategy = ColossalAIStrategy(stage=1, placement_policy='cuda') |
||||
elif args.strategy == 'colossalai_zero1_cpu': |
||||
strategy = ColossalAIStrategy(stage=1, placement_policy='cpu') |
||||
else: |
||||
raise ValueError(f'Unsupported strategy "{args.strategy}"') |
||||
|
||||
model_config = get_gpt_config(args.model) |
||||
|
||||
with strategy.model_init_context(): |
||||
actor = GPTActor(config=model_config).cuda() |
||||
critic = GPTCritic(config=model_config).cuda() |
||||
|
||||
initial_model = deepcopy(actor).cuda() |
||||
reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda() |
||||
|
||||
actor_numel = get_model_numel(actor, strategy) |
||||
critic_numel = get_model_numel(critic, strategy) |
||||
initial_model_numel = get_model_numel(initial_model, strategy) |
||||
reward_model_numel = get_model_numel(reward_model, strategy) |
||||
print_model_numel({ |
||||
'Actor': actor_numel, |
||||
'Critic': critic_numel, |
||||
'Initial model': initial_model_numel, |
||||
'Reward model': reward_model_numel |
||||
}) |
||||
performance_evaluator = PerformanceEvaluator(actor_numel, |
||||
critic_numel, |
||||
initial_model_numel, |
||||
reward_model_numel, |
||||
enable_grad_checkpoint=False, |
||||
ignore_episodes=1) |
||||
|
||||
if args.strategy.startswith('colossalai'): |
||||
actor_optim = HybridAdam(actor.parameters(), lr=5e-6) |
||||
critic_optim = HybridAdam(critic.parameters(), lr=5e-6) |
||||
else: |
||||
actor_optim = Adam(actor.parameters(), lr=5e-6) |
||||
critic_optim = Adam(critic.parameters(), lr=5e-6) |
||||
|
||||
tokenizer = GPT2Tokenizer.from_pretrained('gpt2') |
||||
tokenizer.pad_token = tokenizer.eos_token |
||||
|
||||
trainer = PPOTrainer(strategy, |
||||
actor, |
||||
critic, |
||||
reward_model, |
||||
initial_model, |
||||
actor_optim, |
||||
critic_optim, |
||||
max_epochs=args.max_epochs, |
||||
train_batch_size=args.train_batch_size, |
||||
experience_batch_size=args.experience_batch_size, |
||||
tokenizer=preprocess_batch, |
||||
max_length=512, |
||||
do_sample=True, |
||||
temperature=1.0, |
||||
top_k=50, |
||||
pad_token_id=tokenizer.pad_token_id, |
||||
eos_token_id=tokenizer.eos_token_id, |
||||
prepare_inputs_fn=gpt_prepare_inputs_fn, |
||||
update_model_kwargs_fn=update_model_kwargs_fn, |
||||
callbacks=[performance_evaluator]) |
||||
|
||||
random_prompts = torch.randint(tokenizer.vocab_size, (1000, 400), device=torch.cuda.current_device()) |
||||
trainer.fit(random_prompts, |
||||
num_episodes=args.num_episodes, |
||||
max_timesteps=args.max_timesteps, |
||||
update_timesteps=args.update_timesteps) |
||||
|
||||
print_rank_0(f'Peak CUDA mem: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB') |
||||
|
||||
|
||||
if __name__ == '__main__': |
||||
parser = argparse.ArgumentParser() |
||||
parser.add_argument('--model', default='s') |
||||
parser.add_argument('--strategy', |
||||
choices=[ |
||||
'ddp', 'colossalai_gemini', 'colossalai_gemini_cpu', 'colossalai_zero2', |
||||
'colossalai_zero2_cpu', 'colossalai_zero1', 'colossalai_zero1_cpu' |
||||
], |
||||
default='ddp') |
||||
parser.add_argument('--num_episodes', type=int, default=3) |
||||
parser.add_argument('--max_timesteps', type=int, default=8) |
||||
parser.add_argument('--update_timesteps', type=int, default=8) |
||||
parser.add_argument('--max_epochs', type=int, default=3) |
||||
parser.add_argument('--train_batch_size', type=int, default=8) |
||||
parser.add_argument('--experience_batch_size', type=int, default=8) |
||||
args = parser.parse_args() |
||||
main(args) |
||||
@ -0,0 +1,45 @@
|
||||
#!/usr/bin/env bash |
||||
# Usage: $0 <?number-of-gpus> <?strategy> <?model> |
||||
set -xu |
||||
|
||||
BASE=$(realpath $(dirname $0)) |
||||
|
||||
|
||||
PY_SCRIPT=${BASE}/benchmark_gpt_dummy.py |
||||
export OMP_NUM_THREADS=8 |
||||
|
||||
function tune_batch_size() { |
||||
# we found when experience batch size is equal to train batch size |
||||
# peak CUDA memory usage of making experience phase is less than or equal to that of training phase |
||||
# thus, experience batch size can be larger than or equal to train batch size |
||||
for bs in 1 2 4 8 16 32 64 128 256; do |
||||
torchrun --standalone --nproc_per_node $1 $PY_SCRIPT --model $2 --strategy $3 --experience_batch_size $bs --train_batch_size $bs || return 1 |
||||
done |
||||
} |
||||
|
||||
if [ $# -eq 0 ]; then |
||||
num_gpus=(1 2 4 8) |
||||
else |
||||
num_gpus=($1) |
||||
fi |
||||
|
||||
if [ $# -le 1 ]; then |
||||
strategies=("ddp" "colossalai_zero2" "colossalai_gemini" "colossalai_zero2_cpu" "colossalai_gemini_cpu") |
||||
else |
||||
strategies=($2) |
||||
fi |
||||
|
||||
if [ $# -le 2 ]; then |
||||
models=("s" "m" "l" "xl" "2b" "4b" "6b" "8b" "10b") |
||||
else |
||||
models=($3) |
||||
fi |
||||
|
||||
|
||||
for num_gpu in ${num_gpus[@]}; do |
||||
for strategy in ${strategies[@]}; do |
||||
for model in ${models[@]}; do |
||||
tune_batch_size $num_gpu $model $strategy || break |
||||
done |
||||
done |
||||
done |
||||
@ -0,0 +1,178 @@
|
||||
import argparse |
||||
from copy import deepcopy |
||||
|
||||
import torch |
||||
import torch.distributed as dist |
||||
import torch.nn as nn |
||||
from chatgpt.nn import OPTActor, OPTCritic, RewardModel |
||||
from chatgpt.nn.generation_utils import opt_prepare_inputs_fn, update_model_kwargs_fn |
||||
from chatgpt.trainer import PPOTrainer |
||||
from chatgpt.trainer.callbacks import PerformanceEvaluator |
||||
from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, Strategy |
||||
from torch.optim import Adam |
||||
from transformers import AutoTokenizer |
||||
from transformers.models.opt.configuration_opt import OPTConfig |
||||
|
||||
from colossalai.nn.optimizer import HybridAdam |
||||
|
||||
|
||||
def get_model_numel(model: nn.Module, strategy: Strategy) -> int: |
||||
numel = sum(p.numel() for p in model.parameters()) |
||||
if isinstance(strategy, ColossalAIStrategy) and strategy.stage == 3 and strategy.shard_init: |
||||
numel *= dist.get_world_size() |
||||
return numel |
||||
|
||||
|
||||
def preprocess_batch(samples) -> dict: |
||||
input_ids = torch.stack(samples) |
||||
attention_mask = torch.ones_like(input_ids, dtype=torch.long) |
||||
return {'input_ids': input_ids, 'attention_mask': attention_mask} |
||||
|
||||
|
||||
def print_rank_0(*args, **kwargs) -> None: |
||||
if dist.get_rank() == 0: |
||||
print(*args, **kwargs) |
||||
|
||||
|
||||
def print_model_numel(model_dict: dict) -> None: |
||||
B = 1024**3 |
||||
M = 1024**2 |
||||
K = 1024 |
||||
outputs = '' |
||||
for name, numel in model_dict.items(): |
||||
outputs += f'{name}: ' |
||||
if numel >= B: |
||||
outputs += f'{numel / B:.2f} B\n' |
||||
elif numel >= M: |
||||
outputs += f'{numel / M:.2f} M\n' |
||||
elif numel >= K: |
||||
outputs += f'{numel / K:.2f} K\n' |
||||
else: |
||||
outputs += f'{numel}\n' |
||||
print_rank_0(outputs) |
||||
|
||||
|
||||
def get_gpt_config(model_name: str) -> OPTConfig: |
||||
model_map = { |
||||
'125m': OPTConfig.from_pretrained('facebook/opt-125m'), |
||||
'350m': OPTConfig(hidden_size=1024, ffn_dim=4096, num_hidden_layers=24, num_attention_heads=16), |
||||
'700m': OPTConfig(hidden_size=1280, ffn_dim=5120, num_hidden_layers=36, num_attention_heads=20), |
||||
'1.3b': OPTConfig.from_pretrained('facebook/opt-1.3b'), |
||||
'2.7b': OPTConfig.from_pretrained('facebook/opt-2.7b'), |
||||
'3.5b': OPTConfig(hidden_size=3072, ffn_dim=12288, num_hidden_layers=32, num_attention_heads=32), |
||||
'5.5b': OPTConfig(hidden_size=3840, ffn_dim=15360, num_hidden_layers=32, num_attention_heads=32), |
||||
'6.7b': OPTConfig.from_pretrained('facebook/opt-6.7b'), |
||||
'10b': OPTConfig(hidden_size=5120, ffn_dim=20480, num_hidden_layers=32, num_attention_heads=32), |
||||
'13b': OPTConfig.from_pretrained('facebook/opt-13b'), |
||||
} |
||||
try: |
||||
return model_map[model_name] |
||||
except KeyError: |
||||
raise ValueError(f'Unknown model "{model_name}"') |
||||
|
||||
|
||||
def main(args): |
||||
if args.strategy == 'ddp': |
||||
strategy = DDPStrategy() |
||||
elif args.strategy == 'colossalai_gemini': |
||||
strategy = ColossalAIStrategy(stage=3, placement_policy='cuda', initial_scale=2**5) |
||||
elif args.strategy == 'colossalai_gemini_cpu': |
||||
strategy = ColossalAIStrategy(stage=3, placement_policy='cpu', initial_scale=2**5) |
||||
elif args.strategy == 'colossalai_zero2': |
||||
strategy = ColossalAIStrategy(stage=2, placement_policy='cuda') |
||||
elif args.strategy == 'colossalai_zero2_cpu': |
||||
strategy = ColossalAIStrategy(stage=2, placement_policy='cpu') |
||||
elif args.strategy == 'colossalai_zero1': |
||||
strategy = ColossalAIStrategy(stage=1, placement_policy='cuda') |
||||
elif args.strategy == 'colossalai_zero1_cpu': |
||||
strategy = ColossalAIStrategy(stage=1, placement_policy='cpu') |
||||
else: |
||||
raise ValueError(f'Unsupported strategy "{args.strategy}"') |
||||
|
||||
torch.cuda.set_per_process_memory_fraction(args.cuda_mem_frac) |
||||
|
||||
model_config = get_gpt_config(args.model) |
||||
|
||||
with strategy.model_init_context(): |
||||
actor = OPTActor(config=model_config, lora_rank=args.lora_rank).cuda() |
||||
critic = OPTCritic(config=model_config, lora_rank=args.lora_rank).cuda() |
||||
|
||||
initial_model = deepcopy(actor).cuda() |
||||
reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda() |
||||
|
||||
actor_numel = get_model_numel(actor, strategy) |
||||
critic_numel = get_model_numel(critic, strategy) |
||||
initial_model_numel = get_model_numel(initial_model, strategy) |
||||
reward_model_numel = get_model_numel(reward_model, strategy) |
||||
print_model_numel({ |
||||
'Actor': actor_numel, |
||||
'Critic': critic_numel, |
||||
'Initial model': initial_model_numel, |
||||
'Reward model': reward_model_numel |
||||
}) |
||||
performance_evaluator = PerformanceEvaluator(actor_numel, |
||||
critic_numel, |
||||
initial_model_numel, |
||||
reward_model_numel, |
||||
enable_grad_checkpoint=False, |
||||
ignore_episodes=1) |
||||
|
||||
if args.strategy.startswith('colossalai'): |
||||
actor_optim = HybridAdam(actor.parameters(), lr=5e-6) |
||||
critic_optim = HybridAdam(critic.parameters(), lr=5e-6) |
||||
else: |
||||
actor_optim = Adam(actor.parameters(), lr=5e-6) |
||||
critic_optim = Adam(critic.parameters(), lr=5e-6) |
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained('facebook/opt-350m') |
||||
tokenizer.pad_token = tokenizer.eos_token |
||||
|
||||
trainer = PPOTrainer(strategy, |
||||
actor, |
||||
critic, |
||||
reward_model, |
||||
initial_model, |
||||
actor_optim, |
||||
critic_optim, |
||||
max_epochs=args.max_epochs, |
||||
train_batch_size=args.train_batch_size, |
||||
experience_batch_size=args.experience_batch_size, |
||||
tokenizer=preprocess_batch, |
||||
max_length=512, |
||||
do_sample=True, |
||||
temperature=1.0, |
||||
top_k=50, |
||||
pad_token_id=tokenizer.pad_token_id, |
||||
eos_token_id=tokenizer.eos_token_id, |
||||
prepare_inputs_fn=opt_prepare_inputs_fn, |
||||
update_model_kwargs_fn=update_model_kwargs_fn, |
||||
callbacks=[performance_evaluator]) |
||||
|
||||
random_prompts = torch.randint(tokenizer.vocab_size, (1000, 400), device=torch.cuda.current_device()) |
||||
trainer.fit(random_prompts, |
||||
num_episodes=args.num_episodes, |
||||
max_timesteps=args.max_timesteps, |
||||
update_timesteps=args.update_timesteps) |
||||
|
||||
print_rank_0(f'Peak CUDA mem: {torch.cuda.max_memory_allocated()/1024**3:.2f} GB') |
||||
|
||||
|
||||
if __name__ == '__main__': |
||||
parser = argparse.ArgumentParser() |
||||
parser.add_argument('--model', default='125m') |
||||
parser.add_argument('--strategy', |
||||
choices=[ |
||||
'ddp', 'colossalai_gemini', 'colossalai_gemini_cpu', 'colossalai_zero2', |
||||
'colossalai_zero2_cpu', 'colossalai_zero1', 'colossalai_zero1_cpu' |
||||
], |
||||
default='ddp') |
||||
parser.add_argument('--num_episodes', type=int, default=3) |
||||
parser.add_argument('--max_timesteps', type=int, default=8) |
||||
parser.add_argument('--update_timesteps', type=int, default=8) |
||||
parser.add_argument('--max_epochs', type=int, default=3) |
||||
parser.add_argument('--train_batch_size', type=int, default=8) |
||||
parser.add_argument('--experience_batch_size', type=int, default=8) |
||||
parser.add_argument('--lora_rank', type=int, default=4) |
||||
parser.add_argument('--cuda_mem_frac', type=float, default=1.0) |
||||
args = parser.parse_args() |
||||
main(args) |
||||
@ -0,0 +1,3 @@
|
||||
from .reward_dataset import RewardDataset |
||||
|
||||
__all__ = ['RewardDataset'] |
||||
@ -0,0 +1,52 @@
|
||||
from typing import Callable |
||||
|
||||
from torch.utils.data import Dataset |
||||
from tqdm import tqdm |
||||
|
||||
|
||||
class RewardDataset(Dataset): |
||||
""" |
||||
Dataset for reward model |
||||
|
||||
Args: |
||||
dataset: dataset for reward model |
||||
tokenizer: tokenizer for reward model |
||||
max_length: max length of input |
||||
""" |
||||
|
||||
def __init__(self, dataset, tokenizer: Callable, max_length: int) -> None: |
||||
super().__init__() |
||||
self.chosen = [] |
||||
self.reject = [] |
||||
for data in tqdm(dataset): |
||||
prompt = data['prompt'] |
||||
|
||||
chosen = prompt + data['chosen'] + "<|endoftext|>" |
||||
chosen_token = tokenizer(chosen, |
||||
max_length=max_length, |
||||
padding="max_length", |
||||
truncation=True, |
||||
return_tensors="pt") |
||||
self.chosen.append({ |
||||
"input_ids": chosen_token['input_ids'], |
||||
"attention_mask": chosen_token['attention_mask'] |
||||
}) |
||||
|
||||
reject = prompt + data['rejected'] + "<|endoftext|>" |
||||
reject_token = tokenizer(reject, |
||||
max_length=max_length, |
||||
padding="max_length", |
||||
truncation=True, |
||||
return_tensors="pt") |
||||
self.reject.append({ |
||||
"input_ids": reject_token['input_ids'], |
||||
"attention_mask": reject_token['attention_mask'] |
||||
}) |
||||
|
||||
def __len__(self): |
||||
length = len(self.chosen) |
||||
return length |
||||
|
||||
def __getitem__(self, idx): |
||||
return self.chosen[idx]["input_ids"], self.chosen[idx]["attention_mask"], self.reject[idx][ |
||||
"input_ids"], self.reject[idx]["attention_mask"] |
||||
@ -0,0 +1,4 @@
|
||||
from .base import Experience, ExperienceMaker |
||||
from .naive import NaiveExperienceMaker |
||||
|
||||
__all__ = ['Experience', 'ExperienceMaker', 'NaiveExperienceMaker'] |
||||
@ -0,0 +1,77 @@
|
||||
from abc import ABC, abstractmethod |
||||
from dataclasses import dataclass |
||||
from typing import Optional |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
from chatgpt.nn.actor import Actor |
||||
|
||||
|
||||
@dataclass |
||||
class Experience: |
||||
"""Experience is a batch of data. |
||||
These data should have the the sequence length and number of actions. |
||||
Left padding for sequences is applied. |
||||
|
||||
Shapes of each tensor: |
||||
sequences: (B, S) |
||||
action_log_probs: (B, A) |
||||
values: (B) |
||||
reward: (B) |
||||
advatanges: (B) |
||||
attention_mask: (B, S) |
||||
action_mask: (B, A) |
||||
|
||||
"A" is the number of actions. |
||||
""" |
||||
sequences: torch.Tensor |
||||
action_log_probs: torch.Tensor |
||||
values: torch.Tensor |
||||
reward: torch.Tensor |
||||
advantages: torch.Tensor |
||||
attention_mask: Optional[torch.LongTensor] |
||||
action_mask: Optional[torch.BoolTensor] |
||||
|
||||
@torch.no_grad() |
||||
def to_device(self, device: torch.device) -> None: |
||||
self.sequences = self.sequences.to(device) |
||||
self.action_log_probs = self.action_log_probs.to(device) |
||||
self.values = self.values.to(device) |
||||
self.reward = self.reward.to(device) |
||||
self.advantages = self.advantages.to(device) |
||||
if self.attention_mask is not None: |
||||
self.attention_mask = self.attention_mask.to(device) |
||||
if self.action_mask is not None: |
||||
self.action_mask = self.action_mask.to(device) |
||||
|
||||
def pin_memory(self): |
||||
self.sequences = self.sequences.pin_memory() |
||||
self.action_log_probs = self.action_log_probs.pin_memory() |
||||
self.values = self.values.pin_memory() |
||||
self.reward = self.reward.pin_memory() |
||||
self.advantages = self.advantages.pin_memory() |
||||
if self.attention_mask is not None: |
||||
self.attention_mask = self.attention_mask.pin_memory() |
||||
if self.action_mask is not None: |
||||
self.action_mask = self.action_mask.pin_memory() |
||||
return self |
||||
|
||||
|
||||
class ExperienceMaker(ABC): |
||||
|
||||
def __init__(self, |
||||
actor: Actor, |
||||
critic: nn.Module, |
||||
reward_model: nn.Module, |
||||
initial_model: Actor, |
||||
kl_coef: float = 0.1) -> None: |
||||
super().__init__() |
||||
self.actor = actor |
||||
self.critic = critic |
||||
self.reward_model = reward_model |
||||
self.initial_model = initial_model |
||||
self.kl_coef = kl_coef |
||||
|
||||
@abstractmethod |
||||
def make_experience(self, input_ids: torch.Tensor, **generate_kwargs) -> Experience: |
||||
pass |
||||
@ -0,0 +1,36 @@
|
||||
import torch |
||||
from chatgpt.nn.utils import compute_reward, normalize |
||||
|
||||
from .base import Experience, ExperienceMaker |
||||
|
||||
|
||||
class NaiveExperienceMaker(ExperienceMaker): |
||||
""" |
||||
Naive experience maker. |
||||
""" |
||||
|
||||
@torch.no_grad() |
||||
def make_experience(self, input_ids: torch.Tensor, **generate_kwargs) -> Experience: |
||||
self.actor.eval() |
||||
self.critic.eval() |
||||
self.initial_model.eval() |
||||
self.reward_model.eval() |
||||
|
||||
sequences, attention_mask, action_mask = self.actor.generate(input_ids, |
||||
return_action_mask=True, |
||||
**generate_kwargs) |
||||
num_actions = action_mask.size(1) |
||||
|
||||
action_log_probs = self.actor(sequences, num_actions, attention_mask) |
||||
base_action_log_probs = self.initial_model(sequences, num_actions, attention_mask) |
||||
value = self.critic(sequences, action_mask, attention_mask) |
||||
r = self.reward_model(sequences, attention_mask) |
||||
|
||||
reward = compute_reward(r, self.kl_coef, action_log_probs, base_action_log_probs, action_mask=action_mask) |
||||
|
||||
advantage = reward - value |
||||
# TODO(ver217): maybe normalize adv |
||||
if advantage.ndim == 1: |
||||
advantage = advantage.unsqueeze(-1) |
||||
|
||||
return Experience(sequences, action_log_probs, value, reward, advantage, attention_mask, action_mask) |
||||
@ -0,0 +1,18 @@
|
||||
from .actor import Actor |
||||
from .bloom_actor import BLOOMActor |
||||
from .bloom_critic import BLOOMCritic |
||||
from .bloom_rm import BLOOMRM |
||||
from .critic import Critic |
||||
from .gpt_actor import GPTActor |
||||
from .gpt_critic import GPTCritic |
||||
from .gpt_rm import GPTRM |
||||
from .loss import PairWiseLoss, PolicyLoss, PPOPtxActorLoss, ValueLoss |
||||
from .opt_actor import OPTActor |
||||
from .opt_critic import OPTCritic |
||||
from .opt_rm import OPTRM |
||||
from .reward_model import RewardModel |
||||
|
||||
__all__ = [ |
||||
'Actor', 'Critic', 'RewardModel', 'PolicyLoss', 'ValueLoss', 'PPOPtxActorLoss', 'PairWiseLoss', 'GPTActor', |
||||
'GPTCritic', 'GPTRM', 'BLOOMActor', 'BLOOMCritic', 'BLOOMRM', 'OPTActor', 'OPTCritic', 'OPTRM' |
||||
] |
||||
@ -0,0 +1,62 @@
|
||||
from typing import Optional, Tuple, Union |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
import torch.nn.functional as F |
||||
|
||||
from .generation import generate |
||||
from .lora import LoRAModule |
||||
from .utils import log_probs_from_logits |
||||
|
||||
|
||||
class Actor(LoRAModule): |
||||
""" |
||||
Actor model base class. |
||||
|
||||
Args: |
||||
model (nn.Module): Actor Model. |
||||
lora_rank (int): LoRA rank. |
||||
lora_train_bias (str): LoRA bias training mode. |
||||
""" |
||||
|
||||
def __init__(self, model: nn.Module, lora_rank: int = 0, lora_train_bias: str = 'none') -> None: |
||||
super().__init__(lora_rank=lora_rank, lora_train_bias=lora_train_bias) |
||||
self.model = model |
||||
self.convert_to_lora() |
||||
|
||||
@torch.no_grad() |
||||
def generate( |
||||
self, |
||||
input_ids: torch.Tensor, |
||||
return_action_mask: bool = True, |
||||
**kwargs |
||||
) -> Union[Tuple[torch.LongTensor, torch.LongTensor], Tuple[torch.LongTensor, torch.LongTensor, torch.BoolTensor]]: |
||||
sequences = generate(self.model, input_ids, **kwargs) |
||||
attention_mask = None |
||||
pad_token_id = kwargs.get('pad_token_id', None) |
||||
if pad_token_id is not None: |
||||
attention_mask = sequences.not_equal(pad_token_id).to(dtype=torch.long, device=sequences.device) |
||||
if not return_action_mask: |
||||
return sequences, attention_mask |
||||
input_len = input_ids.size(1) |
||||
eos_token_id = kwargs.get('eos_token_id', None) |
||||
if eos_token_id is None: |
||||
action_mask = torch.ones_like(sequences, dtype=torch.bool) |
||||
else: |
||||
# left padding may be applied, only mask action |
||||
action_mask = (sequences[:, input_len:] == eos_token_id).cumsum(dim=-1) == 0 |
||||
action_mask = F.pad(action_mask, (1 + input_len, -1), value=True) # include eos token and input |
||||
action_mask[:, :input_len] = False |
||||
action_mask = action_mask[:, 1:] |
||||
return sequences, attention_mask, action_mask[:, -(sequences.size(1) - input_len):] |
||||
|
||||
def forward(self, |
||||
sequences: torch.LongTensor, |
||||
num_actions: int, |
||||
attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
||||
"""Returns action log probs |
||||
""" |
||||
output = self.model(sequences, attention_mask=attention_mask) |
||||
logits = output['logits'] |
||||
log_probs = log_probs_from_logits(logits[:, :-1, :], sequences[:, 1:]) |
||||
return log_probs[:, -num_actions:] |
||||
@ -0,0 +1,35 @@
|
||||
from typing import Optional |
||||
|
||||
import torch |
||||
from transformers import BloomConfig, BloomForCausalLM, BloomModel |
||||
|
||||
from .actor import Actor |
||||
|
||||
|
||||
class BLOOMActor(Actor): |
||||
""" |
||||
BLOOM Actor model. |
||||
|
||||
Args: |
||||
pretrained (str): Pretrained model name or path. |
||||
config (BloomConfig): Model config. |
||||
checkpoint (bool): Enable gradient checkpointing. |
||||
lora_rank (int): LoRA rank. |
||||
lora_train_bias (str): LoRA bias training mode. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
pretrained: str = None, |
||||
config: Optional[BloomConfig] = None, |
||||
checkpoint: bool = False, |
||||
lora_rank: int = 0, |
||||
lora_train_bias: str = 'none') -> None: |
||||
if pretrained is not None: |
||||
model = BloomForCausalLM.from_pretrained(pretrained) |
||||
elif config is not None: |
||||
model = BloomForCausalLM(config) |
||||
else: |
||||
model = BloomForCausalLM(BloomConfig()) |
||||
if checkpoint: |
||||
model.gradient_checkpointing_enable() |
||||
super().__init__(model, lora_rank, lora_train_bias) |
||||
@ -0,0 +1,37 @@
|
||||
from typing import Optional |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
from transformers import BloomConfig, BloomForCausalLM, BloomModel |
||||
|
||||
from .critic import Critic |
||||
|
||||
|
||||
class BLOOMCritic(Critic): |
||||
""" |
||||
BLOOM Critic model. |
||||
|
||||
Args: |
||||
pretrained (str): Pretrained model name or path. |
||||
config (BloomConfig): Model config. |
||||
checkpoint (bool): Enable gradient checkpointing. |
||||
lora_rank (int): LoRA rank. |
||||
lora_train_bias (str): LoRA bias training mode. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
pretrained: str = None, |
||||
config: Optional[BloomConfig] = None, |
||||
checkpoint: bool = False, |
||||
lora_rank: int = 0, |
||||
lora_train_bias: str = 'none') -> None: |
||||
if pretrained is not None: |
||||
model = BloomModel.from_pretrained(pretrained) |
||||
elif config is not None: |
||||
model = BloomModel(config) |
||||
else: |
||||
model = BloomModel(BloomConfig()) |
||||
if checkpoint: |
||||
model.gradient_checkpointing_enable() |
||||
value_head = nn.Linear(model.config.hidden_size, 1) |
||||
super().__init__(model, value_head, lora_rank, lora_train_bias) |
||||
@ -0,0 +1,37 @@
|
||||
from typing import Optional |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
from transformers import BloomConfig, BloomForCausalLM, BloomModel |
||||
|
||||
from .reward_model import RewardModel |
||||
|
||||
|
||||
class BLOOMRM(RewardModel): |
||||
""" |
||||
BLOOM Reward model. |
||||
|
||||
Args: |
||||
pretrained (str): Pretrained model name or path. |
||||
config (BloomConfig): Model config. |
||||
checkpoint (bool): Enable gradient checkpointing. |
||||
lora_rank (int): LoRA rank. |
||||
lora_train_bias (str): LoRA bias training mode. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
pretrained: str = None, |
||||
config: Optional[BloomConfig] = None, |
||||
checkpoint: bool = False, |
||||
lora_rank: int = 0, |
||||
lora_train_bias: str = 'none') -> None: |
||||
if pretrained is not None: |
||||
model = BloomModel.from_pretrained(pretrained) |
||||
elif config is not None: |
||||
model = BloomModel(config) |
||||
else: |
||||
model = BloomModel(BloomConfig()) |
||||
if checkpoint: |
||||
model.gradient_checkpointing_enable() |
||||
value_head = nn.Linear(model.config.hidden_size, 1) |
||||
super().__init__(model, value_head, lora_rank, lora_train_bias) |
||||
@ -0,0 +1,47 @@
|
||||
from typing import Optional |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
|
||||
from .lora import LoRAModule |
||||
from .utils import masked_mean |
||||
|
||||
|
||||
class Critic(LoRAModule): |
||||
""" |
||||
Critic model base class. |
||||
|
||||
Args: |
||||
model (nn.Module): Critic model. |
||||
value_head (nn.Module): Value head to get value. |
||||
lora_rank (int): LoRA rank. |
||||
lora_train_bias (str): LoRA bias training mode. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
model: nn.Module, |
||||
value_head: nn.Module, |
||||
lora_rank: int = 0, |
||||
lora_train_bias: str = 'none') -> None: |
||||
|
||||
super().__init__(lora_rank=lora_rank, lora_train_bias=lora_train_bias) |
||||
self.model = model |
||||
self.value_head = value_head |
||||
self.convert_to_lora() |
||||
|
||||
def forward(self, |
||||
sequences: torch.LongTensor, |
||||
action_mask: Optional[torch.Tensor] = None, |
||||
attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
||||
outputs = self.model(sequences, attention_mask=attention_mask) |
||||
last_hidden_states = outputs['last_hidden_state'] |
||||
|
||||
values = self.value_head(last_hidden_states).squeeze(-1)[:, :-1] |
||||
|
||||
if action_mask is not None: |
||||
num_actions = action_mask.size(1) |
||||
values = values[:, -num_actions:] |
||||
value = masked_mean(values, action_mask, dim=1) |
||||
return value |
||||
value = values.mean(dim=1).squeeze(1) |
||||
return value |
||||
@ -0,0 +1,137 @@
|
||||
from typing import Any, Callable, Optional |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
|
||||
try: |
||||
from transformers.generation_logits_process import ( |
||||
LogitsProcessorList, |
||||
TemperatureLogitsWarper, |
||||
TopKLogitsWarper, |
||||
TopPLogitsWarper, |
||||
) |
||||
except ImportError: |
||||
from transformers.generation import LogitsProcessorList, TemperatureLogitsWarper, TopKLogitsWarper, TopPLogitsWarper |
||||
|
||||
|
||||
def prepare_logits_processor(top_k: Optional[int] = None, |
||||
top_p: Optional[float] = None, |
||||
temperature: Optional[float] = None) -> LogitsProcessorList: |
||||
processor_list = LogitsProcessorList() |
||||
if temperature is not None and temperature != 1.0: |
||||
processor_list.append(TemperatureLogitsWarper(temperature)) |
||||
if top_k is not None and top_k != 0: |
||||
processor_list.append(TopKLogitsWarper(top_k)) |
||||
if top_p is not None and top_p < 1.0: |
||||
processor_list.append(TopPLogitsWarper(top_p)) |
||||
return processor_list |
||||
|
||||
|
||||
def sample(model: nn.Module, |
||||
input_ids: torch.Tensor, |
||||
max_length: int, |
||||
early_stopping: bool = False, |
||||
eos_token_id: Optional[int] = None, |
||||
pad_token_id: Optional[int] = None, |
||||
top_k: Optional[int] = None, |
||||
top_p: Optional[float] = None, |
||||
temperature: Optional[float] = None, |
||||
prepare_inputs_fn: Optional[Callable[[torch.Tensor, Any], dict]] = None, |
||||
update_model_kwargs_fn: Optional[Callable[[dict, Any], dict]] = None, |
||||
**model_kwargs) -> torch.Tensor: |
||||
if input_ids.size(1) >= max_length: |
||||
return input_ids |
||||
|
||||
logits_processor = prepare_logits_processor(top_k, top_p, temperature) |
||||
unfinished_sequences = input_ids.new(input_ids.shape[0]).fill_(1) |
||||
|
||||
for _ in range(input_ids.size(1), max_length): |
||||
model_inputs = prepare_inputs_fn(input_ids, **model_kwargs) if prepare_inputs_fn is not None else { |
||||
'input_ids': input_ids |
||||
} |
||||
outputs = model(**model_inputs) |
||||
|
||||
next_token_logits = outputs['logits'][:, -1, :] |
||||
# pre-process distribution |
||||
next_token_logits = logits_processor(input_ids, next_token_logits) |
||||
# sample |
||||
probs = torch.softmax(next_token_logits, dim=-1, dtype=torch.float) |
||||
next_tokens = torch.multinomial(probs, num_samples=1).squeeze(1) |
||||
|
||||
# finished sentences should have their next token be a padding token |
||||
if eos_token_id is not None: |
||||
if pad_token_id is None: |
||||
raise ValueError("If `eos_token_id` is defined, make sure that `pad_token_id` is defined.") |
||||
next_tokens = next_tokens * unfinished_sequences + pad_token_id * (1 - unfinished_sequences) |
||||
|
||||
# update generated ids, model inputs for next step |
||||
input_ids = torch.cat([input_ids, next_tokens[:, None]], dim=-1) |
||||
if update_model_kwargs_fn is not None: |
||||
model_kwargs = update_model_kwargs_fn(outputs, **model_kwargs) |
||||
|
||||
# if eos_token was found in one sentence, set sentence to finished |
||||
if eos_token_id is not None: |
||||
unfinished_sequences = unfinished_sequences.mul((next_tokens != eos_token_id).long()) |
||||
|
||||
# stop when each sentence is finished if early_stopping=True |
||||
if early_stopping and unfinished_sequences.max() == 0: |
||||
break |
||||
|
||||
return input_ids |
||||
|
||||
|
||||
def generate(model: nn.Module, |
||||
input_ids: torch.Tensor, |
||||
max_length: int, |
||||
num_beams: int = 1, |
||||
do_sample: bool = True, |
||||
early_stopping: bool = False, |
||||
eos_token_id: Optional[int] = None, |
||||
pad_token_id: Optional[int] = None, |
||||
top_k: Optional[int] = None, |
||||
top_p: Optional[float] = None, |
||||
temperature: Optional[float] = None, |
||||
prepare_inputs_fn: Optional[Callable[[torch.Tensor, Any], dict]] = None, |
||||
update_model_kwargs_fn: Optional[Callable[[dict, Any], dict]] = None, |
||||
**model_kwargs) -> torch.Tensor: |
||||
"""Generate token sequence. The returned sequence is input_ids + generated_tokens. |
||||
|
||||
Args: |
||||
model (nn.Module): model |
||||
input_ids (torch.Tensor): input sequence |
||||
max_length (int): max length of the returned sequence |
||||
num_beams (int, optional): number of beams. Defaults to 1. |
||||
do_sample (bool, optional): whether to do sample. Defaults to True. |
||||
early_stopping (bool, optional): if True, the sequence length may be smaller than max_length due to finding eos. Defaults to False. |
||||
eos_token_id (Optional[int], optional): end of sequence token id. Defaults to None. |
||||
pad_token_id (Optional[int], optional): pad token id. Defaults to None. |
||||
top_k (Optional[int], optional): the number of highest probability vocabulary tokens to keep for top-k-filtering. Defaults to None. |
||||
top_p (Optional[float], optional): If set to float < 1, only the smallest set of most probable tokens with probabilities that add up to top_p or higher are kept for generation. Defaults to None. |
||||
temperature (Optional[float], optional): The value used to module the next token probabilities. Defaults to None. |
||||
prepare_inputs_fn (Optional[Callable[[torch.Tensor, Any], dict]], optional): Function to preprocess model inputs. Arguments of this function should be input_ids and model_kwargs. Defaults to None. |
||||
update_model_kwargs_fn (Optional[Callable[[dict, Any], dict]], optional): Function to update model_kwargs based on outputs. Arguments of this function should be outputs and model_kwargs. Defaults to None. |
||||
""" |
||||
is_greedy_gen_mode = ((num_beams == 1) and do_sample is False) |
||||
is_sample_gen_mode = ((num_beams == 1) and do_sample is True) |
||||
is_beam_gen_mode = ((num_beams > 1) and do_sample is False) |
||||
if is_greedy_gen_mode: |
||||
# run greedy search |
||||
raise NotImplementedError |
||||
elif is_sample_gen_mode: |
||||
# run sample |
||||
return sample(model, |
||||
input_ids, |
||||
max_length, |
||||
early_stopping=early_stopping, |
||||
eos_token_id=eos_token_id, |
||||
pad_token_id=pad_token_id, |
||||
top_k=top_k, |
||||
top_p=top_p, |
||||
temperature=temperature, |
||||
prepare_inputs_fn=prepare_inputs_fn, |
||||
update_model_kwargs_fn=update_model_kwargs_fn, |
||||
**model_kwargs) |
||||
elif is_beam_gen_mode: |
||||
raise NotImplementedError |
||||
else: |
||||
raise ValueError("Unsupported generation mode") |
||||
@ -0,0 +1,92 @@
|
||||
from typing import Optional |
||||
|
||||
import torch |
||||
|
||||
|
||||
def gpt_prepare_inputs_fn(input_ids: torch.Tensor, past: Optional[torch.Tensor] = None, **kwargs) -> dict: |
||||
token_type_ids = kwargs.get("token_type_ids", None) |
||||
# only last token for inputs_ids if past is defined in kwargs |
||||
if past: |
||||
input_ids = input_ids[:, -1].unsqueeze(-1) |
||||
if token_type_ids is not None: |
||||
token_type_ids = token_type_ids[:, -1].unsqueeze(-1) |
||||
|
||||
attention_mask = kwargs.get("attention_mask", None) |
||||
position_ids = kwargs.get("position_ids", None) |
||||
|
||||
if attention_mask is not None and position_ids is None: |
||||
# create position_ids on the fly for batch generation |
||||
position_ids = attention_mask.long().cumsum(-1) - 1 |
||||
position_ids.masked_fill_(attention_mask == 0, 1) |
||||
if past: |
||||
position_ids = position_ids[:, -1].unsqueeze(-1) |
||||
else: |
||||
position_ids = None |
||||
return { |
||||
"input_ids": input_ids, |
||||
"past_key_values": past, |
||||
"use_cache": kwargs.get("use_cache"), |
||||
"position_ids": position_ids, |
||||
"attention_mask": attention_mask, |
||||
"token_type_ids": token_type_ids, |
||||
} |
||||
|
||||
|
||||
def update_model_kwargs_fn(outputs: dict, **model_kwargs) -> dict: |
||||
if "past_key_values" in outputs: |
||||
model_kwargs["past"] = outputs["past_key_values"] |
||||
else: |
||||
model_kwargs["past"] = None |
||||
|
||||
# update token_type_ids with last value |
||||
if "token_type_ids" in model_kwargs: |
||||
token_type_ids = model_kwargs["token_type_ids"] |
||||
model_kwargs["token_type_ids"] = torch.cat([token_type_ids, token_type_ids[:, -1].unsqueeze(-1)], dim=-1) |
||||
|
||||
# update attention mask |
||||
if "attention_mask" in model_kwargs: |
||||
attention_mask = model_kwargs["attention_mask"] |
||||
model_kwargs["attention_mask"] = torch.cat( |
||||
[attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1) |
||||
|
||||
return model_kwargs |
||||
|
||||
|
||||
def opt_prepare_inputs_fn(input_ids: torch.Tensor, |
||||
past: Optional[torch.Tensor] = None, |
||||
attention_mask: Optional[torch.Tensor] = None, |
||||
use_cache: Optional[bool] = None, |
||||
**kwargs) -> dict: |
||||
# if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly |
||||
if attention_mask is None: |
||||
attention_mask = input_ids.new_ones(input_ids.shape) |
||||
|
||||
if past: |
||||
input_ids = input_ids[:, -1:] |
||||
# first step, decoder_cached_states are empty |
||||
return { |
||||
"input_ids": input_ids, # encoder_outputs is defined. input_ids not needed |
||||
"attention_mask": attention_mask, |
||||
"past_key_values": past, |
||||
"use_cache": use_cache, |
||||
} |
||||
|
||||
|
||||
def bloom_prepare_inputs_fn(input_ids: torch.Tensor, |
||||
past: Optional[torch.Tensor] = None, |
||||
attention_mask: Optional[torch.Tensor] = None, |
||||
use_cache: Optional[bool] = None, |
||||
**kwargs) -> dict: |
||||
# if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly |
||||
if attention_mask is None: |
||||
attention_mask = input_ids.new_ones(input_ids.shape) |
||||
|
||||
if past: |
||||
input_ids = input_ids[:, -1:] |
||||
# first step, decoder_cached_states are empty |
||||
return { |
||||
"input_ids": input_ids, # encoder_outputs is defined. input_ids not needed |
||||
"attention_mask": attention_mask, |
||||
"past_key_values": past, |
||||
"use_cache": use_cache, |
||||
} |
||||
@ -0,0 +1,31 @@
|
||||
from typing import Optional |
||||
|
||||
from transformers.models.gpt2.configuration_gpt2 import GPT2Config |
||||
from transformers.models.gpt2.modeling_gpt2 import GPT2LMHeadModel |
||||
|
||||
from .actor import Actor |
||||
|
||||
|
||||
class GPTActor(Actor): |
||||
""" |
||||
GPT Actor model. |
||||
|
||||
Args: |
||||
pretrained (str): Pretrained model name or path. |
||||
config (GPT2Config): Model config. |
||||
checkpoint (bool): Enable gradient checkpointing. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
pretrained: Optional[str] = None, |
||||
config: Optional[GPT2Config] = None, |
||||
checkpoint: bool = False) -> None: |
||||
if pretrained is not None: |
||||
model = GPT2LMHeadModel.from_pretrained(pretrained) |
||||
elif config is not None: |
||||
model = GPT2LMHeadModel(config) |
||||
else: |
||||
model = GPT2LMHeadModel(GPT2Config()) |
||||
if checkpoint: |
||||
model.gradient_checkpointing_enable() |
||||
super().__init__(model) |
||||
@ -0,0 +1,33 @@
|
||||
from typing import Optional |
||||
|
||||
import torch.nn as nn |
||||
from transformers.models.gpt2.configuration_gpt2 import GPT2Config |
||||
from transformers.models.gpt2.modeling_gpt2 import GPT2Model |
||||
|
||||
from .critic import Critic |
||||
|
||||
|
||||
class GPTCritic(Critic): |
||||
""" |
||||
GPT Critic model. |
||||
|
||||
Args: |
||||
pretrained (str): Pretrained model name or path. |
||||
config (GPT2Config): Model config. |
||||
checkpoint (bool): Enable gradient checkpointing. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
pretrained: Optional[str] = None, |
||||
config: Optional[GPT2Config] = None, |
||||
checkpoint: bool = False) -> None: |
||||
if pretrained is not None: |
||||
model = GPT2Model.from_pretrained(pretrained) |
||||
elif config is not None: |
||||
model = GPT2Model(config) |
||||
else: |
||||
model = GPT2Model(GPT2Config()) |
||||
if checkpoint: |
||||
model.gradient_checkpointing_enable() |
||||
value_head = nn.Linear(model.config.n_embd, 1) |
||||
super().__init__(model, value_head) |
||||
@ -0,0 +1,33 @@
|
||||
from typing import Optional |
||||
|
||||
import torch.nn as nn |
||||
from transformers.models.gpt2.configuration_gpt2 import GPT2Config |
||||
from transformers.models.gpt2.modeling_gpt2 import GPT2Model |
||||
|
||||
from .reward_model import RewardModel |
||||
|
||||
|
||||
class GPTRM(RewardModel): |
||||
""" |
||||
GPT Reward model. |
||||
|
||||
Args: |
||||
pretrained (str): Pretrained model name or path. |
||||
config (GPT2Config): Model config. |
||||
checkpoint (bool): Enable gradient checkpointing. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
pretrained: Optional[str] = None, |
||||
config: Optional[GPT2Config] = None, |
||||
checkpoint: bool = False) -> None: |
||||
if pretrained is not None: |
||||
model = GPT2Model.from_pretrained(pretrained) |
||||
elif config is not None: |
||||
model = GPT2Model(config) |
||||
else: |
||||
model = GPT2Model(GPT2Config()) |
||||
if checkpoint: |
||||
model.gradient_checkpointing_enable() |
||||
value_head = nn.Linear(model.config.n_embd, 1) |
||||
super().__init__(model, value_head) |
||||
@ -0,0 +1,127 @@
|
||||
import math |
||||
from typing import Optional |
||||
|
||||
import loralib as lora |
||||
import torch |
||||
import torch.nn as nn |
||||
import torch.nn.functional as F |
||||
|
||||
|
||||
class LoraLinear(lora.LoRALayer, nn.Module): |
||||
"""Replace in-place ops to out-of-place ops to fit gemini. Convert a torch.nn.Linear to LoraLinear. |
||||
""" |
||||
|
||||
def __init__( |
||||
self, |
||||
weight: nn.Parameter, |
||||
bias: Optional[nn.Parameter], |
||||
r: int = 0, |
||||
lora_alpha: int = 1, |
||||
lora_dropout: float = 0., |
||||
fan_in_fan_out: bool = False, # Set this to True if the layer to replace stores weight like (fan_in, fan_out) |
||||
merge_weights: bool = True, |
||||
): |
||||
nn.Module.__init__(self) |
||||
lora.LoRALayer.__init__(self, |
||||
r=r, |
||||
lora_alpha=lora_alpha, |
||||
lora_dropout=lora_dropout, |
||||
merge_weights=merge_weights) |
||||
self.weight = weight |
||||
self.bias = bias |
||||
|
||||
out_features, in_features = weight.shape |
||||
self.in_features = in_features |
||||
self.out_features = out_features |
||||
|
||||
self.fan_in_fan_out = fan_in_fan_out |
||||
# Actual trainable parameters |
||||
if r > 0: |
||||
self.lora_A = nn.Parameter(self.weight.new_zeros((r, in_features))) |
||||
self.lora_B = nn.Parameter(self.weight.new_zeros((out_features, r))) |
||||
self.scaling = self.lora_alpha / self.r |
||||
# Freezing the pre-trained weight matrix |
||||
self.weight.requires_grad = False |
||||
self.reset_parameters() |
||||
if fan_in_fan_out: |
||||
self.weight.data = self.weight.data.T |
||||
|
||||
def reset_parameters(self): |
||||
if hasattr(self, 'lora_A'): |
||||
# initialize A the same way as the default for nn.Linear and B to zero |
||||
nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) |
||||
nn.init.zeros_(self.lora_B) |
||||
|
||||
def train(self, mode: bool = True): |
||||
|
||||
def T(w): |
||||
return w.T if self.fan_in_fan_out else w |
||||
|
||||
nn.Module.train(self, mode) |
||||
if self.merge_weights and self.merged: |
||||
# Make sure that the weights are not merged |
||||
if self.r > 0: |
||||
self.weight.data -= T(self.lora_B @ self.lora_A) * self.scaling |
||||
self.merged = False |
||||
|
||||
def eval(self): |
||||
|
||||
def T(w): |
||||
return w.T if self.fan_in_fan_out else w |
||||
|
||||
nn.Module.eval(self) |
||||
if self.merge_weights and not self.merged: |
||||
# Merge the weights and mark it |
||||
if self.r > 0: |
||||
self.weight.data += T(self.lora_B @ self.lora_A) * self.scaling |
||||
self.merged = True |
||||
|
||||
def forward(self, x: torch.Tensor): |
||||
|
||||
def T(w): |
||||
return w.T if self.fan_in_fan_out else w |
||||
|
||||
if self.r > 0 and not self.merged: |
||||
result = F.linear(x, T(self.weight), bias=self.bias) |
||||
if self.r > 0: |
||||
result = result + (self.lora_dropout(x) @ self.lora_A.t() @ self.lora_B.t()) * self.scaling |
||||
return result |
||||
else: |
||||
return F.linear(x, T(self.weight), bias=self.bias) |
||||
|
||||
|
||||
def lora_linear_wrapper(linear: nn.Linear, lora_rank: int) -> LoraLinear: |
||||
assert lora_rank <= linear.in_features, f'LoRA rank ({lora_rank}) must be less than or equal to in features ({linear.in_features})' |
||||
lora_linear = LoraLinear(linear.weight, linear.bias, r=lora_rank, merge_weights=False) |
||||
return lora_linear |
||||
|
||||
|
||||
def convert_to_lora_recursively(module: nn.Module, lora_rank: int) -> None: |
||||
for name, child in module.named_children(): |
||||
if isinstance(child, nn.Linear): |
||||
setattr(module, name, lora_linear_wrapper(child, lora_rank)) |
||||
else: |
||||
convert_to_lora_recursively(child, lora_rank) |
||||
|
||||
|
||||
class LoRAModule(nn.Module): |
||||
"""A LoRA module base class. All derived classes should call `convert_to_lora()` at the bottom of `__init__()`. |
||||
This calss will convert all torch.nn.Linear layer to LoraLinear layer. |
||||
|
||||
Args: |
||||
lora_rank (int, optional): LoRA rank. 0 means LoRA is not applied. Defaults to 0. |
||||
lora_train_bias (str, optional): Whether LoRA train biases. |
||||
'none' means it doesn't train biases. 'all' means it trains all biases. 'lora_only' means it only trains biases of LoRA layers. |
||||
Defaults to 'none'. |
||||
""" |
||||
|
||||
def __init__(self, lora_rank: int = 0, lora_train_bias: str = 'none') -> None: |
||||
super().__init__() |
||||
self.lora_rank = lora_rank |
||||
self.lora_train_bias = lora_train_bias |
||||
|
||||
def convert_to_lora(self) -> None: |
||||
if self.lora_rank <= 0: |
||||
return |
||||
convert_to_lora_recursively(self, self.lora_rank) |
||||
lora.mark_only_lora_as_trainable(self, self.lora_train_bias) |
||||
@ -0,0 +1,105 @@
|
||||
from typing import Optional |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
|
||||
from .utils import masked_mean |
||||
|
||||
|
||||
class GPTLMLoss(nn.Module): |
||||
""" |
||||
GPT Language Model Loss |
||||
""" |
||||
|
||||
def __init__(self): |
||||
super().__init__() |
||||
self.loss = nn.CrossEntropyLoss() |
||||
|
||||
def forward(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: |
||||
shift_logits = logits[..., :-1, :].contiguous() |
||||
shift_labels = labels[..., 1:].contiguous() |
||||
# Flatten the tokens |
||||
return self.loss(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1)) |
||||
|
||||
|
||||
class PolicyLoss(nn.Module): |
||||
""" |
||||
Policy Loss for PPO |
||||
""" |
||||
|
||||
def __init__(self, clip_eps: float = 0.2) -> None: |
||||
super().__init__() |
||||
self.clip_eps = clip_eps |
||||
|
||||
def forward(self, |
||||
log_probs: torch.Tensor, |
||||
old_log_probs: torch.Tensor, |
||||
advantages: torch.Tensor, |
||||
action_mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
||||
ratio = (log_probs - old_log_probs).exp() |
||||
surr1 = ratio * advantages |
||||
surr2 = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * advantages |
||||
loss = -torch.min(surr1, surr2) |
||||
if action_mask is not None: |
||||
loss = masked_mean(loss, action_mask) |
||||
loss = loss.mean() |
||||
return loss |
||||
|
||||
|
||||
class ValueLoss(nn.Module): |
||||
""" |
||||
Value Loss for PPO |
||||
""" |
||||
|
||||
def __init__(self, clip_eps: float = 0.4) -> None: |
||||
super().__init__() |
||||
self.clip_eps = clip_eps |
||||
|
||||
def forward(self, |
||||
values: torch.Tensor, |
||||
old_values: torch.Tensor, |
||||
reward: torch.Tensor, |
||||
action_mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
||||
values_clipped = old_values + (values - old_values).clamp(-self.clip_eps, self.clip_eps) |
||||
surr1 = (values_clipped - reward)**2 |
||||
surr2 = (values - reward)**2 |
||||
loss = torch.max(surr1, surr2) |
||||
loss = loss.mean() |
||||
return loss |
||||
|
||||
|
||||
class PPOPtxActorLoss(nn.Module): |
||||
""" |
||||
To Do: |
||||
|
||||
PPO-ptx Actor Loss |
||||
""" |
||||
|
||||
def __init__(self, policy_clip_eps: float = 0.2, pretrain_coef: float = 0.0, pretrain_loss_fn=GPTLMLoss()) -> None: |
||||
super().__init__() |
||||
self.pretrain_coef = pretrain_coef |
||||
self.policy_loss_fn = PolicyLoss(clip_eps=policy_clip_eps) |
||||
self.pretrain_loss_fn = pretrain_loss_fn |
||||
|
||||
def forward(self, |
||||
log_probs: torch.Tensor, |
||||
old_log_probs: torch.Tensor, |
||||
advantages: torch.Tensor, |
||||
lm_logits: torch.Tensor, |
||||
lm_input_ids: torch.Tensor, |
||||
action_mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
||||
policy_loss = self.policy_loss_fn(log_probs, old_log_probs, advantages, action_mask=action_mask) |
||||
lm_loss = self.pretrain_loss_fn(lm_logits, lm_input_ids) |
||||
return policy_loss + self.pretrain_coef * lm_loss |
||||
|
||||
|
||||
class PairWiseLoss(nn.Module): |
||||
""" |
||||
Pairwise Loss for Reward Model |
||||
""" |
||||
|
||||
def forward(self, chosen_reward: torch.Tensor, reject_reward: torch.Tensor) -> torch.Tensor: |
||||
probs = torch.sigmoid(chosen_reward - reject_reward) |
||||
log_probs = torch.log(probs) |
||||
loss = -log_probs.mean() |
||||
return loss |
||||
@ -0,0 +1,35 @@
|
||||
from typing import Optional |
||||
|
||||
from transformers.models.opt.configuration_opt import OPTConfig |
||||
from transformers.models.opt.modeling_opt import OPTForCausalLM |
||||
|
||||
from .actor import Actor |
||||
|
||||
|
||||
class OPTActor(Actor): |
||||
""" |
||||
OPT Actor model. |
||||
|
||||
Args: |
||||
pretrained (str): Pretrained model name or path. |
||||
config (OPTConfig): Model config. |
||||
checkpoint (bool): Enable gradient checkpointing. |
||||
lora_rank (int): Rank of the low-rank approximation. |
||||
lora_train_bias (str): LoRA bias training mode. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
pretrained: Optional[str] = None, |
||||
config: Optional[OPTConfig] = None, |
||||
checkpoint: bool = False, |
||||
lora_rank: int = 0, |
||||
lora_train_bias: str = 'none') -> None: |
||||
if pretrained is not None: |
||||
model = OPTForCausalLM.from_pretrained(pretrained) |
||||
elif config is not None: |
||||
model = OPTForCausalLM(config) |
||||
else: |
||||
model = OPTForCausalLM(OPTConfig()) |
||||
if checkpoint: |
||||
model.gradient_checkpointing_enable() |
||||
super().__init__(model, lora_rank, lora_train_bias) |
||||
@ -0,0 +1,37 @@
|
||||
from typing import Optional |
||||
|
||||
import torch.nn as nn |
||||
from transformers.models.opt.configuration_opt import OPTConfig |
||||
from transformers.models.opt.modeling_opt import OPTModel |
||||
|
||||
from .critic import Critic |
||||
|
||||
|
||||
class OPTCritic(Critic): |
||||
""" |
||||
OPT Critic model. |
||||
|
||||
Args: |
||||
pretrained (str): Pretrained model name or path. |
||||
config (OPTConfig): Model config. |
||||
checkpoint (bool): Enable gradient checkpointing. |
||||
lora_rank (int): Rank of the low-rank approximation. |
||||
lora_train_bias (str): LoRA bias training mode. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
pretrained: Optional[str] = None, |
||||
config: Optional[OPTConfig] = None, |
||||
checkpoint: bool = False, |
||||
lora_rank: int = 0, |
||||
lora_train_bias: str = 'none') -> None: |
||||
if pretrained is not None: |
||||
model = OPTModel.from_pretrained(pretrained) |
||||
elif config is not None: |
||||
model = OPTModel(config) |
||||
else: |
||||
model = OPTModel(OPTConfig()) |
||||
if checkpoint: |
||||
model.gradient_checkpointing_enable() |
||||
value_head = nn.Linear(model.config.hidden_size, 1) |
||||
super().__init__(model, value_head, lora_rank, lora_train_bias) |
||||
@ -0,0 +1,33 @@
|
||||
from typing import Optional |
||||
|
||||
import torch.nn as nn |
||||
from transformers.models.opt.configuration_opt import OPTConfig |
||||
from transformers.models.opt.modeling_opt import OPTModel |
||||
|
||||
from .reward_model import RewardModel |
||||
|
||||
|
||||
class OPTRM(RewardModel): |
||||
""" |
||||
OPT Reward model. |
||||
|
||||
Args: |
||||
pretrained (str): Pretrained model name or path. |
||||
config (OPTConfig): Model config. |
||||
lora_rank (int): Rank of the low-rank approximation. |
||||
lora_train_bias (str): LoRA bias training mode. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
pretrained: Optional[str] = None, |
||||
config: Optional[OPTConfig] = None, |
||||
lora_rank: int = 0, |
||||
lora_train_bias: str = 'none') -> None: |
||||
if pretrained is not None: |
||||
model = OPTModel.from_pretrained(pretrained) |
||||
elif config is not None: |
||||
model = OPTModel(config) |
||||
else: |
||||
model = OPTModel(OPTConfig()) |
||||
value_head = nn.Linear(model.config.hidden_size, 1) |
||||
super().__init__(model, value_head, lora_rank, lora_train_bias) |
||||
@ -0,0 +1,41 @@
|
||||
from typing import Optional |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
|
||||
from .lora import LoRAModule |
||||
|
||||
|
||||
class RewardModel(LoRAModule): |
||||
""" |
||||
Reward model base class. |
||||
|
||||
Args: |
||||
model (nn.Module): Reward model. |
||||
value_head (nn.Module): Value head to get reward score. |
||||
lora_rank (int): LoRA rank. |
||||
lora_train_bias (str): LoRA bias training mode. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
model: nn.Module, |
||||
value_head: Optional[nn.Module] = None, |
||||
lora_rank: int = 0, |
||||
lora_train_bias: str = 'none') -> None: |
||||
super().__init__(lora_rank=lora_rank, lora_train_bias=lora_train_bias) |
||||
self.model = model |
||||
if value_head is not None: |
||||
if value_head.out_features != 1: |
||||
raise ValueError("The value head of reward model's output dim should be 1!") |
||||
self.value_head = value_head |
||||
|
||||
else: |
||||
self.value_head = nn.Linear(model.config.n_embd, 1) |
||||
self.convert_to_lora() |
||||
|
||||
def forward(self, sequences: torch.LongTensor, attention_mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
||||
outputs = self.model(sequences, attention_mask=attention_mask) |
||||
last_hidden_states = outputs['last_hidden_state'] |
||||
values = self.value_head(last_hidden_states)[:, :-1] |
||||
value = values.mean(dim=1).squeeze(1) # ensure shape is (B) |
||||
return value |
||||
@ -0,0 +1,92 @@
|
||||
from typing import Optional, Union |
||||
|
||||
import loralib as lora |
||||
import torch |
||||
import torch.nn as nn |
||||
import torch.nn.functional as F |
||||
|
||||
|
||||
def compute_approx_kl(log_probs: torch.Tensor, |
||||
log_probs_base: torch.Tensor, |
||||
action_mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
||||
""" |
||||
Compute the approximate KL divergence between two distributions. |
||||
Schulman blog: http://joschu.net/blog/kl-approx.html |
||||
|
||||
Args: |
||||
log_probs: Log probabilities of the new distribution. |
||||
log_probs_base: Log probabilities of the base distribution. |
||||
action_mask: Mask for actions. |
||||
""" |
||||
|
||||
log_ratio = log_probs - log_probs_base |
||||
approx_kl = (log_ratio.exp() - 1) - log_ratio |
||||
if action_mask is not None: |
||||
approx_kl = masked_mean(approx_kl, action_mask, dim=1) |
||||
return approx_kl |
||||
approx_kl = approx_kl.mean(dim=1) |
||||
return approx_kl |
||||
|
||||
|
||||
def compute_reward(r: Union[torch.Tensor, float], |
||||
kl_coef: float, |
||||
log_probs: torch.Tensor, |
||||
log_probs_base: torch.Tensor, |
||||
action_mask: Optional[torch.Tensor] = None) -> torch.Tensor: |
||||
if kl_coef <= 0.0: |
||||
return r |
||||
kl = compute_approx_kl(log_probs, log_probs_base, action_mask=action_mask) |
||||
reward = r - kl_coef * kl |
||||
return reward |
||||
|
||||
|
||||
def log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: |
||||
log_probs = F.log_softmax(logits, dim=-1) |
||||
log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)) |
||||
return log_probs_labels.squeeze(-1) |
||||
|
||||
|
||||
def masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor: |
||||
tensor = tensor * mask |
||||
tensor = tensor.sum(dim=dim) |
||||
mask_sum = mask.sum(dim=dim) |
||||
mean = tensor / (mask_sum + 1e-8) |
||||
return mean |
||||
|
||||
|
||||
def masked_normalize(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1, eps: float = 1e-8) -> torch.Tensor: |
||||
tensor = tensor * mask |
||||
mean = masked_mean(tensor, mask, dim=dim) |
||||
mean_centered = tensor - mean |
||||
var = masked_mean(mean_centered**2, mask, dim=dim) |
||||
return mean_centered * var.clamp(min=eps).rsqrt() |
||||
|
||||
|
||||
def normalize(tensor: torch.Tensor, dim: int = 0, eps: float = 1e-8) -> torch.Tensor: |
||||
mean = tensor.mean(dim) |
||||
mean_centered = tensor - mean |
||||
var = (mean_centered**2).mean(dim) |
||||
norm = mean_centered * var.clamp(min=eps).rsqrt() |
||||
return norm |
||||
|
||||
|
||||
def convert_to_lora(model: nn.Module, |
||||
input_size: int, |
||||
output_size: int, |
||||
lora_rank: int = 16, |
||||
lora_alpha: int = 1, |
||||
lora_dropout: float = 0., |
||||
fan_in_fan_out: bool = False, |
||||
merge_weights: bool = True): |
||||
if lora_rank > min(input_size, output_size): |
||||
raise ValueError(f"LoRA rank {lora_rank} must be less or equal than {min(input_size, output_size)}") |
||||
|
||||
for name, module in model.named_modules(): |
||||
if isinstance(module, nn.Linear): |
||||
module._modules[name] = lora.Linear(input_size, |
||||
output_size, |
||||
r=lora_rank, |
||||
lora_alpha=lora_alpha, |
||||
lora_dropout=lora_dropout, |
||||
fan_in_fan_out=fan_in_fan_out, |
||||
merge_weights=merge_weights) |
||||
@ -0,0 +1,4 @@
|
||||
from .base import ReplayBuffer |
||||
from .naive import NaiveReplayBuffer |
||||
|
||||
__all__ = ['ReplayBuffer', 'NaiveReplayBuffer'] |
||||
@ -0,0 +1,43 @@
|
||||
from abc import ABC, abstractmethod |
||||
from typing import Any |
||||
|
||||
from chatgpt.experience_maker.base import Experience |
||||
|
||||
|
||||
class ReplayBuffer(ABC): |
||||
"""Replay buffer base class. It stores experience. |
||||
|
||||
Args: |
||||
sample_batch_size (int): Batch size when sampling. |
||||
limit (int, optional): Limit of number of experience samples. A number <= 0 means unlimited. Defaults to 0. |
||||
""" |
||||
|
||||
def __init__(self, sample_batch_size: int, limit: int = 0) -> None: |
||||
super().__init__() |
||||
self.sample_batch_size = sample_batch_size |
||||
# limit <= 0 means unlimited |
||||
self.limit = limit |
||||
|
||||
@abstractmethod |
||||
def append(self, experience: Experience) -> None: |
||||
pass |
||||
|
||||
@abstractmethod |
||||
def clear(self) -> None: |
||||
pass |
||||
|
||||
@abstractmethod |
||||
def sample(self) -> Experience: |
||||
pass |
||||
|
||||
@abstractmethod |
||||
def __len__(self) -> int: |
||||
pass |
||||
|
||||
@abstractmethod |
||||
def __getitem__(self, idx: int) -> Any: |
||||
pass |
||||
|
||||
@abstractmethod |
||||
def collate_fn(self, batch: Any) -> Experience: |
||||
pass |
||||
@ -0,0 +1,57 @@
|
||||
import random |
||||
from typing import List |
||||
|
||||
import torch |
||||
from chatgpt.experience_maker.base import Experience |
||||
|
||||
from .base import ReplayBuffer |
||||
from .utils import BufferItem, make_experience_batch, split_experience_batch |
||||
|
||||
|
||||
class NaiveReplayBuffer(ReplayBuffer): |
||||
"""Naive replay buffer class. It stores experience. |
||||
|
||||
Args: |
||||
sample_batch_size (int): Batch size when sampling. |
||||
limit (int, optional): Limit of number of experience samples. A number <= 0 means unlimited. Defaults to 0. |
||||
cpu_offload (bool, optional): Whether to offload experience to cpu when sampling. Defaults to True. |
||||
""" |
||||
|
||||
def __init__(self, sample_batch_size: int, limit: int = 0, cpu_offload: bool = True) -> None: |
||||
super().__init__(sample_batch_size, limit) |
||||
self.cpu_offload = cpu_offload |
||||
self.target_device = torch.device(f'cuda:{torch.cuda.current_device()}') |
||||
# TODO(ver217): add prefetch |
||||
self.items: List[BufferItem] = [] |
||||
|
||||
@torch.no_grad() |
||||
def append(self, experience: Experience) -> None: |
||||
if self.cpu_offload: |
||||
experience.to_device(torch.device('cpu')) |
||||
items = split_experience_batch(experience) |
||||
self.items.extend(items) |
||||
if self.limit > 0: |
||||
samples_to_remove = len(self.items) - self.limit |
||||
if samples_to_remove > 0: |
||||
self.items = self.items[samples_to_remove:] |
||||
|
||||
def clear(self) -> None: |
||||
self.items.clear() |
||||
|
||||
@torch.no_grad() |
||||
def sample(self) -> Experience: |
||||
items = random.sample(self.items, self.sample_batch_size) |
||||
experience = make_experience_batch(items) |
||||
if self.cpu_offload: |
||||
experience.to_device(self.target_device) |
||||
return experience |
||||
|
||||
def __len__(self) -> int: |
||||
return len(self.items) |
||||
|
||||
def __getitem__(self, idx: int) -> BufferItem: |
||||
return self.items[idx] |
||||
|
||||
def collate_fn(self, batch) -> Experience: |
||||
experience = make_experience_batch(batch) |
||||
return experience |
||||
@ -0,0 +1,73 @@
|
||||
from dataclasses import dataclass |
||||
from typing import List, Optional |
||||
|
||||
import torch |
||||
import torch.nn.functional as F |
||||
from chatgpt.experience_maker.base import Experience |
||||
|
||||
|
||||
@dataclass |
||||
class BufferItem: |
||||
"""BufferItem is an item of experience data. |
||||
|
||||
Shapes of each tensor: |
||||
sequences: (S) |
||||
action_log_probs: (A) |
||||
values: (1) |
||||
reward: (1) |
||||
advatanges: (1) |
||||
attention_mask: (S) |
||||
action_mask: (A) |
||||
|
||||
"A" is the number of actions. |
||||
""" |
||||
sequences: torch.Tensor |
||||
action_log_probs: torch.Tensor |
||||
values: torch.Tensor |
||||
reward: torch.Tensor |
||||
advantages: torch.Tensor |
||||
attention_mask: Optional[torch.LongTensor] |
||||
action_mask: Optional[torch.BoolTensor] |
||||
|
||||
|
||||
def split_experience_batch(experience: Experience) -> List[BufferItem]: |
||||
batch_size = experience.sequences.size(0) |
||||
batch_kwargs = [{} for _ in range(batch_size)] |
||||
keys = ('sequences', 'action_log_probs', 'values', 'reward', 'advantages', 'attention_mask', 'action_mask') |
||||
for key in keys: |
||||
value = getattr(experience, key) |
||||
if isinstance(value, torch.Tensor): |
||||
vals = torch.unbind(value) |
||||
else: |
||||
# None |
||||
vals = [value for _ in range(batch_size)] |
||||
assert batch_size == len(vals) |
||||
for i, v in enumerate(vals): |
||||
batch_kwargs[i][key] = v |
||||
items = [BufferItem(**kwargs) for kwargs in batch_kwargs] |
||||
return items |
||||
|
||||
|
||||
def zero_pad_sequences(sequences: List[torch.Tensor], side: str = 'left') -> torch.Tensor: |
||||
assert side in ('left', 'right') |
||||
max_len = max(seq.size(0) for seq in sequences) |
||||
padded_sequences = [] |
||||
for seq in sequences: |
||||
pad_len = max_len - seq.size(0) |
||||
padding = (pad_len, 0) if side == 'left' else (0, pad_len) |
||||
padded_sequences.append(F.pad(seq, padding)) |
||||
return torch.stack(padded_sequences, dim=0) |
||||
|
||||
|
||||
def make_experience_batch(items: List[BufferItem]) -> Experience: |
||||
kwargs = {} |
||||
to_pad_keys = set(('action_log_probs', 'action_mask')) |
||||
keys = ('sequences', 'action_log_probs', 'values', 'reward', 'advantages', 'attention_mask', 'action_mask') |
||||
for key in keys: |
||||
vals = [getattr(item, key) for item in items] |
||||
if key in to_pad_keys: |
||||
batch_data = zero_pad_sequences(vals) |
||||
else: |
||||
batch_data = torch.stack(vals, dim=0) |
||||
kwargs[key] = batch_data |
||||
return Experience(**kwargs) |
||||
@ -0,0 +1,5 @@
|
||||
from .base import Trainer |
||||
from .ppo import PPOTrainer |
||||
from .rm import RewardModelTrainer |
||||
|
||||
__all__ = ['Trainer', 'PPOTrainer', 'RewardModelTrainer'] |
||||
@ -0,0 +1,162 @@
|
||||
import random |
||||
from abc import ABC, abstractmethod |
||||
from typing import Any, Callable, Dict, List, Optional, Union |
||||
|
||||
import torch |
||||
from chatgpt.experience_maker import Experience, ExperienceMaker |
||||
from chatgpt.replay_buffer import ReplayBuffer |
||||
from torch import Tensor |
||||
from torch.utils.data import DistributedSampler |
||||
from tqdm import tqdm |
||||
|
||||
from .callbacks import Callback |
||||
from .strategies import Strategy |
||||
from .utils import is_rank_0 |
||||
|
||||
|
||||
class Trainer(ABC): |
||||
""" |
||||
Base class for rlhf trainers. |
||||
|
||||
Args: |
||||
strategy (Strategy):the strategy to use for training |
||||
experience_maker (ExperienceMaker): the experience maker to use for produce experience to fullfill replay buffer |
||||
replay_buffer (ReplayBuffer): the replay buffer to use for training |
||||
experience_batch_size (int, defaults to 8): the batch size to use for experience generation |
||||
max_epochs (int, defaults to 1): the number of epochs of training process |
||||
tokenizer (Callable, optional): the tokenizer to use for tokenizing the input |
||||
sample_replay_buffer (bool, defaults to False): whether to sample from replay buffer |
||||
data_loader_pin_memory (bool, defaults to True): whether to pin memory for data loader |
||||
callbacks (List[Callback], defaults to []): the callbacks to call during training process |
||||
generate_kwargs (dict, optional): the kwargs to use while model generating |
||||
""" |
||||
|
||||
def __init__(self, |
||||
strategy: Strategy, |
||||
experience_maker: ExperienceMaker, |
||||
replay_buffer: ReplayBuffer, |
||||
experience_batch_size: int = 8, |
||||
max_epochs: int = 1, |
||||
tokenizer: Optional[Callable[[Any], dict]] = None, |
||||
sample_replay_buffer: bool = False, |
||||
dataloader_pin_memory: bool = True, |
||||
callbacks: List[Callback] = [], |
||||
**generate_kwargs) -> None: |
||||
super().__init__() |
||||
self.strategy = strategy |
||||
self.experience_maker = experience_maker |
||||
self.replay_buffer = replay_buffer |
||||
self.experience_batch_size = experience_batch_size |
||||
self.max_epochs = max_epochs |
||||
self.tokenizer = tokenizer |
||||
self.generate_kwargs = generate_kwargs |
||||
self.sample_replay_buffer = sample_replay_buffer |
||||
self.dataloader_pin_memory = dataloader_pin_memory |
||||
self.callbacks = callbacks |
||||
|
||||
@abstractmethod |
||||
def training_step(self, experience: Experience) -> Dict[str, Any]: |
||||
pass |
||||
|
||||
def _make_experience(self, inputs: Union[Tensor, Dict[str, Tensor]]) -> Experience: |
||||
if isinstance(inputs, Tensor): |
||||
return self.experience_maker.make_experience(inputs, **self.generate_kwargs) |
||||
elif isinstance(inputs, dict): |
||||
return self.experience_maker.make_experience(**inputs, **self.generate_kwargs) |
||||
else: |
||||
raise ValueError(f'Unsupported input type "{type(inputs)}"') |
||||
|
||||
def _sample_prompts(self, prompts) -> list: |
||||
indices = list(range(len(prompts))) |
||||
sampled_indices = random.sample(indices, self.experience_batch_size) |
||||
return [prompts[i] for i in sampled_indices] |
||||
|
||||
def _learn(self): |
||||
# replay buffer may be empty at first, we should rebuild at each training |
||||
if not self.sample_replay_buffer: |
||||
dataloader = self.strategy.setup_dataloader(self.replay_buffer, self.dataloader_pin_memory) |
||||
device = torch.cuda.current_device() |
||||
if self.sample_replay_buffer: |
||||
pbar = tqdm(range(self.max_epochs), desc='Train epoch', disable=not is_rank_0()) |
||||
for _ in pbar: |
||||
experience = self.replay_buffer.sample() |
||||
metrics = self.training_step(experience) |
||||
pbar.set_postfix(metrics) |
||||
else: |
||||
for epoch in range(self.max_epochs): |
||||
self._on_learn_epoch_start(epoch) |
||||
if isinstance(dataloader.sampler, DistributedSampler): |
||||
dataloader.sampler.set_epoch(epoch) |
||||
pbar = tqdm(dataloader, desc=f'Train epoch [{epoch+1}/{self.max_epochs}]', disable=not is_rank_0()) |
||||
for experience in pbar: |
||||
self._on_learn_batch_start() |
||||
experience.to_device(device) |
||||
metrics = self.training_step(experience) |
||||
self._on_learn_batch_end(metrics, experience) |
||||
pbar.set_postfix(metrics) |
||||
self._on_learn_epoch_end(epoch) |
||||
|
||||
def fit(self, prompts, num_episodes: int = 50000, max_timesteps: int = 500, update_timesteps: int = 5000) -> None: |
||||
time = 0 |
||||
self._on_fit_start() |
||||
for episode in range(num_episodes): |
||||
self._on_episode_start(episode) |
||||
for timestep in tqdm(range(max_timesteps), |
||||
desc=f'Episode [{episode+1}/{num_episodes}]', |
||||
disable=not is_rank_0()): |
||||
time += 1 |
||||
rand_prompts = self._sample_prompts(prompts) |
||||
if self.tokenizer is not None: |
||||
inputs = self.tokenizer(rand_prompts) |
||||
else: |
||||
inputs = rand_prompts |
||||
self._on_make_experience_start() |
||||
experience = self._make_experience(inputs) |
||||
self._on_make_experience_end(experience) |
||||
self.replay_buffer.append(experience) |
||||
if time % update_timesteps == 0: |
||||
self._learn() |
||||
self.replay_buffer.clear() |
||||
self._on_episode_end(episode) |
||||
self._on_fit_end() |
||||
|
||||
# TODO(ver217): maybe simplify these code using context |
||||
def _on_fit_start(self) -> None: |
||||
for callback in self.callbacks: |
||||
callback.on_fit_start() |
||||
|
||||
def _on_fit_end(self) -> None: |
||||
for callback in self.callbacks: |
||||
callback.on_fit_end() |
||||
|
||||
def _on_episode_start(self, episode: int) -> None: |
||||
for callback in self.callbacks: |
||||
callback.on_episode_start(episode) |
||||
|
||||
def _on_episode_end(self, episode: int) -> None: |
||||
for callback in self.callbacks: |
||||
callback.on_episode_end(episode) |
||||
|
||||
def _on_make_experience_start(self) -> None: |
||||
for callback in self.callbacks: |
||||
callback.on_make_experience_start() |
||||
|
||||
def _on_make_experience_end(self, experience: Experience) -> None: |
||||
for callback in self.callbacks: |
||||
callback.on_make_experience_end(experience) |
||||
|
||||
def _on_learn_epoch_start(self, epoch: int) -> None: |
||||
for callback in self.callbacks: |
||||
callback.on_learn_epoch_start(epoch) |
||||
|
||||
def _on_learn_epoch_end(self, epoch: int) -> None: |
||||
for callback in self.callbacks: |
||||
callback.on_learn_epoch_end(epoch) |
||||
|
||||
def _on_learn_batch_start(self) -> None: |
||||
for callback in self.callbacks: |
||||
callback.on_learn_batch_start() |
||||
|
||||
def _on_learn_batch_end(self, metrics: dict, experience: Experience) -> None: |
||||
for callback in self.callbacks: |
||||
callback.on_learn_batch_end(metrics, experience) |
||||
@ -0,0 +1,4 @@
|
||||
from .base import Callback |
||||
from .performance_evaluator import PerformanceEvaluator |
||||
|
||||
__all__ = ['Callback', 'PerformanceEvaluator'] |
||||
@ -0,0 +1,39 @@
|
||||
from abc import ABC |
||||
|
||||
from chatgpt.experience_maker import Experience |
||||
|
||||
|
||||
class Callback(ABC): |
||||
""" |
||||
Base callback class. It defines the interface for callbacks. |
||||
""" |
||||
|
||||
def on_fit_start(self) -> None: |
||||
pass |
||||
|
||||
def on_fit_end(self) -> None: |
||||
pass |
||||
|
||||
def on_episode_start(self, episode: int) -> None: |
||||
pass |
||||
|
||||
def on_episode_end(self, episode: int) -> None: |
||||
pass |
||||
|
||||
def on_make_experience_start(self) -> None: |
||||
pass |
||||
|
||||
def on_make_experience_end(self, experience: Experience) -> None: |
||||
pass |
||||
|
||||
def on_learn_epoch_start(self, epoch: int) -> None: |
||||
pass |
||||
|
||||
def on_learn_epoch_end(self, epoch: int) -> None: |
||||
pass |
||||
|
||||
def on_learn_batch_start(self) -> None: |
||||
pass |
||||
|
||||
def on_learn_batch_end(self, metrics: dict, experience: Experience) -> None: |
||||
pass |
||||
@ -0,0 +1,133 @@
|
||||
from time import time |
||||
from typing import Optional |
||||
|
||||
import torch |
||||
import torch.distributed as dist |
||||
from chatgpt.experience_maker import Experience |
||||
|
||||
from .base import Callback |
||||
|
||||
|
||||
def get_world_size() -> int: |
||||
if dist.is_initialized(): |
||||
return dist.get_world_size() |
||||
return 1 |
||||
|
||||
|
||||
def print_rank_0(*args, **kwargs) -> None: |
||||
if not dist.is_initialized() or dist.get_rank() == 0: |
||||
print(*args, **kwargs) |
||||
|
||||
|
||||
@torch.no_grad() |
||||
def all_reduce_mean(x: float, world_size: int) -> float: |
||||
if world_size == 1: |
||||
return x |
||||
tensor = torch.tensor([x], device=torch.cuda.current_device()) |
||||
dist.all_reduce(tensor) |
||||
tensor = tensor / world_size |
||||
return tensor.item() |
||||
|
||||
|
||||
class PerformanceEvaluator(Callback): |
||||
""" |
||||
Callback for valuate the performance of the model. |
||||
Args: |
||||
actor_num_params: The number of parameters of the actor model. |
||||
critic_num_params: The number of parameters of the critic model. |
||||
initial_model_num_params: The number of parameters of the initial model. |
||||
reward_model_num_params: The number of parameters of the reward model. |
||||
enable_grad_checkpoint: Whether to enable gradient checkpointing. |
||||
ignore_episodes: The number of episodes to ignore when calculating the performance. |
||||
""" |
||||
|
||||
def __init__(self, |
||||
actor_num_params: int, |
||||
critic_num_params: int, |
||||
initial_model_num_params: int, |
||||
reward_model_num_params: int, |
||||
enable_grad_checkpoint: bool = False, |
||||
ignore_episodes: int = 0) -> None: |
||||
super().__init__() |
||||
self.world_size = get_world_size() |
||||
self.actor_num_params = actor_num_params |
||||
self.critic_num_params = critic_num_params |
||||
self.initial_model_num_params = initial_model_num_params |
||||
self.reward_model_num_params = reward_model_num_params |
||||
self.enable_grad_checkpoint = enable_grad_checkpoint |
||||
self.ignore_episodes = ignore_episodes |
||||
self.disable: bool = False |
||||
|
||||
self.make_experience_duration: float = 0. |
||||
self.make_experience_start_time: Optional[float] = None |
||||
self.make_experience_num_samples: int = 0 |
||||
self.make_experience_flop: int = 0 |
||||
self.learn_duration: float = 0. |
||||
self.learn_start_time: Optional[float] = None |
||||
self.learn_num_samples: int = 0 |
||||
self.learn_flop: int = 0 |
||||
|
||||
def on_episode_start(self, episode: int) -> None: |
||||
self.disable = self.ignore_episodes > 0 and episode < self.ignore_episodes |
||||
|
||||
def on_make_experience_start(self) -> None: |
||||
if self.disable: |
||||
return |
||||
self.make_experience_start_time = time() |
||||
|
||||
def on_make_experience_end(self, experience: Experience) -> None: |
||||
if self.disable: |
||||
return |
||||
self.make_experience_duration += time() - self.make_experience_start_time |
||||
|
||||
batch_size, seq_len = experience.sequences.shape |
||||
|
||||
self.make_experience_num_samples += batch_size |
||||
|
||||
# actor generate |
||||
num_actions = experience.action_mask.size(1) |
||||
input_len = seq_len - num_actions |
||||
total_seq_len = (input_len + seq_len - 1) * num_actions / 2 |
||||
self.make_experience_flop += self.actor_num_params * batch_size * total_seq_len * 2 |
||||
# actor forward |
||||
self.make_experience_flop += self.actor_num_params * batch_size * seq_len * 2 |
||||
# critic forward |
||||
self.make_experience_flop += self.critic_num_params * batch_size * seq_len * 2 |
||||
# initial model forward |
||||
self.make_experience_flop += self.initial_model_num_params * batch_size * seq_len * 2 |
||||
# reward model forward |
||||
self.make_experience_flop += self.reward_model_num_params * batch_size * seq_len * 2 |
||||
|
||||
def on_learn_batch_start(self) -> None: |
||||
if self.disable: |
||||
return |
||||
self.learn_start_time = time() |
||||
|
||||
def on_learn_batch_end(self, metrics: dict, experience: Experience) -> None: |
||||
if self.disable: |
||||
return |
||||
self.learn_duration += time() - self.learn_start_time |
||||
|
||||
batch_size, seq_len = experience.sequences.shape |
||||
|
||||
self.learn_num_samples += batch_size |
||||
|
||||
# actor forward-backward, 3 means forward(1) + backward(2) |
||||
self.learn_flop += self.actor_num_params * batch_size * seq_len * 2 * (3 + int(self.enable_grad_checkpoint)) |
||||
# critic foward-backward |
||||
self.learn_flop += self.critic_num_params * batch_size * seq_len * 2 * (3 + int(self.enable_grad_checkpoint)) |
||||
|
||||
def on_fit_end(self) -> None: |
||||
avg_make_experience_duration = all_reduce_mean(self.make_experience_duration, self.world_size) |
||||
avg_learn_duration = all_reduce_mean(self.learn_duration, self.world_size) |
||||
|
||||
avg_make_experience_throughput = self.make_experience_num_samples / (avg_make_experience_duration + 1e-12) |
||||
avg_make_experience_tflops = self.make_experience_flop / 1e12 / (avg_make_experience_duration + 1e-12) |
||||
|
||||
avg_learn_throughput = self.learn_num_samples / (avg_learn_duration + 1e-12) |
||||
avg_learn_tflops = self.learn_flop / 1e12 / (avg_learn_duration + 1e-12) |
||||
|
||||
print_rank_0( |
||||
f'Making experience throughput: {avg_make_experience_throughput:.3f} samples/sec, TFLOPS: {avg_make_experience_tflops:.3f}' |
||||
) |
||||
print_rank_0(f'Learning throughput: {avg_learn_throughput:.3f} samples/sec, TFLOPS: {avg_learn_tflops:.3f}') |
||||
@ -0,0 +1,104 @@
|
||||
from typing import Any, Callable, Dict, List, Optional |
||||
|
||||
import torch.nn as nn |
||||
from chatgpt.experience_maker import Experience, NaiveExperienceMaker |
||||
from chatgpt.nn import Actor, Critic, PolicyLoss, ValueLoss |
||||
from chatgpt.replay_buffer import NaiveReplayBuffer |
||||
from torch.optim import Optimizer |
||||
|
||||
from .base import Trainer |
||||
from .callbacks import Callback |
||||
from .strategies import Strategy |
||||
|
||||
|
||||
class PPOTrainer(Trainer): |
||||
""" |
||||
Trainer for PPO algorithm. |
||||
|
||||
Args: |
||||
strategy (Strategy): the strategy to use for training |
||||
actor (Actor): the actor model in ppo algorithm |
||||
critic (Critic): the critic model in ppo algorithm |
||||
reward_model (nn.Module): the reward model in rlhf algorithm to make reward of sentences |
||||
initial_model (Actor): the initial model in rlhf algorithm to generate reference logits to limit the update of actor |
||||
actor_optim (Optimizer): the optimizer to use for actor model |
||||
critic_optim (Optimizer): the optimizer to use for critic model |
||||
kl_coef (float, defaults to 0.1): the coefficient of kl divergence loss |
||||
train_batch_size (int, defaults to 8): the batch size to use for training |
||||
buffer_limit (int, defaults to 0): the max_size limitaiton of replay buffer |
||||
buffer_cpu_offload (bool, defaults to True): whether to offload replay buffer to cpu |
||||
eps_clip (float, defaults to 0.2): the clip coefficient of policy loss |
||||
value_clip (float, defaults to 0.4): the clip coefficient of value loss |
||||
experience_batch_size (int, defaults to 8): the batch size to use for experience generation |
||||
max_epochs (int, defaults to 1): the number of epochs of training process |
||||
tokenier (Callable, optional): the tokenizer to use for tokenizing the input |
||||
sample_replay_buffer (bool, defaults to False): whether to sample from replay buffer |
||||
dataloader_pin_memory (bool, defaults to True): whether to pin memory for data loader |
||||
callbacks (List[Callback], defaults to []): the callbacks to call during training process |
||||
generate_kwargs (dict, optional): the kwargs to use while model generating |
||||
""" |
||||
|
||||
def __init__(self, |
||||
strategy: Strategy, |
||||
actor: Actor, |
||||
critic: Critic, |
||||
reward_model: nn.Module, |
||||
initial_model: Actor, |
||||
actor_optim: Optimizer, |
||||
critic_optim: Optimizer, |
||||
kl_coef: float = 0.1, |
||||
train_batch_size: int = 8, |
||||
buffer_limit: int = 0, |
||||
buffer_cpu_offload: bool = True, |
||||
eps_clip: float = 0.2, |
||||
value_clip: float = 0.4, |
||||
experience_batch_size: int = 8, |
||||
max_epochs: int = 1, |
||||
tokenizer: Optional[Callable[[Any], dict]] = None, |
||||
sample_replay_buffer: bool = False, |
||||
dataloader_pin_memory: bool = True, |
||||
callbacks: List[Callback] = [], |
||||
**generate_kwargs) -> None: |
||||
actor = Actor(strategy.setup_model(actor.model)) |
||||
critic = strategy.setup_model(critic) |
||||
reward_model = strategy.setup_model(reward_model) |
||||
initial_model = Actor(strategy.setup_model(initial_model.model)) |
||||
experience_maker = NaiveExperienceMaker(actor, critic, reward_model, initial_model, kl_coef) |
||||
replay_buffer = NaiveReplayBuffer(train_batch_size, buffer_limit, buffer_cpu_offload) |
||||
super().__init__(strategy, experience_maker, replay_buffer, experience_batch_size, max_epochs, tokenizer, |
||||
sample_replay_buffer, dataloader_pin_memory, callbacks, **generate_kwargs) |
||||
self.actor = actor |
||||
self.critic = critic |
||||
|
||||
self.actor_loss_fn = PolicyLoss(eps_clip) |
||||
self.critic_loss_fn = ValueLoss(value_clip) |
||||
|
||||
self.actor_optim = strategy.setup_optimizer(actor_optim, self.actor.model) |
||||
self.critic_optim = strategy.setup_optimizer(critic_optim, self.critic) |
||||
|
||||
def training_step(self, experience: Experience) -> Dict[str, float]: |
||||
self.actor.train() |
||||
self.critic.train() |
||||
|
||||
num_actions = experience.action_mask.size(1) |
||||
action_log_probs = self.actor(experience.sequences, num_actions, attention_mask=experience.attention_mask) |
||||
actor_loss = self.actor_loss_fn(action_log_probs, |
||||
experience.action_log_probs, |
||||
experience.advantages, |
||||
action_mask=experience.action_mask) |
||||
self.strategy.backward(actor_loss, self.actor, self.actor_optim) |
||||
self.strategy.optimizer_step(self.actor_optim) |
||||
self.actor_optim.zero_grad() |
||||
|
||||
values = self.critic(experience.sequences, |
||||
action_mask=experience.action_mask, |
||||
attention_mask=experience.attention_mask) |
||||
critic_loss = self.critic_loss_fn(values, |
||||
experience.values, |
||||
experience.reward, |
||||
action_mask=experience.action_mask) |
||||
self.strategy.backward(critic_loss, self.critic, self.critic_optim) |
||||
self.strategy.optimizer_step(self.critic_optim) |
||||
self.critic_optim.zero_grad() |
||||
|
||||
return {'actor_loss': actor_loss.item(), 'critic_loss': critic_loss.item()} |
||||
@ -0,0 +1,77 @@
|
||||
from abc import ABC |
||||
|
||||
import loralib as lora |
||||
from chatgpt.dataset import RewardDataset |
||||
from chatgpt.nn import PairWiseLoss |
||||
from torch.optim import Adam |
||||
from torch.utils.data import DataLoader |
||||
from tqdm import tqdm |
||||
|
||||
|
||||
class RewardModelTrainer(ABC): |
||||
""" |
||||
Trainer to use while training reward model. |
||||
|
||||
Args: |
||||
model (torch.nn.Module): the model to train |
||||
train_dataset (RewardDataset): the dataset to use for training |
||||
eval_dataset (RewardDataset): the dataset to use for evaluation |
||||
batch_size (int, defaults to 1): the batch size while training |
||||
num_epochs (int, defaults to 2): the number of epochs to train |
||||
optim_kwargs (dict, defaults to {'lr':1e-4}): the kwargs to use while initializing optimizer |
||||
""" |
||||
|
||||
def __init__(self, |
||||
model, |
||||
train_dataset: RewardDataset, |
||||
eval_dataset: RewardDataset, |
||||
batch_size: int = 1, |
||||
num_epochs: int = 2, |
||||
optim_kwargs: dict = {'lr': 1e-4}) -> None: |
||||
super().__init__() |
||||
self.model = model |
||||
self.train_dataloader = DataLoader(train_dataset, batch_size=batch_size) |
||||
self.eval_dataloader = DataLoader(eval_dataset, batch_size=batch_size) |
||||
self.loss_fn = PairWiseLoss() |
||||
self.optimizer = Adam(self.model.parameters(), **optim_kwargs) |
||||
self.epochs = num_epochs |
||||
|
||||
def fit(self, use_lora): |
||||
epoch_bar = tqdm(range(self.epochs), desc='Train epoch') |
||||
for epoch in range(self.epochs): |
||||
step_bar = tqdm(range(self.train_dataloader.__len__()), desc='Train step of epoch %d' % epoch) |
||||
# train |
||||
if use_lora > 0: |
||||
print("Using Lora") |
||||
lora.mark_only_lora_as_trainable(self.model) |
||||
else: |
||||
self.model.train() |
||||
for chosen_ids, c_mask, reject_ids, r_mask in self.train_dataloader: |
||||
chosen_ids = chosen_ids.squeeze(1).cuda() |
||||
c_mask = c_mask.squeeze(1).cuda() |
||||
reject_ids = reject_ids.squeeze(1).cuda() |
||||
r_mask = r_mask.squeeze(1).cuda() |
||||
chosen_reward = self.model(chosen_ids, attention_mask=c_mask) |
||||
reject_reward = self.model(reject_ids, attention_mask=r_mask) |
||||
loss = self.loss_fn(chosen_reward, reject_reward) |
||||
loss.backward() |
||||
self.optimizer.step() |
||||
self.optimizer.zero_grad() |
||||
step_bar.update() |
||||
step_bar.set_postfix({'loss': loss.item()}) |
||||
|
||||
# eval |
||||
self.model.eval() |
||||
for chosen_ids, c_mask, reject_ids, r_mask in self.eval_dataloader: |
||||
dist = 0 |
||||
chosen_ids = chosen_ids.squeeze(1).cuda() |
||||
c_mask = c_mask.squeeze(1).cuda() |
||||
reject_ids = reject_ids.squeeze(1).cuda() |
||||
r_mask = r_mask.squeeze(1).cuda() |
||||
chosen_reward = self.model(chosen_ids, attention_mask=c_mask) |
||||
reject_reward = self.model(reject_ids, attention_mask=r_mask) |
||||
dist += (chosen_reward - reject_reward) |
||||
dist_mean = dist / self.eval_dataloader.__len__() |
||||
epoch_bar.update() |
||||
step_bar.set_postfix({'loss': loss.item(), 'dist_mean': dist_mean.item()}) |
||||
step_bar.close() |
||||
@ -0,0 +1,6 @@
|
||||
from .base import Strategy |
||||
from .colossalai import ColossalAIStrategy |
||||
from .ddp import DDPStrategy |
||||
from .naive import NaiveStrategy |
||||
|
||||
__all__ = ['Strategy', 'NaiveStrategy', 'DDPStrategy', 'ColossalAIStrategy'] |
||||
@ -0,0 +1,45 @@
|
||||
from abc import ABC, abstractmethod |
||||
from contextlib import nullcontext |
||||
|
||||
import torch |
||||
import torch.nn as nn |
||||
import torch.optim as optim |
||||
from chatgpt.replay_buffer import ReplayBuffer |
||||
from torch.utils.data import DataLoader |
||||
|
||||
|
||||
class Strategy(ABC): |
||||
""" |
||||
Base class for training strategies. |
||||
""" |
||||
|
||||
def __init__(self) -> None: |
||||
super().__init__() |
||||
self.setup_distributed() |
||||
|
||||
@abstractmethod |
||||
def backward(self, loss: torch.Tensor, model: nn.Module, optimizer: optim.Optimizer, **kwargs) -> None: |
||||
pass |
||||
|
||||
@abstractmethod |
||||
def optimizer_step(self, optimizer: optim.Optimizer, **kwargs) -> None: |
||||
pass |
||||
|
||||
@abstractmethod |
||||
def setup_distributed(self) -> None: |
||||
pass |
||||
|
||||
@abstractmethod |
||||
def setup_model(self, model: nn.Module) -> nn.Module: |
||||
pass |
||||
|
||||
@abstractmethod |
||||
def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Module) -> optim.Optimizer: |
||||
pass |
||||
|
||||
@abstractmethod |
||||
def setup_dataloader(self, replay_buffer: ReplayBuffer, pin_memory: bool = False) -> DataLoader: |
||||
pass |
||||
|
||||
def model_init_context(self): |
||||
return nullcontext() |
||||
@ -0,0 +1,125 @@
|
||||
from typing import Optional |
||||
|
||||
import torch |
||||
import torch.distributed as dist |
||||
import torch.nn as nn |
||||
import torch.optim as optim |
||||
|
||||
import colossalai |
||||
from colossalai.nn.optimizer import CPUAdam, HybridAdam |
||||
from colossalai.nn.parallel import zero_model_wrapper, zero_optim_wrapper |
||||
from colossalai.tensor import ProcessGroup, ShardSpec |
||||
from colossalai.utils import get_current_device |
||||
from colossalai.utils.model.colo_init_context import ColoInitContext |
||||
|
||||
from .ddp import DDPStrategy |
||||
|
||||
|
||||
class ColossalAIStrategy(DDPStrategy): |
||||
""" |
||||
The strategy for training with ColossalAI. |
||||
|
||||
Args: |
||||
stage(int): The stage to use in ZeRO. Choose in (1, 2, 3) |
||||
seed(int): The seed for the random number generator. |
||||
shard_init(bool): Whether to shard the model parameters during initialization. Only for ZeRO-3. |
||||
placement_policy(str): The placement policy for gemini. Choose in ('cpu', 'cuda') |
||||
If it is “cpu”, parameters, gradients and optimizer states will be offloaded to CPU, |
||||
If it is “cuda”, they will not be offloaded, which means max CUDA memory will be used. It is the fastest. |
||||
pin_memory(bool): Whether to pin the memory for the data loader. Only for ZeRO-3. |
||||
force_outputs_fp32(bool): Whether to force the outputs to be fp32. Only for ZeRO-3. |
||||
search_range_mb(int): The search range in MB for the chunk size. Only for ZeRO-3. |
||||
hidden_dim(optional, int): The hidden dimension for the gemini. Only for ZeRO-3. |
||||
min_chunk_size_mb(float): The minimum chunk size in MB. Only for ZeRO-3. |
||||
gpu_margin_mem_ratio(float): The margin memory ratio for the GPU. Only for ZeRO-3. |
||||
reduce_bugket_size(int): The reduce bucket size in bytes. Only for ZeRO-1 and ZeRO-2. |
||||
overlap_communication(bool): Whether to overlap communication and computation. Only for ZeRO-1 and ZeRO-2. |
||||
initial_scale(float): The initial scale for the optimizer. |
||||
growth_factor(float): The growth factor for the optimizer. |
||||
backoff_factor(float): The backoff factor for the optimizer. |
||||
growth_interval(int): The growth interval for the optimizer. |
||||
hysteresis(int): The hysteresis for the optimizer. |
||||
min_scale(float): The minimum scale for the optimizer. |
||||
max_scale(float): The maximum scale for the optimizer. |
||||
max_norm(float): The maximum norm for the optimizer. |
||||
norm_type(float): The norm type for the optimizer. |
||||
|
||||
""" |
||||
|
||||
def __init__( |
||||
self, |
||||
stage: int = 3, |
||||
seed: int = 42, |
||||
shard_init: bool = True, # only for stage 3 |
||||
placement_policy: str = 'cuda', |
||||
pin_memory: bool = True, # only for stage 3 |
||||
force_outputs_fp32: bool = False, # only for stage 3 |
||||
search_range_mb: int = 32, # only for stage 3 |
||||
hidden_dim: Optional[int] = None, # only for stage 3 |
||||
min_chunk_size_mb: float = 32, # only for stage 3 |
||||
gpu_margin_mem_ratio: float = 0.0, # only for stage 3 |
||||
reduce_bucket_size: int = 12 * 1024**2, # only for stage 1&2 |
||||
overlap_communication: bool = True, # only for stage 1&2 |
||||
initial_scale: float = 2**16, |
||||
growth_factor: float = 2, |
||||
backoff_factor: float = 0.5, |
||||
growth_interval: int = 1000, |
||||
hysteresis: int = 2, |
||||
min_scale: float = 1, |
||||
max_scale: float = 2**32, |
||||
max_norm: float = 0.0, |
||||
norm_type: float = 2.0) -> None: |
||||
super().__init__(seed) |
||||
assert placement_policy in ('cpu', 'cuda'), f'Unsupported placement policy "{placement_policy}"' |
||||
self.stage = stage |
||||
self.shard_init = shard_init |
||||
self.gemini_config = dict(device=get_current_device(), |
||||
placement_policy=placement_policy, |
||||
pin_memory=pin_memory, |
||||
force_outputs_fp32=force_outputs_fp32, |
||||
strict_ddp_mode=shard_init, |
||||
search_range_mb=search_range_mb, |
||||
hidden_dim=hidden_dim, |
||||
min_chunk_size_mb=min_chunk_size_mb) |
||||
if stage == 3: |
||||
self.zero_optim_config = dict(gpu_margin_mem_ratio=gpu_margin_mem_ratio) |
||||
else: |
||||
self.zero_optim_config = dict(reduce_bucket_size=reduce_bucket_size, |
||||
overlap_communication=overlap_communication, |
||||
cpu_offload=(placement_policy == 'cpu')) |
||||
self.optim_kwargs = dict(initial_scale=initial_scale, |
||||
growth_factor=growth_factor, |
||||
backoff_factor=backoff_factor, |
||||
growth_interval=growth_interval, |
||||
hysteresis=hysteresis, |
||||
min_scale=min_scale, |
||||
max_scale=max_scale, |
||||
max_norm=max_norm, |
||||
norm_type=norm_type) |
||||
|
||||
def setup_distributed(self) -> None: |
||||
colossalai.launch_from_torch({}, seed=self.seed) |
||||
|
||||
def model_init_context(self): |
||||
if self.stage == 3: |
||||
world_size = dist.get_world_size() |
||||
shard_pg = ProcessGroup(tp_degree=world_size) if self.shard_init else None |
||||
default_dist_spec = ShardSpec([-1], [world_size]) if self.shard_init else None |
||||
return ColoInitContext(device=get_current_device(), |
||||
dtype=torch.half, |
||||
default_pg=shard_pg, |
||||
default_dist_spec=default_dist_spec) |
||||
return super().model_init_context() |
||||
|
||||
def setup_model(self, model: nn.Module) -> nn.Module: |
||||
return zero_model_wrapper(model, zero_stage=self.stage, gemini_config=self.gemini_config) |
||||
|
||||
def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Module) -> optim.Optimizer: |
||||
assert isinstance(optimizer, (CPUAdam, HybridAdam)), f'Unsupported optimizer {type(optimizer)}' |
||||
return zero_optim_wrapper(model, optimizer, optim_config=self.zero_optim_config, **self.optim_kwargs) |
||||
|
||||
def backward(self, loss: torch.Tensor, model: nn.Module, optimizer: optim.Optimizer, **kwargs) -> None: |
||||
optimizer.backward(loss) |
||||
|
||||
def optimizer_step(self, optimizer: optim.Optimizer, **kwargs) -> None: |
||||
optimizer.step() |
||||
@ -0,0 +1,59 @@
|
||||
import os |
||||
import random |
||||
|
||||
import numpy as np |
||||
import torch |
||||
import torch.distributed as dist |
||||
import torch.nn as nn |
||||
from chatgpt.replay_buffer import ReplayBuffer |
||||
from torch.nn.parallel import DistributedDataParallel as DDP |
||||
from torch.utils.data import DataLoader, DistributedSampler |
||||
|
||||
from .naive import NaiveStrategy |
||||
|
||||
|
||||
class DDPStrategy(NaiveStrategy): |
||||
""" |
||||
Strategy for distributed training using torch.distributed. |
||||
""" |
||||
|
||||
def __init__(self, seed: int = 42) -> None: |
||||
self.seed = seed |
||||
super().__init__() |
||||
|
||||
def setup_distributed(self) -> None: |
||||
try: |
||||
rank = int(os.environ['RANK']) |
||||
local_rank = int(os.environ['LOCAL_RANK']) |
||||
world_size = int(os.environ['WORLD_SIZE']) |
||||
host = os.environ['MASTER_ADDR'] |
||||
port = int(os.environ['MASTER_PORT']) |
||||
except KeyError as e: |
||||
raise RuntimeError( |
||||
f"Could not find {e} in the torch environment, visit https://www.colossalai.org/ for more information on launching with torch" |
||||
) |
||||
dist.init_process_group('nccl', init_method=f'tcp://[{host}]:{port}', world_size=world_size, rank=rank) |
||||
self.set_seed(self.seed) |
||||
torch.cuda.set_device(local_rank) |
||||
|
||||
def set_seed(self, seed: int) -> None: |
||||
random.seed(seed) |
||||
np.random.seed(seed) |
||||
torch.manual_seed(seed) |
||||
|
||||
def setup_model(self, model: nn.Module) -> nn.Module: |
||||
device = torch.cuda.current_device() |
||||
return DDP(model, device_ids=[device]) |
||||
|
||||
def setup_dataloader(self, replay_buffer: ReplayBuffer, pin_memory: bool = False) -> DataLoader: |
||||
sampler = DistributedSampler(replay_buffer, |
||||
num_replicas=dist.get_world_size(), |
||||
rank=dist.get_rank(), |
||||
shuffle=True, |
||||
seed=self.seed, |
||||
drop_last=True) |
||||
return DataLoader(replay_buffer, |
||||
batch_size=replay_buffer.sample_batch_size, |
||||
sampler=sampler, |
||||
pin_memory=pin_memory, |
||||
collate_fn=replay_buffer.collate_fn) |
||||
@ -0,0 +1,36 @@
|
||||
import torch |
||||
import torch.nn as nn |
||||
import torch.optim as optim |
||||
from chatgpt.replay_buffer import ReplayBuffer |
||||
from torch.utils.data import DataLoader |
||||
|
||||
from .base import Strategy |
||||
|
||||
|
||||
class NaiveStrategy(Strategy): |
||||
""" |
||||
Strategy for single GPU. No parallelism is used. |
||||
""" |
||||
|
||||
def backward(self, loss: torch.Tensor, model: nn.Module, optimizer: optim.Optimizer, **kwargs) -> None: |
||||
loss.backward() |
||||
|
||||
def optimizer_step(self, optimizer: optim.Optimizer, **kwargs) -> None: |
||||
optimizer.step() |
||||
|
||||
def setup_distributed(self) -> None: |
||||
pass |
||||
|
||||
def setup_model(self, model: nn.Module) -> nn.Module: |
||||
return model |
||||
|
||||
def setup_optimizer(self, optimizer: optim.Optimizer, model: nn.Module) -> optim.Optimizer: |
||||
return optimizer |
||||
|
||||
def setup_dataloader(self, replay_buffer: ReplayBuffer, pin_memory: bool = False) -> DataLoader: |
||||
return DataLoader(replay_buffer, |
||||
batch_size=replay_buffer.sample_batch_size, |
||||
shuffle=True, |
||||
drop_last=True, |
||||
pin_memory=pin_memory, |
||||
collate_fn=replay_buffer.collate_fn) |
||||
@ -0,0 +1,5 @@
|
||||
import torch.distributed as dist |
||||
|
||||
|
||||
def is_rank_0() -> bool: |
||||
return not dist.is_initialized() or dist.get_rank() == 0 |
||||
@ -0,0 +1,105 @@
|
||||
# Examples |
||||
|
||||
## Install requirements |
||||
|
||||
```shell |
||||
pip install -r requirements.txt |
||||
``` |
||||
|
||||
## Train with dummy prompt data |
||||
|
||||
This script supports 3 strategies: |
||||
|
||||
- naive |
||||
- ddp |
||||
- colossalai |
||||
|
||||
It uses random generated prompt data. |
||||
|
||||
Naive strategy only support single GPU training: |
||||
|
||||
```shell |
||||
python train_dummy.py --strategy naive |
||||
# display cli help |
||||
python train_dummy.py -h |
||||
``` |
||||
|
||||
DDP strategy and ColossalAI strategy support multi GPUs training: |
||||
|
||||
```shell |
||||
# run DDP on 2 GPUs |
||||
torchrun --standalone --nproc_per_node=2 train_dummy.py --strategy ddp |
||||
# run ColossalAI on 2 GPUs |
||||
torchrun --standalone --nproc_per_node=2 train_dummy.py --strategy colossalai |
||||
``` |
||||
|
||||
## Train with real prompt data |
||||
|
||||
We use [awesome-chatgpt-prompts](https://huggingface.co/datasets/fka/awesome-chatgpt-prompts) as example dataset. It is a small dataset with hundreds of prompts. |
||||
|
||||
You should download `prompts.csv` first. |
||||
|
||||
This script also supports 3 strategies. |
||||
|
||||
```shell |
||||
# display cli help |
||||
python train_dummy.py -h |
||||
# run naive on 1 GPU |
||||
python train_prompts.py prompts.csv --strategy naive |
||||
# run DDP on 2 GPUs |
||||
torchrun --standalone --nproc_per_node=2 train_prompts.py prompts.csv --strategy ddp |
||||
# run ColossalAI on 2 GPUs |
||||
torchrun --standalone --nproc_per_node=2 train_prompts.py prompts.csv --strategy colossalai |
||||
``` |
||||
|
||||
## Train the reward model |
||||
We use [rm-static](https://huggingface.co/datasets/Dahoas/rm-static) as dataset to train our reward model. It is a dataset of chosen & rejected response of the same prompt. |
||||
|
||||
You can download the dataset from huggingface automatically. |
||||
|
||||
Use these code to train your reward model. |
||||
|
||||
```shell |
||||
# Naive reward model training |
||||
python train_reward_model.py --pretrain <your model path> |
||||
# if to use LoRA |
||||
python train_reward_model.py --pretrain <your model path> --lora_rank 16 |
||||
``` |
||||
|
||||
## Support Model |
||||
|
||||
### GPT |
||||
- [ ] GPT2-S (s) |
||||
- [ ] GPT2-M (m) |
||||
- [ ] GPT2-L (l) |
||||
- [ ] GPT2-XL (xl) |
||||
- [ ] GPT2-4B (4b) |
||||
- [ ] GPT2-6B (6b) |
||||
- [ ] GPT2-8B (8b) |
||||
- [ ] GPT2-10B (10b) |
||||
- [ ] GPT2-12B (12b) |
||||
- [ ] GPT2-15B (15b) |
||||
- [ ] GPT2-18B (18b) |
||||
- [ ] GPT2-20B (20b) |
||||
- [ ] GPT2-24B (24b) |
||||
- [ ] GPT2-28B (28b) |
||||
- [ ] GPT2-32B (32b) |
||||
- [ ] GPT2-36B (36b) |
||||
- [ ] GPT2-40B (40b) |
||||
- [ ] GPT3 (175b) |
||||
|
||||
### BLOOM |
||||
- [x] [BLOOM-560m](https://huggingface.co/bigscience/bloom-560m) |
||||
- [x] [BLOOM-1b1](https://huggingface.co/bigscience/bloom-1b1) |
||||
- [ ] [BLOOM-3b](https://huggingface.co/bigscience/bloom-3b) |
||||
- [ ] [BLOOM-7b](https://huggingface.co/bigscience/bloomz-7b1) |
||||
- [ ] BLOOM-175b |
||||
|
||||
### OPT |
||||
- [x] [OPT-125M](https://huggingface.co/facebook/opt-125m) |
||||
- [x] [OPT-350M](https://huggingface.co/facebook/opt-350m) |
||||
- [ ] [OPT-1.3B](https://huggingface.co/facebook/opt-1.3b) |
||||
- [ ] [OPT-2.7B](https://huggingface.co/facebook/opt-2.7b) |
||||
- [ ] [OPT-6.7B](https://huggingface.co/facebook/opt-6.7b) |
||||
- [ ] [OPT-13B](https://huggingface.co/facebook/opt-13b) |
||||
- [ ] [OPT-30B](https://huggingface.co/facebook/opt-30b) |
||||
@ -0,0 +1,27 @@
|
||||
#!/usr/bin/env bash |
||||
|
||||
set -xue |
||||
|
||||
if [ -z "$PROMPT_PATH" ]; then |
||||
echo "Please set \$PROMPT_PATH to the path to prompts csv." |
||||
exit 1 |
||||
fi |
||||
|
||||
BASE=$(realpath $(dirname $0)) |
||||
|
||||
export OMP_NUM_THREADS=8 |
||||
|
||||
# install requirements |
||||
pip install -r ${BASE}/requirements.txt |
||||
|
||||
# train dummy |
||||
python ${BASE}/train_dummy.py --strategy naive --num_episodes 3 --max_timesteps 3 --update_timesteps 3 --max_epochs 3 --train_batch_size 2 |
||||
for strategy in ddp colossalai_gemini colossalai_zero2; do |
||||
torchrun --standalone --nproc_per_node=2 ${BASE}/train_dummy.py --strategy ${strategy} --num_episodes 3 --max_timesteps 3 --update_timesteps 3 --max_epochs 3 --train_batch_size 2 |
||||
done |
||||
|
||||
# train prompts |
||||
python ${BASE}/train_prompts.py $PROMPT_PATH --strategy naive --num_episodes 3 --max_timesteps 3 --update_timesteps 3 --max_epochs 3 |
||||
for strategy in ddp colossalai_gemini colossalai_zero2; do |
||||
torchrun --standalone --nproc_per_node=2 ${BASE}/train_prompts.py $PROMPT_PATH --strategy ${strategy} --num_episodes 3 --max_timesteps 3 --update_timesteps 3 --max_epochs 3 --train_batch_size 2 |
||||
done |
||||
@ -0,0 +1,121 @@
|
||||
import argparse |
||||
from copy import deepcopy |
||||
|
||||
import torch |
||||
from chatgpt.nn import BLOOMActor, BLOOMCritic, GPTActor, GPTCritic, OPTActor, OPTCritic, RewardModel |
||||
from chatgpt.nn.generation_utils import ( |
||||
bloom_prepare_inputs_fn, |
||||
gpt_prepare_inputs_fn, |
||||
opt_prepare_inputs_fn, |
||||
update_model_kwargs_fn, |
||||
) |
||||
from chatgpt.trainer import PPOTrainer |
||||
from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy |
||||
from torch.optim import Adam |
||||
from transformers import AutoTokenizer, BloomTokenizerFast |
||||
from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer |
||||
|
||||
from colossalai.nn.optimizer import HybridAdam |
||||
|
||||
|
||||
def preprocess_batch(samples): |
||||
input_ids = torch.stack(samples) |
||||
attention_mask = torch.ones_like(input_ids, dtype=torch.long) |
||||
return {'input_ids': input_ids, 'attention_mask': attention_mask} |
||||
|
||||
|
||||
def main(args): |
||||
# configure strategy |
||||
if args.strategy == 'naive': |
||||
strategy = NaiveStrategy() |
||||
elif args.strategy == 'ddp': |
||||
strategy = DDPStrategy() |
||||
elif args.strategy == 'colossalai_gemini': |
||||
strategy = ColossalAIStrategy(stage=3, placement_policy='cuda') |
||||
elif args.strategy == 'colossalai_zero2': |
||||
strategy = ColossalAIStrategy(stage=2, placement_policy='cuda') |
||||
else: |
||||
raise ValueError(f'Unsupported strategy "{args.strategy}"') |
||||
|
||||
# configure model |
||||
with strategy.model_init_context(): |
||||
if args.model == 'gpt2': |
||||
actor = GPTActor().cuda() |
||||
critic = GPTCritic().cuda() |
||||
elif args.model == 'bloom': |
||||
actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda() |
||||
critic = BLOOMCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda() |
||||
elif args.model == 'opt': |
||||
actor = OPTActor().cuda() |
||||
critic = OPTCritic().cuda() |
||||
else: |
||||
raise ValueError(f'Unsupported model "{args.model}"') |
||||
|
||||
initial_model = deepcopy(actor).cuda() |
||||
reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda() |
||||
|
||||
# configure optimizer |
||||
if args.strategy.startswith('colossalai'): |
||||
actor_optim = HybridAdam(actor.parameters(), lr=5e-6) |
||||
critic_optim = HybridAdam(critic.parameters(), lr=5e-6) |
||||
else: |
||||
actor_optim = Adam(actor.parameters(), lr=5e-6) |
||||
critic_optim = Adam(critic.parameters(), lr=5e-6) |
||||
|
||||
# configure tokenizer |
||||
if args.model == 'gpt2': |
||||
tokenizer = GPT2Tokenizer.from_pretrained('gpt2') |
||||
tokenizer.pad_token = tokenizer.eos_token |
||||
prepare_inputs_fn = gpt_prepare_inputs_fn |
||||
elif args.model == 'bloom': |
||||
tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain) |
||||
tokenizer.pad_token = tokenizer.eos_token |
||||
prepare_inputs_fn = bloom_prepare_inputs_fn |
||||
elif args.model == 'opt': |
||||
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") |
||||
prepare_inputs_fn = opt_prepare_inputs_fn |
||||
else: |
||||
raise ValueError(f'Unsupported model "{args.model}"') |
||||
|
||||
# configure trainer |
||||
trainer = PPOTrainer(strategy, |
||||
actor, |
||||
critic, |
||||
reward_model, |
||||
initial_model, |
||||
actor_optim, |
||||
critic_optim, |
||||
max_epochs=args.max_epochs, |
||||
train_batch_size=args.train_batch_size, |
||||
tokenizer=preprocess_batch, |
||||
max_length=128, |
||||
do_sample=True, |
||||
temperature=1.0, |
||||
top_k=50, |
||||
pad_token_id=tokenizer.pad_token_id, |
||||
eos_token_id=tokenizer.eos_token_id, |
||||
prepare_inputs_fn=prepare_inputs_fn, |
||||
update_model_kwargs_fn=update_model_kwargs_fn) |
||||
|
||||
random_prompts = torch.randint(tokenizer.vocab_size, (1000, 64), device=torch.cuda.current_device()) |
||||
trainer.fit(random_prompts, |
||||
num_episodes=args.num_episodes, |
||||
max_timesteps=args.max_timesteps, |
||||
update_timesteps=args.update_timesteps) |
||||
|
||||
|
||||
if __name__ == '__main__': |
||||
parser = argparse.ArgumentParser() |
||||
parser.add_argument('--strategy', |
||||
choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'], |
||||
default='naive') |
||||
parser.add_argument('--model', type=str, default='gpt2', choices=['gpt2', 'bloom', 'opt']) |
||||
parser.add_argument('--pretrain', type=str, default=None) |
||||
parser.add_argument('--num_episodes', type=int, default=50) |
||||
parser.add_argument('--max_timesteps', type=int, default=10) |
||||
parser.add_argument('--update_timesteps', type=int, default=10) |
||||
parser.add_argument('--max_epochs', type=int, default=5) |
||||
parser.add_argument('--train_batch_size', type=int, default=8) |
||||
parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") |
||||
args = parser.parse_args() |
||||
main(args) |
||||
@ -0,0 +1,18 @@
|
||||
set_n_least_used_CUDA_VISIBLE_DEVICES() { |
||||
local n=${1:-"9999"} |
||||
echo "GPU Memory Usage:" |
||||
local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv \ |
||||
| tail -n +2 \ |
||||
| nl -v 0 \ |
||||
| tee /dev/tty \ |
||||
| sort -g -k 2 \ |
||||
| awk '{print $1}' \ |
||||
| head -n $n) |
||||
export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g') |
||||
echo "Now CUDA_VISIBLE_DEVICES is set to:" |
||||
echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" |
||||
} |
||||
|
||||
set_n_least_used_CUDA_VISIBLE_DEVICES 1 |
||||
|
||||
python train_dummy.py --model bloom --pretrain '/data2/users/lczht/bloom-560m' --lora_rank 16 |
||||
@ -0,0 +1,113 @@
|
||||
import argparse |
||||
from copy import deepcopy |
||||
|
||||
import pandas as pd |
||||
from chatgpt.nn import BLOOMActor, BLOOMCritic, GPTActor, GPTCritic, OPTActor, OPTCritic, RewardModel |
||||
from chatgpt.nn.generation_utils import gpt_prepare_inputs_fn, update_model_kwargs_fn |
||||
from chatgpt.trainer import PPOTrainer |
||||
from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy, NaiveStrategy |
||||
from torch.optim import Adam |
||||
from transformers import AutoTokenizer, BloomTokenizerFast |
||||
from transformers.models.gpt2.tokenization_gpt2 import GPT2Tokenizer |
||||
|
||||
from colossalai.nn.optimizer import HybridAdam |
||||
|
||||
|
||||
def main(args): |
||||
# configure strategy |
||||
if args.strategy == 'naive': |
||||
strategy = NaiveStrategy() |
||||
elif args.strategy == 'ddp': |
||||
strategy = DDPStrategy() |
||||
elif args.strategy == 'colossalai_gemini': |
||||
strategy = ColossalAIStrategy(stage=3, placement_policy='cuda') |
||||
elif args.strategy == 'colossalai_zero2': |
||||
strategy = ColossalAIStrategy(stage=2, placement_policy='cuda') |
||||
else: |
||||
raise ValueError(f'Unsupported strategy "{args.strategy}"') |
||||
|
||||
# configure model |
||||
with strategy.model_init_context(): |
||||
if args.model == 'gpt2': |
||||
actor = GPTActor().cuda() |
||||
critic = GPTCritic().cuda() |
||||
elif args.model == 'bloom': |
||||
actor = BLOOMActor(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda() |
||||
critic = BLOOMCritic(pretrained=args.pretrain, lora_rank=args.lora_rank).cuda() |
||||
elif args.model == 'opt': |
||||
actor = OPTActor(lora_rank=args.lora_rank).cuda() |
||||
critic = OPTCritic(lora_rank=args.lora_rank).cuda() |
||||
else: |
||||
raise ValueError(f'Unsupported model "{args.model}"') |
||||
|
||||
initial_model = deepcopy(actor) |
||||
reward_model = RewardModel(deepcopy(critic.model), deepcopy(critic.value_head)).cuda() |
||||
|
||||
# configure optimizer |
||||
if args.strategy.startswith('colossalai'): |
||||
actor_optim = HybridAdam(actor.parameters(), lr=5e-6) |
||||
critic_optim = HybridAdam(critic.parameters(), lr=5e-6) |
||||
else: |
||||
actor_optim = Adam(actor.parameters(), lr=5e-6) |
||||
critic_optim = Adam(critic.parameters(), lr=5e-6) |
||||
|
||||
# configure tokenizer |
||||
if args.model == 'gpt2': |
||||
tokenizer = GPT2Tokenizer.from_pretrained('gpt2') |
||||
tokenizer.pad_token = tokenizer.eos_token |
||||
elif args.model == 'bloom': |
||||
tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain) |
||||
tokenizer.pad_token = tokenizer.eos_token |
||||
elif args.model == 'opt': |
||||
tokenizer = AutoTokenizer.from_pretrained("facebook/opt-350m") |
||||
else: |
||||
raise ValueError(f'Unsupported model "{args.model}"') |
||||
|
||||
dataset = pd.read_csv(args.prompt_path)['prompt'] |
||||
|
||||
def tokenize_fn(texts): |
||||
batch = tokenizer(texts, return_tensors='pt', max_length=96, padding=True, truncation=True) |
||||
return {k: v.cuda() for k, v in batch.items()} |
||||
|
||||
# configure trainer |
||||
trainer = PPOTrainer(strategy, |
||||
actor, |
||||
critic, |
||||
reward_model, |
||||
initial_model, |
||||
actor_optim, |
||||
critic_optim, |
||||
max_epochs=args.max_epochs, |
||||
train_batch_size=args.train_batch_size, |
||||
tokenizer=tokenize_fn, |
||||
max_length=128, |
||||
do_sample=True, |
||||
temperature=1.0, |
||||
top_k=50, |
||||
pad_token_id=tokenizer.pad_token_id, |
||||
eos_token_id=tokenizer.eos_token_id, |
||||
prepare_inputs_fn=gpt_prepare_inputs_fn, |
||||
update_model_kwargs_fn=update_model_kwargs_fn) |
||||
|
||||
trainer.fit(dataset, |
||||
num_episodes=args.num_episodes, |
||||
max_timesteps=args.max_timesteps, |
||||
update_timesteps=args.update_timesteps) |
||||
|
||||
|
||||
if __name__ == '__main__': |
||||
parser = argparse.ArgumentParser() |
||||
parser.add_argument('prompt_path') |
||||
parser.add_argument('--strategy', |
||||
choices=['naive', 'ddp', 'colossalai_gemini', 'colossalai_zero2'], |
||||
default='naive') |
||||
parser.add_argument('--model', default='gpt2', choices=['gpt2', 'bloom', 'opt']) |
||||
parser.add_argument('--pretrain', type=str, default=None) |
||||
parser.add_argument('--num_episodes', type=int, default=10) |
||||
parser.add_argument('--max_timesteps', type=int, default=10) |
||||
parser.add_argument('--update_timesteps', type=int, default=10) |
||||
parser.add_argument('--max_epochs', type=int, default=5) |
||||
parser.add_argument('--train_batch_size', type=int, default=8) |
||||
parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") |
||||
args = parser.parse_args() |
||||
main(args) |
||||
@ -0,0 +1,18 @@
|
||||
set_n_least_used_CUDA_VISIBLE_DEVICES() { |
||||
local n=${1:-"9999"} |
||||
echo "GPU Memory Usage:" |
||||
local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv \ |
||||
| tail -n +2 \ |
||||
| nl -v 0 \ |
||||
| tee /dev/tty \ |
||||
| sort -g -k 2 \ |
||||
| awk '{print $1}' \ |
||||
| head -n $n) |
||||
export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g') |
||||
echo "Now CUDA_VISIBLE_DEVICES is set to:" |
||||
echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" |
||||
} |
||||
|
||||
set_n_least_used_CUDA_VISIBLE_DEVICES 1 |
||||
|
||||
python train_prompts.py prompts.csv --pretrain '/data2/users/lczht/bloom-560m' --lora_rank 16 |
||||
@ -0,0 +1,53 @@
|
||||
import argparse |
||||
|
||||
import loralib as lora |
||||
import torch |
||||
from chatgpt.dataset import RewardDataset |
||||
from chatgpt.nn import BLOOMRM |
||||
from chatgpt.trainer import RewardModelTrainer |
||||
from datasets import load_dataset |
||||
from transformers import BloomTokenizerFast |
||||
|
||||
|
||||
def train(args): |
||||
tokenizer = BloomTokenizerFast.from_pretrained(args.pretrain) |
||||
tokenizer.pad_token = tokenizer.eos_token |
||||
model = BLOOMRM(pretrained=args.pretrain) |
||||
|
||||
model.cuda() |
||||
|
||||
max_len = 1024 |
||||
|
||||
# prepare for data and dataset |
||||
data = load_dataset(args.dataset) |
||||
train_data = data["train"] |
||||
eval_data = data['test'] |
||||
train_dataset = RewardDataset(train_data, tokenizer, max_len) |
||||
eval_dataset = RewardDataset(eval_data, tokenizer, max_len) |
||||
|
||||
# batch_size here is expected to be C(k,2), k means # response of each prompt |
||||
# be limited with the format of dataset 'Dahoas/rm-static', we'd better use batch_size as 1 |
||||
trainer = RewardModelTrainer(model=model, |
||||
train_dataset=train_dataset, |
||||
eval_dataset=eval_dataset, |
||||
batch_size=args.batch_size, |
||||
num_epochs=args.max_epochs) |
||||
|
||||
trainer.fit(use_lora=args.lora_rank) |
||||
|
||||
if args.lora_rank > 0: |
||||
torch.save({'model_state_dict': lora.lora_state_dict(trainer.model)}, args.save_path) |
||||
else: |
||||
torch.save(trainer.model, args.save_path) |
||||
|
||||
|
||||
if __name__ == '__main__': |
||||
parser = argparse.ArgumentParser() |
||||
parser.add_argument('--pretrain', type=str, default=None) |
||||
parser.add_argument('--dataset', type=str, default='Dahoas/rm-static') |
||||
parser.add_argument('--save_path', type=str, default='rm_ckpt.pth') |
||||
parser.add_argument('--max_epochs', type=int, default=2) |
||||
parser.add_argument('--batch_size', type=int, default=1) |
||||
parser.add_argument('--lora_rank', type=int, default=0, help="low-rank adaptation matrices rank") |
||||
args = parser.parse_args() |
||||
train(args) |
||||
@ -0,0 +1,18 @@
|
||||
set_n_least_used_CUDA_VISIBLE_DEVICES() { |
||||
local n=${1:-"9999"} |
||||
echo "GPU Memory Usage:" |
||||
local FIRST_N_GPU_IDS=$(nvidia-smi --query-gpu=memory.used --format=csv \ |
||||
| tail -n +2 \ |
||||
| nl -v 0 \ |
||||
| tee /dev/tty \ |
||||
| sort -g -k 2 \ |
||||
| awk '{print $1}' \ |
||||
| head -n $n) |
||||
export CUDA_VISIBLE_DEVICES=$(echo $FIRST_N_GPU_IDS | sed 's/ /,/g') |
||||
echo "Now CUDA_VISIBLE_DEVICES is set to:" |
||||
echo "CUDA_VISIBLE_DEVICES=$CUDA_VISIBLE_DEVICES" |
||||
} |
||||
|
||||
set_n_least_used_CUDA_VISIBLE_DEVICES 1 |
||||
|
||||
python train_reward_model.py --pretrain '/data2/users/lczht/bloom-560m' --lora_rank 16 |
||||
@ -0,0 +1,6 @@
|
||||
[pytest] |
||||
markers = |
||||
cpu: tests which can run on CPU |
||||
gpu: tests which requires a single GPU |
||||
dist: tests which are run in a multi-GPU or multi-machine environment |
||||
experiment: tests for experimental features |
||||
@ -0,0 +1,6 @@
|
||||
transformers>=4.20.1 |
||||
tqdm |
||||
datasets |
||||
loralib |
||||
colossalai>=0.2.4 |
||||
torch |
||||
@ -0,0 +1,42 @@
|
||||
from setuptools import find_packages, setup |
||||
|
||||
|
||||
def fetch_requirements(path): |
||||
with open(path, 'r') as fd: |
||||
return [r.strip() for r in fd.readlines()] |
||||
|
||||
|
||||
def fetch_readme(): |
||||
with open('README.md', encoding='utf-8') as f: |
||||
return f.read() |
||||
|
||||
|
||||
def fetch_version(): |
||||
with open('version.txt', 'r') as f: |
||||
return f.read().strip() |
||||
|
||||
|
||||
setup( |
||||
name='chatgpt', |
||||
version=fetch_version(), |
||||
packages=find_packages(exclude=( |
||||
'tests', |
||||
'benchmarks', |
||||
'requirements', |
||||
'*.egg-info', |
||||
)), |
||||
description='A RLFH implementation (ChatGPT) powered by ColossalAI', |
||||
long_description=fetch_readme(), |
||||
long_description_content_type='text/markdown', |
||||
license='Apache Software License 2.0', |
||||
url='https://github.com/hpcaitech/ChatGPT', |
||||
install_requires=fetch_requirements('requirements/requirements.txt'), |
||||
python_requires='>=3.6', |
||||
classifiers=[ |
||||
'Programming Language :: Python :: 3', |
||||
'License :: OSI Approved :: Apache Software License', |
||||
'Environment :: GPU :: NVIDIA CUDA', |
||||
'Topic :: Scientific/Engineering :: Artificial Intelligence', |
||||
'Topic :: System :: Distributed Computing', |
||||
], |
||||
) |
||||
@ -0,0 +1,117 @@
|
||||
import os |
||||
from copy import deepcopy |
||||
from functools import partial |
||||
|
||||
import pytest |
||||
import torch |
||||
import torch.distributed as dist |
||||
import torch.multiprocessing as mp |
||||
from chatgpt.experience_maker import NaiveExperienceMaker |
||||
from chatgpt.nn import GPTActor, GPTCritic, RewardModel |
||||
from chatgpt.replay_buffer import NaiveReplayBuffer |
||||
from chatgpt.trainer.strategies import ColossalAIStrategy, DDPStrategy |
||||
|
||||
from colossalai.testing import rerun_if_address_is_in_use |
||||
from colossalai.utils import free_port |
||||
|
||||
|
||||
def get_data(batch_size: int, seq_len: int = 10) -> dict: |
||||
input_ids = torch.randint(0, 50257, (batch_size, seq_len), device='cuda') |
||||
attention_mask = torch.ones_like(input_ids) |
||||
return dict(input_ids=input_ids, attention_mask=attention_mask) |
||||
|
||||
|
||||
def gather_and_equal(tensor: torch.Tensor) -> bool: |
||||
world_size = dist.get_world_size() |
||||
outputs = [torch.empty_like(tensor) for _ in range(world_size)] |
||||
dist.all_gather(outputs, tensor.contiguous()) |
||||
for t in outputs[1:]: |
||||
if not torch.equal(outputs[0], t): |
||||
return False |
||||
return True |
||||
|
||||
|
||||
def run_test_data(strategy): |
||||
EXPERINCE_BATCH_SIZE = 4 |
||||
SAMPLE_BATCH_SIZE = 2 |
||||
|
||||
if strategy == 'ddp': |
||||
strategy = DDPStrategy() |
||||
elif strategy == 'colossalai': |
||||
strategy = ColossalAIStrategy(placement_policy='cuda') |
||||
else: |
||||
raise ValueError(f'Unsupported strategy "{strategy}"') |
||||
|
||||
actor = GPTActor().cuda() |
||||
critic = GPTCritic().cuda() |
||||
|
||||
initial_model = deepcopy(actor) |
||||
reward_model = RewardModel(deepcopy(critic.model)).cuda() |
||||
|
||||
experience_maker = NaiveExperienceMaker(actor, critic, reward_model, initial_model) |
||||
replay_buffer = NaiveReplayBuffer(SAMPLE_BATCH_SIZE, cpu_offload=False) |
||||
|
||||
# experience of all ranks should be the same |
||||
for _ in range(2): |
||||
data = get_data(EXPERINCE_BATCH_SIZE) |
||||
assert gather_and_equal(data['input_ids']) |
||||
assert gather_and_equal(data['attention_mask']) |
||||
experience = experience_maker.make_experience(**data, |
||||
do_sample=True, |
||||
max_length=16, |
||||
eos_token_id=50256, |
||||
pad_token_id=50256) |
||||
assert gather_and_equal(experience.sequences) |
||||
assert gather_and_equal(experience.action_log_probs) |
||||
assert gather_and_equal(experience.values) |
||||
assert gather_and_equal(experience.reward) |
||||
assert gather_and_equal(experience.advantages) |
||||
assert gather_and_equal(experience.action_mask) |
||||
assert gather_and_equal(experience.attention_mask) |
||||
replay_buffer.append(experience) |
||||
|
||||
# replay buffer's data should be the same |
||||
buffer_size = torch.tensor([len(replay_buffer)], device='cuda') |
||||
assert gather_and_equal(buffer_size) |
||||
for item in replay_buffer.items: |
||||
assert gather_and_equal(item.sequences) |
||||
assert gather_and_equal(item.action_log_probs) |
||||
assert gather_and_equal(item.values) |
||||
assert gather_and_equal(item.reward) |
||||
assert gather_and_equal(item.advantages) |
||||
assert gather_and_equal(item.action_mask) |
||||
assert gather_and_equal(item.attention_mask) |
||||
|
||||
# dataloader of each rank should have the same size and different batch |
||||
dataloader = strategy.setup_dataloader(replay_buffer) |
||||
dataloader_size = torch.tensor([len(dataloader)], device='cuda') |
||||
assert gather_and_equal(dataloader_size) |
||||
for experience in dataloader: |
||||
assert not gather_and_equal(experience.sequences) |
||||
assert not gather_and_equal(experience.action_log_probs) |
||||
assert not gather_and_equal(experience.values) |
||||
assert not gather_and_equal(experience.reward) |
||||
assert not gather_and_equal(experience.advantages) |
||||
# action mask and attention mask may be same |
||||
|
||||
|
||||
def run_dist(rank, world_size, port, strategy): |
||||
os.environ['RANK'] = str(rank) |
||||
os.environ['LOCAL_RANK'] = str(rank) |
||||
os.environ['WORLD_SIZE'] = str(world_size) |
||||
os.environ['MASTER_ADDR'] = 'localhost' |
||||
os.environ['MASTER_PORT'] = str(port) |
||||
run_test_data(strategy) |
||||
|
||||
|
||||
@pytest.mark.dist |
||||
@pytest.mark.parametrize('world_size', [2]) |
||||
@pytest.mark.parametrize('strategy', ['ddp', 'colossalai']) |
||||
@rerun_if_address_is_in_use() |
||||
def test_data(world_size, strategy): |
||||
run_func = partial(run_dist, world_size=world_size, port=free_port(), strategy=strategy) |
||||
mp.spawn(run_func, nprocs=world_size) |
||||
|
||||
|
||||
if __name__ == '__main__': |
||||
test_data(2, 'colossalai') |
||||
Loading…
Reference in new issue