Merge pull request #81 from codefuse-ai/support_coba_loss

chencyudel · web-flow · commit a4cd0b906252 · 2024-10-29T17:35:52.000+08:00
Support coba loss
diff --git a/README.md b/README.md
@@ -46,6 +46,8 @@
 
 
 ## News
+🔥🔥🔥 [2024/10/29] Our paper [CoBa: Convergence Balancer for Multitask Finetuning of Large Language Models](https://arxiv.org/abs/2410.06741) has been accepted by EMNLP-2024, which achieves balanced convergence across various tasks.
+
 🔥🔥🔥 [2024/05/20] We released **MFTCoder v0.4**, mainly for MFTCoder-accelerate. It supports **QLoRA + DeepSpeed Zero3** and **QLoRA + FSDP** as options allowing you training very large models. It now supports new models like Qwen2, Qwen2-MoE, Starcoder2, Gemma, etc.
 
 🔥🔥🔥 [2024/05/20] Our paper [MFTCoder: Boosting Code LLMs with Multitask Fine-Tuning](https://arxiv.org/abs/2311.02303) has been accepted by KDD2024.
diff --git a/README_cn.md b/README_cn.md
@@ -45,6 +45,8 @@
 
 
 ## 新闻
+🔥🔥🔥 [2024/10/29] 我们的论文 [CoBa: Convergence Balancer for Multitask Finetuning of Large Language Models](https://arxiv.org/abs/2410.06741) 已被 EMNLP 2024 接收，可以实现多任务收敛均衡。
+
 🔥🔥🔥 [2024/05/20] **MFTCoder-v0.4**发布。新增支持**QLoRA+ DeepSpeed Zero3**, **QLoRA + FSDP**训练模式，可以更好的支持微调更大的模型，比如Qwen1.5-70B等。新增对Qwen2, Qwen2-MoE, Starcoder2, Gemma等模型的支持。
 
 🔥🔥🔥 [2024/05/20] 我们的论文 [MFTCoder: Boosting Code LLMs with Multitask Fine-Tuning](https://arxiv.org/abs/2311.02303) 已被 KDD 2024 接收.
diff --git a/mftcoder_accelerate/README.md b/mftcoder_accelerate/README.md
@@ -7,6 +7,8 @@
 [[中文]](README_cn.md) [**English**]
 
 ## 1. Updates
+🔥 MFTCoder-accelerate supports latest implementation of CoBa Loss (selfpaced Loss) for better Convergence Balance.
+
 🔥 MFTCoder-accelerate now support these modes: QLoRA/LoRA + DeepSpeed ZeRO2， QLoRA + DeepSpeed ZeRO3, Full-parameter + DeepSpeed ZeRO3, QLoRA + FSDP, Full-parameter + FSDP.
 
 🔥 MFTCoder-accelerate supports QLoRA + DeepSpeed ZeRO3 and QLoRA + FSDP, which both work for larger models;
diff --git a/mftcoder_accelerate/README_cn.md b/mftcoder_accelerate/README_cn.md
@@ -7,6 +7,8 @@
 [**中文**] [[English]](README.md)
 
 ## 1. 更新
+🔥 MFTCoder-accelerate 增加了CoBa Loss的最新实现（原selfpaced Loss）, 让收敛均衡更进一步。
+
 🔥 MFTCoder-accelerate 最新支持的训练模式包括: QLoRA/LoRA + DeepSpeed ZeRO2， QLoRA + DeepSpeed ZeRO3, 全量 + DeepSpeed ZeRO3, QLoRA + FSDP, 全量 + FSDP。
 
 🔥 MFTCoder-accelerate 新增支持QLoRA + DeepSpeed ZeRO3， 支持QLoRA + FSDP, 可以训练更大的模型;
diff --git a/mftcoder_accelerate/src/configs/coba_train_config.json b/mftcoder_accelerate/src/configs/coba_train_config.json
@@ -5,16 +5,17 @@
     "pretrained_model_path": "$MODEL_NAME_OR_PATH",
     "model_type": "$MODEL_TYPE",
     "load_raw_dataset": true,
-    "data_split": "98,2,0",
+    "data_split": "95,5,0",
     "padding_mode": "padding",
     "use_dynamic_padding": true,
     "tokenize_mode": "sft",
     "tokenizer_type": "AutoTokenizer",
-    "weighted_loss_mode": "selfpaced",
-    "selfpaced_interval": 1,
-    "selfpaced_history_length": 100,
-    "selfpaced_sample_valid_num": 1,
-    "selfpaced_scale_factor": 50,
+    "weighted_loss_mode": "coba",
+    "coba_warmup_steps": 100,
+    "coba_history_length": 200,
+    "coba_tau": 5,
+    "coba_update_interval": 1,
+    "coba_sample_valid_num": 1,
     "attn_implementation": "flash_attention_2",
     "seq_length": 4096,
     "seed": 1234,
@@ -23,8 +24,8 @@
     "lora_rank": 96,
     "lora_alpha": 32,
     "lora_dropout": 0.05,
-    "per_device_train_batch_size": 2,
-    "per_device_eval_batch_size": 2,
+    "per_device_train_batch_size": 8,
+    "per_device_eval_batch_size": 8,
     "learning_rate": 5e-5,
     "min_lr": 5e-6,
     "weight_decay": 0.1,
@@ -42,4 +43,4 @@
     "early_stopping": true,
     "early_stopping_stall_num": 5,
     "saving_limit": null
-}
+  }
diff --git a/mftcoder_accelerate/src/mpt/mpt_accelerate.py b/mftcoder_accelerate/src/mpt/mpt_accelerate.py
@@ -209,7 +209,7 @@ def prepare_args():
     # generate TASK2ID, ID2TASK
     generate_task_id(args.data_paths)
 
-    if args.weighted_loss_mode == "selfpaced":
+    if args.weighted_loss_mode == "coba":
         args.task_weights = [1.0] * len(ID2TASK)
     elif args.task_weights is not None:
         args.task_weights = [float(wt) for wt in args.task_weights[1:-1].split(",")]
diff --git a/mftcoder_accelerate/src/mpt/mpt_arguments.py b/mftcoder_accelerate/src/mpt/mpt_arguments.py
@@ -126,14 +126,16 @@ class MptTrainArgs:
     # if dynamic padding
     use_dynamic_padding: bool = True
 
-    # interval of update per task train weight in selfpaced
-    selfpaced_interval: int = 1
-    # history length of sample valid loss used to fit the slope curve in selfpaced
-    selfpaced_history_length: int = 100
-    # the number of mini valid batches sampled at each interval
-    selfpaced_sample_valid_num: int = 1
-    # scale factor before softmax
-    selfpaced_scale_factor: int = 50
+    # warm-up steps for CoBa, recommand the number of valid batches
+    coba_warmup_steps: int = 100
+    # history length of sample valid loss used to fit the slope curve in CoBa
+    coba_history_length: int = 200
+    # temperature for divergence factor in CoBa
+    coba_tau: int = 5
+    # iteration interval of update per task train weight in CoBa
+    coba_update_interval: int = 1
+    # the number of mini valid batches sampled at each updated iteration interval
+    coba_sample_valid_num: int = 1
 
     # ATTENTION_CLASSES = { "eager": Normal Attention, "flash_attention_2": FlashAttention2}
     attn_implementation: str = "flash_attention_2"
diff --git a/mftcoder_accelerate/src/mpt/mpt_trainer.py b/mftcoder_accelerate/src/mpt/mpt_trainer.py
@@ -30,7 +30,7 @@
 
 # sys.path.append("..")
 from utils.common_utils import generate_task_id, TASK2ID, ID2TASK
-from utils.loss_utils import loss_func_mft, SelfpacedStatus, load_balancing_loss_func
+from utils.loss_utils import loss_func_mft, CoBaStatus, load_balancing_loss_func
 
 logger = get_logger(__name__)
 
@@ -239,7 +239,7 @@ def accelerate_monitor(
         reduce_task_loss,
         reduce_task_exist,
         completed_steps,
-        selfpaced_status=None,
+        coba_status=None,
     ):
         """
         gather reduce_loss and reduce_task_loss from all N devices.
@@ -263,27 +263,27 @@ def accelerate_monitor(
             f"[lr={self.lr_scheduler.get_lr()[0]:.4e}, {self.optimizer.param_groups[0]['lr']:.4e}]",
             main_process_only=True,
         )
-        if selfpaced_status is not None:
-            if completed_steps > selfpaced_status.selfpaced_history_length:
-                selfpaced_status.log_per_task_weight = selfpaced_status.log_per_task_weight / torch.sum(
-                    selfpaced_status.log_per_task_weight
+        if coba_status is not None:
+            if completed_steps > coba_status.coba_warmup_steps:
+                coba_status.log_per_task_weight = coba_status.log_per_task_weight / torch.sum(
+                    coba_status.log_per_task_weight
                 )
             else:
-                selfpaced_status.log_per_task_weight = torch.ones(len(ID2TASK)) / len(ID2TASK)
+                coba_status.log_per_task_weight = torch.ones(len(ID2TASK)) / len(ID2TASK)
             logger.info(
-                f"[TRAIN][per_task_train_weight={selfpaced_status.log_per_task_weight}]", main_process_only=True
+                f"[TRAIN][per_task_train_weight={coba_status.log_per_task_weight}]", main_process_only=True
             )
         train_log_dict = {"Loss/train": train_loss}
         for i in range(len(ID2TASK)):
             train_log_dict[f"{ID2TASK[i]}_loss/train"] = train_task_loss[i]
-            if selfpaced_status is not None:
-                train_log_dict[f"{ID2TASK[i]}_selfpaced_weight/train"] = selfpaced_status.log_per_task_weight[i].item()
+            if coba_status is not None:
+                train_log_dict[f"{ID2TASK[i]}_coba_weight/train"] = coba_status.log_per_task_weight[i].item()
 
         if self.accelerator.is_main_process:
             write_tensorboard(self.summary_writer, train_log_dict, completed_steps)
 
-        if selfpaced_status is not None:
-            selfpaced_status.log_per_task_weight = torch.zeros(len(ID2TASK))
+        if coba_status is not None:
+            coba_status.log_per_task_weight = torch.zeros(len(ID2TASK))
 
     def accelerate_evaluate(
         self,
@@ -416,18 +416,29 @@ def accelerate_train(self):
         reduce_task_exist = torch.zeros(len(ID2TASK)).to(self.model.device)
         per_task_weight = self.args.task_weights
 
-        if self.args.weighted_loss_mode == "selfpaced":
-            selfpaced_status = SelfpacedStatus(
-                self.args.selfpaced_scale_factor,
-                self.args.selfpaced_interval,
-                self.args.selfpaced_history_length,
-                self.args.selfpaced_sample_valid_num,
+        if self.args.weighted_loss_mode == "coba":
+            self.model.eval()
+            eval_loss, eval_task_loss, _, _, _ = self.accelerate_evaluate(
+                completed_steps,
+                0,
+                min_eval_loss,
+                stall_num,
+                best_step,
+            )
+            self.model.train()
+            coba_status = CoBaStatus(
+                self.args.coba_warmup_steps,
+                self.args.coba_history_length,
+                self.args.coba_tau,
+                self.args.coba_update_interval,
+                self.args.coba_sample_valid_num,
                 self.valid_dataloader,
             )
-            selfpaced_status.sample_valid_batch(self.model, completed_steps)
-            selfpaced_status.valid_iterator = iter(selfpaced_status.valid_dataloader)
+            coba_status.valid_task_loss_begining = eval_task_loss.clone().to(self.model.device)
+            coba_status.sample_valid_batch(self.model, completed_steps)
+            logger.info(f"valid_task_loss: {coba_status.valid_task_loss_accumulated}", main_process_only=True)
         else:
-            selfpaced_status = None
+            coba_status = None
 
         # Training Loop!
         for epoch in range(starting_epoch, self.args.num_train_epochs):
@@ -463,13 +474,15 @@ def accelerate_train(self):
                     )
 
                     if (
-                        self.args.weighted_loss_mode == "selfpaced"
-                        and step % self.args.gradient_accumulation_steps == 0
-                        and completed_steps % self.args.selfpaced_interval == 0
-                        and completed_steps >= self.args.selfpaced_history_length
+                        self.args.weighted_loss_mode == "coba"
+                        and self.accelerator.sync_gradients
+                        and completed_steps % self.args.coba_update_interval == 0
+                        and completed_steps >= self.args.coba_warmup_steps
                     ):
-                        per_task_weight = selfpaced_status.compute_per_task_weight(completed_steps=completed_steps)
-                        selfpaced_status.log_per_task_weight += per_task_weight
+                        with torch.no_grad():
+                            per_task_weight = coba_status.compute_per_task_weight(completed_steps=completed_steps)
+                            coba_status.log_per_task_weight += per_task_weight
+                            # logger.info(f'per_task_weight: {per_task_weight}', main_process_only=True)
 
                     # loss
                     loss, task_loss, _ = loss_func_mft(
@@ -524,11 +537,12 @@ def accelerate_train(self):
                     # If the accelerator has performed an optimization step behind the scenes, thus a completed_step done.
                     if self.accelerator.sync_gradients:
                         if (
-                            self.args.weighted_loss_mode == "selfpaced"
-                            and completed_steps % self.args.selfpaced_interval == 0
+                            self.args.weighted_loss_mode == "coba"
+                            and completed_steps % self.args.coba_update_interval == 0
                             and completed_steps >= 1
                         ):
-                            selfpaced_status.sample_valid_batch(self.model, completed_steps)
+                            coba_status.sample_valid_batch(self.model, completed_steps)
+                            # logger.info(f"valid_task_loss: {coba_status.valid_task_loss_accumulated}", main_process_only=True)
 
                         # progress_bar.update(1)
                         completed_steps += 1
@@ -542,7 +556,7 @@ def accelerate_train(self):
                                 reduce_task_loss,
                                 reduce_task_exist,
                                 completed_steps,
-                                selfpaced_status,
+                                coba_status,
                             )
                             # reset reduce_loss
                             reduce_loss = torch.tensor(0.0).to(self.model.device)
diff --git a/mftcoder_accelerate/src/pefts/mft_accelerate.py b/mftcoder_accelerate/src/pefts/mft_accelerate.py
@@ -217,7 +217,7 @@ def prepare_args():
     # generate TASK2ID, ID2TASK
     generate_task_id(args.data_paths)
 
-    if args.weighted_loss_mode == "selfpaced":
+    if args.weighted_loss_mode == "coba":
         args.task_weights = [1.0] * len(ID2TASK)
     elif args.task_weights is not None:
         args.task_weights = [float(wt) for wt in args.task_weights[1:-1].split(",")]
diff --git a/mftcoder_accelerate/src/pefts/mft_arguments.py b/mftcoder_accelerate/src/pefts/mft_arguments.py
@@ -141,14 +141,16 @@ class MftTrainArgs:
     # if dynamic padding
     use_dynamic_padding: bool = True
 
-    # interval of update per task train weight in selfpaced
-    selfpaced_interval: int = 1
-    # history length of sample valid loss used to fit the slope curve in selfpaced
-    selfpaced_history_length: int = 100
-    # the number of mini valid batches sampled at each interval
-    selfpaced_sample_valid_num: int = 1
-    # scale factor before softmax
-    selfpaced_scale_factor: int = 50
+    # warm-up steps for CoBa, recommand the number of valid batches
+    coba_warmup_steps: int = 100
+    # history length of sample valid loss used to fit the slope curve in CoBa
+    coba_history_length: int = 200
+    # temperature for divergence factor in CoBa
+    coba_tau: int = 5
+    # iteration interval of update per task train weight in CoBa
+    coba_update_interval: int = 1
+    # the number of mini valid batches sampled at each updated iteration interval
+    coba_sample_valid_num: int = 1
 
     # ATTENTION_CLASSES = { "eager": Normal Attention, "flash_attention_2": FlashAttention2}
     attn_implementation: str = "flash_attention_2"
diff --git a/mftcoder_accelerate/src/pefts/mft_trainer.py b/mftcoder_accelerate/src/pefts/mft_trainer.py
diff --git a/mftcoder_accelerate/src/utils/loss_utils.py b/mftcoder_accelerate/src/utils/loss_utils.py
