From 99209da1bb442897def343a460add8b7ae5ce25e Mon Sep 17 00:00:00 2001 From: Tong Li Date: Tue, 10 Sep 2024 09:27:14 +0000 Subject: [PATCH 01/21] update dpo --- .../ColossalChat/coati/trainer/dpo.py | 64 ++++++++++--------- 1 file changed, 35 insertions(+), 29 deletions(-) diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index faa7a90d92de..bd95e61b769c 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -191,13 +191,7 @@ def _train(self, epoch: int): # DPO Loss loss = losses.mean() - self.booster.backward(loss=loss, optimizer=self.optimizer) - if self.num_train_step % self.accumulation_steps == self.accumulation_steps - 1: - self.optimizer.step() - self.optimizer.zero_grad() - self.actor_scheduler.step() - # sync loss_mean = all_reduce_mean(tensor=loss) chosen_rewards_mean = all_reduce_mean(tensor=chosen_rewards) @@ -208,10 +202,20 @@ def _train(self, epoch: int): self.accumulative_meter.add("loss", loss_mean.to(torch.float16).item()) self.accumulative_meter.add("accuracy", reward_accuracies_mean.to(torch.float16).item()) - if i % self.accumulation_steps == self.accumulation_steps - 1: - self.num_train_step += 1 + if (i + 1) % self.accumulation_steps == 0: + self.optimizer.step() + self.optimizer.zero_grad() + self.actor_scheduler.step() + + step_bar.set_postfix( + { + "train/loss": self.accumulative_meter.get("loss"), + "train/chosen_rewards": self.accumulative_meter.get("chosen_rewards"), + "train/rejected_rewards": self.accumulative_meter.get("rejected_rewards"), + "train/accuracy": self.accumulative_meter.get("accuracy"), + } + ) step_bar.update() - # logging if self.writer and is_rank_0(): self.writer.add_scalar("train/loss", self.accumulative_meter.get("loss"), self.num_train_step) self.writer.add_scalar("train/lr", self.optimizer.param_groups[0]["lr"], self.num_train_step) @@ -233,25 +237,26 @@ def _train(self, epoch: int): self.accumulative_meter.get("accuracy"), self.num_train_step, ) + self.num_train_step += 1 self.accumulative_meter.reset() - if self.save_dir is not None and (self.num_train_step + 1) % self.save_interval == 0: - # save checkpoint - self.coordinator.print_on_master("\nStart saving model checkpoint with running states") - save_checkpoint( - save_dir=self.save_dir, - booster=self.booster, - model=self.model, - optimizer=self.optimizer, - lr_scheduler=self.actor_scheduler, - epoch=epoch, - step=i + 1, - batch_size=batch_size, - coordinator=self.coordinator, - ) - self.coordinator.print_on_master( - f"Saved checkpoint at epoch {epoch} step {self.save_interval} at folder {self.save_dir}" - ) + if self.save_dir is not None and self.num_train_step > 0 and self.num_train_step % self.save_interval == 0: + # save checkpoint + self.coordinator.print_on_master("\nStart saving model checkpoint with running states") + save_checkpoint( + save_dir=self.save_dir, + booster=self.booster, + model=self.model, + optimizer=self.optimizer, + lr_scheduler=self.actor_scheduler, + epoch=epoch, + step=self.num_train_step, + batch_size=batch_size, + coordinator=self.coordinator, + ) + self.coordinator.print_on_master( + f"Saved checkpoint at epoch {epoch} step {self.save_interval} at folder {self.save_dir}" + ) step_bar.close() @@ -356,7 +361,8 @@ def _eval(self, epoch: int): for tag in ["loss", "chosen_rewards", "rejected_rewards", "accuracy", "margin"]: msg = msg + f"{tag}: {self.accumulative_meter.get(tag)}\n" self.coordinator.print_on_master(msg) - os.makedirs(self.save_dir, exist_ok=True) - with open(os.path.join(self.save_dir, f"eval_result_epoch{epoch}.txt"), "w") as f: - f.write(msg) + if self.save_dir is not None: + os.makedirs(self.save_dir, exist_ok=True) + with open(os.path.join(self.save_dir, f"eval_result_epoch{epoch}.txt"), "w") as f: + f.write(msg) step_bar.close() From 34627e2c3b0acea62bbe9fd4c55a130c93838df5 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Tue, 10 Sep 2024 10:09:35 +0000 Subject: [PATCH 02/21] remove unsupport plugin --- .../examples/training_scripts/train_dpo.py | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/applications/ColossalChat/examples/training_scripts/train_dpo.py b/applications/ColossalChat/examples/training_scripts/train_dpo.py index 3b324ee784e0..90527c88b38e 100755 --- a/applications/ColossalChat/examples/training_scripts/train_dpo.py +++ b/applications/ColossalChat/examples/training_scripts/train_dpo.py @@ -13,7 +13,7 @@ import colossalai from colossalai.booster import Booster -from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin +from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin, TorchDDPPlugin from colossalai.cluster import DistCoordinator from colossalai.logging import get_dist_logger from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR @@ -29,8 +29,6 @@ def train(args): # check lora compatibility if "gemini" in args.plugin and lora_config is not None and lora_config.r > 0: raise ValueError("LoRA is not supported in GeminiPlugin. Please use other plugin") - if args.plugin == "gemini_auto" and args.accumulation_steps > 1: - raise ValueError("Gradient accumulation is not supported in GeminiPlugin. Please use other plugin") # ============================== # Initialize Distributed Training @@ -46,7 +44,7 @@ def train(args): Default torch ddp plugin without any acceleration, for debugging purpose acceleration, for debugging purpose """ - plugin = TorchDDPPlugin(find_unused_parameters=True) + plugin = TorchDDPPlugin(find_unused_parameters=True if args.grad_checkpoint is False else False) elif args.plugin == "gemini": plugin = GeminiPlugin( precision=args.mixed_precision, @@ -56,14 +54,6 @@ def train(args): enable_gradient_accumulation=True, enable_flash_attention=args.use_flash_attn, ) - elif args.plugin == "gemini_auto": - plugin = GeminiPlugin( - precision=args.mixed_precision, - placement_policy="auto", - initial_scale=2**16, - max_norm=args.grad_clip, - enable_flash_attention=args.use_flash_attn, - ) elif args.plugin == "zero2": plugin = LowLevelZeroPlugin( stage=2, @@ -312,7 +302,7 @@ def train(args): "--plugin", type=str, default="gemini", - choices=["gemini", "gemini_auto", "zero2", "zero2_cpu", "3d"], + choices=["gemini", "zero2", "zero2_cpu", "3d", "ddp"], help="Choose which plugin to use", ) parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping value") From 31270cb79e53817fd113616859727dfbb53e683b Mon Sep 17 00:00:00 2001 From: Tong Li Date: Tue, 10 Sep 2024 10:10:10 +0000 Subject: [PATCH 03/21] update msg --- applications/ColossalChat/coati/trainer/dpo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index bd95e61b769c..eff5998fc3e7 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -357,7 +357,7 @@ def _eval(self, epoch: int): ) step_bar.update() - msg = "Evaluation Result:\n" + msg = "\nEvaluation Result:\n" for tag in ["loss", "chosen_rewards", "rejected_rewards", "accuracy", "margin"]: msg = msg + f"{tag}: {self.accumulative_meter.get(tag)}\n" self.coordinator.print_on_master(msg) From b59c487a3e1b394ad3ac6d1504730a4f3c9779c4 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Tue, 10 Sep 2024 09:27:14 +0000 Subject: [PATCH 04/21] update dpo --- .../ColossalChat/coati/trainer/dpo.py | 64 ++++++++++--------- 1 file changed, 35 insertions(+), 29 deletions(-) diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index faa7a90d92de..bd95e61b769c 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -191,13 +191,7 @@ def _train(self, epoch: int): # DPO Loss loss = losses.mean() - self.booster.backward(loss=loss, optimizer=self.optimizer) - if self.num_train_step % self.accumulation_steps == self.accumulation_steps - 1: - self.optimizer.step() - self.optimizer.zero_grad() - self.actor_scheduler.step() - # sync loss_mean = all_reduce_mean(tensor=loss) chosen_rewards_mean = all_reduce_mean(tensor=chosen_rewards) @@ -208,10 +202,20 @@ def _train(self, epoch: int): self.accumulative_meter.add("loss", loss_mean.to(torch.float16).item()) self.accumulative_meter.add("accuracy", reward_accuracies_mean.to(torch.float16).item()) - if i % self.accumulation_steps == self.accumulation_steps - 1: - self.num_train_step += 1 + if (i + 1) % self.accumulation_steps == 0: + self.optimizer.step() + self.optimizer.zero_grad() + self.actor_scheduler.step() + + step_bar.set_postfix( + { + "train/loss": self.accumulative_meter.get("loss"), + "train/chosen_rewards": self.accumulative_meter.get("chosen_rewards"), + "train/rejected_rewards": self.accumulative_meter.get("rejected_rewards"), + "train/accuracy": self.accumulative_meter.get("accuracy"), + } + ) step_bar.update() - # logging if self.writer and is_rank_0(): self.writer.add_scalar("train/loss", self.accumulative_meter.get("loss"), self.num_train_step) self.writer.add_scalar("train/lr", self.optimizer.param_groups[0]["lr"], self.num_train_step) @@ -233,25 +237,26 @@ def _train(self, epoch: int): self.accumulative_meter.get("accuracy"), self.num_train_step, ) + self.num_train_step += 1 self.accumulative_meter.reset() - if self.save_dir is not None and (self.num_train_step + 1) % self.save_interval == 0: - # save checkpoint - self.coordinator.print_on_master("\nStart saving model checkpoint with running states") - save_checkpoint( - save_dir=self.save_dir, - booster=self.booster, - model=self.model, - optimizer=self.optimizer, - lr_scheduler=self.actor_scheduler, - epoch=epoch, - step=i + 1, - batch_size=batch_size, - coordinator=self.coordinator, - ) - self.coordinator.print_on_master( - f"Saved checkpoint at epoch {epoch} step {self.save_interval} at folder {self.save_dir}" - ) + if self.save_dir is not None and self.num_train_step > 0 and self.num_train_step % self.save_interval == 0: + # save checkpoint + self.coordinator.print_on_master("\nStart saving model checkpoint with running states") + save_checkpoint( + save_dir=self.save_dir, + booster=self.booster, + model=self.model, + optimizer=self.optimizer, + lr_scheduler=self.actor_scheduler, + epoch=epoch, + step=self.num_train_step, + batch_size=batch_size, + coordinator=self.coordinator, + ) + self.coordinator.print_on_master( + f"Saved checkpoint at epoch {epoch} step {self.save_interval} at folder {self.save_dir}" + ) step_bar.close() @@ -356,7 +361,8 @@ def _eval(self, epoch: int): for tag in ["loss", "chosen_rewards", "rejected_rewards", "accuracy", "margin"]: msg = msg + f"{tag}: {self.accumulative_meter.get(tag)}\n" self.coordinator.print_on_master(msg) - os.makedirs(self.save_dir, exist_ok=True) - with open(os.path.join(self.save_dir, f"eval_result_epoch{epoch}.txt"), "w") as f: - f.write(msg) + if self.save_dir is not None: + os.makedirs(self.save_dir, exist_ok=True) + with open(os.path.join(self.save_dir, f"eval_result_epoch{epoch}.txt"), "w") as f: + f.write(msg) step_bar.close() From a732586e298cccdcd0a0fe740d88412527ac0d10 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Tue, 10 Sep 2024 10:09:35 +0000 Subject: [PATCH 05/21] remove unsupport plugin --- .../examples/training_scripts/train_dpo.py | 16 +++------------- 1 file changed, 3 insertions(+), 13 deletions(-) diff --git a/applications/ColossalChat/examples/training_scripts/train_dpo.py b/applications/ColossalChat/examples/training_scripts/train_dpo.py index 3b324ee784e0..90527c88b38e 100755 --- a/applications/ColossalChat/examples/training_scripts/train_dpo.py +++ b/applications/ColossalChat/examples/training_scripts/train_dpo.py @@ -13,7 +13,7 @@ import colossalai from colossalai.booster import Booster -from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin +from colossalai.booster.plugin import GeminiPlugin, HybridParallelPlugin, LowLevelZeroPlugin, TorchDDPPlugin from colossalai.cluster import DistCoordinator from colossalai.logging import get_dist_logger from colossalai.nn.lr_scheduler import CosineAnnealingWarmupLR @@ -29,8 +29,6 @@ def train(args): # check lora compatibility if "gemini" in args.plugin and lora_config is not None and lora_config.r > 0: raise ValueError("LoRA is not supported in GeminiPlugin. Please use other plugin") - if args.plugin == "gemini_auto" and args.accumulation_steps > 1: - raise ValueError("Gradient accumulation is not supported in GeminiPlugin. Please use other plugin") # ============================== # Initialize Distributed Training @@ -46,7 +44,7 @@ def train(args): Default torch ddp plugin without any acceleration, for debugging purpose acceleration, for debugging purpose """ - plugin = TorchDDPPlugin(find_unused_parameters=True) + plugin = TorchDDPPlugin(find_unused_parameters=True if args.grad_checkpoint is False else False) elif args.plugin == "gemini": plugin = GeminiPlugin( precision=args.mixed_precision, @@ -56,14 +54,6 @@ def train(args): enable_gradient_accumulation=True, enable_flash_attention=args.use_flash_attn, ) - elif args.plugin == "gemini_auto": - plugin = GeminiPlugin( - precision=args.mixed_precision, - placement_policy="auto", - initial_scale=2**16, - max_norm=args.grad_clip, - enable_flash_attention=args.use_flash_attn, - ) elif args.plugin == "zero2": plugin = LowLevelZeroPlugin( stage=2, @@ -312,7 +302,7 @@ def train(args): "--plugin", type=str, default="gemini", - choices=["gemini", "gemini_auto", "zero2", "zero2_cpu", "3d"], + choices=["gemini", "zero2", "zero2_cpu", "3d", "ddp"], help="Choose which plugin to use", ) parser.add_argument("--grad_clip", type=float, default=1.0, help="Gradient clipping value") From 380408eae83982b5dadbef44146a18dc37780c1e Mon Sep 17 00:00:00 2001 From: Tong Li Date: Tue, 10 Sep 2024 10:10:10 +0000 Subject: [PATCH 06/21] update msg --- applications/ColossalChat/coati/trainer/dpo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index bd95e61b769c..eff5998fc3e7 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -357,7 +357,7 @@ def _eval(self, epoch: int): ) step_bar.update() - msg = "Evaluation Result:\n" + msg = "\nEvaluation Result:\n" for tag in ["loss", "chosen_rewards", "rejected_rewards", "accuracy", "margin"]: msg = msg + f"{tag}: {self.accumulative_meter.get(tag)}\n" self.coordinator.print_on_master(msg) From a9d222f93610e625b93d88ca9effa49d625b8a63 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Tue, 10 Sep 2024 11:32:04 +0000 Subject: [PATCH 07/21] update template --- .../ColossalChat/tests/test_templating.sh | 21 ++++++++++--------- 1 file changed, 11 insertions(+), 10 deletions(-) diff --git a/applications/ColossalChat/tests/test_templating.sh b/applications/ColossalChat/tests/test_templating.sh index 6ee10e8bed87..defe6f71bb4c 100755 --- a/applications/ColossalChat/tests/test_templating.sh +++ b/applications/ColossalChat/tests/test_templating.sh @@ -4,7 +4,7 @@ BASE_TEMP_DIR=$BASE_DIR/temp EXAMPLES_DIR=$BASE_DIR/examples TEST_DATA_DIR=$BASE_DIR/tests/test_data DATA_SAVE_PATH=$BASE_TEMP_DIR/tests -CONFIG_DIR=$BASE_DIR/config +CONFIG_DIR=$BASE_DIR/conversation_template # MODELS=("colossal-llama2" "llama2" "mistral" "chatGLM2" "chatGLM3" "deepseek" "Yi" "baichuan") # for local test MODELS=("colossal-llama2" "llama2" "chatGLM2" "chatGLM3" "deepseek" "Yi") @@ -39,23 +39,23 @@ get_pretrain() { get_conversation_template_config() { local model=$1 if [[ $model == "colossal-llama2" ]]; then - echo "$CONFIG_DIR/conversation_template/colossal-llama2.json" + echo "$CONFIG_DIR/colossal-llama2.json" elif [[ $model == "llama2" ]]; then - echo "$CONFIG_DIR/conversation_template/llama2.json" + echo "$CONFIG_DIR/llama2.json" elif [[ $model == "deepseek" ]]; then - echo "$CONFIG_DIR/conversation_template/deepseek-ai_DeepSeek-V2-Lite.json" + echo "$CONFIG_DIR/deepseek-ai_DeepSeek-V2-Lite.json" elif [[ $model == "mistral" ]]; then - echo "$CONFIG_DIR/conversation_template/mistralai_Mixtral-8x7B-Instruct-v0.1.json" + echo "$CONFIG_DIR/mistralai_Mixtral-8x7B-Instruct-v0.1.json" elif [[ $model == "chatGLM2" ]]; then - echo "$CONFIG_DIR/conversation_template/THUDM_chatglm2-6b.json" + echo "$CONFIG_DIR/THUDM_chatglm2-6b.json" elif [[ $model == "chatGLM3" ]]; then - echo "$CONFIG_DIR/conversation_template/THUDM_chatglm3-6b.json" + echo "$CONFIG_DIR/THUDM_chatglm3-6b.json" elif [[ $model == "phi" ]]; then - echo "$CONFIG_DIR/conversation_template/microsoft_phi-2.json" + echo "$CONFIG_DIR/microsoft_phi-2.json" elif [[ $model == "Yi" ]]; then - echo "$CONFIG_DIR/conversation_template/01-ai_Yi-1.5-9B-Chat.json" + echo "$CONFIG_DIR/01-ai_Yi-1.5-9B-Chat.json" elif [[ $model == "baichuan" ]]; then - echo "$CONFIG_DIR/conversation_template/baichuan-inc_Baichuan2-13B-Chat.json" + echo "$CONFIG_DIR/baichuan-inc_Baichuan2-13B-Chat.json" else echo "Unknown model $model" exit 1 @@ -71,6 +71,7 @@ for model in ${MODELS[@]}; do rm -rf $SAVE_DIR/arrow pretrain=$(get_pretrain $model) conversation_template_config=$(get_conversation_template_config $model) + echo $conversation_template_config python $EXAMPLES_DIR/data_preparation_scripts/prepare_dataset.py --type sft --data_input_dirs $TEST_DATA_DIR/sft \ --tokenizer_dir $pretrain \ --conversation_template_config $conversation_template_config \ From 6d89af2e895f88e88e93d468d45e1043ef138635 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Wed, 11 Sep 2024 06:42:21 +0000 Subject: [PATCH 08/21] update dataset --- .../examples/data_preparation_scripts/prepare_dataset.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/applications/ColossalChat/examples/data_preparation_scripts/prepare_dataset.py b/applications/ColossalChat/examples/data_preparation_scripts/prepare_dataset.py index a35f2bf52dfd..b551497b9cce 100644 --- a/applications/ColossalChat/examples/data_preparation_scripts/prepare_dataset.py +++ b/applications/ColossalChat/examples/data_preparation_scripts/prepare_dataset.py @@ -73,8 +73,7 @@ def main(): "--conversation_template_config", type=str, default="conversation_template_config", - help="Path \ - to save conversation template config files.", + help="Path to save conversation template config files.", ) parser.add_argument("--data_cache_dir", type=str, default="cache", help="Data cache directory") parser.add_argument( From 779aff86288a6088c7a33f11dfac16d1be9456d8 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Thu, 12 Sep 2024 11:30:49 +0000 Subject: [PATCH 09/21] add pp for dpo --- .../ColossalChat/coati/models/loss.py | 2 +- .../ColossalChat/coati/trainer/dpo.py | 333 ++++++++++++------ .../examples/training_scripts/train_dpo.py | 49 ++- colossalai/shardformer/modeling/llama.py | 1 + 4 files changed, 250 insertions(+), 135 deletions(-) diff --git a/applications/ColossalChat/coati/models/loss.py b/applications/ColossalChat/coati/models/loss.py index bd0bbd36b9bc..1bc57a2e94b7 100755 --- a/applications/ColossalChat/coati/models/loss.py +++ b/applications/ColossalChat/coati/models/loss.py @@ -156,7 +156,7 @@ def forward( pi_logratios = logprob_actor_chosen.sum(-1) - logprob_actor_reject.sum(-1) logits = pi_logratios - ref_logratios - self.gamma / self.beta losses = -torch.nn.functional.logsigmoid(self.beta * logits) - + losses = losses.mean() # Calculate rewards for logging if logprob_ref_chosen is not None: chosen_rewards = self.beta * (logprob_actor_chosen.sum(-1) - logprob_ref_chosen.sum(-1)).detach() diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index eff5998fc3e7..420279edc380 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -6,6 +6,7 @@ from typing import Any, Optional import torch +import torch.distributed as dist from coati.models.loss import DpoLoss from coati.models.utils import calc_masked_log_probs from coati.trainer.utils import all_reduce_mean @@ -13,10 +14,11 @@ from torch.optim import Optimizer from torch.optim.lr_scheduler import _LRScheduler from torch.utils.data import DataLoader -from tqdm import trange +from tqdm import tqdm, trange from transformers import PreTrainedTokenizerBase from colossalai.booster import Booster, Plugin +from colossalai.booster.plugin import HybridParallelPlugin from colossalai.cluster import DistCoordinator from colossalai.utils import get_current_device @@ -117,128 +119,227 @@ def _train(self, epoch: int): epoch int: the number of current epoch """ self.model.train() - self.accumulative_meter.reset() - step_bar = trange( - len(self.train_dataloader) // self.accumulation_steps, - desc=f"Epoch {epoch + 1}/{self.max_epochs}", - disable=not is_rank_0(), - ) - for i, batch in enumerate(self.train_dataloader): - batch = to_device(batch, self.device) - ( - chosen_input_ids, - chosen_attention_mask, - chosen_loss_mask, - reject_input_ids, - reject_attention_mask, - reject_loss_mask, - ) = ( - batch["chosen_input_ids"], - batch["chosen_attention_mask"], - batch["chosen_loss_mask"], - batch["reject_input_ids"], - batch["reject_attention_mask"], - batch["reject_loss_mask"], - ) - if not self.apply_loss_mask: - chosen_loss_mask = chosen_loss_mask.fill_(1.0) - reject_loss_mask = reject_loss_mask.fill_(1.0) - - batch_size = chosen_input_ids.size()[0] - - actor_all_logits = self.model( - input_ids=torch.cat([chosen_input_ids, reject_input_ids]), - attention_mask=torch.cat([chosen_attention_mask, reject_attention_mask]), - )["logits"] - actor_chosen_logits = actor_all_logits[:batch_size] - actor_reject_logits = actor_all_logits[batch_size:] - logprob_actor_chosen = calc_masked_log_probs( - actor_chosen_logits, chosen_input_ids, chosen_loss_mask[:, 1:], self.length_normalization + if isinstance(self.plugin, HybridParallelPlugin) and self.plugin.pp_size > 1: + step_bar = tqdm( + range(len(self.train_dataloader)), + desc="Step", + disable=not (dist.get_rank() == dist.get_world_size() - 1), ) + for i, batch in enumerate(self.train_dataloader): + batch = to_device(batch, self.device) + ( + chosen_input_ids, + chosen_attention_mask, + chosen_loss_mask, + reject_input_ids, + reject_attention_mask, + reject_loss_mask, + ) = ( + batch["chosen_input_ids"], + batch["chosen_attention_mask"], + batch["chosen_loss_mask"], + batch["reject_input_ids"], + batch["reject_attention_mask"], + batch["reject_loss_mask"], + ) + batch_size = chosen_input_ids.size()[0] + # Calculate logits from reference model. + if self.ref_model is not None: + self.ref_model.eval() + with torch.no_grad(): + ref_all_logits = self.ref_model( + input_ids=torch.cat([chosen_input_ids, reject_input_ids]), + attention_mask=torch.cat([chosen_attention_mask, reject_attention_mask]), + )["logits"] + ref_chosen_logits = ref_all_logits[:batch_size] + ref_reject_logits = ref_all_logits[batch_size:] + logprob_ref_chosen = calc_masked_log_probs( + ref_chosen_logits, chosen_input_ids, chosen_loss_mask[:, 1:], self.length_normalization + ) + logprob_ref_reject = calc_masked_log_probs( + ref_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization + ) + else: + logprob_ref_chosen = None + logprob_ref_reject = None + + # Merge chosen and reject + inputs_ids = torch.stack([item for tup in zip(chosen_input_ids, reject_input_ids) for item in tup]) + attention_mask = torch.stack( + [item for tup in zip(chosen_attention_mask, reject_attention_mask) for item in tup] + ) + loss_mask = torch.stack([item for tup in zip(chosen_loss_mask, reject_loss_mask) for item in tup]) + logprob_ref = torch.stack([item for tup in zip(logprob_ref_chosen, logprob_ref_reject) for item in tup]) + + data_iter = iter( + [ + { + "input_ids": inputs_ids, + "attention_mask": attention_mask, + "loss_mask": loss_mask, + "logprob_ref": logprob_ref, + } + ] + ) - logprob_actor_reject = calc_masked_log_probs( - actor_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization - ) + outputs = self.booster.execute_pipeline( + data_iter, + self.model, + criterion=lambda outputs, inputs: self.actor_loss_fn( + calc_masked_log_probs( + outputs["logits"][0::2], + inputs["input_ids"][0::2], + inputs["loss_mask"][0::2][:, 1:], + self.length_normalization, + ), + calc_masked_log_probs( + outputs["logits"][1::2], + inputs["input_ids"][1::2], + inputs["loss_mask"][1::2][:, 1:], + self.length_normalization, + ), + inputs["logprob_ref"][0::2] if inputs["logprob_ref"] is not None else None, + inputs["logprob_ref"][1::2] if inputs["logprob_ref"] is not None else None, + inputs["loss_mask"][0::2][:, 1:], + inputs["loss_mask"][1::2][:, 1:], + )[0], + optimizer=self.optimizer, + return_loss=True, + ) + loss = outputs["loss"] + if self.booster.plugin.stage_manager.is_last_stage(): + global_loss = all_reduce_mean(loss, self.plugin) + if dist.get_rank() == dist.get_world_size() - 1: + step_bar.set_postfix({"train/loss": global_loss.item()}) + step_bar.update() - if self.ref_model is not None: - self.ref_model.eval() - with torch.no_grad(): - ref_all_logits = self.ref_model( - input_ids=torch.cat([chosen_input_ids, reject_input_ids]), - attention_mask=torch.cat([chosen_attention_mask, reject_attention_mask]), - )["logits"] - ref_chosen_logits = ref_all_logits[:batch_size] - ref_reject_logits = ref_all_logits[batch_size:] - logprob_ref_chosen = calc_masked_log_probs( - ref_chosen_logits, chosen_input_ids, chosen_loss_mask[:, 1:], self.length_normalization - ) - logprob_ref_reject = calc_masked_log_probs( - ref_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization - ) - else: - logprob_ref_chosen = None - logprob_ref_reject = None - - losses, chosen_rewards, rejected_rewards = self.actor_loss_fn( - logprob_actor_chosen, - logprob_actor_reject, - logprob_ref_chosen if logprob_ref_chosen is not None else None, - logprob_ref_reject if logprob_ref_reject is not None else None, - chosen_loss_mask[:, 1:], - reject_loss_mask[:, 1:], - ) - reward_accuracies = (chosen_rewards > rejected_rewards).float().mean() - - # DPO Loss - loss = losses.mean() - self.booster.backward(loss=loss, optimizer=self.optimizer) - # sync - loss_mean = all_reduce_mean(tensor=loss) - chosen_rewards_mean = all_reduce_mean(tensor=chosen_rewards) - rejected_rewards_mean = all_reduce_mean(tensor=rejected_rewards) - reward_accuracies_mean = all_reduce_mean(tensor=reward_accuracies) - self.accumulative_meter.add("chosen_rewards", chosen_rewards_mean.to(torch.float16).mean().item()) - self.accumulative_meter.add("rejected_rewards", rejected_rewards_mean.to(torch.float16).mean().item()) - self.accumulative_meter.add("loss", loss_mean.to(torch.float16).item()) - self.accumulative_meter.add("accuracy", reward_accuracies_mean.to(torch.float16).item()) - - if (i + 1) % self.accumulation_steps == 0: self.optimizer.step() self.optimizer.zero_grad() - self.actor_scheduler.step() - - step_bar.set_postfix( - { - "train/loss": self.accumulative_meter.get("loss"), - "train/chosen_rewards": self.accumulative_meter.get("chosen_rewards"), - "train/rejected_rewards": self.accumulative_meter.get("rejected_rewards"), - "train/accuracy": self.accumulative_meter.get("accuracy"), - } + + else: + self.accumulative_meter.reset() + step_bar = trange( + len(self.train_dataloader) // self.accumulation_steps, + desc=f"Epoch {epoch + 1}/{self.max_epochs}", + disable=not is_rank_0(), + ) + for i, batch in enumerate(self.train_dataloader): + batch = to_device(batch, self.device) + ( + chosen_input_ids, + chosen_attention_mask, + chosen_loss_mask, + reject_input_ids, + reject_attention_mask, + reject_loss_mask, + ) = ( + batch["chosen_input_ids"], + batch["chosen_attention_mask"], + batch["chosen_loss_mask"], + batch["reject_input_ids"], + batch["reject_attention_mask"], + batch["reject_loss_mask"], ) - step_bar.update() - if self.writer and is_rank_0(): - self.writer.add_scalar("train/loss", self.accumulative_meter.get("loss"), self.num_train_step) - self.writer.add_scalar("train/lr", self.optimizer.param_groups[0]["lr"], self.num_train_step) - self.writer.add_scalar( - "train/chosen_rewards", self.accumulative_meter.get("chosen_rewards"), self.num_train_step - ) - self.writer.add_scalar( - "train/rejected_rewards", - self.accumulative_meter.get("rejected_rewards"), - self.num_train_step, - ) - self.writer.add_scalar( - "train/margin", - self.accumulative_meter.get("chosen_rewards") - self.accumulative_meter.get("rejected_rewards"), - self.num_train_step, - ) - self.writer.add_scalar( - "train/accuracy", - self.accumulative_meter.get("accuracy"), - self.num_train_step, + if not self.apply_loss_mask: + chosen_loss_mask = chosen_loss_mask.fill_(1.0) + reject_loss_mask = reject_loss_mask.fill_(1.0) + + batch_size = chosen_input_ids.size()[0] + + actor_all_logits = self.model( + input_ids=torch.cat([chosen_input_ids, reject_input_ids]), + attention_mask=torch.cat([chosen_attention_mask, reject_attention_mask]), + )["logits"] + actor_chosen_logits = actor_all_logits[:batch_size] + actor_reject_logits = actor_all_logits[batch_size:] + logprob_actor_chosen = calc_masked_log_probs( + actor_chosen_logits, chosen_input_ids, chosen_loss_mask[:, 1:], self.length_normalization + ) + + logprob_actor_reject = calc_masked_log_probs( + actor_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization + ) + + if self.ref_model is not None: + self.ref_model.eval() + with torch.no_grad(): + ref_all_logits = self.ref_model( + input_ids=torch.cat([chosen_input_ids, reject_input_ids]), + attention_mask=torch.cat([chosen_attention_mask, reject_attention_mask]), + )["logits"] + ref_chosen_logits = ref_all_logits[:batch_size] + ref_reject_logits = ref_all_logits[batch_size:] + logprob_ref_chosen = calc_masked_log_probs( + ref_chosen_logits, chosen_input_ids, chosen_loss_mask[:, 1:], self.length_normalization + ) + logprob_ref_reject = calc_masked_log_probs( + ref_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization + ) + else: + logprob_ref_chosen = None + logprob_ref_reject = None + + losses, chosen_rewards, rejected_rewards = self.actor_loss_fn( + logprob_actor_chosen, + logprob_actor_reject, + logprob_ref_chosen if logprob_ref_chosen is not None else None, + logprob_ref_reject if logprob_ref_reject is not None else None, + chosen_loss_mask[:, 1:], + reject_loss_mask[:, 1:], + ) + reward_accuracies = (chosen_rewards > rejected_rewards).float().mean() + + # DPO Loss + loss = losses.mean() + self.booster.backward(loss=loss, optimizer=self.optimizer) + # sync + loss_mean = all_reduce_mean(tensor=loss) + chosen_rewards_mean = all_reduce_mean(tensor=chosen_rewards) + rejected_rewards_mean = all_reduce_mean(tensor=rejected_rewards) + reward_accuracies_mean = all_reduce_mean(tensor=reward_accuracies) + self.accumulative_meter.add("chosen_rewards", chosen_rewards_mean.to(torch.float16).mean().item()) + self.accumulative_meter.add("rejected_rewards", rejected_rewards_mean.to(torch.float16).mean().item()) + self.accumulative_meter.add("loss", loss_mean.to(torch.float16).item()) + self.accumulative_meter.add("accuracy", reward_accuracies_mean.to(torch.float16).item()) + + if (i + 1) % self.accumulation_steps == 0: + self.optimizer.step() + self.optimizer.zero_grad() + self.actor_scheduler.step() + + step_bar.set_postfix( + { + "train/loss": self.accumulative_meter.get("loss"), + "train/chosen_rewards": self.accumulative_meter.get("chosen_rewards"), + "train/rejected_rewards": self.accumulative_meter.get("rejected_rewards"), + "train/accuracy": self.accumulative_meter.get("accuracy"), + } ) - self.num_train_step += 1 - self.accumulative_meter.reset() + step_bar.update() + if self.writer and is_rank_0(): + self.writer.add_scalar("train/loss", self.accumulative_meter.get("loss"), self.num_train_step) + self.writer.add_scalar("train/lr", self.optimizer.param_groups[0]["lr"], self.num_train_step) + self.writer.add_scalar( + "train/chosen_rewards", self.accumulative_meter.get("chosen_rewards"), self.num_train_step + ) + self.writer.add_scalar( + "train/rejected_rewards", + self.accumulative_meter.get("rejected_rewards"), + self.num_train_step, + ) + self.writer.add_scalar( + "train/margin", + self.accumulative_meter.get("chosen_rewards") + - self.accumulative_meter.get("rejected_rewards"), + self.num_train_step, + ) + self.writer.add_scalar( + "train/accuracy", + self.accumulative_meter.get("accuracy"), + self.num_train_step, + ) + self.num_train_step += 1 + self.accumulative_meter.reset() if self.save_dir is not None and self.num_train_step > 0 and self.num_train_step % self.save_interval == 0: # save checkpoint diff --git a/applications/ColossalChat/examples/training_scripts/train_dpo.py b/applications/ColossalChat/examples/training_scripts/train_dpo.py index 90527c88b38e..9e0fa92e9f67 100755 --- a/applications/ColossalChat/examples/training_scripts/train_dpo.py +++ b/applications/ColossalChat/examples/training_scripts/train_dpo.py @@ -82,20 +82,24 @@ def train(args): parallel_output=False, max_norm=args.grad_clip, precision=args.mixed_precision, + microbatch_size=args.microbatch_size, ) else: raise ValueError(f"Unknown plugin {args.plugin}") booster = Booster(plugin=plugin) - ref_booster = Booster(plugin=plugin) - # ====================================================== - # Initialize Model, Objective, Optimizer and LR Scheduler - # ====================================================== - # Temp Fix: Disable lazy init due to version conflict - # init_ctx = ( - # LazyInitContext(default_device=get_current_device()) if isinstance(plugin, (GeminiPlugin,)) else nullcontext() - # ) + ref_plugin = HybridParallelPlugin( + tp_size=args.ref_tp, + pp_size=1, + zero_stage=args.zero_stage, + enable_flash_attention=args.use_flash_attn, + cpu_offload=True if args.zero_stage >= 1 and args.zero_cpu_offload else False, + parallel_output=False, + max_norm=args.grad_clip, + precision=args.mixed_precision, + ) + ref_booster = Booster(plugin=ref_plugin) init_ctx = nullcontext() with init_ctx: @@ -120,6 +124,7 @@ def train(args): ref_model = AutoModelForCausalLM.from_pretrained(args.pretrain) else: ref_model = None + if args.lora_config is not None: model = convert_to_lora_module(model, lora_config=lora_config) for name, module in model.named_modules(): @@ -129,7 +134,9 @@ def train(args): disable_dropout(ref_model) if args.grad_checkpoint: - # Note, for some models, lora may not be compatible with gradient checkpointing + # Make sure gradient checkpointing can be activated. + model.train() + # Note, for some models, lora may not be compatible with gradient checkpointing. model.gradient_checkpointing_enable(gradient_checkpointing_kwargs={"use_reentrant": False}) coordinator.print_on_master(msg="Gradient checkpointing enabled successfully") @@ -159,7 +166,7 @@ def train(args): adamw_mode=True, ) - # configure dataset + # Configure dataset coordinator.print_on_master(f"Load dataset: {args.dataset}") mode_map = {"train": "train", "valid": "validation", "test": "test"} train_dataset = load_tokenized_dataset(dataset_paths=args.dataset, mode="train", mode_map=mode_map) @@ -203,14 +210,15 @@ def train(args): default_dtype = torch.float16 if args.mixed_precision == "fp16" else torch.bfloat16 torch.set_default_dtype(default_dtype) + model, optim, _, train_dataloader, lr_scheduler = booster.boost( model=model, optimizer=optim, lr_scheduler=lr_scheduler, dataloader=train_dataloader, ) - if ref_model is not None: - ref_model, _, _, _, _ = ref_booster.boost(model=ref_model, dataloader=train_dataloader) + ref_model, _, _, _, _ = ref_booster.boost(model=ref_model) + torch.set_default_dtype(torch.float) coordinator.print_on_master(f"Booster init max CUDA memory: {torch.cuda.max_memory_allocated() / 1024 ** 2:.2f} MB") @@ -332,12 +340,6 @@ def train(args): parser.add_argument("--max_length", type=int, default=2048, help="Model max length") parser.add_argument("--max_epochs", type=int, default=3) parser.add_argument("--batch_size", type=int, default=4) - parser.add_argument( - "--disable_reference_model", - action="store_true", - default=False, - help="Disable the reference model (enabled by default)", - ) parser.add_argument("--disable_loss_mask", default=False, action="store_true") parser.add_argument("--mixed_precision", type=str, default="fp16", choices=["fp16", "bf16"], help="Mixed precision") parser.add_argument("--lora_config", type=str, default=None, help="low-rank adaptation config file path") @@ -348,6 +350,17 @@ def train(args): parser.add_argument("--use_wandb", default=False, action="store_true") parser.add_argument("--grad_checkpoint", default=False, action="store_true") parser.add_argument("--use_flash_attn", default=False, action="store_true") + parser.add_argument("--microbatch_size", type=int, default=1) + # Parameter for reference model + parser.add_argument( + "--disable_reference_model", + action="store_true", + default=False, + help="Disable the reference model (enabled by default)", + ) + parser.add_argument( + "--ref_tp", type=int, default=1, help="TP size for reference model; used only when reference model is to large." + ) args = parser.parse_args() # fool proof hyperparameter setup diff --git a/colossalai/shardformer/modeling/llama.py b/colossalai/shardformer/modeling/llama.py index 47c17e7494f2..2a5b60287e32 100644 --- a/colossalai/shardformer/modeling/llama.py +++ b/colossalai/shardformer/modeling/llama.py @@ -271,6 +271,7 @@ def llama_for_causal_lm_forward( hidden_states: Optional[torch.FloatTensor] = None, stage_index: Optional[List[int]] = None, shard_config: ShardConfig = None, + **kwargs, ): r""" Args: From 47e347c9f65eef8db6aaa699bfc3b620fd97ac3d Mon Sep 17 00:00:00 2001 From: Tong Li Date: Thu, 12 Sep 2024 11:32:09 +0000 Subject: [PATCH 10/21] update dpo --- .../ColossalChat/examples/training_scripts/train_dpo.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/applications/ColossalChat/examples/training_scripts/train_dpo.py b/applications/ColossalChat/examples/training_scripts/train_dpo.py index 9e0fa92e9f67..eeca080ca0ae 100755 --- a/applications/ColossalChat/examples/training_scripts/train_dpo.py +++ b/applications/ColossalChat/examples/training_scripts/train_dpo.py @@ -359,7 +359,7 @@ def train(args): help="Disable the reference model (enabled by default)", ) parser.add_argument( - "--ref_tp", type=int, default=1, help="TP size for reference model; used only when reference model is to large." + "--ref_tp", type=int, default=1, help="TP size for reference model; used only when reference model is too large." ) args = parser.parse_args() From 0cfb7fb11b7a52ad56255ab4adf27222c3cab94d Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Thu, 12 Sep 2024 11:33:38 +0000 Subject: [PATCH 11/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- .../ColossalChat/examples/training_scripts/train_dpo.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/applications/ColossalChat/examples/training_scripts/train_dpo.py b/applications/ColossalChat/examples/training_scripts/train_dpo.py index eeca080ca0ae..ad5a1c803d8b 100755 --- a/applications/ColossalChat/examples/training_scripts/train_dpo.py +++ b/applications/ColossalChat/examples/training_scripts/train_dpo.py @@ -359,7 +359,10 @@ def train(args): help="Disable the reference model (enabled by default)", ) parser.add_argument( - "--ref_tp", type=int, default=1, help="TP size for reference model; used only when reference model is too large." + "--ref_tp", + type=int, + default=1, + help="TP size for reference model; used only when reference model is too large.", ) args = parser.parse_args() From 41b10ef7bed880180afcecdc8f4888a9db5084fa Mon Sep 17 00:00:00 2001 From: Tong Li Date: Sat, 14 Sep 2024 17:45:19 +0800 Subject: [PATCH 12/21] add dpo fn --- .../ColossalChat/coati/trainer/dpo.py | 27 ++++++++++++++----- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index 420279edc380..725d366656ed 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -181,11 +181,10 @@ def _train(self, epoch: int): } ] ) + rewards = [] - outputs = self.booster.execute_pipeline( - data_iter, - self.model, - criterion=lambda outputs, inputs: self.actor_loss_fn( + def _criterion(outputs, inputs): + loss, chosen_rewards, rejected_rewards = self.actor_loss_fn( calc_masked_log_probs( outputs["logits"][0::2], inputs["input_ids"][0::2], @@ -202,20 +201,34 @@ def _train(self, epoch: int): inputs["logprob_ref"][1::2] if inputs["logprob_ref"] is not None else None, inputs["loss_mask"][0::2][:, 1:], inputs["loss_mask"][1::2][:, 1:], - )[0], + ) + rewards.append(chosen_rewards) + rewards.append(rejected_rewards) + return loss + + outputs = self.booster.execute_pipeline( + data_iter, + self.model, + criterion=_criterion, optimizer=self.optimizer, return_loss=True, ) loss = outputs["loss"] if self.booster.plugin.stage_manager.is_last_stage(): + chosen_rewards, rejected_rewards = rewards[0], rewards[1] global_loss = all_reduce_mean(loss, self.plugin) if dist.get_rank() == dist.get_world_size() - 1: - step_bar.set_postfix({"train/loss": global_loss.item()}) + step_bar.set_postfix( + { + "train/loss": global_loss.item(), + "train/chosen_rewards": chosen_rewards.to(torch.float16).mean().item(), + "train/rejected_rewards": rejected_rewards.to(torch.float16).mean().item(), + } + ) step_bar.update() self.optimizer.step() self.optimizer.zero_grad() - else: self.accumulative_meter.reset() step_bar = trange( From f342573d84a12bc6a50f85e3f293b951bf7572d9 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Sat, 14 Sep 2024 18:07:33 +0800 Subject: [PATCH 13/21] update dpo --- applications/ColossalChat/coati/trainer/dpo.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index 725d366656ed..39e52d364b70 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -221,6 +221,7 @@ def _criterion(outputs, inputs): step_bar.set_postfix( { "train/loss": global_loss.item(), + "train/lr": self.actor_scheduler.get_last_lr()[0], "train/chosen_rewards": chosen_rewards.to(torch.float16).mean().item(), "train/rejected_rewards": rejected_rewards.to(torch.float16).mean().item(), } @@ -229,6 +230,7 @@ def _criterion(outputs, inputs): self.optimizer.step() self.optimizer.zero_grad() + self.actor_scheduler.step() else: self.accumulative_meter.reset() step_bar = trange( @@ -432,9 +434,6 @@ def _eval(self, epoch: int): logprob_actor_reject = calc_masked_log_probs( actor_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization ) - - self.ref_model.eval() - ref_all_logits = self.ref_model( torch.cat([chosen_input_ids, reject_input_ids]), torch.cat([chosen_attention_mask, reject_attention_mask]), From c3812652585747e7c58a3a7118c8a404a7b87f0e Mon Sep 17 00:00:00 2001 From: Tong Li Date: Sat, 14 Sep 2024 18:37:52 +0800 Subject: [PATCH 14/21] update dpo --- .../ColossalChat/coati/trainer/dpo.py | 285 +++++++++++++----- 1 file changed, 206 insertions(+), 79 deletions(-) diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index 39e52d364b70..1fcdefe4bd36 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -386,94 +386,221 @@ def _eval(self, epoch: int): return self.model.eval() self.ref_model.eval() + self.accumulative_meter.reset() self.coordinator.print_on_master("\nStart evaluation...") - step_bar = trange( - len(self.eval_dataloader), - desc=f"Epoch {epoch + 1}/{self.max_epochs}", - disable=not is_rank_0(), - ) - - self.accumulative_meter.reset() + if isinstance(self.plugin, HybridParallelPlugin) and self.plugin.pp_size > 1: + step_bar = tqdm( + range(len(self.eval_dataloader)), + desc="Step", + disable=not (dist.get_rank() == dist.get_world_size() - 1), + ) + with torch.no_grad(): + for _, batch in enumerate(self.eval_dataloader): + batch = to_device(batch, self.device) + ( + chosen_input_ids, + chosen_attention_mask, + chosen_loss_mask, + reject_input_ids, + reject_attention_mask, + reject_loss_mask, + ) = ( + batch["chosen_input_ids"], + batch["chosen_attention_mask"], + batch["chosen_loss_mask"], + batch["reject_input_ids"], + batch["reject_attention_mask"], + batch["reject_loss_mask"], + ) + batch_size = chosen_input_ids.size()[0] + # Calculate logits from reference model. + if self.ref_model is not None: + self.ref_model.eval() + with torch.no_grad(): + ref_all_logits = self.ref_model( + input_ids=torch.cat([chosen_input_ids, reject_input_ids]), + attention_mask=torch.cat([chosen_attention_mask, reject_attention_mask]), + )["logits"] + ref_chosen_logits = ref_all_logits[:batch_size] + ref_reject_logits = ref_all_logits[batch_size:] + logprob_ref_chosen = calc_masked_log_probs( + ref_chosen_logits, chosen_input_ids, chosen_loss_mask[:, 1:], self.length_normalization + ) + logprob_ref_reject = calc_masked_log_probs( + ref_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization + ) + else: + logprob_ref_chosen = None + logprob_ref_reject = None + + # Merge chosen and reject + inputs_ids = torch.stack([item for tup in zip(chosen_input_ids, reject_input_ids) for item in tup]) + attention_mask = torch.stack( + [item for tup in zip(chosen_attention_mask, reject_attention_mask) for item in tup] + ) + loss_mask = torch.stack([item for tup in zip(chosen_loss_mask, reject_loss_mask) for item in tup]) + logprob_ref = torch.stack( + [item for tup in zip(logprob_ref_chosen, logprob_ref_reject) for item in tup] + ) - with torch.no_grad(): - for i, batch in enumerate(self.eval_dataloader): - batch = to_device(batch, self.device) - ( - chosen_input_ids, - chosen_attention_mask, - chosen_loss_mask, - reject_input_ids, - reject_attention_mask, - reject_loss_mask, - ) = ( - batch["chosen_input_ids"], - batch["chosen_attention_mask"], - batch["chosen_loss_mask"], - batch["reject_input_ids"], - batch["reject_attention_mask"], - batch["reject_loss_mask"], - ) - if not self.apply_loss_mask: - chosen_loss_mask = chosen_loss_mask.fill_(1.0) - reject_loss_mask = reject_loss_mask.fill_(1.0) + data_iter = iter( + [ + { + "input_ids": inputs_ids, + "attention_mask": attention_mask, + "loss_mask": loss_mask, + "logprob_ref": logprob_ref, + } + ] + ) + rewards = [] + + def _criterion(outputs, inputs): + loss, chosen_rewards, rejected_rewards = self.actor_loss_fn( + calc_masked_log_probs( + outputs["logits"][0::2], + inputs["input_ids"][0::2], + inputs["loss_mask"][0::2][:, 1:], + self.length_normalization, + ), + calc_masked_log_probs( + outputs["logits"][1::2], + inputs["input_ids"][1::2], + inputs["loss_mask"][1::2][:, 1:], + self.length_normalization, + ), + inputs["logprob_ref"][0::2] if inputs["logprob_ref"] is not None else None, + inputs["logprob_ref"][1::2] if inputs["logprob_ref"] is not None else None, + inputs["loss_mask"][0::2][:, 1:], + inputs["loss_mask"][1::2][:, 1:], + ) + rewards.append(chosen_rewards) + rewards.append(rejected_rewards) + return loss + + outputs = self.booster.execute_pipeline( + data_iter, + self.model, + criterion=_criterion, + optimizer=self.optimizer, + return_loss=True, + ) + loss = outputs["loss"] + if self.booster.plugin.stage_manager.is_last_stage(): + chosen_rewards, rejected_rewards = rewards[0], rewards[1] + global_loss = all_reduce_mean(loss, self.plugin) + chosen_rewards_mean = all_reduce_mean(chosen_rewards, self.plugin) + rejected_rewards_mean = all_reduce_mean(rejected_rewards, self.plugin) + if dist.get_rank() == dist.get_world_size() - 1: + step_bar.set_postfix( + { + "eval/loss": global_loss.item(), + "eval/lr": self.actor_scheduler.get_last_lr()[0], + "eval/chosen_rewards": chosen_rewards.to(torch.float16).mean().item(), + "eval/rejected_rewards": rejected_rewards.to(torch.float16).mean().item(), + } + ) + self.accumulative_meter.add( + "chosen_rewards", chosen_rewards_mean.to(torch.float16).mean().item() + ) + self.accumulative_meter.add( + "rejected_rewards", rejected_rewards_mean.to(torch.float16).mean().item() + ) + self.accumulative_meter.add("loss", global_loss.to(torch.float16).item()) + step_bar.update() + if self.booster.plugin.stage_manager.is_last_stage(): + msg = "\nEvaluation Result:\n" + for tag in ["loss", "chosen_rewards", "rejected_rewards"]: + msg = msg + f"{tag}: {self.accumulative_meter.get(tag)}\n" + if dist.get_rank() == dist.get_world_size() - 1: + print(msg) + else: + step_bar = trange( + len(self.eval_dataloader), + desc=f"Epoch {epoch + 1}/{self.max_epochs}", + disable=not is_rank_0(), + ) + with torch.no_grad(): + for i, batch in enumerate(self.eval_dataloader): + batch = to_device(batch, self.device) + ( + chosen_input_ids, + chosen_attention_mask, + chosen_loss_mask, + reject_input_ids, + reject_attention_mask, + reject_loss_mask, + ) = ( + batch["chosen_input_ids"], + batch["chosen_attention_mask"], + batch["chosen_loss_mask"], + batch["reject_input_ids"], + batch["reject_attention_mask"], + batch["reject_loss_mask"], + ) + if not self.apply_loss_mask: + chosen_loss_mask = chosen_loss_mask.fill_(1.0) + reject_loss_mask = reject_loss_mask.fill_(1.0) - batch_size = chosen_input_ids.size()[0] + batch_size = chosen_input_ids.size()[0] - actor_all_logits = self.model( - torch.cat([chosen_input_ids, reject_input_ids]), - torch.cat([chosen_attention_mask, reject_attention_mask]), - )["logits"] - actor_chosen_logits = actor_all_logits[:batch_size] - actor_reject_logits = actor_all_logits[batch_size:] + actor_all_logits = self.model( + torch.cat([chosen_input_ids, reject_input_ids]), + torch.cat([chosen_attention_mask, reject_attention_mask]), + )["logits"] + actor_chosen_logits = actor_all_logits[:batch_size] + actor_reject_logits = actor_all_logits[batch_size:] - logprob_actor_chosen = calc_masked_log_probs( - actor_chosen_logits, chosen_input_ids, chosen_loss_mask[:, 1:], self.length_normalization - ) + logprob_actor_chosen = calc_masked_log_probs( + actor_chosen_logits, chosen_input_ids, chosen_loss_mask[:, 1:], self.length_normalization + ) - logprob_actor_reject = calc_masked_log_probs( - actor_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization - ) - ref_all_logits = self.ref_model( - torch.cat([chosen_input_ids, reject_input_ids]), - torch.cat([chosen_attention_mask, reject_attention_mask]), - )["logits"] - ref_chosen_logits = ref_all_logits[:batch_size] - ref_reject_logits = ref_all_logits[batch_size:] - logprob_ref_chosen = calc_masked_log_probs( - ref_chosen_logits, chosen_input_ids, chosen_loss_mask[:, 1:], self.length_normalization - ) - logprob_ref_reject = calc_masked_log_probs( - ref_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization - ) + logprob_actor_reject = calc_masked_log_probs( + actor_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization + ) + ref_all_logits = self.ref_model( + torch.cat([chosen_input_ids, reject_input_ids]), + torch.cat([chosen_attention_mask, reject_attention_mask]), + )["logits"] + ref_chosen_logits = ref_all_logits[:batch_size] + ref_reject_logits = ref_all_logits[batch_size:] + logprob_ref_chosen = calc_masked_log_probs( + ref_chosen_logits, chosen_input_ids, chosen_loss_mask[:, 1:], self.length_normalization + ) + logprob_ref_reject = calc_masked_log_probs( + ref_reject_logits, reject_input_ids, reject_loss_mask[:, 1:], self.length_normalization + ) - losses, chosen_rewards, rejected_rewards = self.actor_loss_fn( - logprob_actor_chosen, - logprob_actor_reject, - logprob_ref_chosen if logprob_ref_chosen is not None else None, - logprob_ref_reject if logprob_ref_reject is not None else None, - chosen_loss_mask[:, 1:], - reject_loss_mask[:, 1:], - ) - reward_accuracies = (chosen_rewards > rejected_rewards).float().mean() - loss = losses.mean() - loss_mean = all_reduce_mean(tensor=loss) - chosen_rewards_mean = all_reduce_mean(tensor=chosen_rewards) - rejected_rewards_mean = all_reduce_mean(tensor=rejected_rewards) - reward_accuracies_mean = all_reduce_mean(tensor=reward_accuracies) - self.accumulative_meter.add("chosen_rewards", chosen_rewards_mean.to(torch.float16).mean().item()) - self.accumulative_meter.add("rejected_rewards", rejected_rewards_mean.to(torch.float16).mean().item()) - self.accumulative_meter.add("loss", loss_mean.to(torch.float16).item()) - self.accumulative_meter.add("accuracy", reward_accuracies_mean.to(torch.float16).item()) - self.accumulative_meter.add( - "margin", (chosen_rewards_mean - rejected_rewards_mean).to(torch.float16).mean().item() - ) - step_bar.update() + losses, chosen_rewards, rejected_rewards = self.actor_loss_fn( + logprob_actor_chosen, + logprob_actor_reject, + logprob_ref_chosen if logprob_ref_chosen is not None else None, + logprob_ref_reject if logprob_ref_reject is not None else None, + chosen_loss_mask[:, 1:], + reject_loss_mask[:, 1:], + ) + reward_accuracies = (chosen_rewards > rejected_rewards).float().mean() + loss = losses.mean() + loss_mean = all_reduce_mean(tensor=loss) + chosen_rewards_mean = all_reduce_mean(tensor=chosen_rewards) + rejected_rewards_mean = all_reduce_mean(tensor=rejected_rewards) + reward_accuracies_mean = all_reduce_mean(tensor=reward_accuracies) + self.accumulative_meter.add("chosen_rewards", chosen_rewards_mean.to(torch.float16).mean().item()) + self.accumulative_meter.add( + "rejected_rewards", rejected_rewards_mean.to(torch.float16).mean().item() + ) + self.accumulative_meter.add("loss", loss_mean.to(torch.float16).item()) + self.accumulative_meter.add("accuracy", reward_accuracies_mean.to(torch.float16).item()) + self.accumulative_meter.add( + "margin", (chosen_rewards_mean - rejected_rewards_mean).to(torch.float16).mean().item() + ) + step_bar.update() - msg = "\nEvaluation Result:\n" - for tag in ["loss", "chosen_rewards", "rejected_rewards", "accuracy", "margin"]: - msg = msg + f"{tag}: {self.accumulative_meter.get(tag)}\n" - self.coordinator.print_on_master(msg) + msg = "\nEvaluation Result:\n" + for tag in ["loss", "chosen_rewards", "rejected_rewards", "accuracy", "margin"]: + msg = msg + f"{tag}: {self.accumulative_meter.get(tag)}\n" + self.coordinator.print_on_master(msg) if self.save_dir is not None: os.makedirs(self.save_dir, exist_ok=True) with open(os.path.join(self.save_dir, f"eval_result_epoch{epoch}.txt"), "w") as f: From d853d2d5541ab4f6b01da7c02455f31d64dc23b0 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Thu, 19 Sep 2024 06:53:42 +0000 Subject: [PATCH 15/21] update dpo --- applications/ColossalChat/coati/trainer/dpo.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index 1fcdefe4bd36..5bfddd5468a2 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -595,6 +595,14 @@ def _criterion(outputs, inputs): self.accumulative_meter.add( "margin", (chosen_rewards_mean - rejected_rewards_mean).to(torch.float16).mean().item() ) + step_bar.set_postfix( + { + "eval/loss": self.accumulative_meter.get("loss"), + "eval/chosen_rewards": self.accumulative_meter.get("chosen_rewards"), + "eval/rejected_rewards": self.accumulative_meter.get("rejected_rewards"), + "eval/accuracy": self.accumulative_meter.get("accuracy"), + } + ) step_bar.update() msg = "\nEvaluation Result:\n" From c41525ae0dd9f2d703152dcc498ba99f42baed13 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Thu, 19 Sep 2024 10:58:53 +0000 Subject: [PATCH 16/21] update dpo --- .../ColossalChat/coati/trainer/dpo.py | 32 +++++++++++++++++-- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index 5bfddd5468a2..c696ddc08c6e 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -98,18 +98,25 @@ def _before_fit( self.train_dataloader = train_preference_dataloader self.eval_dataloader = eval_preference_dataloader self.writer = None - if use_wandb and is_rank_0(): + + init_criterion = ( + dist.get_rank() == dist.get_world_size() - 1 + if isinstance(self.plugin, HybridParallelPlugin) and self.plugin.pp_size > 1 + else is_rank_0() + ) + + if use_wandb and init_criterion: assert log_dir is not None, "log_dir must be provided when use_wandb is True" import wandb self.wandb_run = wandb.init(project="Coati-dpo", sync_tensorboard=True) - if log_dir is not None and is_rank_0(): + if log_dir is not None and init_criterion: import os import time from torch.utils.tensorboard import SummaryWriter - log_dir = os.path.join(log_dir, "dpo") + log_dir = os.path.join(log_dir, "DPO") log_dir = os.path.join(log_dir, time.strftime("%Y-%m-%d_%H:%M:%S", time.localtime())) self.writer = SummaryWriter(log_dir=log_dir) @@ -227,6 +234,25 @@ def _criterion(outputs, inputs): } ) step_bar.update() + self.accumulative_meter.add("loss", global_loss.item()) + self.accumulative_meter.add("chosen_rewards", chosen_rewards.to(torch.float16).mean().item()) + self.accumulative_meter.add( + "rejected_rewards", rejected_rewards.to(torch.float16).mean().item() + ) + + self.writer.add_scalar("train/loss", self.accumulative_meter.get("loss"), i) + self.writer.add_scalar("train/chosen_rewards", self.accumulative_meter.get("chosen_rewards"), i) + self.writer.add_scalar( + "train/rejected_rewards", + self.accumulative_meter.get("rejected_rewards"), + i, + ) + self.writer.add_scalar( + "train/margin", + self.accumulative_meter.get("chosen_rewards") + - self.accumulative_meter.get("rejected_rewards"), + i, + ) self.optimizer.step() self.optimizer.zero_grad() From e3e7b8a0f2565851c37350aa0e6eea8e8d7413c1 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Fri, 20 Sep 2024 07:26:55 +0000 Subject: [PATCH 17/21] minor update --- .../ColossalChat/coati/models/loss.py | 1 + .../ColossalChat/coati/models/utils.py | 4 +-- .../ColossalChat/coati/trainer/dpo.py | 32 +++++++++---------- 3 files changed, 18 insertions(+), 19 deletions(-) diff --git a/applications/ColossalChat/coati/models/loss.py b/applications/ColossalChat/coati/models/loss.py index 1bc57a2e94b7..4f7de203457a 100755 --- a/applications/ColossalChat/coati/models/loss.py +++ b/applications/ColossalChat/coati/models/loss.py @@ -153,6 +153,7 @@ def forward( else: # If no reference model is provided ref_logratios = 0.0 + pi_logratios = logprob_actor_chosen.sum(-1) - logprob_actor_reject.sum(-1) logits = pi_logratios - ref_logratios - self.gamma / self.beta losses = -torch.nn.functional.logsigmoid(self.beta * logits) diff --git a/applications/ColossalChat/coati/models/utils.py b/applications/ColossalChat/coati/models/utils.py index c583f057a5ab..fe7ab209822f 100755 --- a/applications/ColossalChat/coati/models/utils.py +++ b/applications/ColossalChat/coati/models/utils.py @@ -50,8 +50,8 @@ def _log_probs_from_logits(logits: torch.Tensor, labels: torch.Tensor) -> torch. torch.Tensor: The log probabilities corresponding to the labels. """ log_probs = F.log_softmax(logits, dim=-1) - log_probs_labels = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)) - return log_probs_labels.squeeze(-1) + per_label_logps = log_probs.gather(dim=-1, index=labels.unsqueeze(-1)) + return per_label_logps.squeeze(-1) def calc_action_log_probs(logits: torch.Tensor, sequences: torch.LongTensor, num_actions: int) -> torch.Tensor: diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index c696ddc08c6e..c72b5de7bb7c 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -239,20 +239,20 @@ def _criterion(outputs, inputs): self.accumulative_meter.add( "rejected_rewards", rejected_rewards.to(torch.float16).mean().item() ) - - self.writer.add_scalar("train/loss", self.accumulative_meter.get("loss"), i) - self.writer.add_scalar("train/chosen_rewards", self.accumulative_meter.get("chosen_rewards"), i) - self.writer.add_scalar( - "train/rejected_rewards", - self.accumulative_meter.get("rejected_rewards"), - i, - ) - self.writer.add_scalar( - "train/margin", - self.accumulative_meter.get("chosen_rewards") - - self.accumulative_meter.get("rejected_rewards"), - i, - ) + if self.writer is not None: + self.writer.add_scalar("train/loss", self.accumulative_meter.get("loss"), i) + self.writer.add_scalar("train/chosen_rewards", self.accumulative_meter.get("chosen_rewards"), i) + self.writer.add_scalar( + "train/rejected_rewards", + self.accumulative_meter.get("rejected_rewards"), + i, + ) + self.writer.add_scalar( + "train/margin", + self.accumulative_meter.get("chosen_rewards") + - self.accumulative_meter.get("rejected_rewards"), + i, + ) self.optimizer.step() self.optimizer.zero_grad() @@ -320,7 +320,7 @@ def _criterion(outputs, inputs): logprob_ref_chosen = None logprob_ref_reject = None - losses, chosen_rewards, rejected_rewards = self.actor_loss_fn( + loss, chosen_rewards, rejected_rewards = self.actor_loss_fn( logprob_actor_chosen, logprob_actor_reject, logprob_ref_chosen if logprob_ref_chosen is not None else None, @@ -330,8 +330,6 @@ def _criterion(outputs, inputs): ) reward_accuracies = (chosen_rewards > rejected_rewards).float().mean() - # DPO Loss - loss = losses.mean() self.booster.backward(loss=loss, optimizer=self.optimizer) # sync loss_mean = all_reduce_mean(tensor=loss) From 69970d6d6804898fc4aeb559d338540001822407 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 20 Sep 2024 07:27:46 +0000 Subject: [PATCH 18/21] [pre-commit.ci] auto fixes from pre-commit.com hooks for more information, see https://pre-commit.ci --- applications/ColossalChat/coati/trainer/dpo.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/applications/ColossalChat/coati/trainer/dpo.py b/applications/ColossalChat/coati/trainer/dpo.py index c72b5de7bb7c..499113e9646f 100755 --- a/applications/ColossalChat/coati/trainer/dpo.py +++ b/applications/ColossalChat/coati/trainer/dpo.py @@ -241,7 +241,9 @@ def _criterion(outputs, inputs): ) if self.writer is not None: self.writer.add_scalar("train/loss", self.accumulative_meter.get("loss"), i) - self.writer.add_scalar("train/chosen_rewards", self.accumulative_meter.get("chosen_rewards"), i) + self.writer.add_scalar( + "train/chosen_rewards", self.accumulative_meter.get("chosen_rewards"), i + ) self.writer.add_scalar( "train/rejected_rewards", self.accumulative_meter.get("rejected_rewards"), From 18cf308a649d50b380cefd2fd31e1e16845dc220 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Fri, 20 Sep 2024 07:36:52 +0000 Subject: [PATCH 19/21] update loss --- applications/ColossalChat/coati/models/loss.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/applications/ColossalChat/coati/models/loss.py b/applications/ColossalChat/coati/models/loss.py index 4f7de203457a..927dfd5a89b6 100755 --- a/applications/ColossalChat/coati/models/loss.py +++ b/applications/ColossalChat/coati/models/loss.py @@ -157,7 +157,7 @@ def forward( pi_logratios = logprob_actor_chosen.sum(-1) - logprob_actor_reject.sum(-1) logits = pi_logratios - ref_logratios - self.gamma / self.beta losses = -torch.nn.functional.logsigmoid(self.beta * logits) - losses = losses.mean() + loss = losses.mean() # Calculate rewards for logging if logprob_ref_chosen is not None: chosen_rewards = self.beta * (logprob_actor_chosen.sum(-1) - logprob_ref_chosen.sum(-1)).detach() @@ -168,7 +168,7 @@ def forward( else: rejected_rewards = self.beta * logprob_actor_reject.sum(-1).detach() - return losses, chosen_rewards, rejected_rewards + return loss, chosen_rewards, rejected_rewards class LogSigLoss(nn.Module): From ff278dcd59c21920c5095beb034a2693fb886c57 Mon Sep 17 00:00:00 2001 From: Tong Li Date: Fri, 20 Sep 2024 07:46:19 +0000 Subject: [PATCH 20/21] update help --- .../ColossalChat/examples/training_scripts/train_dpo.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/applications/ColossalChat/examples/training_scripts/train_dpo.py b/applications/ColossalChat/examples/training_scripts/train_dpo.py index ad5a1c803d8b..351f59092328 100755 --- a/applications/ColossalChat/examples/training_scripts/train_dpo.py +++ b/applications/ColossalChat/examples/training_scripts/train_dpo.py @@ -345,12 +345,17 @@ def train(args): parser.add_argument("--lora_config", type=str, default=None, help="low-rank adaptation config file path") parser.add_argument("--save_interval", type=int, default=1000, help="number of step between two checkpoints") parser.add_argument("--lr", type=float, default=5e-6) - parser.add_argument("--accumulation_steps", type=int, default=8) + parser.add_argument("--accumulation_steps", type=int, default=1) parser.add_argument("--log_dir", default=None, type=str) parser.add_argument("--use_wandb", default=False, action="store_true") parser.add_argument("--grad_checkpoint", default=False, action="store_true") parser.add_argument("--use_flash_attn", default=False, action="store_true") - parser.add_argument("--microbatch_size", type=int, default=1) + parser.add_argument( + "--microbatch_size", + type=int, + default=2, + help="Micro batch size for PP training. To activate PP training for DPO-like algorithm, you must keep size even and the size should be equal or greater than 2.", + ) # Parameter for reference model parser.add_argument( "--disable_reference_model", From d71660217d04935b57fd60862c5082dddb57dc5c Mon Sep 17 00:00:00 2001 From: Tong Li Date: Fri, 11 Oct 2024 05:47:48 +0000 Subject: [PATCH 21/21] polish code --- .../ColossalChat/examples/training_scripts/train_dpo.py | 2 +- .../ColossalChat/examples/training_scripts/train_sft.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/applications/ColossalChat/examples/training_scripts/train_dpo.py b/applications/ColossalChat/examples/training_scripts/train_dpo.py index 351f59092328..ad81db73aa53 100755 --- a/applications/ColossalChat/examples/training_scripts/train_dpo.py +++ b/applications/ColossalChat/examples/training_scripts/train_dpo.py @@ -44,7 +44,7 @@ def train(args): Default torch ddp plugin without any acceleration, for debugging purpose acceleration, for debugging purpose """ - plugin = TorchDDPPlugin(find_unused_parameters=True if args.grad_checkpoint is False else False) + plugin = TorchDDPPlugin(find_unused_parameters=not args.grad_checkpoint) elif args.plugin == "gemini": plugin = GeminiPlugin( precision=args.mixed_precision, diff --git a/applications/ColossalChat/examples/training_scripts/train_sft.py b/applications/ColossalChat/examples/training_scripts/train_sft.py index 62acad32f66a..e319340c3f60 100755 --- a/applications/ColossalChat/examples/training_scripts/train_sft.py +++ b/applications/ColossalChat/examples/training_scripts/train_sft.py @@ -68,7 +68,7 @@ def train(args): Default torch ddp plugin without any acceleration, for debugging purpose acceleration, for debugging purpose """ - plugin = TorchDDPPlugin(find_unused_parameters=True if args.grad_checkpoint is False else False) + plugin = TorchDDPPlugin(find_unused_parameters=not args.grad_checkpoint) elif args.plugin == "gemini": plugin = GeminiPlugin( precision=args.mixed_precision,