diff --git a/.github/workflows/run_chatgpt_examples.yml b/.github/workflows/run_chatgpt_examples.yml index f25a6189fcf2..1194825369ff 100644 --- a/.github/workflows/run_chatgpt_examples.yml +++ b/.github/workflows/run_chatgpt_examples.yml @@ -21,7 +21,7 @@ jobs: container: image: image-cloud.luchentech.com/hpcaitech/pytorch-cuda:2.2.2-12.1.0 options: --gpus all --rm -v /data/scratch/examples-data:/data/scratch/examples-data --shm-size=10.24gb - timeout-minutes: 60 + timeout-minutes: 180 defaults: run: shell: bash @@ -34,6 +34,10 @@ jobs: pip install --no-cache-dir -v -e . - name: Install ChatGPT + env: + CFLAGS: "-O1" + CXXFLAGS: "-O1" + MAX_JOBS: 4 run: | cd applications/ColossalChat pip install --no-cache-dir -v . diff --git a/.github/workflows/run_chatgpt_unit_tests.yml b/.github/workflows/run_chatgpt_unit_tests.yml index 9180ede37a31..5cbcaa7fbbc0 100644 --- a/.github/workflows/run_chatgpt_unit_tests.yml +++ b/.github/workflows/run_chatgpt_unit_tests.yml @@ -21,7 +21,7 @@ jobs: container: image: image-cloud.luchentech.com/hpcaitech/pytorch-cuda:2.2.2-12.1.0 options: --gpus all --rm -v /data/scratch/examples-data:/data/scratch/examples-data - timeout-minutes: 30 + timeout-minutes: 180 defaults: run: shell: bash @@ -30,6 +30,10 @@ jobs: uses: actions/checkout@v2 - name: Install ChatGPT + env: + CFLAGS: "-O1" + CXXFLAGS: "-O1" + MAX_JOBS: 4 run: | cd applications/ColossalChat pip install -v . diff --git a/.gitignore b/.gitignore index 8bc74b4c8c2c..94d952295410 100644 --- a/.gitignore +++ b/.gitignore @@ -163,3 +163,13 @@ coverage.xml # log, test files - ColossalChat applications/ColossalChat/logs applications/ColossalChat/tests/logs +applications/ColossalChat/wandb +applications/ColossalChat/model +applications/ColossalChat/eval +applications/ColossalChat/rollouts +applications/ColossalChat/*.txt +applications/ColossalChat/*.db +applications/ColossalChat/stdin +applications/ColossalChat/*.zip +applications/ColossalChat/*.prof +applications/ColossalChat/*.png diff --git a/applications/ColossalChat/coati/dataset/loader.py b/applications/ColossalChat/coati/dataset/loader.py index fc3a4930c058..16fd385bad30 100755 --- a/applications/ColossalChat/coati/dataset/loader.py +++ b/applications/ColossalChat/coati/dataset/loader.py @@ -352,14 +352,30 @@ def apply_chat_template_and_mask( tokenizer: PreTrainedTokenizer, chat: List[Dict[str, str]], max_length: Optional[int] = None, + system_prompt: str = None, padding: bool = True, truncation: bool = True, ignore_idx: int = -100, ) -> Dict[str, torch.Tensor]: + + if system_prompt is None: + system_prompt = "You are a helpful assistant. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within and tags, respectively, i.e., reasoning process here answer here . Now the user asks you to solve a math problem that involves reasoning. After thinking, when you finally reach a conclusion, clearly output the final answer without explanation within the tags, i.e., 123 .\n\n" + + system_element = { + "role": "system", + "content": system_prompt, + } + + # Format for RL. + if "messages" in chat: + gt_answer = chat.get("gt_answer", None) + test_cases = chat.get("test_cases", None) + chat = [chat["messages"]] + tokens = [] assistant_mask = [] for i, msg in enumerate(chat): - msg_tokens = tokenizer.apply_chat_template([msg], tokenize=True) + msg_tokens = tokenizer.apply_chat_template([system_element, msg], tokenize=True, add_generation_prompt=True) # remove unexpected bos token if i > 0 and msg_tokens[0] == tokenizer.bos_token_id: msg_tokens = msg_tokens[1:] @@ -372,14 +388,10 @@ def apply_chat_template_and_mask( if max_length is not None: if padding and len(tokens) < max_length: to_pad = max_length - len(tokens) - if tokenizer.padding_side == "right": - tokens.extend([tokenizer.pad_token_id] * to_pad) - assistant_mask.extend([False] * to_pad) - attention_mask.extend([0] * to_pad) - else: - tokens = [tokenizer.pad_token_id] * to_pad + tokens - assistant_mask = [False] * to_pad + assistant_mask - attention_mask = [0] * to_pad + attention_mask + # Left padding for generation. + tokens = [tokenizer.pad_token_id] * to_pad + tokens + assistant_mask = [False] * to_pad + assistant_mask + attention_mask = [0] * to_pad + attention_mask if truncation and len(tokens) > max_length: tokens = tokens[:max_length] assistant_mask = assistant_mask[:max_length] @@ -389,6 +401,15 @@ def apply_chat_template_and_mask( labels = input_ids.clone() labels[~torch.tensor(assistant_mask, dtype=torch.bool)] = ignore_idx + if gt_answer is not None: + return {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels, "gt_answer": gt_answer} + elif test_cases is not None: + return { + "input_ids": input_ids, + "attention_mask": attention_mask, + "labels": labels, + "test_cases": test_cases, + } return { "input_ids": input_ids, "attention_mask": attention_mask, @@ -402,7 +423,7 @@ class RawConversationDataset(Dataset): Each instance is a dictionary with fields `system`, `roles`, `messages`, `offset`, `sep_style`, `seps`. """ - def __init__(self, tokenizer: PreTrainedTokenizer, input_file: str, max_length: int) -> None: + def __init__(self, tokenizer: PreTrainedTokenizer, input_file: str, max_length: int, system_prompt: str) -> None: self.tokenizer = tokenizer self.raw_texts = [] with jsonlines.open(input_file) as f: @@ -410,6 +431,7 @@ def __init__(self, tokenizer: PreTrainedTokenizer, input_file: str, max_length: self.raw_texts.append(line) self.tokenized_texts = [None] * len(self.raw_texts) self.max_length = max_length + self.system_prompt = system_prompt def __len__(self) -> int: return len(self.raw_texts) @@ -417,6 +439,23 @@ def __len__(self) -> int: def __getitem__(self, index: int): if self.tokenized_texts[index] is None: message = self.raw_texts[index] - tokens = apply_chat_template_and_mask(self.tokenizer, message, self.max_length) + tokens = apply_chat_template_and_mask(self.tokenizer, message, self.max_length, self.system_prompt) self.tokenized_texts[index] = dict(tokens) return self.tokenized_texts[index] + + +def collate_fn_grpo(batch): + input_ids = [item["input_ids"] for item in batch] + attention_mask = [item["attention_mask"] for item in batch] + labels = [item["labels"] for item in batch] + # Assume input_ids, attention_mask, labels are already of the same length, + # otherwise use pad_sequence(input_ids, batch_first=True, padding_value=tokenizer.pad_token_id) + input_ids = torch.stack(input_ids) + attention_mask = torch.stack(attention_mask) + labels = torch.stack(labels) + ret = {"input_ids": input_ids, "attention_mask": attention_mask, "labels": labels} + if "test_cases" in batch[0]: + ret["test_cases"] = [item["test_cases"] for item in batch] + if "gt_answer" in batch[0]: + ret["gt_answer"] = [item["gt_answer"] for item in batch] + return ret diff --git a/applications/ColossalChat/coati/distributed/README.md b/applications/ColossalChat/coati/distributed/README.md new file mode 100644 index 000000000000..21647a8cc896 --- /dev/null +++ b/applications/ColossalChat/coati/distributed/README.md @@ -0,0 +1,323 @@ +# Distributed RL Framework for Language Model Fine-Tuning + +This repository implements a distributed Reinforcement Learning (RL) training framework designed to fine-tune large language models using algorithms such as **GRPO** and **DAPO**. It supports multi-node and multi-GPU setups, scalable rollout generation, and policy optimization using libraries like VLLM. Currently, we support two Reinforcement Learning with Verifiable Reward (RLVR) tasks: solving math problems and code generation. + +**Please note that we are still under intensive development, stay tuned.** + +--- + +## 🚀 Features + +* **Distributed Training with Ray**: Scalable to multiple machines and GPUs. +* **Support for GRPO and DAPO**: Choose your preferred policy optimization algorithm. +* **Model Backends**: Support `vllm` as inference backends. +* **Rollout and Policy Decoupling**: Efficient generation and consumption of data through parallel inferencer-trainer architecture. +* **Evaluation Integration**: Easily plug in task-specific eval datasets. +* **Checkpoints and Logging**: Configurable intervals and directories. + +--- + +## 🛠 Installation + +### Prepare Develop Environment + +Install Colossalai & ColossalChat +```bash +git clone https://github.com/hpcaitech/ColossalAI.git +git checkout grpo-latest +BUILD_EXT=1 pip install -e . + +cd ./applications/ColossalChat +pip install -e . +``` + +Install vllm +```bash +pip install vllm==0.7.3 +``` + +Install Ray. +```bash +pip install ray +``` + +Install Other Dependencies +```bash +pip install cupy-cuda12x +python -m cupyx.tools.install_library --cuda 12.x --library nccl +``` + +To support long input/output sequence length (e.g., 32K), you may need to manually change the default setting (180 seconds) for the `timeout_s` variable in your ray installation to a larger value as shown in the screenshot below. + +
+

+ +

+
+ +Prepare Model & dataset +```bash +huggingface-cli download --local-dir-use-symlinks False Qwen/Qwen2.5-7B --local-dir /models/Qwen/Qwen2.5-7B +``` + +## Architecture Design + +
+

+ +

