project-llm-trainer 0.7.9__py3-none-any.whl → 0.8.2__py3-none-any.whl

llm_trainer/dpo_trainer.py

@@ -70,12 +70,12 @@ class DPOTrainer(Trainer):
 
         return criterion, None
 
-    def _convert_train_args(self) -> Tuple[dict, dict, dict, bool]:
+    def _convert_train_args(self) -> Tuple[dict, dict, dict]:
         dpo_collate_fn = get_dpo_collate_fn(self.train_config.mask_prompt)
-        parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim = super()._convert_train_args()
+        parallel_kwargs, data_loader_kwargs, sampler_kwargs = super()._convert_train_args()
         data_loader_kwargs.update({"collate_fn": dpo_collate_fn})
 
-        return parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim
+        return parallel_kwargs, data_loader_kwargs, sampler_kwargs
 
     def _create_dataset(self, file_idx) -> Tuple[Dataset, str]:
         file_path = self.train_config.file_dataset[file_idx]

llm_trainer/grpo_trainer.py

@@ -82,11 +82,11 @@ class GRPOTrainer(Trainer):
 
         return criterion, None
 
-    def _convert_train_args(self) -> Tuple[dict, dict, dict, bool]:
-        parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim = super()._convert_train_args()
+    def _convert_train_args(self) -> Tuple[dict, dict, dict]:
+        parallel_kwargs, data_loader_kwargs, sampler_kwargs = super()._convert_train_args()
         data_loader_kwargs.update({"collate_fn": lambda x: x})
 
-        return parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim
+        return parallel_kwargs, data_loader_kwargs, sampler_kwargs
 
     def _create_dataset(self, file_idx) -> Tuple[Dataset, str]:
         file_path = self.train_config.file_dataset[file_idx]

llm_trainer/sft_trainer.py

@@ -23,12 +23,12 @@ class SFTTrainer(Trainer):
         )
         self.packed_sequences = False
 
-    def _convert_train_args(self) -> Tuple[dict, dict, dict, bool]:
+    def _convert_train_args(self) -> Tuple[dict, dict, dict]:
         sft_collate_fn = get_sft_collate_fn(self.train_config.mask_prompt)
-        parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim = super()._convert_train_args()
+        parallel_kwargs, data_loader_kwargs, sampler_kwargs = super()._convert_train_args()
         data_loader_kwargs.update({"collate_fn": sft_collate_fn})
 
-        return parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim
+        return parallel_kwargs, data_loader_kwargs, sampler_kwargs
 
     def _create_dataset(self, file_idx) -> Tuple[Dataset, str]:
         file_path = self.train_config.file_dataset[file_idx]

llm_trainer/train_configs.py

@@ -107,7 +107,8 @@ class DataLoaderConfig:
 
 
 @dataclass(kw_only=True)
-class LrConfig:
+class OptimConfig:
+    optim_type: str = 'adam' # or 'lion'
     enable_lr_scheduler: bool = False
     initial_lr: float
     weight_decay: float = 0.1
@@ -195,8 +196,8 @@ class TrainConfig:
                 grpo训练时不生效该配置！
             eval_batch_interval (`int`, default is 100):
                 每隔多少个batch进行模型eval
-            lr_config (`LrConfig`):
-                lr配置项
+            optim_config (`OptimConfig`):
+                optim配置项
             data_loader_config: (`DataLoaderConfig`):
                 data loader配置项
             kd_config: (`KDConfig`, *Optional*, default is None):
@@ -213,7 +214,7 @@ class TrainConfig:
     image_tags_file_dataset: Optional[FileDataset] = None
 
     loss_config: LossConfig = field(default_factory=LossConfig)
-    lr_config: LrConfig = field(default_factory=LrConfig)
+    optim_config: OptimConfig = field(default_factory=OptimConfig)
 
     ds_config: DsConfig = field(default_factory=DsConfig)
 

llm_trainer/trainer.py

@@ -77,19 +77,15 @@ class Trainer:
         if self.eval_image_tags:
             assert len(self.eval_prompts) == len(self.eval_image_tags)
 
-        parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim = self._convert_train_args()
-        self.parallel_kwargs = parallel_kwargs
-        self.data_loader_kwargs: dict[str, Any] = data_loader_kwargs
-        self.sampler_kwargs: dict[str, Any] = sampler_kwargs
-
+        self.parallel_kwargs, self.data_loader_kwargs, self.sampler_kwargs = self._convert_train_args()
         # initialize a GradScaler. If enabled=False scaler is a no-op
         self.scalar = torch.GradScaler(enabled=TrainerTools().use_amp)
 
         # 注意：学习率要根据GPU的数量进行倍增：
         # 在训练的过程中，损失梯度决定下降的方向，学习率决定下降的步长。如果有两块gpu，前进的综合步长为：平均学习率*2
