project-llm-trainer 0.9.0__py3-none-any.whl → 0.9.2__py3-none-any.whl

llm_trainer/train_configs.py

@@ -1,4 +1,4 @@
-from typing import Optional, Union, Callable, List, Mapping, Any
+from typing import Optional, Union, Callable, List, Mapping, Any, Tuple
 from dataclasses import dataclass, field
 
 import torch
@@ -111,7 +111,8 @@ class OptimConfig:
     optim_type: str = 'adam' # or 'lion'
     enable_lr_scheduler: bool = False
     initial_lr: float
-    weight_decay: float = 0.1
+    weight_decay: Optional[float] = None
+    betas: Optional[Tuple[float, float]] = None
     warmup_iters: Optional[int] = None
     max_lr: Optional[float] = None
     min_lr: Optional[float] = None

llm_trainer/trainer.py

@@ -1,5 +1,7 @@
 from typing import Optional, Tuple, List, Dict, Any
 import copy
+import importlib.metadata
+from packaging import version
 
 import torch
 import torch.distributed as dist
@@ -154,26 +156,30 @@ class Trainer:
 
         model, optim = TrainerTools().parallel.process(
             model=model,
-            optimizer=self._get_optim(model, initial_lr),
+            optimizer=self._config_optim(model, initial_lr),
             kwargs=self.parallel_kwargs
         )
 
         return model, optim
 
-    def _get_optim(self, model, initial_lr):
+    def _config_optim(self, model, initial_lr):
         optimizer = None
+        use_lion_optim = self.train_config.optim_config.optim_type == 'lion'
 
         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!!!!!')
+                    if version.parse(importlib.metadata.version("deepspeed")) >= version.parse('0.17.6'):
+                        optimizer = deepspeed.ops.lion.DeepSpeedCPULion
+                    else:
+                        optimizer = deepspeed.ops.adam.DeepSpeedCPUAdam
+                        use_lion_optim = False
+                        log('When set offload_optimizer, lion optim is unsupported, so set optim to adam!!!!!')
+                else:
+                    optimizer = deepspeed.ops.adam.DeepSpeedCPUAdam
             else:
                 if self.train_config.optim_config.optim_type == 'lion':
                     optimizer = deepspeed.ops.lion.FusedLion
@@ -191,10 +197,26 @@ class Trainer:
             else:
                 optimizer = torch.optim.AdamW
 
+        betas = self.train_config.optim_config.betas
+        weight_decay = self.train_config.optim_config.weight_decay
+
+        if betas is None:
+            if use_lion_optim:
+                betas = (0.95, 0.98)
+            else:
+                betas = (0.9, 0.999)
+
+        if weight_decay is None:
+            if use_lion_optim:
+                weight_decay = 0.015
+            else:
+                weight_decay = 0.01
+
         return optimizer(
             self._get_trainable_params(model),
             lr=initial_lr,
-            weight_decay=self.train_config.optim_config.weight_decay
+            betas=betas,
+            weight_decay=weight_decay
         )
 
     def _init_lr_scheduler(self, initial_lr: float) -> LRScheduler:

{project_llm_trainer-0.9.0.dist-info → project_llm_trainer-0.9.2.dist-info}/METADATA

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

{project_llm_trainer-0.9.0.dist-info → project_llm_trainer-0.9.2.dist-info}/RECORD

@@ -17,17 +17,17 @@ llm_trainer/scheduler.py,sha256=LAI_0VxClsIQkix0bRoduRD4vPfVuIZDhZgTAT_KK8k,4901
 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=u3b3RhZn7S6vm2DKNF5sNRHzp73-eYnActGSlzcIbzs,7638
-llm_trainer/trainer.py,sha256=sz350VaLR7odGm0vYGTGvFWcr5bC7tmwMx7PUY6LA2o,28980
+llm_trainer/train_configs.py,sha256=afXUZ7M_Uoj0B3c2Nwf5xE-Lv7QAZZHTdW8LBw-QeWE,7704
+llm_trainer/trainer.py,sha256=bVghqvQY4bvYAZFPgyh2ywX8WanqAC525Lkg8bNv4FQ,29721
 llm_trainer/utils.py,sha256=xC5plG-8-_Al5yIF5xIU5lroOcBBk98TEhtUJrazZPE,12305
-project_llm_trainer-0.9.0.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.9.0.data/scripts/ddp_train,sha256=eZSud6KYQAoKLsYB5QB-FI2zq5AZm6Apq1azKdupV3o,477
-project_llm_trainer-0.9.0.data/scripts/ds_train,sha256=41q4rOxwbvZDUY0FDdAIpG13PEaUWBpthhvFvww8uOc,388
-project_llm_trainer-0.9.0.data/scripts/plot_loss,sha256=O9ooioAJ-79-X06LosgqF8XOqQe-beRxYm3LsLunmoU,908
-project_llm_trainer-0.9.0.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.9.0.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.9.0.data/scripts/smart_train,sha256=N8dp2n7k6bghGczedBVwOdtf1O66oM_cNPh9QmZt0bM,914
-project_llm_trainer-0.9.0.dist-info/METADATA,sha256=PSaXkuGCDIwRi8gPHD3LEYXmDhmMj9rC_8xlZP5ftxw,195
-project_llm_trainer-0.9.0.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.9.0.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.9.0.dist-info/RECORD,,
+project_llm_trainer-0.9.2.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.9.2.data/scripts/ddp_train,sha256=eZSud6KYQAoKLsYB5QB-FI2zq5AZm6Apq1azKdupV3o,477
+project_llm_trainer-0.9.2.data/scripts/ds_train,sha256=41q4rOxwbvZDUY0FDdAIpG13PEaUWBpthhvFvww8uOc,388
+project_llm_trainer-0.9.2.data/scripts/plot_loss,sha256=O9ooioAJ-79-X06LosgqF8XOqQe-beRxYm3LsLunmoU,908
+project_llm_trainer-0.9.2.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.9.2.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.9.2.data/scripts/smart_train,sha256=N8dp2n7k6bghGczedBVwOdtf1O66oM_cNPh9QmZt0bM,914
+project_llm_trainer-0.9.2.dist-info/METADATA,sha256=hoIO4KbvNU5xaZdzuNljSZcZSb_Iozl_Skp4miE3U6Y,195
+project_llm_trainer-0.9.2.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.9.2.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.9.2.dist-info/RECORD,,