+
+Producer-Consumer Pattern: a classic software design pattern used for managing resources, data, or tasks between two different processes or threads. + +* Producer: inference engine which rollouts out examples and saves them into a shared buffer. +* Consumer: training framework which takes training examples from the shared buffer and train the policy model. + +Key features for Producer-Consumer Pattern: +* Buffer: Acts as a shared queue where the producer adds data and the consumer removes data. +* Concurrency: Rollout and training can work concurrently. + +## 🧠 Data Format + +Samples in the training or evaluation `.jsonl` file should follow the format specific to the type of task. We currently support two RLVR tasks: solving math problems and code generation. + +### Math Data Format +```json +{ + "messages": { + "role": "user", + "content": "Simplify $\\sqrt[3]{1+8} \\cdot \\sqrt[3]{1+\\sqrt[3]{8}}$." + }, + "gt_answer": "3" +} +``` + +### Code Data Format +We support [Prime code dataset format](https://github.com/PRIME-RL/PRIME). Inputs and outputs in test cases should be two lists containing only strings and matching in the number of elements. Your prompt must properly instruct the LLM to generate code to read test cases from stdin and output results to stdout. +```json +{ + "messages": { + "role": "user", + "content": "Solve the following coding problem using the programming language python:\n\nMikhail walks on a Cartesian plane. He starts at the point $(0, 0)$, and in one move he can go to any of eight adjacent points. For example, ..." + }, + "test_cases": { + "inputs": [ + "3\n2 2 3\n4 3 7\n10 1 9\n" + ], + "outputs": [ + "1\n6\n-1\n" + ] + } +} +``` + +--- + +## ⚙️ Hyperparameters & Arguments + +| Argument | Description | Example | +| ---------------- | --------------------------------------- | ----------------- | +| `--model` | Model path or identifier | `/path/to/model` | +| `--dataset` | Path to training `.jsonl` | `/path/to/train_data.jsonl` | +| `--eval-dataset` | JSON of task\:eval\_dataset\_path pairs | `{"eval_1":"/path/to/eval_1.jsonl"}` | +| `--project` | Project name | `Project1` | +| `--num-episodes` | Number of training episodes | `1` | + +### Distributed Training + +| Argument | Description | Example | +| ----------------------------- | ------------------------------------- | ------- | +| `--num-trainers` | Number of trainer processes | `4` | +| `--num-inferencer` | Number of inferencer processes | `4` | +| `--inference-batch-size` | Prompts per inference step | `8` | +| `--inference-microbatch-size` | Per-GPU batch size for inference | `8` | +| `--train-batch-size` | Prompts per trainer step per dp group | `8` | +| `--train-minibatch-size` | Mini-batch size before forward pass | `8` | +| `--train-microbatch-size` | Per-GPU batch size for training | `2` | + +### Sampling + +| Argument | Description | Example | +| --------------------- | --------------------- | -------------- | +| `--backend` | Generation backend, choose from `vllm` | `vllm` | +| `--temperature` | Sampling temperature for generation | `1.0` | +| `--top-k` | Top-K sampling parameter for generation | `None` | +| `--top-p` | Top-P sampling parameter for generation | `1.0` | +| `--system-prompt` | System prompt, optional, default to the default system prompt for each reward types. For more information, refer to the [**reward type**](#-constraints-and-notes) section | `Please reason step by step, and put your final answer within \\boxed{}.` | +| `--max-new-tokens` | Max generation tokens | `3584` | +| `--max-prompt-tokens` | Max prompt tokens | `512` | + +### GRPO Specific + +| Argument | Description | Example | +| ----------------- | ---------------------------- | ------------------- | +| `--algo` | Algorithm (`GRPO` or `DAPO`), for more customization refer to [GRPO Settings](#️-grpo-settings) | `GRPO` | +| `--learning-rate` | Learning rate | `1e-6` | +| `--kl-coeff` | KL penalty coefficient, if nonzero, a reference model will be used | `0.01` | +| `--reward-type` | Reward signal type (choose from 'think_answer_tags', 'boxed', 'code') For more information, refer to the [**reward type**](#-constraints-and-notes) section | `think_answer_tags` | +| `--eval-interval` | Evaluation interval in number of training steps (positive value to enable evaluation) | `10` | + +### Logging and Checkpointing + +| Argument | Description | Example | +| -------------------- | ------------------------- | ------------ | +| `--save-interval` | Training steps between checkpoints | `20` | +| `--save-dir` | Checkpoint directory | `./model` | +| `--eval-save-dir` | Evaluation save path | `./eval` | +| `--rollout-save-dir` | Rollout logs directory | `./rollouts` | + +### Miscellaneous + +| Argument | Description | Example | +| ------------------ | --------------------------------------- | ------- | +| `--ray_dir` | Custom Ray temp dir of a running Ray cluster (optional) | `None` | +| `--master_address` | Master address of a running Ray cluster | `None` | +| `--master_port` | Master port for torch DDP | `29506` | + +--- + +## ⚙️ GRPO Settings + +In addition to the two default training settings provided—`GRPO` and `DAPO`—users can customize their training by changing the following hyperparameters in `grpo_config` in `rl_example.py`. + +| Argument Name | Description | Default | +| ----------------------------- | ---------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------- | +| `filter_range` | Filters out rollout group if the success rate within that group is out of this range.| `[0.01, 0.99]` | +| `dynamic_batching` | Enables dynamic batching as described in the [DAPO paper](https://arxiv.org/abs/2503.14476). | `True` | +| `clip_eps_low` | epsilon_low in DAPO in equation in [DAPO paper](https://arxiv.org/abs/2503.14476) | `0.2` | +| `clip_eps_high` | epsilon_high in DAPO equation in [DAPO paper](https://arxiv.org/abs/2503.14476) | `0.28` | +| `skip_threshold` | If ratio is above this threshold, the sample is skipped to avoid instability. | `20.0` | +| `loss_variation` | Type of loss variation. Supports `"token_level"` for token-wise policy gradient loss and `sample_level` for original GRPO loss. | `"token_level"` | +| `soft_over_length_punishment` | Whether to use soft overlength penalty in [DAPO paper](https://arxiv.org/abs/2503.14476) or not. | `True` | +| `cache_length` | `L_cache` parameter for soft overlength penalty in e.q. 13 in [DAPO paper](https://arxiv.org/abs/2503.14476) | `min(1024, int(args.max_new_tokens / 4))` | +| `filter_truncated_response` | Mask out truncated responses in loss calculation. | `True` | + + + +## 🔄 Constraints and Notes + +* `num_inferencer + num_trainer == NUM_GPUs` +* `num_inferencer % num_trainer == 0` +* `(num_inferencer * inference_batch_size) % (num_trainer * train_batch_size) == 0` +* `train_batch_size >= train_minibatch_size >= train_microbatch_size` +* `inference_batch_size >= inference_microbatch_size` +* Set microbatch sizes based on **VRAM capacity** +* To use tensor parallelism on inferencer + * set backend to `vllm` + * change `tensor_parallel_size` in `inference_model_config` in rl_example.py + * set `num_inferencer = NUM_INFERENCE_GPUs / tensor_parallel_size` +* To set tensor parallelism / pipeline parallelism / zero stage + * change corresponding settings in `plugin_config` in rl_example.py +* Ensure rollout generation rate matches trainer consumption: + + ``` + num_inferencer * inference_batch_size % ( + num_trainer * train_batch_size / + train_pipeline_parallelism_size / + train_tensor_parallelism_size + ) == 0 + ``` +* Model weights sync every: + + ``` + (num_inferencer * inference_batch_size) / + (num_trainer * train_batch_size / + train_pipeline_parallelism_size / + train_tensor_parallelism_size) + ``` +* Reward Type + + We currently support three reward types--- `think_answer_tags`, `boxed`, `code`, each varies in details such as how answer is extracted and the reward calculation process. Please select one from `think_answer_tags`, `boxed` for math problem solving and use `code` for code generation. The default system prompt for each reward type is as follows. Please make sure your system prompt provides information for the answer to be correctly extracted from model responses. + + * think_answer_tags + + Answer extraction: extract the content between the last ``, `` tags. + + ``` + You are a helpful assistant. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within and tags, respectively, i.e., reasoning process here answer here . Now the user asks you to solve a math problem that involves reasoning. After thinking, when you finally reach a conclusion, clearly output the final answer without explanation within the tags, i.e., 123 .\n\n + ``` + * boxed + + Answer extraction: extract the last content marked by `\\boxed{}` + ``` + Please reason step by step, and put your final answer within \\boxed{}. + ``` + * code + + Answer extraction: extract code inside ` ```python\n...``` ` + ``` + You are a helpful assistant. + ``` +--- + +## 🧪 Example: single machine 8-GPU Zero2 Strategy + +```bash +python rl_example.py \ + --dataset /path/to/train_data.jsonl \ + --model /path/to/Qwen2.5-3B/ \ + -t 4 -i 4 \ + -b vllm \ + -ibs 2 -tbs 4 -tMbs 1 -tmbs 4 -imbs 1 \ + -rt boxed \ + -g 4 \ + -ibs 1 \ + -tbs 2 \ + -tMbs 1 \ + -tmbs 2 \ + -imbs 1 \ + -s "Please reason step by step, and put your final answer within \\boxed{}." \ + -tMbs 8 \ + -p GRPO-Train-Align-Debug \ +``` + +## 🧪 Example: multi-machine TP+PP Strategy + +### Create ray cluster on multi-machine +For example, now we have 4 nodes and their IPs are 10.0.0.3, 10.0.0.4, 10.0.0.5, 10.0.0.6. +We use 10.0.0.3 as master node. First we start a ray cluster on 10.0.0.3: +```bash +ray start --head --node-ip-address=10.0.0.3 +``` + +Then, for each slave node (10.0.0.4/10.0.0.5/10.0.0.6), we add to the ray cluster by following code: +```bash +ray start --address='10.0.0.3:6379' +``` + +Modify plugin_config in ./applications/ColossalChat/rl_example.py +```python +plugin_config={ + "tp_size": 4, + "pp_size": 2, + "microbatch_size": max( + 1, args.train_microbatch_size // 2 + ), # microbatch size should be set to train_microbatch_size // pp_size + "zero_stage": 1, + "max_norm": 1.0, + }, # for pp, tp +``` + +```bash +# Hint1: replace /models/Qwen/Qwen2.5-7B to your model path +# replace /datasets/train-alignment.jsonl to your dataset path +python rl_example.py + -m /path/to/Qwen2.5-Math-7B/ \ + -d /path/to/train_data.jsonl \ + --master_address '10.0.0.3' + -t 16 \ + -i 16 \ + -p GRPO-Train-Align-Debug \ + -g 2 \ + -ibs 1 \ + -tbs 2 \ + -tMbs 1 \ + -tmbs 2 \ + -imbs 1 \ + -b vllm \ + -e 2 \ + -rt boxed \ + -s "Please reason step by step, and put your final answer within \\boxed{}." +``` + +## Acknowledgement +Colossal-RL is a distributed version of ColossalChat and inspired by a few awesome open-source projects. We would like to express our gratitude to the following awesome open-source projects and algorithms: GRPO, DAPO, TRL, Verl, OpenRLHF, StreamRL, Qwen, Logic-RL. diff --git a/applications/ColossalChat/coati/distributed/__init__.py b/applications/ColossalChat/coati/distributed/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/applications/ColossalChat/coati/distributed/comm.py b/applications/ColossalChat/coati/distributed/comm.py new file mode 100644 index 000000000000..3824303f55bd --- /dev/null +++ b/applications/ColossalChat/coati/distributed/comm.py @@ -0,0 +1,57 @@ +from typing import Any, Dict + +import ray.util.collective as cc +import torch +import torch.distributed.distributed_c10d as c10d +from packaging.version import Version + + +def ray_broadcast_object(obj: Any, src: int = 0, device=None, group_name: str = "default") -> Any: + rank = cc.get_rank(group_name) + if rank == src: + if Version(torch.__version__) >= Version("2.3.0"): + obj_tensor, size_tensor = c10d._object_to_tensor(obj, device=device, group=None) + elif Version(torch.__version__) >= Version("1.13.0"): + obj_tensor, size_tensor = c10d._object_to_tensor(obj, device=device) + else: + obj_tensor, size_tensor = c10d._object_to_tensor(obj) + obj_tensor = obj_tensor.to(device) + size_tensor = size_tensor.to(device) + else: + size_tensor = torch.empty(1, dtype=torch.int64, device=device) + cc.broadcast(size_tensor, src, group_name) + if rank != src: + obj_tensor = torch.empty(size_tensor.item(), dtype=torch.uint8, device=device) + cc.broadcast(obj_tensor, src, group_name) + if rank != src: + if Version(torch.__version__) >= Version("2.3.0"): + obj = c10d._tensor_to_object(obj_tensor, size_tensor.item(), group=None) + else: + obj = c10d._tensor_to_object(obj, size_tensor.item()) + return obj + + +def ray_broadcast_tensor_dict( + tensor_dict: Dict[str, torch.Tensor], src: int = 0, device=None, group_name: str = "default" +) -> Dict[str, torch.Tensor]: + rank = cc.get_rank(group_name) + if rank == src: + metadata = [] + for k, v in tensor_dict.items(): + metadata.append((k, v.shape, v.dtype)) + else: + metadata = None + metadata = ray_broadcast_object(metadata, src, device, group_name) + if rank != src: + out_dict = {} + for k, shape, dtype in metadata: + if rank == src: + tensor = tensor_dict[k] + else: + tensor = torch.empty(shape, dtype=dtype, device=device) + cc.broadcast(tensor, src, group_name) + if rank != src: + out_dict[k] = tensor + if rank == src: + out_dict = tensor_dict + return out_dict diff --git a/applications/ColossalChat/coati/distributed/consumer.py b/applications/ColossalChat/coati/distributed/consumer.py new file mode 100644 index 000000000000..ba7d882c9dea --- /dev/null +++ b/applications/ColossalChat/coati/distributed/consumer.py @@ -0,0 +1,400 @@ +import os +from contextlib import nullcontext +from typing import Any, Dict, Optional + +import ray +import ray.util.collective as cc +import torch +import torch.distributed as dist +from coati.distributed.profiling_utils import CustomProfiler +from tqdm import tqdm +from transformers import AutoModelForCausalLM + +from colossalai.booster import Booster +from colossalai.booster.plugin import HybridParallelPlugin +from colossalai.initialize import launch +from colossalai.nn.optimizer import HybridAdam +from colossalai.utils import get_current_device + +from .comm import ray_broadcast_tensor_dict +from .utils import bind_batch, post_recv, unbind_batch + + +class BaseConsumer: + def __init__( + self, + num_producers: int, + num_episodes: int, + rank: int, + world_size: int, + master_addr: str, + master_port: int, + num_update_per_episode: int, + num_recv_per_update: int, + batch_size: int, + model_config: Dict[str, Any], + plugin_config: Dict[str, Any], + minibatch_size: int = 1, + save_interval: int = 100, + save_dir: str = "./model", + enable_profiling: bool = False, + n_behind: int = 0, + ): + self.num_producers = num_producers + self.num_episodes = num_episodes + self.rank = rank + self.world_size = world_size + self.master_addr = master_addr + self.master_port = master_port + self.num_update_per_episode = num_update_per_episode + self.num_recv_per_update = num_recv_per_update + self.batch_size = batch_size + self.minibatch_size = minibatch_size + self.save_interval = save_interval + self.save_dir = save_dir + self.enable_profiling = enable_profiling + assert batch_size % minibatch_size == 0, "batch_size should be divisible by microbatch_size" + self.num_microbatches = batch_size // minibatch_size + + self.model_config = model_config + self.plugin_config = plugin_config + + self.device = get_current_device() + self.lr_scheduler = None + self.n_behind = n_behind + + def setup(self) -> None: + launch(self.rank, self.world_size, self.master_addr, self.master_port, local_rank=0) + + plugin_config = dict(tp_size=1, pp_size=1, precision="bf16", zero_stage=2) + if ( + self.plugin_config.get("pp_size", 1) > 1 + and "num_microbatches" not in self.plugin_config + and "microbatch_size" not in self.plugin_config + ): + plugin_config["microbatch_size"] = max(1, self.minibatch_size // plugin_config.get("pp_size", 1)) + plugin_config.update(self.plugin_config) + self.plugin = HybridParallelPlugin(**plugin_config) + self.booster = Booster(plugin=self.plugin) + self.dp_rank = dist.get_rank(self.plugin.dp_group) + self.tp_rank = dist.get_rank(self.plugin.tp_group) + self.pp_rank = dist.get_rank(self.plugin.pp_group) + + self.dp_size = dist.get_world_size(self.plugin.dp_group) + self.tp_size = dist.get_world_size(self.plugin.tp_group) + self.pp_size = dist.get_world_size(self.plugin.pp_group) + + # Init Hybrid ray process group + for i in range(self.num_producers): + cc.init_collective_group(self.world_size + 1, self.rank + 1, group_name=f"sync_data_{i}") + if self.pp_size > 1: + # use hybrid tp + pp + if self.tp_rank == 0 and self.dp_rank == 0: + cc.init_collective_group( + self.num_producers + 1, self.num_producers, group_name=f"sync_model_{self.pp_rank}" + ) + else: + if self.rank == 0: + cc.init_collective_group(self.num_producers + 1, self.num_producers, group_name="sync_model") + + self.buffer = [] + self.recv_cnt = 0 + self.profiler = CustomProfiler(f"C{self.rank}", disabled=not self.enable_profiling) + + def state_dict(self) -> Dict[str, torch.Tensor]: + raise NotImplementedError + + def step(self, step_idx: int, **kwargs) -> Optional[float]: + raise NotImplementedError + + def prepare_mini_batch(self, effective_group_to_raw_group_mapping: Dict[int, int]) -> Dict[str, torch.Tensor]: + """ + Prepare a mini-batch from the effective group to raw group mapping. + This method is used to create a mini-batch for training. + """ + batches = [ + self.buffer[effective_group_to_raw_group_mapping[i]] + for i in range(self.dp_rank * self.minibatch_size, (self.dp_rank + 1) * self.minibatch_size) + ] + # every dp_rank will receive a complete mini-batch, no need to sync within step() later + # each mini-batch use the first self.dp_size * minibatch_size effective samples + raw_mini_batches = self.buffer[ + : effective_group_to_raw_group_mapping[self.dp_size * self.minibatch_size - 1] + 1 + ] # include the last effective sample + raw_mini_batches_metric_dict = { + "raw_train_mini_batch_reward": [t[1] for t in raw_mini_batches], + "raw_train_mini_batch_format_acc": [t[2] for t in raw_mini_batches], + "raw_train_mini_batch_ans_acc": [t[3] for t in raw_mini_batches], + "raw_train_mini_batch_response_len": [t[4] for t in raw_mini_batches], + } + batch = bind_batch([t[0] for t in batches]) + batch = post_recv(batch) + return batch, raw_mini_batches_metric_dict + + def calculate_effective_group_to_raw_group_mapping(self, step): + effective_group_to_raw_group_mapping = {} + for buffer_idx in range(len(self.buffer)): + if self.buffer[buffer_idx][0] is not None: + if self.n_behind == 0: + effective_group_to_raw_group_mapping[len(effective_group_to_raw_group_mapping)] = buffer_idx + else: + if self.buffer[buffer_idx][-1] <= step - self.n_behind: + effective_group_to_raw_group_mapping[len(effective_group_to_raw_group_mapping)] = buffer_idx + return effective_group_to_raw_group_mapping + + def loop(self) -> None: + print( + f"Consumer{self.rank} num_update: {self.num_update_per_episode}, num_recv: {self.num_recv_per_update}, nmb: {self.num_microbatches}" + ) + for episode in range(self.num_episodes): + with tqdm( + range(self.num_update_per_episode), + desc=f"Episode {episode} with rollout step(s)", + disable=self.rank != 0, + ) as pbar: + for step in pbar: + torch.cuda.reset_peak_memory_stats() + i = 0 + + self.profiler.enter(f"rollout_episode_{episode}_step_{step}") + for _ in range(self.num_recv_per_update): + if self.n_behind > 0: + # after sync model, do not wait for more data to arrive as rollout takes time, use buffered data + effective_group_to_raw_group_mapping = self.calculate_effective_group_to_raw_group_mapping( + step=step + ) + while len(effective_group_to_raw_group_mapping) >= self.dp_size * self.minibatch_size: + self.profiler.log( + f"Still have {len(effective_group_to_raw_group_mapping)} effective groups, greater than {self.dp_size * self.minibatch_size}, start training" + ) + batch, raw_mini_batches_metric_dict = self.prepare_mini_batch( + effective_group_to_raw_group_mapping + ) + self.profiler.enter("step") + loss = self.step(i, pbar, **batch, **raw_mini_batches_metric_dict) + self.profiler.exit("step") + self.buffer = self.buffer[ + effective_group_to_raw_group_mapping[self.dp_size * self.minibatch_size - 1] + 1 : + ] + # recalculate the effective group to raw group mapping + effective_group_to_raw_group_mapping_size_before = len( + effective_group_to_raw_group_mapping + ) + effective_group_to_raw_group_mapping = ( + self.calculate_effective_group_to_raw_group_mapping(step=step) + ) + assert ( + len(effective_group_to_raw_group_mapping) + == effective_group_to_raw_group_mapping_size_before + - self.dp_size * self.minibatch_size + ) + if loss is not None: + pbar.set_postfix({"loss": loss}) + i += 1 + + # receive data from producers + for r in range(self.num_producers): + print(f"[T{dist.get_rank()}] Recv data episode {episode} step {step} from {r}") + self.profiler.enter(f"recv_broadcast_data_P{r}") + raw_batch = ray_broadcast_tensor_dict( + None, src=0, device=self.device, group_name=f"sync_data_{r}" + ) + self.profiler.exit(f"recv_broadcast_data_P{r}") + # calculate group reward et al. filtering. As only the filtered group will be used for training (which is incomplete), + # we need to calculate the metrics before filtering here for logging + # [batch_size, num_generations, ...] -> [batch_size * num_generations, ...] + raw_batch = { + k: v.view(-1, self.num_generations, v.size(-1)) if k != "temperature" else v + for k, v in raw_batch.items() + } + # [batch_size, num_generations] -> [batch_size] + reward = raw_batch["reward"][:, :, 0] + format_acc = raw_batch["format_acc"][:, :, 0] + ans_acc = raw_batch["ans_acc"][:, :, 0] + response_len = ( + raw_batch["response_idx"][:, :, 1] - raw_batch["response_idx"][:, :, 0] + 1 + ).type(torch.float32) + effective_group_mask = None + if self.filter_range is not None and self.grpo_config.get("dynamic_batching", True): + # filter the group based on the reward and accuracy + group_ans_acc_mean = ans_acc.mean(dim=1) + effective_group_mask = torch.logical_and( + group_ans_acc_mean > self.filter_range[0], group_ans_acc_mean < self.filter_range[1] + ) + raw_batch = unbind_batch(raw_batch) # List[Dict[str, torch.Tensor]] + for group_idx, group_with_reward in enumerate(raw_batch): + self.buffer.append( + [ + ( + group_with_reward + if effective_group_mask is None or effective_group_mask[group_idx] + else None + ), + reward[group_idx], + format_acc[group_idx], + ans_acc[group_idx], + response_len[group_idx], + step, + ] + ) + if effective_group_mask is not None: + print( + f"[T{dist.get_rank()}] Filter recv data: {len(raw_batch)} -> {torch.sum(effective_group_mask).cpu().item()} effective groups" + ) + # mapping the effective group to the raw group for indexing + effective_group_to_raw_group_mapping = self.calculate_effective_group_to_raw_group_mapping( + step=step + ) + print( + f"[T{dist.get_rank()}] Collect Effective Prompt: {len(effective_group_to_raw_group_mapping)}/{self.dp_size * self.minibatch_size}" + ) + + if self.n_behind == 0: + # If n_behind is 0, we start training after receiving data from producers. + while len(effective_group_to_raw_group_mapping) >= self.dp_size * self.minibatch_size: + self.profiler.log( + f"Collect {len(effective_group_to_raw_group_mapping)} effective groups, greater than {self.dp_size * self.minibatch_size}, start training" + ) + batch, raw_mini_batches_metric_dict = self.prepare_mini_batch( + effective_group_to_raw_group_mapping + ) + self.profiler.enter("step") + loss = self.step(i, pbar, **batch, **raw_mini_batches_metric_dict) + self.profiler.exit("step") + self.buffer = self.buffer[ + effective_group_to_raw_group_mapping[self.dp_size * self.minibatch_size - 1] + 1 : + ] + # recalculate the effective group to raw group mapping + effective_group_to_raw_group_mapping_size_before = len( + effective_group_to_raw_group_mapping + ) + effective_group_to_raw_group_mapping = ( + self.calculate_effective_group_to_raw_group_mapping(step=step) + ) + assert ( + len(effective_group_to_raw_group_mapping) + == effective_group_to_raw_group_mapping_size_before + - self.dp_size * self.minibatch_size + ) + if loss is not None: + pbar.set_postfix({"loss": loss}) + i += 1 + + if self.lr_scheduler is not None: + self.lr_scheduler.step() + if (step + 1) % self.save_interval == 0 or (step + 1) == self.num_update_per_episode: + if self.rank == 0: + print(f"Start saving policy model at step {step + 1}.") + save_path = os.path.join(self.save_dir, f"modeling-episode-{episode}-step-{step + 1}") + self.booster.save_model(self.policy_model, save_path, shard=True) + if self.rank == 0: + print(f"Saved model checkpoint at step {step + 1} in folder {save_path}") + + if (episode != self.num_episodes - 1 or step != self.num_update_per_episode - 1) and ( + episode != 0 or step >= self.n_behind + ): + if self.pp_size > 1: + print( + f"[T{dist.get_rank()}] Sync model PP stage {self.pp_rank} episode {episode} step {step}" + ) + else: + print(f"[T{dist.get_rank()}] Sync model episode {episode} step {step}") + self.profiler.enter("sync_model") + torch.cuda.empty_cache() + state_dict = self.state_dict() + if self.pp_size > 1: + if self.tp_rank == 0 and self.dp_rank == 0: + ray_broadcast_tensor_dict( + state_dict, + src=self.num_producers, + device=self.device, + group_name=f"sync_model_{self.pp_rank}", + ) + else: + if self.rank == 0: + ray_broadcast_tensor_dict( + state_dict, src=self.num_producers, device=self.device, group_name="sync_model" + ) + del state_dict + torch.cuda.empty_cache() + self.profiler.exit("sync_model") + self.profiler.log(f"Peak memory usage: {torch.cuda.max_memory_allocated() / 1024 ** 2:.2f} MB") + self.profiler.exit(f"rollout_episode_{episode}_step_{step}") + + def __del__(self): + if hasattr(self, "profiler"): + self.profiler.close() + + +@ray.remote +class SimpleConsumer(BaseConsumer): + def __init__( + self, + num_producers, + num_episodes, + rank, + world_size, + master_addr, + master_port, + num_update_per_episode, + num_recv_per_update, + batch_size, + model_config, + plugin_config, + minibatch_size=1, + save_interval: int = 100, + save_dir="./model", + ): + super().__init__( + num_producers, + num_episodes, + rank, + world_size, + master_addr, + master_port, + num_update_per_episode, + num_recv_per_update, + batch_size, + model_config, + plugin_config, + minibatch_size, + save_interval, + save_dir, + ) + path = model_config.pop("path") + self.model = AutoModelForCausalLM.from_pretrained(path, **model_config) + self.model.train() + self.model.gradient_checkpointing_enable() + self.optimizer = HybridAdam(self.model.parameters(), lr=1e-3, weight_decay=0.01) + self.accum_loss = torch.zeros(1, device=self.device) + + def setup(self): + super().setup() + self.model, self.optimizer, *_ = self.booster.boost(self.model, self.optimizer) + + def step(self, step_idx: int, pbar: Any, **kwargs) -> Optional[float]: + labels = kwargs["input_ids"].clone() + labels[kwargs["attention_mask"] == 0] = -100 + kwargs["labels"] = labels + assert kwargs.pop("action_mask").shape == kwargs.pop("action_log_probs").shape + + need_update = (step_idx + 1) % self.num_microbatches == 0 + + ctx = nullcontext() if need_update else self.booster.no_sync(self.model, self.optimizer) + with ctx: + out = self.model(**kwargs) + loss = out.loss / self.num_microbatches + self.accum_loss.add_(loss.data) + self.booster.backward(loss, self.optimizer) + if need_update: + self.optimizer.step() + self.optimizer.zero_grad() + loss_scalar = self.accum_loss.item() + self.accum_loss.zero_() + return loss_scalar + + def state_dict(self): + self.model._force_wait_all_gather() + model = self.model.unwrap() + state_dict = model.state_dict() + return state_dict diff --git a/applications/ColossalChat/coati/distributed/grpo_consumer.py b/applications/ColossalChat/coati/distributed/grpo_consumer.py new file mode 100644 index 000000000000..424d460989ef --- /dev/null +++ b/applications/ColossalChat/coati/distributed/grpo_consumer.py @@ -0,0 +1,533 @@ +from contextlib import nullcontext +from typing import Any, Optional + +import ray +import torch +import wandb +from coati.distributed.consumer import BaseConsumer +from coati.distributed.loss import PolicyLoss +from coati.distributed.utils import entropy_from_logits, memory_efficient_logprob +from coati.trainer.utils import all_reduce_mean, all_reduce_sum +from transformers import AutoModelForCausalLM, AutoTokenizer + +from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR +from colossalai.nn.optimizer import HybridAdam + + +@ray.remote +class GRPOConsumer(BaseConsumer): + def __init__( + self, + num_producers, + num_episodes, + rank, + world_size, + master_addr, + master_port, + num_update_per_episode, + num_recv_per_update, + batch_size, + model_config, + plugin_config, + minibatch_size=1, + num_generations=8, + generate_config=None, + grpo_config={}, + save_interval: int = 100, + save_dir="./model", + project_name: str = None, + run_name: str = None, + wandb_group_name: str = None, + enable_profiling: bool = False, + n_behind: int = 0, + ): + print(f"Using GRPO config: {grpo_config}") + if ( + plugin_config.get("pp_size", 1) > 1 + and "num_microbatches" not in plugin_config + and "microbatch_size" not in plugin_config + ): + plugin_config["microbatch_size"] = max( + 1, grpo_config.get("train_microbatch_size") // plugin_config.get("pp_size", 1) + ) + super().__init__( + num_producers, + num_episodes, + rank, + world_size, + master_addr, + master_port, + num_update_per_episode, + num_recv_per_update, + batch_size, + model_config, + plugin_config, + minibatch_size, + save_interval=save_interval, + save_dir=save_dir, + enable_profiling=enable_profiling, + n_behind=n_behind, + ) + path = model_config.pop("path") + self.policy_model = AutoModelForCausalLM.from_pretrained(path, **model_config) + self.policy_model.train() + self.policy_model.gradient_checkpointing_enable() + self.optimizer = HybridAdam( + self.policy_model.parameters(), + lr=grpo_config.get("lr", 1e-6), + weight_decay=grpo_config.get("weight_decay", 0.01), + ) + self.accum_loss = torch.zeros(1, device=self.device) + self.accum_kl = torch.zeros(1, device=self.device) + self.accum_entropy = torch.zeros(1, device=self.device) + self.accum_advantages = torch.zeros(1, device=self.device) + self.raw_train_batch_reward = [] + self.raw_train_batch_format_acc = [] + self.raw_train_batch_ans_acc = [] + self.raw_train_batch_response_len = [] + self.accum_count = 0 + self.generate_config = generate_config + self.grpo_config = grpo_config + self.project_name = project_name + self.effective_sample_count = 0 + self.effective_prompt_count = 0 + self.project_name = project_name + self.run_name = run_name + self.wandb_group_name = wandb_group_name + + self.policy_loss_fn = PolicyLoss( + clip_eps_low=grpo_config.get("clip_eps_low", 0.2), + clip_eps_high=grpo_config.get("clip_eps_high", 0.2), + beta=grpo_config.get("beta", 0.01), + loss_variation=grpo_config.get("loss_variation", "sample_level"), + ) + + # Reference model is initialized from policy model. + if self.policy_loss_fn.beta > 0: + self.reference_model = AutoModelForCausalLM.from_pretrained(path, **model_config) + self.reference_model.eval() + + self.tokenizer = AutoTokenizer.from_pretrained(path) + self.pad_token_id = self.tokenizer.pad_token_id + self.num_generations = num_generations + self.filter_range = grpo_config.get("filter_range", None) + if self.filter_range is not None: + assert len(self.filter_range) == 2, "Filter range should have 2 values." + + self.filter_truncated_response = grpo_config.get("filter_truncated_response", False) + if self.filter_truncated_response: + self.max_length = 0 + if "max_tokens" in self.generate_config: + self.max_length = self.generate_config["max_tokens"] + elif "max_new_tokens" in self.generate_config: + self.max_length = self.generate_config["max_new_tokens"] + else: + raise ValueError( + "either max_tokens (vllm) or max_new_tokens (transformers) must be set in generate_config." + ) + # Initialize verifiable reward. + grpo_config.get("response_format_tags", None) + self.global_step = 0 + + self.lr_scheduler = CosineAnnealingWarmupLR( + optimizer=self.optimizer, + total_steps=min(self.num_episodes, 4) * self.num_update_per_episode, + warmup_steps=0, + eta_min=0.1 * grpo_config.get("lr", 1e-6), + ) + + def setup(self): + super().setup() + if (not self.plugin.pp_size > 1 and self.rank == 0) or ( + self.plugin.pp_size > 1 + and self.booster.plugin.stage_manager.is_last_stage() + and self.tp_rank == 0 + and self.dp_rank == 0 + ): + self.wandb_run = wandb.init( + project=self.project_name, + sync_tensorboard=False, + dir="./wandb", + name=self.run_name, + group=self.wandb_group_name, + ) + + self.policy_model, self.optimizer, _, _, self.lr_scheduler = self.booster.boost( + self.policy_model, self.optimizer, lr_scheduler=self.lr_scheduler + ) + if self.policy_loss_fn.beta > 0: + self.reference_model, *_ = self.booster.boost(self.reference_model) + self.plugin.logger.set_level("ERROR") + + def step(self, step_idx: int, pbar: Any, **kwargs) -> Optional[float]: + """ + Step data from policy model: + [{ + "input_ids": torch.Tensor, + "attention_mask": torch.Tensor, + "action_mask": torch.Tensor, + "action_log_probs": torch.Tensor, + }, + ...] + Format: + [minibatch_size, num_of_generation, prompt_length + response_length] --- ............. + """ + # Reshape to [minibatch_size x num_of_generation, prompt_length + response_length] + data = {k: v.view(-1, v.size(-1)) for k, v in kwargs.items() if "raw_train_mini_batch_" not in k} + self.raw_train_batch_reward.extend(kwargs["raw_train_mini_batch_reward"]) + self.raw_train_batch_format_acc.extend(kwargs["raw_train_mini_batch_format_acc"]) + self.raw_train_batch_ans_acc.extend(kwargs["raw_train_mini_batch_ans_acc"]) + self.raw_train_batch_response_len.extend(kwargs["raw_train_mini_batch_response_len"]) + action_mask = data["action_mask"] + num_action = action_mask.shape[1] + old_action_log_probs = data["action_log_probs"] + response_length = torch.sum(action_mask, dim=1).to(torch.float32) + train_microbatch_size = self.grpo_config.get("train_microbatch_size", data["input_ids"].size(0)) + + reward = data["reward"].view((-1)) + format_acc = data["format_acc"].view((-1)) + ans_acc = data["ans_acc"].view((-1)) + + # [minibatch_size, num_generations] + + group_reward = reward.view(-1, self.num_generations) + reward_mean = group_reward.mean(dim=1) + # [minibatch_size x num_generations] + reward_mean = reward_mean.repeat_interleave(self.num_generations, dim=0) + + reward_std = group_reward.std(dim=1).repeat_interleave(self.num_generations, dim=0) + # [minibatch_size x num_generations] + advantages = ((reward - reward_mean) / (reward_std + 1e-4)).unsqueeze(dim=-1) + + # [minibatch_size x num_of_generation] + loss_mask = torch.ones(action_mask.size(0), device=action_mask.device).bool() + + # filter out overlength samples + if self.filter_truncated_response and action_mask.size(1) == self.max_length: + loss_mask = torch.logical_and( + loss_mask, + action_mask[:, -1] == False, + ) + if self.filter_range is not None and self.grpo_config.get("dynamic_batching", False) == False: + # filter out samples with reward outside the range + # if dynamic batching is enabled, we filter out out of range groups before training + group_ans_acc_mean = ( + ans_acc.view(-1, self.num_generations).mean(dim=1).repeat_interleave(self.num_generations, dim=-1) + ) + loss_mask = torch.logical_and( + loss_mask, + torch.logical_and( + group_ans_acc_mean > self.filter_range[0], + group_ans_acc_mean < self.filter_range[1], + ), + ) + self.effective_prompt_count += group_reward.size(0) * self.dp_size + + mean_kl, mean_loss = [], [] + + if self.grpo_config.get("dynamic_batching", True): + need_update = self.effective_prompt_count >= self.batch_size * self.dp_size + else: + # If dynamic batching is disabled, we need to use all samples for training. + need_update = (step_idx + 1) % self.num_microbatches == 0 + + effective_samples = all_reduce_sum(torch.sum(loss_mask), self.plugin) + effective_tokens_count = torch.sum(action_mask, dim=-1) * loss_mask + total_effective_tokens_count = all_reduce_sum(torch.sum(effective_tokens_count), self.plugin) + self.effective_sample_count += effective_samples.item() + pbar.set_postfix( + { + "Global Step": self.global_step, + "Gradient Accumulation on": f"{self.effective_prompt_count}/{self.batch_size * self.dp_size} effective prompts, {self.effective_sample_count}/{self.batch_size * self.dp_size * self.num_generations} effective samples", + } + ) + + # Gradient must be synchronized if zero2 is enabled. https://github.com/hpcaitech/ColossalAI/blob/44d4053fec005fe0b06b6bc755fdc962463145df/colossalai/booster/plugin/hybrid_parallel_plugin.py#L1500 + ctx = ( + nullcontext() + if need_update or self.booster.plugin.zero_stage == 2 + else self.booster.no_sync(self.policy_model, self.optimizer) + ) + with ctx: + mini_batch_entropies = [] + for forward_micro_batch_start in range(0, data["input_ids"].size(0), train_microbatch_size): + input_ids_forward_micro_batch = data["input_ids"][ + forward_micro_batch_start : forward_micro_batch_start + train_microbatch_size + ] + old_action_log_probs_micro_batch = old_action_log_probs[ + forward_micro_batch_start : forward_micro_batch_start + train_microbatch_size + ] + attention_mask_forward_micro_batch = data["attention_mask"][ + forward_micro_batch_start : forward_micro_batch_start + train_microbatch_size + ] + action_mask_forward_micro_batch = action_mask[ + forward_micro_batch_start : forward_micro_batch_start + train_microbatch_size + ] + loss_mask_forward_micro_batch = ( + loss_mask[forward_micro_batch_start : forward_micro_batch_start + train_microbatch_size] + if loss_mask is not None + else None + ) + advantages_forward_micro_batch = advantages[ + forward_micro_batch_start : forward_micro_batch_start + train_microbatch_size + ] + + if self.plugin.pp_size > 1: + # Support training with PP. + if self.policy_loss_fn.beta > 0: + with torch.no_grad(): + reference_model_outputs = self.booster.execute_pipeline( + iter( + [ + { + "input_ids": input_ids_forward_micro_batch, + "attention_mask": attention_mask_forward_micro_batch, + } + ] + ), + self.reference_model, + criterion=lambda outputs, inputs: torch.tensor( + [0.0], device=action_mask.device + ), # dummy criterion + optimizer=None, + return_loss=False, + return_outputs=True, + ) + + if self.booster.plugin.stage_manager.is_last_stage(): + reference_action_log_probs = memory_efficient_logprob( + reference_model_outputs["outputs"]["logits"] / self.generate_config["temperature"], + input_ids_forward_micro_batch, + num_action, + shard_config=self.plugin.shard_config, + ) + else: + # Dummy reference logprobs for data iterator. + reference_action_log_probs = None + else: + reference_action_log_probs = None + + data_policy_forward = { + "input_ids": input_ids_forward_micro_batch, + "attention_mask": attention_mask_forward_micro_batch, + "action_mask": action_mask_forward_micro_batch, + "advantages": advantages_forward_micro_batch, + "loss_mask": loss_mask_forward_micro_batch, + "old_action_log_probs": old_action_log_probs_micro_batch, + "source": self.rank, + } + if reference_action_log_probs is not None: + data_policy_forward["reference_action_log_probs"] = reference_action_log_probs + + kl = [] + + def _criterion(outputs, inputs): + action_logits = outputs.logits + mini_batch_entropies.append( + ( + ((entropy_from_logits(action_logits[:, -num_action:]) * inputs["action_mask"]).sum(-1)) + / inputs["action_mask"].sum(-1) + ).detach() + ) + action_log_probs = memory_efficient_logprob( + action_logits / self.generate_config["temperature"], + inputs["input_ids"], + num_action, + shard_config=self.plugin.shard_config, + ) + if "reference_action_log_probs" in inputs: + per_token_kl = ( + torch.exp(inputs["reference_action_log_probs"] - action_log_probs) + - (inputs["reference_action_log_probs"] - action_log_probs) + - 1 + ) + appox_kl = torch.sum(per_token_kl * inputs["action_mask"], dim=-1) / torch.sum( + inputs["action_mask"], dim=-1 + ) + kl.append(appox_kl.mean()) + else: + per_token_kl = 0.0 + kl.append(torch.tensor(0.0)) + + loss, _ = self.policy_loss_fn( + action_log_probs, + inputs["old_action_log_probs"], + inputs["advantages"].repeat_interleave(action_log_probs.size(-1), dim=-1), + per_token_kl, + inputs["action_mask"], + loss_mask=inputs["loss_mask"], + total_effective_tokens_in_batch=total_effective_tokens_count, + ) + return loss + + policy_model_outputs = self.booster.execute_pipeline( + iter([data_policy_forward]), + self.policy_model, + criterion=_criterion, + optimizer=self.optimizer, + return_loss=True, + return_outputs=False, + ) + loss = policy_model_outputs["loss"] + + if self.booster.plugin.stage_manager.is_last_stage(): + if len(kl) > 0: + kl = all_reduce_mean(torch.mean(torch.stack(kl)).to(loss.device), self.plugin).data + mean_kl.append(kl) + mean_loss.append(all_reduce_mean(loss, self.plugin).data) + else: + policy_model_logits = self.policy_model( + input_ids=input_ids_forward_micro_batch, + attention_mask=attention_mask_forward_micro_batch, + ).logits + action_log_probs = memory_efficient_logprob( + policy_model_logits / self.generate_config["temperature"], + input_ids_forward_micro_batch, + num_action, + shard_config=self.plugin.shard_config, + ) + + if self.policy_loss_fn.beta > 0: + with torch.no_grad(): + reference_model_logits = self.reference_model( + input_ids=input_ids_forward_micro_batch, + attention_mask=attention_mask_forward_micro_batch, + ).logits + reference_action_log_probs = memory_efficient_logprob( + reference_model_logits / self.generate_config["temperature"], + input_ids_forward_micro_batch, + num_action, + shard_config=self.plugin.shard_config, + ) + per_token_kl = ( + torch.exp(reference_action_log_probs - action_log_probs) + - (reference_action_log_probs - action_log_probs) + - 1 + ) + kl = torch.sum(per_token_kl * action_mask_forward_micro_batch, dim=-1) / torch.sum( + action_mask_forward_micro_batch, dim=-1 + ) + else: + per_token_kl = 0.0 + kl = None + + loss, _ = self.policy_loss_fn( + action_log_probs, + old_action_log_probs_micro_batch, + advantages_forward_micro_batch.repeat_interleave(action_log_probs.size(-1), dim=-1), + per_token_kl, + action_mask_forward_micro_batch, + loss_mask=loss_mask_forward_micro_batch, + total_effective_tokens_in_batch=total_effective_tokens_count, + ) + + self.booster.backward(loss, self.optimizer) + loss = all_reduce_mean(loss, self.plugin) + # Calculate accumulate value. + if kl is not None: + kl = all_reduce_mean(kl.mean(), self.plugin) + mean_kl.append(kl.data) + mean_loss.append(loss.data) + mini_batch_entropies.append( + all_reduce_mean( + ( + ( + ( + entropy_from_logits(policy_model_logits[:, -num_action:]) + * action_mask_forward_micro_batch + ).sum(-1) + ) + / action_mask_forward_micro_batch.sum(-1) + ).detach(), + self.plugin, + ) + ) + if not self.plugin.pp_size > 1 or ( + self.plugin.pp_size > 1 + and self.booster.plugin.stage_manager.is_last_stage() + and self.tp_rank == 0 + and self.dp_rank == 0 + ): + reward = all_reduce_mean(reward.mean(), self.plugin) + format_acc = all_reduce_mean(format_acc.mean(), self.plugin) + ans_acc = all_reduce_mean(ans_acc.mean(), self.plugin) + advantages = all_reduce_mean(advantages.mean(), self.plugin) + response_length = all_reduce_mean(response_length.mean(), self.plugin) + entropy = all_reduce_mean(torch.cat(mini_batch_entropies, dim=0).mean(), self.plugin) + self.accum_loss.add_(sum(mean_loss) / len(mean_loss)) + self.accum_entropy.add_(entropy.data) + if self.policy_loss_fn.beta > 0: + self.accum_kl.add_(sum(mean_kl) / len(mean_kl)) + self.accum_advantages.add_(advantages.data) + self.accum_count += 1 + if need_update: + self.optimizer.step() + self.optimizer.zero_grad() + self.global_step += 1 + # no need to run all reduce as raw_train_batch_* are not splited across dp rank + sample_utilization = self.effective_sample_count / len(self.raw_train_batch_reward) / self.num_generations + self.effective_prompt_count = 0 + self.effective_sample_count = 0 + loss_scalar = self.accum_loss.item() + if not self.plugin.pp_size > 1 or ( + self.plugin.pp_size > 1 and self.booster.plugin.stage_manager.is_last_stage() and self.tp_rank == 0 + ): + if (not self.plugin.pp_size > 1 and self.rank == 0) or ( + self.plugin.pp_size > 1 and self.booster.plugin.stage_manager.is_last_stage() and self.tp_rank == 0 + ): + raw_batch_reward_mean = torch.cat(self.raw_train_batch_reward, dim=0).mean().cpu().item() + raw_batch_format_acc_mean = torch.cat(self.raw_train_batch_format_acc, dim=0).mean().cpu().item() + raw_batch_ans_acc_mean = torch.cat(self.raw_train_batch_ans_acc, dim=0).mean().cpu().item() + raw_batch_response_len = torch.cat(self.raw_train_batch_response_len, dim=0) + raw_batch_response_len_mean = raw_batch_response_len.mean().cpu().item() + overlength_samples_ratio = ( + (raw_batch_response_len >= action_mask.size(-1)).to(float).mean().cpu().item() + ) # not an exact figure, but a close estimate + self.raw_train_batch_reward = [] + self.raw_train_batch_format_acc = [] + self.raw_train_batch_ans_acc = [] + self.raw_train_batch_response_len = [] + to_log_msg = [ + f"Loss: {self.accum_loss.item() / self.accum_count:.4f}", + f"Reward: {raw_batch_reward_mean:.4f}", + f"format Reward: {raw_batch_format_acc_mean:.4f}", + f"Acc Reward: {raw_batch_ans_acc_mean:.4f}", + f"Advantages: {self.accum_advantages.item() / self.accum_count:.4f}", + f"Response Length: {raw_batch_response_len_mean:.4f}", + f"Sample_utilization: {sample_utilization:.4f}", + f"Overlength samples ratio: {overlength_samples_ratio:.4f}", + f"Entropy: {self.accum_entropy.item() / self.accum_count:.4f}", + ] + ([f"KL: {self.accum_kl.item() / self.accum_count:.4f}"] if self.policy_loss_fn.beta > 0 else []) + print("\n".join(to_log_msg)) + metrics = { + "metrics/reward": raw_batch_reward_mean, + "metrics/format_acc": raw_batch_format_acc_mean, + "metrics/ans_acc": raw_batch_ans_acc_mean, + "metrics/response_length": raw_batch_response_len_mean, + "train/loss": self.accum_loss.item() / self.accum_count, + "train/advantages": self.accum_advantages.item() / self.accum_count, + "train/learning_rate": self.lr_scheduler.get_last_lr()[0], + "train/sample_utilization": sample_utilization, + "train/entropy": self.accum_entropy.item() / self.accum_count, + "train/overlength_samples_ratio": overlength_samples_ratio, + "rollout/temperature": data["temperature"].cpu().numpy()[0][0], + } + if self.policy_loss_fn.beta > 0: + metrics["train/kl"] = self.accum_kl.item() / self.accum_count + if self.wandb_run is not None: + self.wandb_run.log(metrics) + self.accum_loss.zero_() + self.accum_kl.zero_() + self.accum_entropy.zero_() + self.accum_advantages.zero_() + self.accum_count = 0 + return loss_scalar + else: + return None + + def state_dict(self): + self.policy_model._force_wait_all_gather() + model = self.policy_model.unwrap() + state_dict = model.state_dict() + state_dict["consumer_global_step"] = torch.tensor([self.global_step], device=self.device) + return state_dict diff --git a/applications/ColossalChat/coati/distributed/inference_backend.py b/applications/ColossalChat/coati/distributed/inference_backend.py new file mode 100644 index 000000000000..34827e4e2cf9 --- /dev/null +++ b/applications/ColossalChat/coati/distributed/inference_backend.py @@ -0,0 +1,287 @@ +from typing import Any, Dict + +import torch +import torch.nn.functional as F +from transformers import AutoConfig, AutoModelForCausalLM, PreTrainedModel, PreTrainedTokenizer + +from colossalai.utils import get_current_device + +from .utils import log_probs_from_logits, update_by_default + +try: + import sglang as sgl +except ImportError: + sgl = None + +try: + from vllm import LLM, SamplingParams +except ImportError: + LLM = None + + +class BaseInferenceBackend: + def __init__(self, model_config: Dict[str, Any], generate_config: Dict[str, Any], tokenizer: PreTrainedTokenizer): + pass + + def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: + """Generate new tokens given input_ids and attention_mask. + + Args: + input_ids (torch.Tensor): shape [B, S] + attention_mask (torch.Tensor): shape [B, S] + + Returns: + Dict[str, torch.Tensor]: containing the + - input_ids (torch.Tensor): shape [B, S+N] + - attention_mask (torch.Tensor): shape [B, S+N] + - action_log_probs (torch.Tensor): shape [B, N] + - action_mask (torch.Tensor): shape [B, N] + where N is the number of generated tokens. And all tensors should be on CUDA. + """ + + def load_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> None: + pass + + +class TransformersInferenceBackend(BaseInferenceBackend): + DEFAULT_MODEL_CONFIG = dict( + trust_remote_code=True, + torch_dtype=torch.bfloat16, + ) + FORCE_MODEL_CONFIG = dict( + device_map="auto", + ) + FORCE_GENERATE_CONFIG = dict(output_logits=True, return_dict_in_generate=True) + + def __init__( + self, + model_config: Dict[str, Any], + generate_config: Dict[str, Any], + tokenizer: PreTrainedTokenizer, + num_generations: int = 8, + ): + model_config = update_by_default(model_config, self.DEFAULT_MODEL_CONFIG) + model_config.update(self.FORCE_MODEL_CONFIG) + path = model_config.pop("path") + self.model: PreTrainedModel = AutoModelForCausalLM.from_pretrained(path, **model_config) + self.generate_config = generate_config.copy() + self.generate_config.update(self.FORCE_GENERATE_CONFIG) + self.tokenizer = tokenizer + self.num_generations = num_generations + + @torch.no_grad() + def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: + micro_batch_size = input_ids.size(0) + input_ids = input_ids.to(get_current_device()) + attention_mask = attention_mask.to(get_current_device()) + gt_answer = kwargs.pop("gt_answer", None) + test_cases = kwargs.pop("test_cases", None) + if self.num_generations > 1: + input_ids = input_ids.repeat_interleave(self.num_generations, dim=0) + attention_mask = attention_mask.repeat_interleave(self.num_generations, dim=0) + out = self.model.generate( + input_ids, attention_mask=attention_mask, **kwargs, **self.generate_config, tokenizer=self.tokenizer + ) + input_len = input_ids.shape[-1] + new_token_ids = out.sequences[:, input_len:] + # get log probs + assert new_token_ids.shape[-1] == len(out.logits) + action_log_probs = [] + for i, logits in enumerate(out.logits): + action_log_probs.append(log_probs_from_logits(logits[:, None, :], new_token_ids[:, i : i + 1])) + action_log_probs = torch.cat(action_log_probs, dim=1) + # get action mask + response_idx = torch.zeros((new_token_ids.size(0), 2), dtype=torch.int).to(get_current_device()) + action_mask = torch.ones_like(new_token_ids, dtype=attention_mask.dtype) + if self.tokenizer.eos_token_id is not None: + for indices in torch.nonzero(new_token_ids == self.tokenizer.eos_token_id): + action_mask[indices[0], indices[1] + 1 :] = 0 + response_idx[:, 0] = input_len + response_idx[:, 1] = input_len + action_mask.sum(dim=1) - 1 + + if attention_mask.size(0) != action_mask.size(0): + assert action_mask.size(0) % attention_mask.size(0) == 0 + attention_mask = attention_mask.repeat_interleave(action_mask.size(0) // attention_mask.size(0), dim=0) + + attention_mask = torch.cat((attention_mask, action_mask), dim=1) + data = { + "input_ids": out.sequences, + "attention_mask": attention_mask, + "action_log_probs": action_log_probs, + "action_mask": action_mask, + "response_idx": response_idx, + } + + data = {k: v.view(micro_batch_size, self.num_generations, v.size(-1)) for k, v in data.items()} + + if gt_answer is not None: + data["gt_answer"] = gt_answer + if test_cases is not None: + data["test_cases"] = test_cases + data = {k: v.to(get_current_device()) for k, v in data.items()} + return data + + def load_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> None: + self.model.load_state_dict(state_dict) + + +class SGLangInferenceBackend(BaseInferenceBackend): + def __init__( + self, + model_config: Dict[str, Any], + generate_config: Dict[str, Any], + tokenizer: PreTrainedTokenizer, + num_generations: int = 8, + ): + if sgl is None: + raise ImportError("sglang is not installed") + path = model_config.pop("path") + defaut_config = dict( + trust_remote_code=True, + skip_tokenizer_init=True, + ) + defaut_config.update(model_config) + self.llm = sgl.Engine(model_path=path, **defaut_config) + self.generate_config = generate_config + self.tokenizer = tokenizer + self.config = AutoConfig.from_pretrained(path) + + @torch.no_grad() + def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: + outputs = self.llm.generate(input_ids=input_ids.tolist(), sampling_params=self.generate_config) + out_tokens = [] + out_len = [] + for out in outputs: + out_tokens.append(out["token_ids"]) + out_len.append(out["meta_info"]["completion_tokens"]) + max_len = max(out_len) + input_len = input_ids.shape[-1] + attention_mask = F.pad(attention_mask, (0, max_len), value=1) + for i in range(len(out_tokens)): + out_tokens[i] = out_tokens[i] + [self.tokenizer.pad_token_id] * (max_len - out_len[i]) + attention_mask[i, input_len + out_len[i] :] = 0 + out = torch.tensor(out_tokens) + out = torch.cat((input_ids, out), dim=1) + labels = out.clone() + labels[..., :input_len] = -100 + for i in range(len(out_len)): + labels[i, input_len + out_len[i] :] = -100 + data = { + "input_ids": out, + "attention_mask": attention_mask, + "labels": labels, + } + data = {k: v.to(get_current_device()) for k, v in data.items()} + return data + + def load_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> None: + if self.config.tie_word_embeddings: + del state_dict["lm_head.weight"] + named_tensors = [(k, v) for k, v in state_dict.items()] + self.llm.update_weights_from_tensor(named_tensors) + + +class VLLMInferenceBackend(BaseInferenceBackend): + DEFAULT_MODEL_CONFIG = dict( + trust_remote_code=True, + enable_sleep_mode=False, + ) + FORCE_GENERATE_CONFIG = dict( + logprobs=0, + ) + + def __init__( + self, + model_config: Dict[str, Any], + generate_config: Dict[str, Any], + tokenizer: PreTrainedTokenizer, + num_generations: int = 8, + ): + if LLM is None: + raise ImportError("vllm is not installed") + model_config = update_by_default(model_config, self.DEFAULT_MODEL_CONFIG) + path = model_config.pop("path") + self.llm = LLM(model=path, **model_config) + generate_config = generate_config.copy() + generate_config.update(self.FORCE_GENERATE_CONFIG) + generate_config.update({"n": num_generations}) + self.generate_config = generate_config + self.sample_params = SamplingParams(**generate_config) + self.model_config = model_config + self.tokenizer = tokenizer + self.num_generations = num_generations + + @torch.no_grad() + def generate(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: + micro_batch_size = input_ids.size(0) + response_start_idx = input_ids.size(1) + first_non_padding_token_idx = (input_ids != self.tokenizer.pad_token_id).int().argmax(dim=1) + micro_batch_input_ids = input_ids.tolist() + micro_batch_input_ids_no_padding = [ + micro_batch_input_ids[i][first_non_padding_token_idx[i] :] for i in range(micro_batch_size) + ] + sample_params = kwargs.get("sample_params", self.sample_params) + outputs = self.llm.generate( + prompt_token_ids=micro_batch_input_ids_no_padding, sampling_params=sample_params, use_tqdm=False + ) + out_tokens = [] + out_len = [] + log_probs = [] + response_idx = [] + for out in outputs: + for output_i in out.outputs: + out_len.append(len(output_i.token_ids)) + out_tokens.append(list(output_i.token_ids)) + response_idx.append((response_start_idx, response_start_idx + len(output_i.token_ids) - 1)) + assert len(output_i.logprobs) == len(output_i.token_ids) + p = [m[t].logprob for m, t in zip(output_i.logprobs, output_i.token_ids)] + log_probs.append(p) + + # pad them + max_len = self.sample_params.max_tokens + action_mask = torch.ones(len(out_tokens), max_len, dtype=attention_mask.dtype) + + for i, new_token_ids in enumerate(out_tokens): + pad_len = max_len - out_len[i] + out_tokens[i] = new_token_ids + [self.tokenizer.pad_token_id] * pad_len + log_probs[i] = log_probs[i] + [0.0] * pad_len + action_mask[i, out_len[i] :] = 0 + + out_tokens = torch.tensor(out_tokens) + log_probs = torch.tensor(log_probs) + response_idx = torch.tensor(response_idx) + + if attention_mask.size(0) != action_mask.size(0): + assert action_mask.size(0) % attention_mask.size(0) == 0 + num_returns = action_mask.size(0) // attention_mask.size(0) + attention_mask = attention_mask.repeat_interleave(num_returns, dim=0) + input_ids = input_ids.repeat_interleave(num_returns, dim=0) + + out_tokens = torch.cat((input_ids, out_tokens), dim=1) + attention_mask = torch.cat((attention_mask, action_mask), dim=1) + + data = { + "input_ids": out_tokens, + "attention_mask": attention_mask, + "action_log_probs": log_probs, + "action_mask": action_mask, + "response_idx": response_idx, + } + + data = {k: v.view(micro_batch_size, -1, v.size(-1)) for k, v in data.items()} + data = {k: v.to(get_current_device()) for k, v in data.items()} + if "gt_answer" in kwargs: + data["gt_answer"] = kwargs["gt_answer"] + if "test_cases" in kwargs: + data["test_cases"] = kwargs["test_cases"] + return data + + def load_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> None: + self.llm.llm_engine.model_executor.driver_worker.model_runner.model.load_weights(state_dict.items()) + + +BACKEND_MAP = { + "transformers": TransformersInferenceBackend, + # "sglang": SGLangInferenceBackend, # sglang backend will stuck the process due to unknown reason + "vllm": VLLMInferenceBackend, +} diff --git a/applications/ColossalChat/coati/distributed/launch.py b/applications/ColossalChat/coati/distributed/launch.py new file mode 100644 index 000000000000..a48246c87526 --- /dev/null +++ b/applications/ColossalChat/coati/distributed/launch.py @@ -0,0 +1,183 @@ +import copy +import os +import uuid +from typing import Any, Dict, Optional + +import ray + +from .consumer import SimpleConsumer +from .grpo_consumer import GRPOConsumer +from .producer import SimpleProducer + +ALGO_MAP = {"Simple": SimpleConsumer, "GRPO": GRPOConsumer, "DAPO": GRPOConsumer} + + +def get_jsonl_size_fast(path: str) -> int: + with open(path) as f: + lines = f.readlines() + lines = [line for line in lines if line.strip()] + return len(lines) + + +def get_dp_size_fast(n_procs: int, plugin_config: Dict[str, Any]) -> int: + tp_size = plugin_config.get("tp_size", 1) + pp_size = plugin_config.get("pp_size", 1) + ep_size = plugin_config.get("ep_size", 1) + sp_size = plugin_config.get("sp_size", 1) + return n_procs // (tp_size * pp_size * ep_size * sp_size) + + +def launch_distributed( + num_producers: int, + num_proc_per_producer: int, + num_consumer_procs: int, + num_episodes: int, + inference_batch_size: int, + inference_microbatch_size: int, + train_batch_size: int, + train_minibatch_size: int, + train_dataset_config: Dict[str, Any], + inference_model_config: Dict[str, Any], + generate_config: Dict[str, Any], + train_model_config: Dict[str, Any], + grpo_config: Dict[str, Any], + plugin_config: Dict[str, Any], + tokenizer_config: Optional[Dict[str, Any]] = None, + inference_backend: str = "transformers", + num_generations: int = 8, + master_addr: str = "localhost", + master_port: int = 29500, + core_algo: str = "GRPO", + project_name: Optional[str] = None, + save_interval: int = 100, + save_dir: str = "./model", + eval_dataset_config: Optional[Dict[str, Any]] = None, + eval_interval: int = 100, + eval_save_dir: Optional[str] = None, + eval_generation_config: Optional[Dict[str, Any]] = None, + log_rollout_interval: int = 20, + rollout_save_dir: str = "./rollout", + enable_profiling: bool = False, + n_behind: int = 0, +): + if core_algo not in ALGO_MAP: + raise NotImplementedError(f"{core_algo} is not supported yet.") + else: + core_consumer = ALGO_MAP.get(core_algo, SimpleConsumer) + + train_dp_size = get_dp_size_fast(num_consumer_procs, plugin_config) + assert (inference_batch_size * num_producers) % (train_batch_size * train_dp_size) == 0 + + dataset_path = train_dataset_config["path"] + num_samples = get_jsonl_size_fast(dataset_path) + global_inference_batch_size = inference_batch_size * num_producers + num_update_per_episode = num_samples // global_inference_batch_size + num_recv_per_update = inference_batch_size // inference_microbatch_size + + run_name = f"{inference_backend}_bs_{train_batch_size * train_dp_size}_temp_{generate_config['temperature']:.01f}_top_p_{generate_config['top_p']:.02f}" + wandb_group_name = str(uuid.uuid4()) + rollout_log_file = os.path.join( + rollout_save_dir, + f"{project_name.replace(' ','_')}_run_{wandb_group_name}.jsonl", + ) + + # Attention: Ray use complex schedualing method that consider various factors including load-balancing. + # when requesting resources, it is not guaranteed that the resource comes from a node with lower node it + # this go against the design principle of our implementation, and we need to manually force the schedualing, + # allocating the producer to nodes with lower node id and the consumer to the resouces from nodes with higher + # node id. See the reference here: https://docs.ray.io/en/latest/ray-core/scheduling/index.html#nodeaffinityschedulingstrategy + nodes = ray.nodes() + node_info = { + node["NodeID"]: { + "num_gpus": node["Resources"].get("GPU", 0), + "address": node["NodeManagerAddress"], + } # Default to 0 if no GPUs are available + for node in nodes + } + gpu_to_node_id = [] + gpu_to_ip_address = [] + for node_id in node_info: + for idx in range(int(node_info[node_id]["num_gpus"])): + gpu_to_node_id.append(node_id) + gpu_to_ip_address.append(node_info[node_id]["address"]) + print(node_info) + + producer_procs = [] + for i in range(num_producers): + node_id = gpu_to_node_id[0] + producer_ip_address = gpu_to_ip_address[0] + for _ in range(num_proc_per_producer): + gpu_to_node_id.pop(0) + gpu_to_ip_address.pop(0) + print(f"Schedual Producer P[{i}] which requires {num_proc_per_producer} GPUs on node {producer_ip_address}") + producer = SimpleProducer.options(num_gpus=num_proc_per_producer).remote( + producer_idx=i, + num_producers=num_producers, + num_consumer_procs=num_consumer_procs, + num_episodes=num_episodes, + batch_size=inference_batch_size, + train_dataset_config=train_dataset_config, + model_config=inference_model_config, + generate_config=generate_config, + tokenizer_config=tokenizer_config, + microbatch_size=inference_microbatch_size, + backend=inference_backend, + num_generations=num_generations, + consumer_plugin_config=plugin_config, + eval_dataset_config=eval_dataset_config, + eval_interval=eval_interval, + grpo_config=grpo_config, + eval_save_dir=eval_save_dir, + eval_generation_config=eval_generation_config, + project_name=project_name, + run_name=run_name, + wandb_group_name=wandb_group_name, + log_rollout_interval=log_rollout_interval, + rollout_log_file=rollout_log_file, + enable_profiling=enable_profiling, + n_behind=n_behind, + ) + producer_procs.append(producer) + ray.get([p.setup.remote() for p in producer_procs]) + generate_config_consumer = copy.deepcopy(generate_config) + generate_config_consumer.update( + dict( + backend=inference_backend, + ) + ) + consumer_master_ip_address = gpu_to_ip_address[0] + print(f"Use {consumer_master_ip_address} as master address for torch DDP.") + consumer_procs = [] + for i in range(num_consumer_procs): + node_id = gpu_to_node_id[0] + consumer_ip_address = gpu_to_ip_address[0] + gpu_to_node_id.pop(0) + gpu_to_ip_address.pop(0) + print(f"Schedual Consumer T[{i}] which requires 1 GPUs on node {consumer_ip_address}") + consumer = core_consumer.options(num_gpus=1).remote( + num_producers=num_producers, + num_episodes=num_episodes, + rank=i, + world_size=num_consumer_procs, + master_addr=consumer_master_ip_address, + master_port=master_port, + num_update_per_episode=num_update_per_episode, + num_recv_per_update=num_recv_per_update, + batch_size=train_batch_size, + model_config=train_model_config, + plugin_config=plugin_config, + minibatch_size=train_minibatch_size, + generate_config=generate_config_consumer, + grpo_config=grpo_config, + num_generations=num_generations, + save_interval=save_interval, + save_dir=save_dir, + project_name=project_name, + run_name=run_name, + wandb_group_name=wandb_group_name, + enable_profiling=enable_profiling, + n_behind=n_behind, + ) + consumer_procs.append(consumer) + ray.get([p.setup.remote() for p in consumer_procs]) + ray.get([p.loop.remote() for p in (producer_procs + consumer_procs)]) diff --git a/applications/ColossalChat/coati/distributed/loss.py b/applications/ColossalChat/coati/distributed/loss.py new file mode 100644 index 000000000000..36057b24faf5 --- /dev/null +++ b/applications/ColossalChat/coati/distributed/loss.py @@ -0,0 +1,68 @@ +from typing import Optional + +import torch +import torch.nn as nn +from coati.distributed.utils import masked_mean, masked_sum + + +class PolicyLoss(nn.Module): + """ + Policy Loss for PPO + """ + + def __init__( + self, + clip_eps_low: float = 0.2, + clip_eps_high: float = 0.2, + beta: float = 0.01, + loss_variation: str = "sample_level", + ) -> None: + super().__init__() + self.clip_eps_low = clip_eps_low + self.clip_eps_high = clip_eps_high + self.beta = beta + self.loss_variation = loss_variation + assert loss_variation in ["sample_level", "token_level"], f"Unsupported loss variation: {loss_variation}" + + def forward( + self, + log_probs: torch.Tensor, + old_log_probs: torch.Tensor, + advantages: torch.Tensor, + per_token_kl: torch.Tensor, + action_mask: Optional[torch.Tensor] = None, + loss_mask: Optional[torch.Tensor] = None, + total_effective_tokens_in_batch: torch.Tensor = None, + ) -> torch.Tensor: + if action_mask is None: + ratio = (log_probs - log_probs.detach()).exp() + else: + ratio = ((log_probs - log_probs.detach()) * action_mask).exp() + + surr1 = ratio * advantages + surr2 = ratio.clamp(1 - self.clip_eps_low, 1 + self.clip_eps_high) * advantages + if self.beta == 0: + # skip kl term if kl coefficient is zero + per_token_kl = 0.0 + loss = -torch.min(surr1, surr2) + self.beta * per_token_kl + + if self.loss_variation == "sample_level": + if action_mask is not None: + loss = masked_mean(loss, action_mask) + else: + loss = loss.mean(dim=1) + if loss_mask is not None: + loss = loss * loss_mask + loss = loss.mean() + elif self.loss_variation == "token_level": + if action_mask is not None: + loss = masked_sum(loss, action_mask) + else: + loss = loss.sum(dim=1) + if loss_mask is not None: + loss = loss * loss_mask + loss = loss.sum() / (total_effective_tokens_in_batch + 1e-8) + else: + raise ValueError(f"Unsupported loss variation: {self.loss_variation}") + + return loss, ratio.max() diff --git a/applications/ColossalChat/coati/distributed/producer.py b/applications/ColossalChat/coati/distributed/producer.py new file mode 100644 index 000000000000..11fb5d3aa3a9 --- /dev/null +++ b/applications/ColossalChat/coati/distributed/producer.py @@ -0,0 +1,477 @@ +import copy +import json +import os +from typing import Any, Dict, Optional + +import ray +import ray.util.collective as cc +import torch +import tqdm +import wandb +from coati.dataset.loader import RawConversationDataset, collate_fn_grpo +from coati.distributed.profiling_utils import CustomProfiler +from coati.distributed.reward.reward_fn import boxed_math_reward_fn, code_reward_fn, math_reward_fn +from coati.distributed.reward.verifiable_reward import VerifiableReward +from ray.util.collective import allreduce +from ray.util.collective.types import Backend, ReduceOp +from torch.utils.data import DataLoader, DistributedSampler +from transformers import AutoTokenizer + +from colossalai.utils import get_current_device + +from .comm import ray_broadcast_tensor_dict +from .inference_backend import BACKEND_MAP +from .utils import pre_send, safe_append_to_jsonl_file + +try: + from vllm import SamplingParams +except ImportError: + LLM = None + + +class BaseProducer: + def __init__( + self, + producer_idx: int, + num_producers: int, + num_consumer_procs: int, + num_episodes: int, + batch_size: int, + train_dataset_config: Dict[str, Any], + model_config: Dict[str, Any], + generate_config: Dict[str, Any], + tokenizer_config: Optional[Dict[str, Any]] = None, + microbatch_size: int = 1, + backend: str = "transformers", + consumer_plugin_config: Dict[str, Any] = None, + eval_dataset_config=None, + eval_interval=-1, # disable evaluation + grpo_config: Dict[str, Any] = None, + eval_save_dir: str = "./eval", + project_name: str = None, + run_name: str = None, + wandb_group_name: str = None, + log_rollout_interval: int = 20, + rollout_log_file: str = "./rollout_log.jsonl", + enable_profiling: bool = False, + n_behind: int = 0, + ): + self.producer_idx = producer_idx + self.num_producers = num_producers + self.num_consumer_procs = num_consumer_procs + self.num_episodes = num_episodes + self.batch_size = batch_size + self.microbatch_size = microbatch_size + assert batch_size % microbatch_size == 0 + self.num_microbatches = batch_size // microbatch_size + self.latest_eval_step = -1 + self.profiler = CustomProfiler(f"P{self.producer_idx}", disabled=not enable_profiling) + + self.train_dataset_config = train_dataset_config + self.model_config = model_config + self.generate_config = generate_config + self.tokenizer_config = tokenizer_config + self.consumer_plugin_config = consumer_plugin_config + self.eval_interval = eval_interval + self.eval_save_dir = eval_save_dir + self.consumer_global_step = 0 + self.eval_mode = False + self.log_rollout_interval = log_rollout_interval + self.latest_rollout_log_step = -1 + self.grpo_config = grpo_config + self.n_behind = n_behind + reward_model_kwargs = { + k: v + for k, v in grpo_config.items() + if k in ["soft_over_length_punishment", "max_new_tokens", "cache_length", "code_verifier_api_url"] + } + self.response_format_tags = grpo_config.get("response_format_tags", None) + if producer_idx == 0: + if os.path.exists(rollout_log_file): + raise ValueError( + f"Rollout log file {rollout_log_file} already exists. Please delete it or change the name." + ) + else: + os.makedirs(os.path.dirname(rollout_log_file), exist_ok=True) + self.rollout_log_file = open(rollout_log_file, "w", encoding="utf8") + if self.producer_idx == 0: + self.wandb_run = wandb.init( + project=project_name, + sync_tensorboard=False, + dir="./wandb", + name=run_name + "_eval", + group=wandb_group_name, + ) + + if os.path.exists(self.eval_save_dir) and self.eval_interval > 0: + raise ValueError(f"Eval save dir {self.eval_save_dir} already exists. Please delete it or change the name.") + + # init tokenizer + if tokenizer_config is None: + tokenizer_path = model_config["path"] + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path) + else: + tokenizer_path = tokenizer_config.pop("path") + self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, **tokenizer_config) + self.tokenizer.padding_side = "left" + + # init dataloader + train_dataset_path = train_dataset_config.pop("path") + self.train_dataset = RawConversationDataset(self.tokenizer, train_dataset_path, **train_dataset_config) + self.train_dataloader = DataLoader( + self.train_dataset, + batch_size=microbatch_size, + sampler=DistributedSampler( + self.train_dataset, + num_replicas=num_producers, + rank=producer_idx, + shuffle=True, + drop_last=True, + seed=42, + ), + num_workers=4, + drop_last=True, + collate_fn=collate_fn_grpo, + ) + if grpo_config["reward_fn_type"] == "think_answer_tags": + self.evaluation_function = math_reward_fn + elif grpo_config["reward_fn_type"] == "boxed": + self.evaluation_function = boxed_math_reward_fn + elif grpo_config["reward_fn_type"] == "code": + self.evaluation_function = code_reward_fn + else: + raise ValueError(f"Unknown evaluation function type {grpo_config['reward_fn_type']}") + + self.eval_dataset_config = eval_dataset_config + if self.eval_dataset_config is not None: + self.eval_dataloaders = {} + for eval_task_name in self.eval_dataset_config: + eval_dataset_path = eval_dataset_config[eval_task_name].pop("path") + eval_dataset = RawConversationDataset( + self.tokenizer, eval_dataset_path, **eval_dataset_config[eval_task_name] + ) + print(f"[P{self.producer_idx}] eval dataset {eval_task_name} size: {len(eval_dataset)}") + self.eval_dataloaders[eval_task_name] = DataLoader( + eval_dataset, + batch_size=microbatch_size, + sampler=DistributedSampler( + eval_dataset, + num_replicas=num_producers, + rank=producer_idx, + shuffle=False, + drop_last=False, + seed=42, + ), + collate_fn=collate_fn_grpo, + ) + else: + print("No eval dataset provided, skip eval") + self.device = get_current_device() + self.reward_model = VerifiableReward( + reward_fns=[self.evaluation_function], # multiple reward functions can be added here + tokenizer=self.tokenizer, + tags=self.response_format_tags, + **reward_model_kwargs, + ) + + # init backend + if backend in BACKEND_MAP: + self.backend_cls = BACKEND_MAP[backend] + else: + raise ValueError(f"Unexpected backend {backend}") + + self.consumer_pp_size = consumer_plugin_config.get("pp_size", 1) # consumer pp size + + def setup(self) -> None: + cc.init_collective_group( + world_size=self.num_producers, + rank=self.producer_idx, + backend=Backend.NCCL, + group_name="producer_group", + ) + cc.init_collective_group(1 + self.num_consumer_procs, 0, group_name=f"sync_data_{self.producer_idx}") + if self.consumer_pp_size > 1: + for i in range(self.consumer_pp_size): + cc.init_collective_group(self.num_producers + 1, self.producer_idx, group_name=f"sync_model_{i}") + else: + cc.init_collective_group(self.num_producers + 1, self.producer_idx, group_name="sync_model") + + def rollout(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, **kwargs) -> Dict[str, torch.Tensor]: + raise NotImplementedError + + def load_state_dict(self, state_dict: Dict[str, torch.Tensor]) -> None: + raise NotImplementedError + + def loop(self) -> None: + num_update_per_episode = len(self.train_dataloader) // self.num_microbatches + num_valid_microbatches = num_update_per_episode * self.num_microbatches + + print( + f"[P{self.producer_idx}] num_valid_microbatches {num_valid_microbatches}, nmb: {self.num_microbatches}, dl: {len(self.train_dataloader)}" + ) + for episode in range(self.num_episodes): + self.train_dataloader.sampler.set_epoch(episode) + for i, batch in enumerate(self.train_dataloader): + if i >= num_valid_microbatches: + break + if self.eval_interval > 0 and self.eval_dataset_config is not None: + if ( + self.consumer_global_step - self.latest_eval_step >= self.eval_interval + and self.consumer_global_step > self.latest_eval_step + ) or self.latest_eval_step == -1: + to_log_msg = {} + self.eval_mode = True + for eval_task_name in self.eval_dataloaders: + if self.producer_idx == 0: + print( + f"[P{self.producer_idx}] Evaluate model at training step {self.consumer_global_step} on task {eval_task_name}" + ) + eval_results = [] + eval_statistics_tensor = torch.zeros((2,), dtype=torch.float32).to(self.device) + for eval_batch in tqdm.tqdm( + self.eval_dataloaders[eval_task_name], disable=self.producer_idx != 0 + ): + eval_outputs = self.rollout(**eval_batch, sample_params=self.eval_sample_params) + eval_results = eval_results + [ + self.evaluation_function( + eval_outputs["input_ids"][m][n], + eval_outputs[ + ( + "test_cases" + if self.grpo_config["reward_fn_type"] == "code" + else "gt_answer" + ) + ][m], + eval_outputs["response_idx"][m][n], + tokenizer=self.tokenizer, + eval_mode=True, + tags=self.response_format_tags, + ) + for m in range(eval_outputs["input_ids"].size(0)) + for n in range(eval_outputs["input_ids"].size(1)) + ] + eval_statistics_tensor[0] += sum([max(0, res["ans_valid"]) for res in eval_results]) + eval_statistics_tensor[1] += len(eval_results) + allreduce(eval_statistics_tensor, op=ReduceOp.SUM, group_name="producer_group") + to_log_msg[f"eval/{eval_task_name}"] = ( + eval_statistics_tensor[0].item() / eval_statistics_tensor[1].item() + ) + if self.producer_idx == 0: + print( + f"[P{self.producer_idx}]: Accuracy on {eval_task_name}: {to_log_msg[f'eval/{eval_task_name}']}" + ) + # save eval results + safe_append_to_jsonl_file( + os.path.join( + self.eval_save_dir, + f"{eval_task_name}_training_step_{self.consumer_global_step}.jsonl", + ), + eval_results, + ) + + if self.producer_idx == 0: + self.wandb_run.log(to_log_msg, step=self.consumer_global_step) + self.eval_mode = False + self.latest_eval_step = self.consumer_global_step + self.profiler.enter("rollout") + outputs = self.rollout(**batch) + self.profiler.exit("rollout") + outputs["temperature"] = torch.tensor( + [self.model.generate_config["temperature"]] * outputs["input_ids"].size(0) + ).to(outputs["input_ids"].device) + bs, num_gen = outputs["input_ids"].size(0), outputs["input_ids"].size(1) + self.profiler.enter("calculate_reward") + if self.grpo_config["reward_fn_type"] == "code": + test_cases = [] + for prompt_id in range(bs): + test_cases.extend([outputs["test_cases"][prompt_id]] * num_gen) + reward_model_output = self.reward_model( + outputs["input_ids"].view((-1, outputs["input_ids"].size(-1))), + test_cases=test_cases, + response_idx=outputs["response_idx"].view((-1, 2)), + ) + else: + gt_answer = [] + for prompt_id in range(bs): + gt_answer.extend([outputs["gt_answer"][prompt_id]] * num_gen) + reward_model_output = self.reward_model( + outputs["input_ids"].view((-1, outputs["input_ids"].size(-1))), + gt_answer=gt_answer, + response_idx=outputs["response_idx"].view((-1, 2)), + ) + outputs["reward"] = ( + torch.tensor([value[0] for value in reward_model_output]) + .to(outputs["input_ids"].device) + .view((bs, num_gen, 1)) + ) + outputs["format_acc"] = ( + torch.tensor([value[1] for value in reward_model_output]) + .to(outputs["input_ids"].device) + .view((bs, num_gen, 1)) + ) + outputs["ans_acc"] = ( + torch.tensor([value[2] for value in reward_model_output]) + .to(outputs["input_ids"].device) + .view((bs, num_gen, 1)) + ) + if "gt_answer" in outputs: + outputs.pop("gt_answer") + if "test_cases" in outputs: + outputs.pop("test_cases") + self.profiler.exit("calculate_reward") + + print(f"[P{self.producer_idx}] Send data {[(k, v.shape) for k, v in outputs.items()]}") + outputs = pre_send(outputs) + self.profiler.enter("send_broadcast_data") + ray_broadcast_tensor_dict( + outputs, src=0, device=self.device, group_name=f"sync_data_{self.producer_idx}" + ) + self.profiler.exit("send_broadcast_data") + if ( + (i + 1) % self.num_microbatches == 0 + and (episode != self.num_episodes - 1 or i != num_valid_microbatches - 1) + and (episode != 0 or (i + 1) > self.n_behind * self.num_microbatches) + ): + if isinstance(self.model, BACKEND_MAP["vllm"]) and self.model.model_config.get( + "enable_sleep_mode", False + ): + self.model.llm.sleep() # revict KV_cache to avoid OOM + # don't sync model for last iteration + torch.cuda.empty_cache() + self.profiler.enter("sync_model") + if self.consumer_pp_size > 1: + for pp_idx in range(self.consumer_pp_size): + print( + f"[P{self.producer_idx}] Sync model PP stage {pp_idx} episode {episode} step {(i + 1) // self.num_microbatches - 1}" + ) + state_dict = ray_broadcast_tensor_dict( + None, self.num_producers, device=self.device, group_name=f"sync_model_{pp_idx}" + ) + if "consumer_global_step" in state_dict: + self.consumer_global_step = state_dict.pop("consumer_global_step").item() + self.load_state_dict(state_dict) + else: + print( + f"[P{self.producer_idx}] Sync model episode {episode} step {(i + 1) // self.num_microbatches - 1}" + ) + state_dict = ray_broadcast_tensor_dict( + None, self.num_producers, device=self.device, group_name="sync_model" + ) + if "consumer_global_step" in state_dict: + self.consumer_global_step = state_dict.pop("consumer_global_step").item() + self.load_state_dict(state_dict) + self.profiler.exit("sync_model") + del state_dict + torch.cuda.empty_cache() + if isinstance(self.model, BACKEND_MAP["vllm"]) and self.model.model_config.get( + "enable_sleep_mode", False + ): + self.model.llm.wake_up() + # linear annealing for 1 episode, temperature from initial to 0.9 + if episode <= 0: + ratio = 1 - (len(self.train_dataloader) - i) / len(self.train_dataloader) + self.model.generate_config["temperature"] = (1 - ratio) * self.generate_config[ + "temperature" + ] + ratio * 0.9 + if isinstance(self.model, BACKEND_MAP["vllm"]): + self.model.sample_params.temperature = (1 - ratio) * self.generate_config[ + "temperature" + ] + ratio * 0.9 + + def __del__(self): + self.profiler.close() + + +@ray.remote +class SimpleProducer(BaseProducer): + def __init__( + self, + producer_idx, + num_producers, + num_consumer_procs, + num_episodes, + batch_size, + train_dataset_config, + model_config, + generate_config, + tokenizer_config=None, + microbatch_size=1, + backend="transformers", + num_generations: int = 8, + consumer_plugin_config=None, + eval_dataset_config=None, + eval_interval=-1, # disable evaluation + grpo_config: Dict[str, Any] = None, + eval_save_dir: str = "./eval", + eval_generation_config={}, + project_name: str = None, + run_name: str = None, + wandb_group_name: str = None, + log_rollout_interval: int = 20, + rollout_log_file: str = "./rollout_log.jsonl", + enable_profiling: bool = False, + n_behind: int = 0, + ): + super().__init__( + producer_idx, + num_producers, + num_consumer_procs, + num_episodes, + batch_size, + train_dataset_config, + model_config, + generate_config, + tokenizer_config, + microbatch_size, + backend, + consumer_plugin_config, + eval_dataset_config=eval_dataset_config, + eval_interval=eval_interval, + grpo_config=grpo_config, + eval_save_dir=eval_save_dir, + project_name=project_name, + run_name=run_name, + wandb_group_name=wandb_group_name, + log_rollout_interval=log_rollout_interval, + rollout_log_file=rollout_log_file, + enable_profiling=enable_profiling, + n_behind=n_behind, + ) + self.model = self.backend_cls(model_config, generate_config, self.tokenizer, num_generations) + self.eval_generation_config = copy.deepcopy(self.model.generate_config) + self.eval_generation_config["n"] = 1 # use 1 generation for evaluation + self.eval_generation_config.update(eval_generation_config) + self.eval_sample_params = SamplingParams(**self.eval_generation_config) + + @torch.no_grad() + def rollout(self, input_ids, attention_mask, **kwargs): + rollouts = self.model.generate(input_ids, attention_mask, **kwargs) + if self.producer_idx == 0 and not self.eval_mode: + if ( + self.consumer_global_step - self.latest_rollout_log_step >= self.log_rollout_interval + or self.latest_rollout_log_step == -1 + ): + new_record = ( + json.dumps( + { + "train_step": self.consumer_global_step, + "rollout": self.tokenizer.batch_decode( + rollouts["input_ids"][:, 0], skip_special_tokens=True + ), + } + ) + + "\n" + ) + self.rollout_log_file.write(new_record) + self.rollout_log_file.flush() + self.latest_rollout_log_step = self.consumer_global_step + return rollouts + + def __del__(self): + if self.producer_idx == 0: + self.wandb_run.finish() + if hasattr(self, "rollout_log_file"): + self.rollout_log_file.close() + + def load_state_dict(self, state_dict): + self.model.load_state_dict(state_dict) diff --git a/applications/ColossalChat/coati/distributed/profiling_utils.py b/applications/ColossalChat/coati/distributed/profiling_utils.py new file mode 100644 index 000000000000..1c1169b50ada --- /dev/null +++ b/applications/ColossalChat/coati/distributed/profiling_utils.py @@ -0,0 +1,37 @@ +import os +import time + + +class CustomProfiler: + def __init__(self, name, disabled=True): + self.disabled = disabled + if not disabled: + self.name = name + self.pid = os.getpid() + self.file = open(f"{name}.prof", "w") + + def _log(self, message): + if self.disabled: + return + current_time = time.time() + self.file.write(f"{current_time} {self.name} {self.pid}:: {message}\n") + self.file.flush() + + def log(self, message): + if self.disabled: + return + current_time = time.time() + self.file.write(f"[Log]: {current_time} {self.name} {self.pid}:: {message}\n") + self.file.flush() + + def enter(self, event_name): + self._log(f"Enter {event_name}") + + def exit(self, event_name): + self._log(f"Exit {event_name}") + + def close(self): + if self.disabled: + return + self.file.close() + print(f"Profiler data written to {self.name}.prof") diff --git a/applications/ColossalChat/coati/distributed/reward/code_reward/testing_util.py b/applications/ColossalChat/coati/distributed/reward/code_reward/testing_util.py new file mode 100644 index 000000000000..19be101f66b7 --- /dev/null +++ b/applications/ColossalChat/coati/distributed/reward/code_reward/testing_util.py @@ -0,0 +1,674 @@ +# Code from the verl Project (https://github.com/agentica-project/rllm), +# which itself is adapted from Prime (https://github.com/PRIME-RL/PRIME) +# +# Copyright 2024 ByteDance Group + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + + +import ast +import faulthandler +import json +import platform + +# to run the solution files we're using a timing based approach +import signal +import sys +import traceback + +# used for debugging to time steps +from datetime import datetime +from enum import Enum + +# for capturing the stdout +from io import StringIO + +# used for testing the code that reads from input +from unittest.mock import mock_open, patch + +import numpy as np +from pyext import RuntimeModule + + +def truncatefn(s, length=300): + assert isinstance(s, str) + if len(s) <= length: + return s + + return s[: length // 2] + "...(truncated) ..." + s[-length // 2 :] + + +class CODE_TYPE(Enum): + call_based = 0 + standard_input = 1 + + +# used to capture stdout as a list +# from https://stackoverflow.com/a/16571630/6416660 +# alternative use redirect_stdout() from contextlib +class Capturing(list): + def __enter__(self): + self._stdout = sys.stdout + sys.stdout = self._stringio = StringIO() + # Make closing the StringIO a no-op + self._stringio.close = lambda x: 1 + return self + + def __exit__(self, *args): + self.append(self._stringio.getvalue()) + del self._stringio # free up some memory + sys.stdout = self._stdout + + +def only_int_check(val): + return isinstance(val, int) + + +def string_int_check(val): + return isinstance(val, str) and val.isdigit() + + +def combined_int_check(val): + return only_int_check(val) or string_int_check(val) + + +def clean_traceback(error_traceback): + file_start = error_traceback.find('File ""') + # print(file_start) + error_traceback = "Traceback (most recent call last):\n " + error_traceback[file_start:] + return error_traceback + + +def run_test(in_outs, test=None, debug=False, timeout=15, run_all_tests=False): + """ + if test(generated_code) is not None it'll try to run the code. + otherwise it'll just return an input and output pair. + """ + # Disable functionalities that can make destructive changes to the test. + reliability_guard() + + if debug: + print(f"start = {datetime.now().time()}") + + if in_outs: + if in_outs.get("fn_name") is None: + which_type = CODE_TYPE.standard_input # Standard input + method_name = None + else: + which_type = CODE_TYPE.call_based # Call-based + method_name = in_outs["fn_name"] + + if debug: + print(f"loaded input_output = {datetime.now().time()}") + + if test is None: + raise AssertionError("should not happen: test code is none") + elif test is not None: + results = [] + sol = "from string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\nfrom copy import *\nfrom math import *\nfrom random import *\nfrom statistics import *\nfrom itertools import *\nfrom functools import *\nfrom operator import *\nfrom io import *\nfrom sys import *\nfrom json import *\nfrom builtins import *\nfrom typing import *\nimport string\nimport re\nimport datetime\nimport collections\nimport heapq\nimport bisect\nimport copy\nimport math\nimport random\nimport statistics\nimport itertools\nimport functools\nimport operator\nimport io\nimport sys\nimport json\nsys.setrecursionlimit(6*10**5)\n" # noqa: E501 + if debug: + print(f"loading test code = {datetime.now().time()}") + + if which_type == CODE_TYPE.call_based: + sol += test + if debug: + print(f"sol = {sol}") + signal.alarm(timeout) + try: + tmp_sol = RuntimeModule.from_string("tmp_sol", "", sol) + tmp = tmp_sol if "class Solution" not in test else tmp_sol.Solution() + signal.alarm(0) + except Exception as e: + signal.alarm(0) + error_traceback = traceback.format_exc() + if debug: + print(f"type 0 compilation error = {e}") + results.append(-2) + return results, { + "error": repr(e), + # "error_code": -1, + # "error_message": "Compilation Error", + "traceback": clean_traceback(error_traceback), + } + signal.alarm(0) + + elif which_type == CODE_TYPE.standard_input: + # sol + # if code has if __name__ == "__main__": then remove it + try: + astree = ast.parse(test) + last_block = astree.body[-1] + if isinstance(last_block, ast.If): + condition = last_block.test + if ast.unparse(condition).strip() == "__name__ == '__main__'": + test = ast.unparse(astree.body[:-1]) + "\n" + ast.unparse(last_block.body) + except Exception: + pass + + tmp_test = test.split("\n") + + new_test = [] + for x in tmp_test: + if (not x.startswith("from ")) and (not x.startswith("import ")): + new_test.append("\t" + x + "\n") + else: + new_test.append(x + "\n") + tmp_test = new_test + + new_test = "" + started = False + for i in tmp_test: + if i.startswith("\t") and not started: + new_test += "stdin = sys.stdin\nstdout = sys.stdout\n" + new_test += "def code():\n" + new_test += i + started = True + elif started and ((i.startswith("from ")) or (i.startswith("import "))): + new_test += "\t" + i + else: + new_test += i + tmp_test = new_test + + sol += tmp_test + method_name = "code" + signal.alarm(timeout) + try: + tmp_sol = RuntimeModule.from_string("tmp_sol", "", sol) + tmp = tmp_sol + signal.alarm(0) + except Exception as e: + signal.alarm(0) + error_traceback = traceback.format_exc() + if debug: + print(f"type 1 compilation error = {e}") + results.append(-2) + return results, { + "error": repr(e), + # "error_code": -1, + # "error_message": "Compilation Error", + "traceback": clean_traceback(error_traceback), + } + signal.alarm(0) + if debug: + print(f"get method {method_name} = {datetime.now().time()}") + try: + method = getattr(tmp, method_name) # get_attr second arg must be str + except Exception: + signal.alarm(0) + error_traceback = traceback.format_exc() + error_info = sys.exc_info() + print(f"unable to get function error = {error_info}") + results.append(-2) + return results, { + "error": repr(error_info), + # "error_code": -1, + # "error_message": "Unable to extract code", + "traceback": clean_traceback(error_traceback), + } + + for index, inputs in enumerate(in_outs["inputs"]): + raw_inputs = inputs + raw_outputs = in_outs["outputs"][index] + if which_type == CODE_TYPE.call_based: + inputs = [json.loads(line) for line in inputs.split("\n")] + in_outs["outputs"][index] = json.loads(in_outs["outputs"][index]) + + truncate_line_size = 300 // (raw_inputs.count("\n") + 1) + raw_inputs = "\n".join( + [truncatefn(line, truncate_line_size) for line in raw_inputs.strip().split("\n")] + ) + raw_outputs = truncatefn(raw_outputs, 200) + else: + raw_inputs = truncatefn(raw_inputs) + raw_outputs = truncatefn(raw_outputs, 200) + # JSON forces dictionaries to have string keys; this undoes this (assuming a singleton list) + try: + if isinstance(inputs[0], dict): + inputs = [{int(k): v for k, v in inputs[0].items()}] + except Exception: + pass + try: + if isinstance(in_outs["outputs"][index], dict): + in_outs["outputs"][index] = [{int(k): v for k, v in in_outs["outputs"][index].items()}] + except Exception: + pass + try: + if isinstance(in_outs["outputs"][index][0], dict): + in_outs["outputs"][index] = [{int(k): v for k, v in in_outs["outputs"][index][0].items()}] + except Exception: + pass + + if debug: + print( + f"time: {datetime.now().time()} testing index = {index} inputs = {inputs}, {type(inputs)}. type = {which_type}" + ) + if which_type == CODE_TYPE.call_based: # Call-based + signal.alarm(timeout) + faulthandler.enable() + try: + output = method(*inputs) + raw_true_output = output + + raw_true_output_copy = json.dumps(output) + raw_true_output_copy = truncatefn(raw_true_output_copy, 200) + + # ground truth sequences are not tuples + if isinstance(output, tuple): + output = list(output) + + tmp_result = output == in_outs["outputs"][index] + if isinstance(in_outs["outputs"][index], list) and in_outs["outputs"][index]: + tmp_result = tmp_result or (output == in_outs["outputs"][index][0]) + + # ground truth sequences are not tuples + try: + if isinstance(output[0], tuple): + tmp_result = tmp_result or ([list(x) for x in output] == in_outs["outputs"][index][0]) + except Exception: + pass + results.append(tmp_result) + if tmp_result is not True: + return results, { + "output": raw_true_output_copy, + "expected": raw_outputs, + "inputs": raw_inputs, + # "error_code": -2, + "error_message": "Wrong Answer", + } + # reset the alarm + signal.alarm(0) + except Exception as e: + signal.alarm(0) + error_traceback = traceback.format_exc() + faulthandler.disable() + if debug: + print(f"Standard input runtime error or time limit exceeded error = {e}") + results.append(-1) + return results, { + "error": repr(e), + "traceback": clean_traceback(error_traceback), + } + faulthandler.disable() + signal.alarm(0) + if debug: + print( + f"outputs = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + elif which_type == CODE_TYPE.standard_input: # Standard input + faulthandler.enable() + passed = False + + if isinstance(inputs, list): + inputs = "\n".join(inputs) + if isinstance(in_outs["outputs"][index], list): + in_outs["outputs"][index] = "\n".join(in_outs["outputs"][index]) + + signal.alarm(timeout) + with Capturing() as output: + try: + call_method(method, inputs) + # reset the alarm + signal.alarm(0) + passed = True + except Exception as e: + # runtime error or took too long + signal.alarm(0) + error_traceback = traceback.format_exc() + print(f"Call-based runtime error or time limit exceeded error = {repr(e)}{e}") + results.append(-1) + signal.alarm(0) + if run_all_tests: + continue + return results, { + "error": repr(e), + "traceback": clean_traceback(error_traceback), + } + signal.alarm(0) + raw_true_output = output[0] + raw_true_output_copy = truncatefn(raw_true_output, 200) + output = raw_true_output.splitlines() + if not passed: + if debug: + nl = "\n" + if not isinstance(inputs, list): + print( + f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl, ' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + else: + print( + f"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + continue + + if passed and debug: + print(f"==> output = {output}, test outputs = {in_outs['outputs'][index]}") + + if custom_compare_(output, in_outs["outputs"][index]): + tmp_result = True + results.append(tmp_result) + continue + + # ground truth sequences are expressed as lists not tuples + if isinstance(output, tuple): + output = list(output) + + tmp_result = False + try: + tmp_result = output == [in_outs["outputs"][index]] + if isinstance(in_outs["outputs"][index], list): + tmp_result = tmp_result or (output == in_outs["outputs"][index]) + if isinstance(output[0], str): + tmp_result = tmp_result or ([e.strip() for e in output] == in_outs["outputs"][index]) + except Exception as e: + if debug: + print(f"Failed check1 exception = {e}") + + if tmp_result is True: + results.append(tmp_result) + continue + + # try one more time without \n + if isinstance(in_outs["outputs"][index], list): + for tmp_index, i in enumerate(in_outs["outputs"][index]): + in_outs["outputs"][index][tmp_index] = i.split("\n") + in_outs["outputs"][index][tmp_index] = [ + x.strip() for x in in_outs["outputs"][index][tmp_index] if x + ] + else: + in_outs["outputs"][index] = in_outs["outputs"][index].split("\n") + in_outs["outputs"][index] = list(filter(len, in_outs["outputs"][index])) + in_outs["outputs"][index] = list(map(lambda x: x.strip(), in_outs["outputs"][index])) + + try: + tmp_result = output == [in_outs["outputs"][index]] + if isinstance(in_outs["outputs"][index], list): + tmp_result = tmp_result or (output == in_outs["outputs"][index]) + except Exception as e: + if debug: + print(f"Failed check2 exception = {e}") + + if tmp_result is True: + results.append(tmp_result) + continue + + # try by converting the output into a split up list too + if isinstance(output, list): + output = list(filter(len, output)) + + if debug: + nl = "\n" + if not isinstance(inputs, list): + print( + f"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl, ' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]} {tmp_result=}" + ) + else: + print( + f"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]} {tmp_result=}" + ) + + if debug: + print(f"{tmp_result=} @a") + + try: + tmp_result = output == [in_outs["outputs"][index]] + if isinstance(in_outs["outputs"][index], list): + tmp_result = tmp_result or (output == in_outs["outputs"][index]) + except Exception as e: + if debug: + print(f"Failed check3 exception = {e}") + + if debug: + print(f"{tmp_result=} @b") + + try: + all_ints = all( + combined_int_check(e1) and combined_int_check(e2) + for e1, e2 in zip(output, in_outs["outputs"][index]) + ) + if not all_ints: + if debug: + print( + [ + combined_int_check(e1) and combined_int_check(e2) + for e1, e2 in zip(output, in_outs["outputs"][index]) + ] + ) + output_float = [float(e) for e in output] + gt_float = [float(e) for e in in_outs["outputs"][index]] + tmp_result = tmp_result or ( + (len(output_float) == len(gt_float)) and np.allclose(output_float, gt_float) + ) + except Exception: + pass + + if debug: + print(f"{tmp_result=} @c") + + try: + if isinstance(output[0], list): + all_ints = all( + combined_int_check(e1) and combined_int_check(e2) + for e1, e2 in zip(output[0], in_outs["outputs"][index]) + ) + if not all_ints: + output_float = [float(e) for e in output[0]] + gt_float = [float(e) for e in in_outs["outputs"][index][0]] + tmp_result = tmp_result or ( + (len(output_float) == len(gt_float)) and np.allclose(output_float, gt_float) + ) + except Exception: + pass + + if tmp_result is True: + results.append(tmp_result) + continue + + if debug: + print(f"{tmp_result=} @d") + # try by converting the stuff into split up list + if isinstance(in_outs["outputs"][index], list): + for tmp_index, i in enumerate(in_outs["outputs"][index]): + in_outs["outputs"][index][tmp_index] = set(i.split()) + else: + in_outs["outputs"][index] = set(in_outs["outputs"][index].split()) + + if debug: + print(f"{tmp_result=} @e") + + try: + tmp_result = output == in_outs["outputs"][index] + except Exception as e: + if debug: + print(f"Failed check4 exception = {e}") + continue + + if tmp_result is True: + results.append(tmp_result) + continue + + if debug: + print(f"{tmp_result=} @f") + + # try by converting the output into a split up list too + if isinstance(output, list): + for tmp_index, i in enumerate(output): + output[tmp_index] = i.split() + output = list(filter(len, output)) + for tmp_index, i in enumerate(output): + output[tmp_index] = set(i) + else: + output = output.split() + output = list(filter(len, output)) + output = set(output) + + if debug: + print(f"{tmp_result=} @g") + + if tmp_result is True and debug: + print("PASSED") + + results.append(tmp_result) + if tmp_result is not True: + if debug: + print("final result:", results) + if run_all_tests: + continue + return results, { + "output": raw_true_output_copy, + "expected": raw_outputs, + "inputs": raw_inputs, + # "error_code": -2, + "error_message": "Wrong Answer", + } + + if debug: + nl = "\n" + if not isinstance(inputs, list): + print( + f"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs.replace(nl, ' new-line ')}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + else: + print( + f"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}" + ) + + print(f"results = {results}") + if debug: + print("final results", results) + return results, {} + + +def custom_compare_(output, ground_truth): + if isinstance(output, list): + output_1 = "\n".join(output) + if stripped_string_compare(output_1, ground_truth): + return True + + if isinstance(output, list): + output_2 = [o.lstrip().rstrip() for o in output] + output_2 = "\n".join(output_2) + if stripped_string_compare(output_2, ground_truth): + return True + + return False + + +def stripped_string_compare(s1, s2): + s1 = s1.lstrip().rstrip() + s2 = s2.lstrip().rstrip() + return s1 == s2 + + +def call_method(method, inputs): + if isinstance(inputs, list): + inputs = "\n".join(inputs) + + inputs_line_iterator = iter(inputs.split("\n")) + + # sys.setrecursionlimit(10000) + + # @patch('builtins.input', side_effect=inputs.split("\n")) + @patch("builtins.open", mock_open(read_data=inputs)) + @patch("sys.stdin", StringIO(inputs)) + @patch("sys.stdin.readline", lambda *args: next(inputs_line_iterator)) + @patch("sys.stdin.readlines", lambda *args: inputs.split("\n")) + @patch("sys.stdin.read", lambda *args: inputs) + # @patch('sys.stdout.write', print) + def _inner_call_method(_method): + try: + return _method() + except SystemExit: + pass + finally: + pass + + return _inner_call_method(method) + + +def reliability_guard(maximum_memory_bytes=None): + """ + This disables various destructive functions and prevents the generated code + from interfering with the test (e.g. fork bomb, killing other processes, + removing filesystem files, etc.) + WARNING + This function is NOT a security sandbox. Untrusted code, including, model- + generated code, should not be blindly executed outside of one. See the + Codex paper for more information about OpenAI's code sandbox, and proceed + with caution. + """ + + if maximum_memory_bytes is not None: + import resource + + resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)) + resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)) + if platform.uname().system != "Darwin": + resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)) + + faulthandler.disable() + + import builtins + + builtins.exit = None + builtins.quit = None + + import os + + os.environ["OMP_NUM_THREADS"] = "1" + + os.kill = None + os.system = None # 防止干扰repl评测 + os.putenv = None + os.remove = None + os.removedirs = None + os.rmdir = None + os.fchdir = None + os.setuid = None + os.fork = None + os.forkpty = None + os.killpg = None + os.rename = None + os.renames = None + os.truncate = None + os.replace = None + os.unlink = None + os.fchmod = None + os.fchown = None + os.chmod = None + os.chown = None + os.chroot = None + os.lchflags = None + os.lchmod = None + os.lchown = None + os.getcwd = None + os.chdir = None + + import shutil + + shutil.rmtree = None + shutil.move = None + shutil.chown = None + + import subprocess + + subprocess.Popen = None # type: ignore + + __builtins__["help"] = None + + import sys + + sys.modules["ipdb"] = None + sys.modules["joblib"] = None + sys.modules["resource"] = None + sys.modules["psutil"] = None + sys.modules["tkinter"] = None diff --git a/applications/ColossalChat/coati/distributed/reward/code_reward/utils.py b/applications/ColossalChat/coati/distributed/reward/code_reward/utils.py new file mode 100644 index 000000000000..dbb3b2149780 --- /dev/null +++ b/applications/ColossalChat/coati/distributed/reward/code_reward/utils.py @@ -0,0 +1,71 @@ +# Code from the verl Project (https://github.com/agentica-project/rllm), +# which itself is adapted from Prime (https://github.com/PRIME-RL/PRIME) +# +# Copyright 2024 ByteDance Group + +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at + +# http://www.apache.org/licenses/LICENSE-2.0 + +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import multiprocessing +import traceback +from typing import Optional + +import requests + +from .testing_util import run_test + + +def _temp_run(sample, generation, debug, result, metadata_list, timeout): + try: + res, metadata = run_test(in_outs=sample, test=generation, debug=debug, timeout=timeout) + result.append(res) + metadata_list.append(metadata) + except Exception: + # print(e) # some tracebacks are extremely long. + traceback.print_exc(10) + result.append([-1 for i in range(len(sample["inputs"]))]) + metadata_list.append({}) + + +def check_correctness(in_outs: Optional[dict], generation, timeout=10, debug=True): + """Check correctness of code generation with a global timeout. + The global timeout is to catch some extreme/rare cases not handled by the timeouts + inside `run_test`""" + + manager = multiprocessing.Manager() + result = manager.list() + metadata_list = manager.list() + p = multiprocessing.Process(target=_temp_run, args=(in_outs, generation, debug, result, metadata_list, timeout)) + p.start() + p.join(timeout=600) # Global timeout of 10 minutes that's for all test cases combined + if p.is_alive(): + p.kill() + # p.terminate() + if not result: + # consider that all tests failed + result = [[-1 for i in range(len(in_outs["inputs"]))]] + if debug: + print("global timeout") + return result[0], metadata_list + + +def check_correctness_code_api( + in_outs: Optional[dict], generation, timeout=10, debug=True, url="http://localhost:8000/check_correctness" +): + payload = {"in_outs": in_outs, "generation": generation, "timeout": timeout, "debug": debug} + response = requests.post(url, json=payload) + if response.status_code == 200: + results = response.json() + return results["result"], results["metadata"] + else: + print(f"Error: {response.status_code} - {response.text}") + return [-1 for i in range(len(in_outs["inputs"]))], {} diff --git a/applications/ColossalChat/coati/distributed/reward/reward_fn.py b/applications/ColossalChat/coati/distributed/reward/reward_fn.py new file mode 100644 index 000000000000..f7a2fb89cadb --- /dev/null +++ b/applications/ColossalChat/coati/distributed/reward/reward_fn.py @@ -0,0 +1,306 @@ +# Copyright 2024 ByteDance Group + +# 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. +""" +Some functions in this file are adapted from the verl project +under the Apache License 2.0: +https://github.com/volcengine/verl +""" + + +import json + +import torch +from latex2sympy2_extended import NormalizationConfig +from math_verify import ExprExtractionConfig, LatexExtractionConfig, parse, verify + +from .code_reward.utils import check_correctness_code_api as check_correctness_code +from .reward_utils import extract_boxed_solution, extract_solution, validate_response_structure + +CANNOT_PARSE_GT_ANSWER = -1 +CANNOT_PARSE_PREDICTION = -2 +SUCCESS = 1 +MATCHING_FAIL = 0 + + +def verify_math_representation(completion, gt_answer): + """ + Verify if the completion is a valid math representation of the gt_answer. + """ + if not completion.startswith("\\boxed{"): + completion = "\\boxed{" + completion + "}" + if not gt_answer.startswith("\\boxed{"): + gt_answer = "\\boxed{" + gt_answer + "}" + target = ( + ExprExtractionConfig(), + LatexExtractionConfig( + normalization_config=NormalizationConfig( + nits=False, + malformed_operators=False, + basic_latex=True, + boxed="all", + units=True, + ), + boxed_match_priority=0, + ), + ) + if not isinstance(gt_answer, str) or len(gt_answer) == 0: + raise ValueError("gt_answer should be a string, please verify your training data.") + if not isinstance(completion, str) or len(completion) == 0: + return MATCHING_FAIL + try: + parsed_gt_answer = parse(gt_answer, extraction_config=target) + if len(parsed_gt_answer) == 0: + return CANNOT_PARSE_GT_ANSWER + parsed_completion = parse(completion, extraction_config=target) + if len(parsed_completion) == 0: + return CANNOT_PARSE_PREDICTION + if verify(parsed_gt_answer, parsed_completion): + return SUCCESS + else: + return MATCHING_FAIL + except Exception: + return MATCHING_FAIL + + +def verify_model_answer(decoded_final_answer, gt_answer, ans_acc, acc_score, reward): + math_verify_result = verify_math_representation(decoded_final_answer, gt_answer) + exact_match_result = ( + SUCCESS + if decoded_final_answer.strip().replace(" ", "").replace("{", "").replace("}", "").replace(",", "") + == gt_answer.strip().replace(" ", "").replace("{", "").replace("}", "").replace(",", "") + else MATCHING_FAIL + ) + if math_verify_result == SUCCESS: + ans_acc += 1 + reward += acc_score + elif exact_match_result == SUCCESS: + # sometimes for answers that's not a (valid) math expression, math_verify will fail + ans_acc += 1 + if math_verify_result == CANNOT_PARSE_PREDICTION: + reward += ( + acc_score / 2 + ) # not a valid latex math representation, but the answer is correct, receive half of the score + else: + reward += acc_score + return reward, ans_acc + + +def math_reward_fn(input_ids, gt_answer, response_idx, **kwargs): + tokenizer = kwargs["tokenizer"] + eval_mode = kwargs.get("eval_mode", False) + soft_over_length_punishment = kwargs.get("soft_over_length_punishment", False) + acc_score = 10.0 + reward = torch.tensor(0.0) + format_acc = torch.tensor(0.0) + ans_acc = torch.tensor(0.0) + s, e = response_idx[0], response_idx[1] + + length_reward = 0.0 + res_length = e.item() - s.item() + 1 + if not eval_mode: + max_new_tokens = kwargs["max_new_tokens"] + else: + max_new_tokens = -1 # for eval mode, we don't need to check the length + if not eval_mode and soft_over_length_punishment: + cache_length = kwargs["cache_length"] + if max_new_tokens - cache_length < res_length < max_new_tokens: + length_reward = ((max_new_tokens - cache_length) - res_length) / cache_length * acc_score + + if gt_answer is None: + raise ValueError("no gt_answer is provided, please check your training dataset.") + + decoded_final_answer = tokenizer.decode(input_ids[s : e + 1], skip_special_tokens=True) + + final_answer, processed_str = extract_solution(decoded_final_answer) + + format_valid = validate_response_structure(processed_str, kwargs["tags"]) + + # Check answer accuracy, answer is considered correct if the answer is correct and the format is valid + if final_answer is not None: + if eval_mode or format_valid: + reward, ans_acc = verify_model_answer(final_answer, gt_answer, ans_acc, acc_score, reward) + if not eval_mode: + reward = reward + length_reward + + # Check format accuracy + if format_valid: + format_acc += 1 + + # Check if the sequence is over length + if not eval_mode and res_length >= max_new_tokens: + reward *= 0.0 + + if not eval_mode: + return torch.tensor([reward, format_acc, ans_acc]).to(input_ids.device) + else: + prompt = tokenizer.decode(input_ids[:s], skip_special_tokens=True) + return { + "prompt": prompt, + "prediction": decoded_final_answer, + "gold": gt_answer, + "parsed": final_answer, + "format_valid": format_acc.item(), + "ans_valid": ans_acc.item(), + "response_length": res_length, + "reward": reward.item(), + } + + +def boxed_math_reward_fn(input_ids, gt_answer, response_idx, **kwargs): + tokenizer = kwargs["tokenizer"] + eval_mode = kwargs.get("eval_mode", False) + soft_over_length_punishment = kwargs.get("soft_over_length_punishment", False) + acc_score = 10.0 + reward = torch.tensor(0.0) + format_acc = torch.tensor(0.0) + ans_acc = torch.tensor(0.0) + s, e = response_idx[0], response_idx[1] + + length_reward = 0.0 + res_length = e.item() - s.item() + 1 + if not eval_mode: + max_new_tokens = kwargs["max_new_tokens"] + else: + max_new_tokens = -1 # for eval mode, we don't need to check the length + if not eval_mode and soft_over_length_punishment: + cache_length = kwargs["cache_length"] + if max_new_tokens - cache_length < res_length < max_new_tokens: + length_reward = ((max_new_tokens - cache_length) - res_length) / cache_length * acc_score + + if gt_answer is None: + raise ValueError("no gt_answer is provided, please check your training dataset.") + + decoded_final_answer = tokenizer.decode(input_ids[s : e + 1], skip_special_tokens=True) + + final_answer = extract_boxed_solution(decoded_final_answer) + format_valid = final_answer is not None + if "tags" in kwargs and kwargs["tags"]: + tags = kwargs["tags"] + format_valid = format_valid and all( + [decoded_final_answer.count(tags[tag]["text"]) == tags[tag]["num_occur"] for tag in tags] + ) + + # Check answer accuracy, answer is considered correct if the answer is correct and the format is valid + if final_answer is not None: + if eval_mode or format_valid: + reward, ans_acc = verify_model_answer(final_answer, gt_answer, ans_acc, acc_score, reward) + if not eval_mode: + reward = reward + length_reward + + # Check format accuracy + if format_valid: + format_acc += 1 + + # Check if the sequence is over length + if not eval_mode and res_length >= max_new_tokens: + reward *= 0.0 + + if not eval_mode: + return torch.tensor([reward, format_acc, ans_acc]).to(input_ids.device) + else: + prompt = tokenizer.decode(input_ids[:s], skip_special_tokens=True) + return { + "prompt": prompt, + "prediction": decoded_final_answer, + "gold": gt_answer, + "parsed": final_answer, + "format_valid": format_acc.item(), + "ans_valid": ans_acc.item(), + "response_length": res_length, + "reward": reward.item(), + } + + +def code_reward_fn(input_ids, test_cases, response_idx, **kwargs): + url = kwargs.get("url", "http://localhost:8000/check_correctness") + tokenizer = kwargs["tokenizer"] + eval_mode = kwargs.get("eval_mode", False) + soft_over_length_punishment = kwargs.get("soft_over_length_punishment", False) + acc_score = 10.0 + reward = torch.tensor(0.0) + format_acc = torch.tensor(0.0) + ans_acc = torch.tensor(0.0) + s, e = response_idx[0], response_idx[1] + + length_reward = 0.0 + res_length = e.item() - s.item() + 1 + if not eval_mode: + max_new_tokens = kwargs["max_new_tokens"] + else: + max_new_tokens = -1 # for eval mode, we don't need to check the length + if not eval_mode and soft_over_length_punishment: + cache_length = kwargs["cache_length"] + if max_new_tokens - cache_length < res_length < max_new_tokens: + length_reward = ((max_new_tokens - cache_length) - res_length) / cache_length * acc_score + + # try to get code solution from completion. if the completion is pure code, this will not take effect. + decoded_final_answer = tokenizer.decode(input_ids[s : e + 1], skip_special_tokens=True) + + solution = decoded_final_answer.split("```python")[-1].split("```")[0] + format_valid = False + if "```python" in decoded_final_answer: + format_valid = solution is not None + + # Check format accuracy + if format_valid: + format_acc += 1 + + res = [] + metadata = [] + + try: + try: + if not isinstance(test_cases, dict): + test_cases = json.loads(test_cases) + except Exception as e: + print(f"Error {e}: Cannot parse test cases.") + raise e + # Complete check on all in-out pairs first. If there is no failure, per-sample test can be skipped. + try: + res, metadata = check_correctness_code( + in_outs=test_cases, generation=solution, timeout=10, debug=False, url=url + ) + metadata = dict(enumerate(metadata))[0] + success = all(map(lambda x: x == 1, res)) + if success: + ans_acc += 1 + if eval_mode or format_valid: + reward += acc_score + if not eval_mode: + reward = reward + length_reward + + except Exception: + pass + + # Check if the sequence is over length + if not eval_mode and res_length >= max_new_tokens: + reward *= 0.0 + except Exception: + pass + if not eval_mode: + return torch.tensor([reward, format_acc, ans_acc]).to(input_ids.device) + else: + prompt = tokenizer.decode(input_ids[:s], skip_special_tokens=True) + return { + "prompt": prompt, + "prediction": decoded_final_answer, + "test_cases": test_cases, + "test_results": res, + "test_metadata": metadata, + "parsed": solution, + "format_valid": format_acc.item(), + "ans_valid": ans_acc.item(), + "response_length": res_length, + "reward": reward.item(), + } diff --git a/applications/ColossalChat/coati/distributed/reward/reward_utils.py b/applications/ColossalChat/coati/distributed/reward/reward_utils.py new file mode 100644 index 000000000000..ffc220846ea7 --- /dev/null +++ b/applications/ColossalChat/coati/distributed/reward/reward_utils.py @@ -0,0 +1,124 @@ +# Copyright Unakar +# Modified from https://github.com/Unakar/Logic-RL/blob/086373176ac198c97277ff50f4b6e7e1bfe669d3/verl/utils/reward_score/kk.py#L99 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import re +from typing import Dict, Optional, Tuple + + +def validate_response_structure(processed_str: str, tags: Dict = None) -> bool: + """Performs comprehensive validation of response structure. + + Args: + processed_str: Processed response string from the model + + Returns: + Boolean indicating whether all formatting requirements are met + """ + validation_passed = True + # Check required tags + if tags is None: + tags = { + "think_start": {"text": "", "num_occur": 1}, + "think_end": {"text": "", "num_occur": 1}, + "answer_start": {"text": "", "num_occur": 1}, + "answer_end": {"text": "", "num_occur": 1}, + } + positions = {} + for tag_name, tag_info in tags.items(): + tag_str = tag_info["text"] + expected_count = tag_info["num_occur"] + count = processed_str.count(tag_str) + positions[tag_name] = pos = processed_str.find(tag_str) + if count != expected_count: + validation_passed = False + # Verify tag order + if ( + positions["think_start"] > positions["think_end"] + or positions["think_end"] > positions["answer_start"] + or positions["answer_start"] > positions["answer_end"] + ): + validation_passed = False + if len(processed_str) - positions["answer_end"] != len(tags["answer_end"]["text"]): + validation_passed = False + return validation_passed + + +def extract_solution(solution_str: str) -> Tuple[Optional[str], str]: + """Extracts the final answer from the model's response string. + + Args: + solution_str: Raw response string from the language model + + Returns: + Tuple containing (extracted_answer, processed_string) + """ + + # Extract final answer using XML-style tags + answer_pattern = r"(.*?)" + matches = list(re.finditer(answer_pattern, solution_str, re.DOTALL)) + + if not matches: + return None, solution_str + + final_answer = matches[-1].group(1).strip() + return final_answer, solution_str + + +def extract_boxed_solution(text: str) -> Optional[str]: + """ + Modified from: https://gist.github.com/lewtun/9c2ce1937b741404090a3dc4c7c022b3 + Retrieves the content from the last occurrence of `\boxed{}` in a LaTeX-like string. + + Args: + text (str): A string potentially containing LaTeX-style boxed expressions. + + Returns: + Optional[str]: The text inside the final `\boxed{}` if successfully extracted; + returns `None` if no properly closed box is found. + + Examples: + >>> extract_boxed_solution("The answer is \\boxed{42}.") + '42' + >>> extract_boxed_solution("Here is an unmatched \\boxed{42") + None + """ + try: + # Find the last occurrence of "\boxed{" + start_idx = text.rindex("\\boxed{") + # Move past "\boxed{" to find the start of the content + content_start = start_idx + len("\\boxed{") + open_braces = 1 + pos = content_start + + # Traverse the string to find the matching closing brace + while open_braces > 0 and pos < len(text): + if text[pos] == "{": + open_braces += 1 + elif text[pos] == "}": + open_braces -= 1 + pos += 1 + + # If all braces are matched, extract and return the content + if open_braces == 0: + return text[content_start : pos - 1].strip() + else: + return None + + except ValueError: + # "\boxed{" not found + return None + except Exception: + # Any other unexpected error + return None diff --git a/applications/ColossalChat/coati/distributed/reward/verifiable_reward.py b/applications/ColossalChat/coati/distributed/reward/verifiable_reward.py new file mode 100644 index 000000000000..2f0b3778c153 --- /dev/null +++ b/applications/ColossalChat/coati/distributed/reward/verifiable_reward.py @@ -0,0 +1,71 @@ +""" +Function-based reward verification module. +""" + +import inspect +from typing import Any, Dict, List + +import torch + + +class VerifiableReward: + def __init__(self, reward_fns: List[callable], **kwargs: List[Dict[str, Any]]): + self.reward_fns = reward_fns + self.kwargs = kwargs + + def __call__( + self, + input_ids: torch.LongTensor, + gt_answer: List[str] = None, + test_cases: List[str] = None, + response_idx: List[torch.Tensor] = None, + ) -> torch.Tensor: + # Get batch size + bs = input_ids.size(0) + # Initialize reward + rewards = torch.zeros((bs, 3), device=input_ids.device) + + # Loop through reward functions + for reward_fn in self.reward_fns: + # Apply the reward function to the entire batch at once + if "gt_answer" in inspect.getfullargspec(reward_fn).args: + reward_batch = torch.stack( + [ + reward_fn( + input_ids[i], + gt_answer=gt_answer[i], + response_idx=response_idx[i], + **self.kwargs, + ) + for i in range(bs) + ], + dim=0, + ) + elif "test_cases" in inspect.getfullargspec(reward_fn).args: + reward_batch = torch.stack( + [ + reward_fn( + input_ids[i], + test_cases=test_cases[i], + response_idx=response_idx[i], + **self.kwargs, + ) + for i in range(bs) + ], + dim=0, + ) + else: + reward_batch = torch.stack( + [ + reward_fn( + input_ids[i], + response_idx=response_idx[i], + **self.kwargs, + ) + for i in range(bs) + ], + dim=0, + ) + + rewards += reward_batch + return rewards diff --git a/applications/ColossalChat/coati/distributed/utils.py b/applications/ColossalChat/coati/distributed/utils.py new file mode 100644 index 000000000000..466914cc0d4d --- /dev/null +++ b/applications/ColossalChat/coati/distributed/utils.py @@ -0,0 +1,167 @@ +import json +import os +from typing import Any, Dict, List + +import torch +from filelock import FileLock + +from colossalai.shardformer.layer.loss import dist_log_prob + + +def unbind_batch(batch: Dict[str, torch.Tensor]) -> List[Dict[str, torch.Tensor]]: + batches = [] + for k, v in batch.items(): + if len(batches) == 0: + unbinded_tensors = v.unbind(0) + batches = [{k: tensor} for tensor in unbinded_tensors] + else: + unbinded_tensors = v.unbind(0) + assert len(batches) == len(unbinded_tensors) + for i, tensor in enumerate(unbinded_tensors): + batches[i][k] = tensor + return batches + + +def bind_batch(batches: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]: + batch = {} + for k in batches[0].keys(): + batch[k] = torch.stack([batch[k] for batch in batches], dim=0) + return batch + + +def pre_send(batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + # compress mask to save bandwidth + if "attention_mask" in batch: + batch["attention_mask"] = batch["attention_mask"].to(torch.bool) + if "action_mask" in batch: + batch["action_mask"] = batch["action_mask"].to(torch.bool) + return batch + + +def post_recv(batch: Dict[str, torch.Tensor]) -> Dict[str, torch.Tensor]: + # decompress mask + if "attention_mask" in batch: + batch["attention_mask"] = batch["attention_mask"].to(torch.int) + if "action_mask" in batch: + batch["action_mask"] = batch["action_mask"].to(torch.int) + return batch + + +def update_by_default(data: Dict[str, Any], default: Dict[str, Any]) -> Dict[str, Any]: + data = data.copy() + for k, v in default.items(): + if k not in data: + data[k] = v + return data + + +def log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + """ + Compute the log probabilities from logits for the given labels. + + Args: + logits (torch.Tensor): The input logits. + labels (torch.Tensor): The target labels. + + Returns: + torch.Tensor: The log probabilities corresponding to the labels. + """ + log_probs = torch.log_softmax(logits, dim=-1) + per_label_logps = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)) + return per_label_logps.squeeze(-1) + + +def memory_efficient_logprob( + logits: torch.Tensor, + inputs: torch.Tensor, + num_action: int, + chunk_size: int = 2048, + shard_config: Any = None, + vocab_size: int = None, +) -> torch.Tensor: + """ + Calculate action log probs in a memory-efficient way by processing in chunks. + Args: + logits (torch.Tensor): Output tensor of Actor.forward.logits. + inputs (torch.LongTensor): Input sequences. + num_action (int): Number of actions. + chunk_size (int, optional): Size of each chunk to process. Default is 2048. + shard_config: Shard configuration for distributed computation. + vocab_size (int, optional): Vocabulary size. Default is None. + Returns: + torch.Tensor: Action log probs. + """ + action_log_probs = torch.zeros((logits.size(0), num_action), device=logits.device, dtype=logits.dtype) + context_length = logits.size(1) - num_action + for i in range(action_log_probs.size(0)): + # loop over each sample in the micro-batch + for start in range(context_length, logits.size(1), chunk_size): + end = min(start + chunk_size, logits.size(1)) + # calculate log probs in chunks to save memory + log_probs = dist_log_prob( + inputs[i : i + 1, start - 1 : end], + logits[i : i + 1, start - 1 : end], + shard_config, + vocab_size, + logits.dtype, + ) # [1, chunk_size, 1] + log_probs = log_probs.squeeze(-1) + action_log_probs[i, start - context_length : end - context_length] += log_probs[0] + return action_log_probs + + +def entropy_from_logits(logits: torch.Tensor) -> torch.Tensor: + """ + Calculate entropy + Reference: https://github.com/volcengine/verl/blob/96b730bbed80292a439f0c0057d3920ab8b28d52/verl/utils/torch_functional.py#L145 + """ + p = torch.nn.functional.softmax(logits, dim=-1) + entropy = torch.logsumexp(logits, dim=-1) - torch.sum(p * logits, dim=-1) + return entropy + + +def masked_mean(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor: + """ + Compute the masked mean of a tensor along a specified dimension. + + Args: + tensor (torch.Tensor): The input tensor. + mask (torch.Tensor): The mask tensor with the same shape as the input tensor. + dim (int, optional): The dimension along which to compute the mean. Default is 1. + + Returns: + torch.Tensor: The masked mean 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_sum(tensor: torch.Tensor, mask: torch.Tensor, dim: int = 1) -> torch.Tensor: + """ + Compute the masked sum of a tensor along a specified dimension. + + Args: + tensor (torch.Tensor): The input tensor. + mask (torch.Tensor): The mask tensor with the same shape as the input tensor. + dim (int, optional): The dimension along which to compute the sum. Default is 1. + + Returns: + torch.Tensor: The masked sum tensor. + + """ + tensor = tensor * mask + return tensor.sum(dim=dim) + + +def safe_append_to_jsonl_file(file_path, data): + with FileLock(file_path + ".lock"): + # Ensure file exists + os.makedirs(os.path.dirname(file_path), exist_ok=True) + with open(file_path, "a", encoding="utf8") as f: + for entry in data: + json_line = json.dumps(entry, ensure_ascii=False) + f.write(json_line + "\n") diff --git a/applications/ColossalChat/coati/trainer/utils.py b/applications/ColossalChat/coati/trainer/utils.py index 22a5f492ebd2..25615d19184f 100755 --- a/applications/ColossalChat/coati/trainer/utils.py +++ b/applications/ColossalChat/coati/trainer/utils.py @@ -128,7 +128,7 @@ def all_reduce_mean(tensor: torch.Tensor, plugin: Plugin = None) -> torch.Tensor return tensor -def all_reduce_sum(tensor: torch.Tensor) -> torch.Tensor: +def all_reduce_sum(tensor: torch.Tensor, plugin: Plugin = None) -> torch.Tensor: """ Performs an all-reduce operation to sum the values of the given tensor across all processes. @@ -138,5 +138,35 @@ def all_reduce_sum(tensor: torch.Tensor) -> torch.Tensor: Returns: torch.Tensor: The reduced tensor with the sum of values across all processes. """ - dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM) + # All reduce sum across DP group + if plugin is not None: + dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM, group=plugin.dp_group) + else: + dist.all_reduce(tensor=tensor, op=dist.ReduceOp.SUM) return tensor + + +def all_gather_tensors(local_tensor_list: torch.Tensor, plugin: Plugin = None) -> torch.Tensor: + """ + Gathers tensors from all processes and concatenates them along the first dimension. + + Args: + tensor (torch.Tensor): The input tensor to be gathered. + + Returns: + torch.Tensor: The gathered tensor. + """ + # Gather tensors across DP group + if plugin is not None: + all_tensor_lists = [None] * plugin.dp_size + dist.all_gather_object(all_tensor_lists, local_tensor_list, group=plugin.dp_group) + gathered_tensor_list = [] + for tensors in all_tensor_lists: + gathered_tensor_list.extend(tensors) + else: + all_tensor_lists = [None] * dist.get_world_size() + dist.all_gather_object(all_tensor_lists, local_tensor_list) + gathered_tensor_list = [] + for tensors in all_tensor_lists: + gathered_tensor_list.extend(tensors) + return gathered_tensor_list diff --git a/applications/ColossalChat/examples/requirements.txt b/applications/ColossalChat/examples/requirements.txt index 31eef5256e74..3e4cc1a95f75 100644 --- a/applications/ColossalChat/examples/requirements.txt +++ b/applications/ColossalChat/examples/requirements.txt @@ -1,4 +1,4 @@ pandas>=1.4.1 sentencepiece -colossalai==0.4.7 +colossalai>=0.4.7 prompt_toolkit diff --git a/applications/ColossalChat/profiling.sh b/applications/ColossalChat/profiling.sh new file mode 100755 index 000000000000..d9f3d9a931df --- /dev/null +++ b/applications/ColossalChat/profiling.sh @@ -0,0 +1,13 @@ +export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1 + +# 8K context length +# rm -rf *.prof +# MAX_NEW_TOKENS=$((8192-512)) +# python rl_example.py --dataset /mnt/nfs/yeanbang/experiments/RLHF/grpo/train-alignment-samll.jsonl --model /home/share/data/model/Qwen2.5-Math-7B/ -t 4 -i 4 -b vllm -a GRPO -ibs 16 -tbs 16 -e 1 -rt boxed -si 100 -s "Please reason step by step, and put your final answer within \\boxed{}." -tmbs 4 -p GRPO-Math-Profile -ei -5 -zero 1 -pp 2 -mnt $MAX_NEW_TOKENS -nb 1 --enable_profiling 2>&1| tee ibs_64_tbs_32_tmbs_2_pp_2_ptp_2_8192_GRPO_profile.txt +# python profile_grpo.py --visualization actor_timelines_ibs_64_tbs_32_tmbs_2_pp_2_ptp_2_8192_GRPO_profile.png + +# 4K context length +rm -rf *.prof +MAX_NEW_TOKENS=$((4096-512)) +python rl_example.py --dataset /mnt/nfs/yeanbang/experiments/RLHF/grpo/train-alignment-samll.jsonl --model /home/share/data/model/Qwen2.5-Math-7B/ -t 4 -i 4 -b vllm -a GRPO -ibs 8 -tbs 8 -e 1 -rt boxed -si 100 -s "Please reason step by step, and put your final answer within \\boxed{}." -tMbs 4 -tmbs 4 -p GRPO-Math-Profile -ei -5 -zero 2 -mnt $MAX_NEW_TOKENS -nb 1 --enable_profiling 2>&1| tee ibs_32_tbs_32_tmbs_4_zero_2_4096_GRPO_profile.txt +python visualization.py --visualization actor_timelines_ibs_32_tbs_32_tmbs_4_zero_2_4096_GRPO_profile.png diff --git a/applications/ColossalChat/requirements.txt b/applications/ColossalChat/requirements.txt index 472080101b9b..ecd876ef391e 100755 --- a/applications/ColossalChat/requirements.txt +++ b/applications/ColossalChat/requirements.txt @@ -1,4 +1,4 @@ -transformers==4.39.3 +transformers>=4.39.3 tqdm datasets==2.14.7 loralib @@ -22,3 +22,6 @@ sentencepiece==0.1.99 flash-attn tiktoken jsonlines +math_verify +latex2sympy2_extended +pyext diff --git a/applications/ColossalChat/rl_example.py b/applications/ColossalChat/rl_example.py new file mode 100644 index 000000000000..58148b67e4a8 --- /dev/null +++ b/applications/ColossalChat/rl_example.py @@ -0,0 +1,392 @@ +import argparse +import json +import os + +import ray +import torch +from coati.distributed.launch import launch_distributed + +DEFAUT_SYSTEM_PROMPT = { + "think_answer_tags": "You are a helpful assistant. The assistant first thinks about the reasoning process in the mind and then provides the user with the answer. The reasoning process and answer are enclosed within and tags, respectively, i.e., reasoning process here answer here . Now the user asks you to solve a math problem that involves reasoning. After thinking, when you finally reach a conclusion, clearly output the final answer without explanation within the tags, i.e., 123 .\n\n", + "boxed": "Please reason step by step, and put your final answer within \\boxed{}.", + "code": "You are a helpful assistant.", +} + +# bypass the proxy for local addresses +os.environ["no_proxy"] = "127.0.0.1,localhost" + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + parser.add_argument("-m", "--model", type=str, default="Qwen/Qwen2.5-7B") + parser.add_argument("-d", "--dataset", type=str, default="data.jsonl") + parser.add_argument( + "-ed", + "--eval-dataset", + type=str, + default=None, + help="Evaluation dataset for each task, please use json format to specify the dataset for each task. \ + For example: {'task1':'data_eval_task1.jsonl', 'task2':'data_eval_task2.jsonl'}, the jsonl file should be in the same format as the training dataset. \ + The key is the task name, and the value is the path to the jsonl file", + ) + parser.add_argument("-p", "--project", type=str, default="GRPO", help="Project name.") + parser.add_argument("-e", "--num-episodes", type=int, default=1, help="Number of episodes to train.") + + # Distributed training parameters + parser.add_argument("-t", "--num-trainers", type=int, default=2) + parser.add_argument("-i", "--num-inferencer", type=int, default=2) + parser.add_argument("-g", "--num-generations", type=int, default=8, help="Number of generations per prompt.") + parser.add_argument( + "-ibs", + "--inference-batch-size", + type=int, + default=64, + help="Number of prompts to generate per inference step. It should be divisible by tbs, and the weights on the inference backend will be synced every ibs/tbs training steps of the policy model.", + ) + parser.add_argument( + "-imbs", + "--inference-microbatch-size", + type=int, + default=8, + help="Effective batch size for the inference backend to run generation. Please select based on memory constraint.", + ) + parser.add_argument( + "-tbs", + "--train-batch-size", + type=int, + default=32, + help="Number of unique prompts to update policy model per step per dp group. Gradient is accumulated across tbs * dp_size unique prompts, equivalently tbs * g * dp_size samples", + ) + parser.add_argument( + "-tMbs", + "--train-minibatch-size", + type=int, + default=8, + help="Number of unique prompts in each training batch per dp group. The inference backend must generate tMbs * g * dp_size samples before forwarding. Satisfy tMbs * g >= tmbs", + ) + parser.add_argument( + "-tmbs", + "--train-microbatch-size", + type=int, + default=2, + help="Effective batch size per dp group for forwarding and backwarding. Please select based on the availiable memory.", + ) + parser.add_argument( + "-tp", + "--tensor-parallel-size", + type=int, + default=1, + help="Tensor parallel size for the trainer (consumer). Please check the generation arguments documentation for your backend.", + ) + parser.add_argument( + "-pp", + "--pipeline-parallel-size", + type=int, + default=1, + help="Pipeline parallel size for the trainer (consumer). Please check the generation arguments documentation for your backend.", + ) + parser.add_argument( + "-zero", + "--zero-stage", + type=int, + default=0, + help="Zero stage for the trainer (consumer). Please check the generation arguments documentation for your backend.", + ) + parser.add_argument( + "--ray_dir", type=str, default=None, help="Custom temperary directory for storing ray cluster data, Optional" + ) + parser.add_argument( + "--master_address", type=str, default=None, help="Master address for multi-node distributed training, Optional" + ) + parser.add_argument( + "--master_port", type=int, default=29506, help="Master port for multi-node distributed training, Optional" + ) + + # Sampling parameters + parser.add_argument("-b", "--backend", type=str, default="transformers", choices=["transformers", "vllm"]) + parser.add_argument("-temp", "--temperature", type=float, default=1.0, help="Temperature for sampling.") + parser.add_argument( + "-topk", + "--top-k", + type=int, + default=None, + help="Top k for sampling. Please check the generation arguments documentation for your backend.", + ) + parser.add_argument( + "-topp", + "--top-p", + type=float, + default=1.0, + help="Top p for sampling. Please check the generation arguments documentation for your backend.", + ) + parser.add_argument("-s", "--system-prompt", type=str, default=None, help="System prompt for data construction.") + parser.add_argument("-mnt", "--max-new-tokens", type=int, default=1024 * 4 - 512, help="Max length for generation.") + parser.add_argument("-mpt", "--max-prompt-tokens", type=int, default=512, help="Max length for prompt.") + parser.add_argument( + "-ptp", + "--producer-tensor-parallel-size", + type=int, + default=1, + help="Tensor parallel size for the producer. Please check the generation arguments documentation for your backend.", + ) + + # GRPO parameters + parser.add_argument("-a", "--algo", type=str, default="GRPO", choices=["DAPO", "GRPO"]) + parser.add_argument("-lr", "--learning-rate", type=float, default=1e-6, help="Learning rate for GRPO.") + parser.add_argument("-kl", "--kl-coeff", type=float, default=0.01, help="KL penalty coefficient for GRPO.") + parser.add_argument( + "-rt", + "--reward-type", + type=str, + default="think_answer_tags", + choices=["think_answer_tags", "boxed", "code"], + help="Reward type for GRPO.", + ) + parser.add_argument( + "-cv", + "--code-verifier-api-url", + type=str, + default=None, + help="API URL for code verifier. If not provided, the code verifier will be disabled.", + ) + parser.add_argument( + "-ei", + "--eval-interval", + type=int, + default=-1, + help="Interval for evaluation. Evaluate every ei training steps.", + ) + parser.add_argument( + "-nb", + "--n-behind", + type=int, + default=0, + help="Number of producer batches to rollout to fill the data buffer before trainer starts to decrease bubble time", + ) + + # Logging/Checkpointing parameters + parser.add_argument("-si", "--save-interval", type=int, default=100, help="Interval for saving checkpoints.") + parser.add_argument("-sd", "--save-dir", type=str, default="./model", help="Directory for saving checkpoints.") + parser.add_argument( + "-esd", "--eval-save-dir", type=str, default="./eval", help="Directory for saving evaluation results." + ) + parser.add_argument( + "-rsd", "--rollout-save-dir", type=str, default="./rollouts", help="Directory for saving rollout loggings." + ) + parser.add_argument( + "--enable_profiling", action="store_true", default=False, help="Enable profiling for the training process." + ) + + args = parser.parse_args() + + if args.train_minibatch_size is None: + # Default settings: Using train batch size as mini batch size + args.train_minibatch_size = args.train_batch_size + if args.inference_batch_size is None: + # Default settings: Using train batch size as inference batch size, sync every inference model every train step + args.inference_batch_size = args.train_batch_size + assert ( + args.train_minibatch_size * args.num_generations >= args.train_microbatch_size + and args.train_microbatch_size > 0 + ), "Train micro batch size must be greater than 0 less than train mini batch size * num generations" + assert ( + args.train_minibatch_size <= args.train_batch_size and args.train_batch_size % args.train_minibatch_size == 0 + ), "Train mini batch size must be less than or equals to train batch size and train batch size must be divisible by train mini batch size" + + if args.master_address is None: + # Default settings: Using single machine + ray.init( + address="local", + namespace="ray-example", + runtime_env={ + "env_vars": { + # "RAY_DEBUG_POST_MORTEM": "1", # enable post-mortem debugging with ray + "TOKENIZERS_PARALLELISM": "false" + }, + }, + ) + else: + # For ray distributed multi-machine training, Please change _node_ip_address to your IP address of your master node + ray.init( + _node_ip_address=args.master_address, + namespace="ray-example", + _temp_dir=args.ray_dir, + runtime_env={ + "env_vars": { + # "RAY_DEBUG_POST_MORTEM": "1", # enable post-mortem debugging with ray + "TOKENIZERS_PARALLELISM": "false" + }, + }, + ) + + if args.top_k is None: + if args.backend == "transformers": + args.top_k = 50 + elif args.backend == "vllm": + args.top_k = -1 + + os.environ["TOKENIZERS_PARALLELISM"] = "false" # Disable tokenizers parallelism to avoid deadlock + + inference_model_config = dict(path=args.model) + train_model_config = dict(path=args.model, use_flash_attention_2=True, use_cache=False) + generate_config = dict(top_k=args.top_k, top_p=args.top_p, temperature=args.temperature) + + if args.backend == "transformers": + inference_model_config.update( + dict( + use_flash_attention_2=True, + torch_dtype=torch.bfloat16, + ) + ) + generate_config.update( + dict( + max_length=args.max_new_tokens + args.max_prompt_tokens, + do_sample=True, + max_new_tokens=None, + early_stopping=False if args.reward_type == "think_answer_tags" else True, + stop_strings=[""] if args.reward_type == "think_answer_tags" else None, + ) + ) + eval_generation_config = {"temperature": 0.6} # used to update generation config for evaluation + elif args.backend == "vllm": + inference_model_config.update( + dict( + gpu_memory_utilization=0.7, + enforce_eager=True, + enable_chunked_prefill=True, + max_model_len=args.max_new_tokens + args.max_prompt_tokens, + tensor_parallel_size=args.producer_tensor_parallel_size, + ) + ) + if args.enable_profiling: + # If profiling is enabled, we force model to generate to max_new_tokens + generate_config.update( + dict( + max_tokens=args.max_new_tokens, # max new tokens + ignore_eos=True, + include_stop_str_in_output=True, + stop=None, + ) + ) + else: + generate_config.update( + dict( + max_tokens=args.max_new_tokens, # max new tokens + ignore_eos=True if args.reward_type == "think_answer_tags" else False, + include_stop_str_in_output=True, + stop=[""] if args.reward_type == "think_answer_tags" else None, + ) + ) + eval_generation_config = {"temperature": 0.6} # used to update generation config for evaluation + else: + raise ValueError(f"Unsupported backend: {args.backend}") + + if args.algo == "GRPO": + # Default Settings + grpo_config = { + "lr": args.learning_rate, + "train_microbatch_size": args.train_microbatch_size, + "beta": args.kl_coeff, # KL penalty coefficient + "loss_variation": "sample_level", + "reward_fn_type": args.reward_type, + "max_length": args.max_new_tokens + args.max_prompt_tokens, + "max_new_tokens": args.max_new_tokens, + "response_format_tags": ( + { + "think_start": {"text": "", "num_occur": 1}, + "think_end": {"text": "", "num_occur": 1}, + "answer_start": {"text": "", "num_occur": 1}, + "answer_end": {"text": "", "num_occur": 1}, + } + if args.reward_type == "think_answer_tags" + else None + ), + } + elif args.algo == "DAPO": + # DAPO variant settings + grpo_config = { + "filter_range": [0.01, 0.99], # only filter out all zero batch and all one batch + "lr": args.learning_rate, + "train_microbatch_size": args.train_microbatch_size, + "dynamic_batching": True, + "clip_eps_low": 0.2, + "clip_eps_high": 0.28, + "skip_threshold": 20.0, + "beta": 0, # no KL penalty for DAPO + "loss_variation": "token_level", + "soft_over_length_punishment": True, + "max_length": args.max_new_tokens + args.max_prompt_tokens, + "max_new_tokens": args.max_new_tokens, + "cache_length": min(1024, int(args.max_new_tokens / 4)), + "filter_truncated_response": True, + "reward_fn_type": args.reward_type, + "response_format_tags": ( + { + "think_start": {"text": "", "num_occur": 1}, + "think_end": {"text": "", "num_occur": 1}, + "answer_start": {"text": "", "num_occur": 1}, + "answer_end": {"text": "", "num_occur": 1}, + } + if args.reward_type == "think_answer_tags" + else None + ), + } + else: + raise ValueError(f"Unsupported algorithm: {args.algo}") + if args.reward_type == "code": + assert args.code_verifier_api_url is not None, "Please provide a code verifier API URL for code reward type." + grpo_config.update({"code_verifier_api_url": args.code_verifier_api_url}) + if args.system_prompt is None: + # Default system prompt + args.system_prompt = DEFAUT_SYSTEM_PROMPT[args.reward_type] + + launch_distributed( + num_producers=args.num_inferencer, + num_proc_per_producer=inference_model_config.get("tensor_parallel_size", args.producer_tensor_parallel_size), + num_consumer_procs=args.num_trainers, + num_episodes=args.num_episodes, + inference_batch_size=args.inference_batch_size, + inference_microbatch_size=args.inference_microbatch_size, + train_batch_size=args.train_batch_size, + train_minibatch_size=args.train_minibatch_size, + train_dataset_config={ + "path": args.dataset, + "max_length": args.max_prompt_tokens, + "system_prompt": args.system_prompt, + }, + inference_model_config=inference_model_config, + generate_config=generate_config, + num_generations=args.num_generations, + train_model_config=train_model_config, + grpo_config=grpo_config, + plugin_config={ + "tp_size": args.tensor_parallel_size, + "pp_size": args.pipeline_parallel_size, + "microbatch_size": max( + 1, args.train_microbatch_size // args.pipeline_parallel_size + ), # microbatch size should be set to train_microbatch_size // pp_size + "zero_stage": args.zero_stage, + "max_norm": 1.0, + }, # for pp, tp + inference_backend=args.backend, + master_addr="localhost", + master_port=args.master_port, + core_algo=args.algo, + project_name=args.project, + save_interval=args.save_interval, + save_dir=os.path.join(args.save_dir, args.project.replace(" ", "_")), + eval_dataset_config=( + { + k: {"path": v, "max_length": args.max_prompt_tokens, "system_prompt": args.system_prompt} + for k, v in json.loads(args.eval_dataset).items() + } + if args.eval_dataset + else None + ), + eval_interval=args.eval_interval, + eval_save_dir=os.path.join(args.eval_save_dir, args.project.replace(" ", "_")), + eval_generation_config=eval_generation_config, + log_rollout_interval=20, + rollout_save_dir=args.rollout_save_dir, + enable_profiling=args.enable_profiling, + n_behind=args.n_behind, + ) diff --git a/applications/ColossalChat/start_code_verifier.py b/applications/ColossalChat/start_code_verifier.py new file mode 100644 index 000000000000..d1924f610698 --- /dev/null +++ b/applications/ColossalChat/start_code_verifier.py @@ -0,0 +1,35 @@ +from typing import List, Optional + +from coati.distributed.reward.code_reward.utils import check_correctness # Assuming utils.py is in the same directory +from fastapi import FastAPI, HTTPException +from pydantic import BaseModel + +app = FastAPI() + + +class CheckCorrectnessRequest(BaseModel): + in_outs: Optional[dict] + generation: str + timeout: int = 10 + debug: bool = True + eval_mode: bool = False + + +class CheckCorrectnessResponse(BaseModel): + result: List[int] + metadata: List[dict] + + +@app.post("/check_correctness", response_model=CheckCorrectnessResponse) +def check_correctness_api(request: CheckCorrectnessRequest): + try: + result, metadata = check_correctness( + in_outs=request.in_outs, + generation=request.generation, + timeout=request.timeout, + debug=request.debug, + eval_mode=request.eval_mode, + ) + return CheckCorrectnessResponse(result=result, metadata=metadata) + except Exception as e: + raise HTTPException(status_code=500, detail=str(e)) diff --git a/applications/ColossalChat/visualization.py b/applications/ColossalChat/visualization.py new file mode 100644 index 000000000000..afb729884ad3 --- /dev/null +++ b/applications/ColossalChat/visualization.py @@ -0,0 +1,100 @@ +# Re-import required libraries due to kernel reset +import argparse +from collections import defaultdict + +import matplotlib.cm as cm +import matplotlib.pyplot as plt + +# Argument parser for command line arguments +parser = argparse.ArgumentParser(description="Process profiling logs and generate a timeline plot.") +parser.add_argument("--visualization", type=str, default="actor_timelines.png", help="Path to the visualization file.") +args = parser.parse_args() + +# Raw log lines +log_lines = [] + +import glob + +files = glob.glob("*.prof") +for file in files: + with open(file, "r") as f: + log_lines += f.readlines() + +# Parse logs and collect function intervals grouped by actor +actors = defaultdict(lambda: defaultdict(list)) +current_entries = {} + +# First, collect all timestamps to find the minimum +all_timestamps = [] +parsed_lines = [] + +for line in log_lines: + if line.startswith("[Log]"): + continue + parts = line.split() + timestamp = float(parts[0]) + actor = parts[1] + action = parts[3] + func_name = parts[4] + parsed_lines.append((timestamp, actor, action, func_name)) + all_timestamps.append(timestamp) + +if not all_timestamps: + raise ValueError("No valid log entries found.") + +min_timestamp = min(all_timestamps) + +for timestamp, actor, action, func_name in parsed_lines: + rel_timestamp = timestamp - min_timestamp + key = (actor, func_name) + if action == "Enter": + current_entries[key] = rel_timestamp + elif action == "Exit": + start_time = current_entries.pop(key, None) + if start_time is not None: + actors[actor][func_name].append((start_time, rel_timestamp)) + +# Plotting setup +fig, ax = plt.subplots(figsize=(12, 6)) +colors = cm.get_cmap("tab10", len(actors)) + +actor_offsets = {} +base_offset = 0 +function_spacing = 0.9 + +yticks = [] +yticklabels = [] + +for idx, (actor, func_dict) in enumerate(actors.items()): + actor_offsets[actor] = base_offset + color = colors(idx) + for j, (func, intervals) in enumerate(func_dict.items()): + print(actor, func, intervals) + y_val = base_offset + j * function_spacing + yticks.append(y_val) + yticklabels.append(f"{actor}:{func}") + for start, end in intervals: + if end - start < 1: + end = start + 1 # Ensure all lines are at least 3 units long + ax.plot( + [start, end], + [y_val, y_val], + color=color, + linewidth=2, + label=actor if j == 0 else "", + ) + base_offset += len(func_dict) * function_spacing + 1 + +# Formatting +ax.set_yticks(yticks) +ax.set_yticklabels(yticklabels) +ax.set_xlabel("Time") +ax.set_title("Timeline per Actor") +# Remove duplicate labels in legend +handles, labels = ax.get_legend_handles_labels() +unique = dict(zip(labels, handles)) +ax.legend(unique.values(), unique.keys()) +plt.tight_layout() +plt.grid(True) +plt.savefig(args.visualization, dpi=600) # Increase dpi for higher resolution +print(f"Plot saved as {args.visualization}") diff --git a/colossalai/booster/plugin/hybrid_parallel_plugin.py b/colossalai/booster/plugin/hybrid_parallel_plugin.py index 1e0f7be240f6..e22d157d3626 100644 --- a/colossalai/booster/plugin/hybrid_parallel_plugin.py +++ b/colossalai/booster/plugin/hybrid_parallel_plugin.py @@ -1412,8 +1412,10 @@ def execute_pipeline( ) # run with gradients accumulation - if model.require_grad_sync == False or ( - isinstance(optimizer, HybridParallelZeroOptimizer) and optimizer.require_grad_sync == False + if ( + not torch.is_grad_enabled() + or model.require_grad_sync == False + or (isinstance(optimizer, HybridParallelZeroOptimizer) and optimizer.require_grad_sync == False) ): return outputs diff --git a/colossalai/shardformer/layer/__init__.py b/colossalai/shardformer/layer/__init__.py index 0bd1b60923e9..a1b80bf56b63 100644 --- a/colossalai/shardformer/layer/__init__.py +++ b/colossalai/shardformer/layer/__init__.py @@ -3,7 +3,7 @@ from .dropout import DropoutForParallelInput, DropoutForReplicatedInput from .embedding import Embedding1D, PaddingEmbedding, VocabParallelEmbedding1D from .linear import Linear1D_Col, Linear1D_Row, LinearWithGradAccum, PaddingLMHead, VocabParallelLMHead1D -from .loss import cross_entropy_1d, dist_cross_entropy +from .loss import cross_entropy_1d, dist_cross_entropy, dist_log_prob, dist_log_prob_1d from .normalization import FusedLayerNorm, FusedRMSNorm, LayerNorm, RMSNorm from .parallel_module import ParallelModule from .qkv_fused_linear import ( @@ -28,6 +28,8 @@ "DropoutForReplicatedInput", "cross_entropy_1d", "dist_cross_entropy", + "dist_log_prob_1d", + "dist_log_prob", "BaseLayerNorm", "LayerNorm", "RMSNorm", diff --git a/colossalai/shardformer/layer/loss.py b/colossalai/shardformer/layer/loss.py index 0e2241af9fc9..a9bb76fc7d6b 100644 --- a/colossalai/shardformer/layer/loss.py +++ b/colossalai/shardformer/layer/loss.py @@ -3,13 +3,21 @@ from torch.autograd import Function from torch.distributed import ProcessGroup from torch.nn import CrossEntropyLoss +from torch.nn.functional import log_softmax from colossalai.shardformer.layer._operation import reduce_forward from colossalai.shardformer.shard import ShardConfig from .utils import is_share_sp_tp -__all__ = ["DistCrossEntropy", "cross_entropy_1d", "dist_cross_entropy"] +__all__ = [ + "DistCrossEntropy", + "cross_entropy_1d", + "dist_cross_entropy", + "DistLogProb", + "dist_log_prob_1d", + "dist_log_prob", +] _IGNORE_IDX = -100 @@ -137,6 +145,98 @@ def backward(ctx, grad_output): return grad_logits, None, None, None, None, None, None +class DistLogProb(Function): + r""" + Overwrite the forward and backward function to calculate the log prob before gather + + Args: + Function (:class:`torch.autograd.Function`): default + """ + + @staticmethod + def forward( + ctx, + vocab_logits: torch.Tensor, + target: torch.Tensor, + process_group: ProcessGroup, + vocab_size: int, + dtype=torch.float32, + ): + + ################## + # Step1:Find the global maximum value of logits + ################## + logits_max = torch.max(vocab_logits, dim=-1)[0] + handle = dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=process_group, async_op=True) + + ################## + # Step2:Find the local mask. local mask will be use to select log_probs value in Step 4. + # For accleration, we overlap Step 2 and Step 3 + ################## + rank = dist.get_rank(group=process_group) + world_size = dist.get_world_size(group=process_group) + if vocab_size is None: + partition_vocab_size = vocab_logits.size()[-1] + global_vocab_size = partition_vocab_size * world_size + else: + global_vocab_size = vocab_size + partition_vocab_size = global_vocab_size // world_size + # down and up threshold for local logits + delta = (global_vocab_size + world_size - 1) // world_size + down_threshold = rank * delta + up_threshold = down_threshold + delta + if up_threshold > global_vocab_size: + up_threshold = global_vocab_size + # mask + mask = (target < down_threshold) | (target >= up_threshold) + masked_target = target.clone() - down_threshold + masked_target[mask] = 0 + masked_target_1d = masked_target.view(-1).contiguous() + handle.wait() + + ################## + # Step3:Calculate global summation exp logits + ################## + vocab_logits = vocab_logits - logits_max.unsqueeze(dim=-1) + exp_logits = torch.exp(vocab_logits) + sum_exp_logits = torch.sum(exp_logits, dim=-1, dtype=torch.float32) # local summation exp logits + dist.all_reduce(sum_exp_logits, op=dist.ReduceOp.SUM, group=process_group) + + ################## + # Step4:Calculate local prob. We first cal log_softmax, then select log probs via local mask + ################## + log_probs = vocab_logits - torch.log(sum_exp_logits.unsqueeze(dim=-1)) # cal log_softmax + log_probs = log_probs.gather(dim=-1, index=masked_target.unsqueeze(-1)) + log_probs[mask.unsqueeze(-1)] = 0 # set masked val to zero + dist.all_reduce(log_probs, op=dist.ReduceOp.SUM, group=process_group) + + ctx.save_for_backward(exp_logits, mask, masked_target_1d, sum_exp_logits) + ctx.dtype = dtype + return log_probs + + @staticmethod + def backward(ctx, grad_output): + exp_logits, mask, masked_target_1d, sum_exp_logits = ctx.saved_tensors + ################## + # Step1:Find the global sofmax value + ################## + softmax_logits = exp_logits / sum_exp_logits.unsqueeze(dim=-1) + + ################## + # Step2:Update softmax value based on local target index + ################## + partion_vocab_size = softmax_logits.shape[-1] + softmax_logits_2d = softmax_logits.view(-1, partion_vocab_size) + update = 1.0 - mask.view(-1).float().to(ctx.dtype) + softmax_logits_2d[torch.arange(0, softmax_logits_2d.shape[0]), masked_target_1d] -= update + + ################## + # Step3:Calculate grad_output, which is the gradient of the loss function with respect to the output of logsoftmax + ################## + grad_logits = -softmax_logits.mul_(grad_output) + return grad_logits, None, None, None, None, None, None + + def cross_entropy_1d( vocab_logits: torch.Tensor, labels: torch.Tensor, @@ -149,6 +249,16 @@ def cross_entropy_1d( return DistCrossEntropy.apply(vocab_logits, labels, ignore_index, process_group, vocab_size, dtype, mode) +def dist_log_prob_1d( + vocab_logits: torch.Tensor, + labels: torch.Tensor, + process_group: ProcessGroup = None, + vocab_size: int = None, + dtype: torch.dtype = None, +) -> torch.Tensor: + return DistLogProb.apply(vocab_logits, labels, process_group, vocab_size, dtype) + + def dist_cross_entropy( labels: torch.Tensor, # [B, S] or [B, S, Vocab_size] logits: torch.Tensor, # [B, S, Vocab_size] @@ -243,3 +353,41 @@ def dist_cross_entropy( loss, num_nonzero = loss[0], loss[1].detach() loss = (loss / num_nonzero).squeeze() return loss + + +def dist_log_prob( + labels: torch.Tensor, # [B, S] or [B, S, Vocab_size] + logits: torch.Tensor, # [B, S, Vocab_size] + shard_config: ShardConfig, + vocab_size: int, + dtype: torch.dtype, + seq_dim: int = 1, +) -> torch.Tensor: + """ + Helper to compute log prob for most shardformer models supporting PP, TP. + """ + # Split labels if not gather output + parallel_output = shard_config.parallel_output + is_tp = shard_config.enable_tensor_parallelism + + # TODO:support sp + labels = labels[..., 1:] + logits = logits[..., :-1, :] + labels = labels.contiguous() + logits = logits.contiguous() + assert labels.shape == logits.shape[:-1], f"label shape {labels.shape} does not match logit shape {logits.shape}" + + # Flatten the tokens + if is_tp and parallel_output: + log_prob = dist_log_prob_1d( + logits, + labels, + process_group=shard_config.tensor_parallel_process_group, + vocab_size=vocab_size, + dtype=dtype, + ) + else: + log_prob = log_softmax(logits, dim=-1) + log_prob = log_prob.gather(dim=-1, index=labels.unsqueeze(-1)) + + return log_prob diff --git a/colossalai/shardformer/modeling/qwen2.py b/colossalai/shardformer/modeling/qwen2.py index 5417bf4ebb3f..0c437e93e0aa 100644 --- a/colossalai/shardformer/modeling/qwen2.py +++ b/colossalai/shardformer/modeling/qwen2.py @@ -275,6 +275,7 @@ def qwen2_for_causal_lm_forward( hidden_states: Optional[torch.FloatTensor] = None, stage_index: Optional[List[int]] = None, shard_config: ShardConfig = None, + **kwargs, ): r""" Args: @@ -823,7 +824,6 @@ def forward( loss = None if labels is not None: loss = dist_cross_entropy(labels, logits, shard_config, self.lm_head.out_features, logits.dtype) - if not return_dict: output = (logits,) + outputs[1:] return (loss,) + output if loss is not None else output diff --git a/colossalai/shardformer/modeling/qwen3.py b/colossalai/shardformer/modeling/qwen3.py index 5e8c0762c9fa..5f96f5f4927b 100644 --- a/colossalai/shardformer/modeling/qwen3.py +++ b/colossalai/shardformer/modeling/qwen3.py @@ -273,6 +273,7 @@ def qwen3_for_causal_lm_forward( hidden_states: Optional[torch.FloatTensor] = None, stage_index: Optional[List[int]] = None, shard_config: ShardConfig = None, + **kwargs, ): r""" Args: diff --git a/colossalai/shardformer/policies/qwen2.py b/colossalai/shardformer/policies/qwen2.py index 78b3bf5285a8..32f7d37b190f 100644 --- a/colossalai/shardformer/policies/qwen2.py +++ b/colossalai/shardformer/policies/qwen2.py @@ -20,6 +20,7 @@ PaddingEmbedding, RMSNorm, VocabParallelEmbedding1D, + VocabParallelLMHead1D, ) from ..modeling.qwen2 import ( @@ -411,8 +412,12 @@ def module_policy(self): sub_module_replacement=[ SubModuleReplacementDescription( suffix="lm_head", - target_module=Linear1D_Col, - kwargs=dict(fp8_communication=self.shard_config.fp8_communication, use_zbv=use_zbv), + target_module=VocabParallelLMHead1D, + kwargs=dict( + gather_output=not self.shard_config.parallel_output, + fp8_communication=self.shard_config.fp8_communication, + use_zbv=use_zbv, + ), ) ], method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, @@ -428,7 +433,16 @@ def module_policy(self): suffix="lm_head", target_module=LinearWithGradAccum, kwargs=dict(fp8_communication=self.shard_config.fp8_communication, use_zbv=use_zbv), - ) + ), + SubModuleReplacementDescription( + suffix="lm_head", + target_module=VocabParallelLMHead1D, + kwargs={ + "gather_output": not self.shard_config.parallel_output, + "make_vocab_size_divisible_by": self.shard_config.make_vocab_size_divisible_by, + "fp8_communication": self.shard_config.fp8_communication, + }, + ), ], method_replacement={"forward": get_lm_forward_with_dist_cross_entropy(self.shard_config)}, ) diff --git a/tests/test_shardformer/test_layer/test_dist_log_prob.py b/tests/test_shardformer/test_layer/test_dist_log_prob.py new file mode 100644 index 000000000000..05a6a5d4766f --- /dev/null +++ b/tests/test_shardformer/test_layer/test_dist_log_prob.py @@ -0,0 +1,52 @@ +import pytest +import torch +from coati.distributed.utils import log_probs_from_logits + +import colossalai +from colossalai.logging import disable_existing_loggers +from colossalai.shardformer.layer import dist_log_prob_1d +from colossalai.testing import rerun_if_address_is_in_use, spawn + +CONFIG = dict( + parallel=dict(data=1, pipeline=1, tensor=dict(size=2, mode="1d")), +) + + +def check_dist_log_prob(rank, world_size, port): + disable_existing_loggers() + colossalai.launch(rank=rank, world_size=world_size, port=port, host="localhost", backend="nccl") + + # prepare data + pred = torch.randn(2, 4, 8, requires_grad=True).cuda() + labels = torch.randint(8, (2, 4)).cuda() + + logprob = log_probs_from_logits(pred, labels) + + pred.retain_grad() + logprob.mean().backward() + + dist_pred = pred.clone().chunk(world_size, -1)[rank].detach() + dist_pred.requires_grad = True + dist_logprob = dist_log_prob_1d(dist_pred, labels) + + dist_pred.retain_grad() + dist_logprob.squeeze(-1).mean().backward() + + assert torch.allclose( + logprob, dist_logprob.squeeze(-1), atol=1e-5 + ), f"dist cross entropy logprob is not equal to orgin logprob\n{logprob}\n{dist_logprob.squeeze(-1)}" + + pred_grad_partial = pred.grad.clone().chunk(world_size, -1)[rank].detach() + assert torch.allclose( + pred_grad_partial, dist_pred.grad + ), f"dist grad is not equal to orgin grad\n{pred.grad}\n{dist_pred.grad}" + + +@pytest.mark.dist +@rerun_if_address_is_in_use() +def test_dist_log_prob(): + spawn(check_dist_log_prob, 2) + + +if __name__ == "__main__": + test_dist_log_prob()