-        initial_lr = train_config.lr_config.initial_lr
+        initial_lr = train_config.optim_config.initial_lr
 
-        self.train_model, self.optimizer = self._init_train_model_and_optim(initial_lr, parallel_kwargs, use_ds_optim)
+        self.train_model, self.optimizer = self._init_train_model_and_optim(initial_lr)
         self.lr_scheduler = self._init_lr_scheduler(initial_lr)
 
         self.criterion, self.kd_loss = self._init_loss()
@@ -127,12 +123,7 @@ class Trainer:
         freeze_llm_model = self.train_config.freeze_llm_model
         return model.parameters() if not freeze_llm_model else filter(lambda p: p.requires_grad, model.parameters())
 
-    def _init_train_model_and_optim(
-            self,
-            initial_lr: float,
-            parallel_kwargs: dict,
-            use_ds_optim: bool
-    ):
+    def _init_train_model_and_optim(self, initial_lr: float):
         model = self._new_model(self.train_config)
 
         if self.train_config.init_state_dict:
@@ -161,34 +152,58 @@ class Trainer:
             total_size_mb = total_size_bytes / (1024 * 1024)
             log(f"Total size of the model: {total_size_mb:.2f} MB")
 
-        if use_ds_optim:
-            import deepspeed
-            origin_optim = deepspeed.ops.adam.DeepSpeedCPUAdam(
-                self._get_trainable_params(model),
-                lr=initial_lr,
-                weight_decay=self.train_config.lr_config.weight_decay
-            )
-        else:
-            origin_optim = torch.optim.AdamW(
-                self._get_trainable_params(model),
-                lr=initial_lr,
-                weight_decay=self.train_config.lr_config.weight_decay
-            )
         model, optim = TrainerTools().parallel.process(
             model=model,
-            optimizer=origin_optim,
-            kwargs=parallel_kwargs
+            optimizer=self._get_optim(model, initial_lr),
+            kwargs=self.parallel_kwargs
         )
 
         return model, optim
 
+    def _get_optim(self, model, initial_lr):
+        optimizer = None
+
+        if isinstance(TrainerTools().parallel, DsParallel) and self.parallel_kwargs:
+            import deepspeed
+            if ('zero_optimization' in self.parallel_kwargs
+                    and 'offload_optimizer' in self.parallel_kwargs['zero_optimization']
+                    and self.parallel_kwargs['zero_optimization']['offload_optimizer']['device'] == 'cpu'):
+                # offline optimizer to cpu
+                # 不能使用 deepspeed.ops.lion.cpu_lion.DeepSpeedCPULion???
+                # 所以，这里忽略lion判断
+                optimizer = deepspeed.ops.adam.DeepSpeedCPUAdam
+                if self.train_config.optim_config.optim_type == 'lion':
+                    log('When set offload_optimizer, lion optim is unsupported, so set optim to adam!!!!!')
+            else:
+                if self.train_config.optim_config.optim_type == 'lion':
+                    optimizer = deepspeed.ops.lion.FusedLion
+                else:
+                    optimizer = deepspeed.ops.adam.FusedAdam
+
+        if not optimizer:
+            if self.train_config.optim_config.optim_type == 'lion':
+                try:
+                    import lion_pytorch
+                except:
+                    raise Exception('lion is not detected, please use `pip3 install lion_pytorch` to install or set optim_type to adam')
+
+                optimizer = lion_pytorch.Lion
+            else:
+                optimizer = torch.optim.AdamW
+
+        return optimizer(
+            self._get_trainable_params(model),
+            lr=initial_lr,
+            weight_decay=self.train_config.optim_config.weight_decay
+        )
+
     def _init_lr_scheduler(self, initial_lr: float) -> LRScheduler:
-        if self.train_config.lr_config.enable_lr_scheduler:
-            warmup_iters = self.train_config.lr_config.warmup_iters
-            min_lr = self.train_config.lr_config.min_lr
-            max_lr = self.train_config.lr_config.max_lr
-            cosine_annealing_period = self.train_config.lr_config.cosine_annealing_period
-            cosine_annealing_period_mul = self.train_config.lr_config.cosine_annealing_period_mul
+        if self.train_config.optim_config.enable_lr_scheduler:
+            warmup_iters = self.train_config.optim_config.warmup_iters
+            min_lr = self.train_config.optim_config.min_lr
+            max_lr = self.train_config.optim_config.max_lr
+            cosine_annealing_period = self.train_config.optim_config.cosine_annealing_period
+            cosine_annealing_period_mul = self.train_config.optim_config.cosine_annealing_period_mul
 
             return WarmupCosineAnnealingLRScheduler(
                 optimizer=self.optimizer,
@@ -220,9 +235,8 @@ class Trainer:
 
         return criterion, kd_loss
 
-    def _convert_train_args(self) -> Tuple[dict, dict, dict, bool]:
+    def _convert_train_args(self) -> Tuple[dict, dict, dict]:
         parallel_kwargs: Optional[Dict[str, Any]] = None
-        use_ds_optim: bool = False
         if isinstance(TrainerTools().parallel, DsParallel) and self.train_config.ds_config:
             parallel_kwargs = {
                 'gradient_accumulation_steps': 1,
@@ -253,7 +267,6 @@ class Trainer:
                             "device": zero_config.offload_optimizer.device,
                             "pin_memory": zero_config.offload_optimizer.pin_memory
                         }
-                        use_ds_optim = True
                     if zero_config.offload_param is not None:
                         zero_optimization['offload_param'] = {
                             "device": zero_config.offload_param.device,
@@ -328,10 +341,10 @@ class Trainer:
             "drop_last": dataloader_args.data_loader_drop_last,
         }
 
-        return parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim
+        return parallel_kwargs, data_loader_kwargs, sampler_kwargs
 
     def _init_ref_model_args(self) -> dict:
-        parallel_kwargs = copy.deepcopy(self.parallel_kwargs)
+        parallel_kwargs = copy.deepcopy(self.parallel_kwargs) if self.parallel_kwargs else None
 
         if parallel_kwargs and isinstance(TrainerTools().parallel, DsParallel):
             # reference to https://github.com/huggingface/trl/blob/main/trl/models/utils.py:prepare_deepspeed
@@ -435,7 +448,7 @@ class Trainer:
         exception_file = e.__traceback__.tb_frame.f_globals["__file__"]
         exception_line = e.__traceback__.tb_lineno
         log_msg = f"epoch: {epoch}, batch: {batch}, {e} at {exception_file} line {exception_line}\n"
-        log(log_msg, f'{log_dir}log.txt')
+        log(log_msg, f'{log_dir}exception.txt')
 
         raise e
 

project_llm_trainer-0.8.2.data/scripts/ddp_train

@@ -0,0 +1,21 @@
+#!python
+
+if __name__ == '__main__':
+    import os, sys
+    arguments = sys.argv[1:]
+    # file_name
+    run_file_name = arguments[0]
+
+    extra_args = ''
+    if len(arguments) > 1:
+        extra_args = f"{' '.join(arguments[1:])} "
+
+    os.environ['PARALLEL_TYPE'] = 'ddp'
+
+    if len(extra_args) == 0:
+        extra_args = '--standalone --nproc_per_node=gpu '
+
+    command = f'torchrun {extra_args}{run_file_name}'
+
+    print(f'run command {command}')
+    os.system(command)

project_llm_trainer-0.8.2.data/scripts/ds_train

@@ -0,0 +1,17 @@
+#!python
+
+if __name__ == '__main__':
+    import os, sys
+    arguments = sys.argv[1:]
+    # file_name
+    run_file_name = arguments[0]
+
+    extra_args = ''
+    if len(arguments) > 1:
+        extra_args = f"{' '.join(arguments[1:])} "
+
+    os.environ['PARALLEL_TYPE'] = 'ds'
+    command = f'deepspeed {extra_args}{run_file_name}'
+
+    print(f'run command {command}')
+    os.system(command)

{project_llm_trainer-0.7.9.data → project_llm_trainer-0.8.2.data}/scripts/smart_train

@@ -7,18 +7,9 @@ if __name__ == '__main__':
     # file name
     run_file_name = arguments[0]
 
-    # cuda_visible_devive
+    extra_args = ''
     if len(arguments) > 1:
-        # 0,1,2,3
-        cuda_visible_devive = arguments[1]
-    else:
-        cuda_visible_devive = None
-
-    # cuda location
-    if len(arguments) > 2:
-        cuda_loc = arguments[2]
-    else:
-        cuda_loc = 'localhost'
+        extra_args = f"{' '.join(arguments[1:])} "
 
     try:
         import deepspeed
@@ -33,12 +24,12 @@ if __name__ == '__main__':
     os.environ['PARALLEL_TYPE'] = parallel_type
 
     if parallel_type == 'ds':
-        cuda_ctrl = f' --include {cuda_loc}:{cuda_visible_devive}' if cuda_visible_devive else ''
-        command = f'deepspeed{cuda_ctrl} {run_file_name}'
+        command = f'deepspeed {extra_args}{run_file_name}'
     elif parallel_type == 'ddp':
-        if cuda_visible_devive:
-            os.environ['CUDA_VISIBLE_DEVICES'] = cuda_visible_devive
-        command = f'torchrun --standalone --nproc_per_node=gpu {run_file_name}'
+        if len(extra_args) == 0:
+            extra_args = '--standalone --nproc_per_node=gpu '
+
+        command = f'torchrun {extra_args}{run_file_name}'
     else:
         command = f'python3 {run_file_name}'
 

{project_llm_trainer-0.7.9.dist-info → project_llm_trainer-0.8.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: project_llm_trainer
-Version: 0.7.9
+Version: 0.8.2
 Summary: LLM and VLM trainer
 Author: qibin
 Author-email: qibin0506@gmail.com

{project_llm_trainer-0.7.9.dist-info → project_llm_trainer-0.8.2.dist-info}/RECORD

@@ -1,11 +1,11 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
 llm_trainer/checkpoint.py,sha256=X5ZeUtJlxVz7pnWQLaS-y7UIZOaOAnZTt2L8rSAPzUs,4428
 llm_trainer/dataset.py,sha256=UL3fGeM4XSlyNQRZH-139u3LujqAQx3YyaxNRewk6LE,8935
-llm_trainer/dpo_trainer.py,sha256=Bgds18UWFhzf_UNCFN-iBCdhKf9pcXJBFPEc32oJeXA,13354
+llm_trainer/dpo_trainer.py,sha256=Qi7WKhFO4fdnj9W8BNIF_so6-F8g_YKUoPU9sNjWK_M,13320
 llm_trainer/ds_checkpoint.py,sha256=X2IWgpgi0yOtogph7n6DEwvK_0Ceb7juu1WMutv3HSk,2270
 llm_trainer/eval.py,sha256=ZyUfSo2Q8P-lrCdPEnGkoo5pGubd0AabREK5eMISRII,1109
 llm_trainer/generate_utils.py,sha256=8K3YFbp7IF_lCkmkzjHhqTW26EBFb2AilQmarVcfMvs,15001
-llm_trainer/grpo_trainer.py,sha256=MXnP8Kc9CQJw0CB3uMbHxIYwvpuujai4hgbbpUut_K4,16808
+llm_trainer/grpo_trainer.py,sha256=3CcV-cuyV4ZUTymN9vz3au4uf3gZdyo8SGgSj2NEofs,16774
 llm_trainer/log.py,sha256=XwychwKF6gvFPhthCIZCAEUZ0G3DY3fiQrOHqPWsxz0,463
 llm_trainer/loss.py,sha256=RhTxftLMj1Tqc5pkUvJiZumfbMEPWL8GBGxdTfQggmk,6744
 llm_trainer/parallel.py,sha256=yjStV21DJ26yM8-0O6GTMxdFAcyShY5GsQWSZmbI7HU,4543
@@ -14,20 +14,20 @@ llm_trainer/parallel_ds.py,sha256=oy8RRxHud3rACWubFlJqqd0pjPEQhKeAPGPQUSdJX2c,11
 llm_trainer/parallel_none.py,sha256=TG6Pm829Dg-yQu-97O-EHV3FCARBlNcP47KkGFAs16E,676
 llm_trainer/partition_utils.py,sha256=eEYNhfEIF4hGzZ3OLa6sEBIECz261drptEz_n7fZYtk,8396
 llm_trainer/scheduler.py,sha256=LAI_0VxClsIQkix0bRoduRD4vPfVuIZDhZgTAT_KK8k,4901
-llm_trainer/sft_trainer.py,sha256=LudTRIaqLQYy6ym6jjMX7v9xtFBJelrR3nnPCwb48nM,1821
+llm_trainer/sft_trainer.py,sha256=rSOGZx53jMgOuJdztfxQASYJ62uD0dVaih4IAnSwGBc,1787
 llm_trainer/tokenizer.py,sha256=0-xQCMz1xiPTDAZiYsVsiECSoZ_1eIvW9XsZOoFfakQ,7250
 llm_trainer/tools.py,sha256=5op5qrjjkK-Lr9oes5VxIVnOVYOYGoAdlIJq9mPUf64,2637
-llm_trainer/train_configs.py,sha256=N3ykM1uaLHcSNRC8ErYIxp9VYhSP7voJyAP-2D4ZJe0,7574
-llm_trainer/trainer.py,sha256=jS31zEXIIj9BoPTPlmaGYq61x72HGCjKfS2u3_gOkDk,27924
+llm_trainer/train_configs.py,sha256=pPZkbliRdTnWSv3TUuTM23x9RDdMhGSPrxbNAyzDklY,7636
+llm_trainer/trainer.py,sha256=diP-1suOf2U5dY_R8QH5arAx4MgBrKW-GBQ2_ScGNM8,28799
 llm_trainer/utils.py,sha256=xC5plG-8-_Al5yIF5xIU5lroOcBBk98TEhtUJrazZPE,12305
-project_llm_trainer-0.7.9.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.7.9.data/scripts/ddp_train,sha256=Z-309mM56CN0m3bxoeC5us4LUuwuNnoiOm3-fDdLMjQ,566
-project_llm_trainer-0.7.9.data/scripts/ds_train,sha256=tME0xmMdX1D9XuVo07D9dilW5VIWavBS3UK9DoY67WI,709
-project_llm_trainer-0.7.9.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.7.9.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.7.9.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.7.9.data/scripts/smart_train,sha256=3oLIDuuqb4U4TU1lXy9V8lw_0gIf7i8tGsxlQ_s6bro,1220
-project_llm_trainer-0.7.9.dist-info/METADATA,sha256=mDGLc1BjmIlOPz85JYB5bFnlXJgJ5VaNesW4z0HDZCA,195
-project_llm_trainer-0.7.9.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.7.9.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.7.9.dist-info/RECORD,,
+project_llm_trainer-0.8.2.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.8.2.data/scripts/ddp_train,sha256=eZSud6KYQAoKLsYB5QB-FI2zq5AZm6Apq1azKdupV3o,477
+project_llm_trainer-0.8.2.data/scripts/ds_train,sha256=41q4rOxwbvZDUY0FDdAIpG13PEaUWBpthhvFvww8uOc,388
+project_llm_trainer-0.8.2.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.8.2.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.8.2.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.8.2.data/scripts/smart_train,sha256=N8dp2n7k6bghGczedBVwOdtf1O66oM_cNPh9QmZt0bM,914
+project_llm_trainer-0.8.2.dist-info/METADATA,sha256=XlNe-d24OrjYkzrJMiQCjiZPT70QOFRd4K2XrVDWZiY,195
+project_llm_trainer-0.8.2.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.8.2.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.8.2.dist-info/RECORD,,

project_llm_trainer-0.7.9.data/scripts/ddp_train

@@ -1,24 +0,0 @@
-#!python
-
-if __name__ == '__main__':
-    import os, sys
-    arguments = sys.argv[1:]
-    # file_name
-    run_file_name = arguments[0]
-
-    # cuda_visible_devive
-    if len(arguments) > 1:
-        # 0,1,2,3
-        cuda_visible_devive = arguments[1]
-    else:
-        cuda_visible_devive = None
-
-    os.environ['PARALLEL_TYPE'] = 'ddp'
-
-    if cuda_visible_devive:
-        os.environ['CUDA_VISIBLE_DEVICES'] = cuda_visible_devive
-
-    command = f'torchrun --standalone --nproc_per_node=gpu {run_file_name}'
-
-    print(f'run command {command}')
-    os.system(command)

project_llm_trainer-0.7.9.data/scripts/ds_train

@@ -1,30 +0,0 @@
-#!python
-
-if __name__ == '__main__':
-    import os, sys
-    arguments = sys.argv[1:]
-    # file_name
-    run_file_name = arguments[0]
-
-    # cuda_visible_devive
-    if len(arguments) > 1:
-        # 0,1,2,3
-        cuda_visible_devive = arguments[1]
-
-        # cuda location
-        if len(arguments) > 2:
-            cuda_loc = arguments[2]
-        else:
-            cuda_loc = 'localhost'
-    else:
-        cuda_visible_devive = None
-        cuda_loc = None
-
-    os.environ['PARALLEL_TYPE'] = 'ds'
-
-    cuda_ctrl = f' --include {cuda_loc}:{cuda_visible_devive}' if cuda_visible_devive else ''
-
-    command = f'deepspeed{cuda_ctrl} {run_file_name}'
-
-    print(f'run command {command}')
-    os.system(command)