project-llm-trainer 0.9.1__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

@@ -156,14 +156,15 @@ 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
@@ -175,6 +176,7 @@ class Trainer:
                         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
@@ -195,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.1.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.1
+Version: 0.9.2
 Summary: LLM and VLM trainer
 Author: qibin
 Author-email: qibin0506@gmail.com

{project_llm_trainer-0.9.1.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=Wf1JLWYhmxKKef49asxkaR4ZocKO-jlupxksv-BNfZI,29170
+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.1.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.9.1.data/scripts/ddp_train,sha256=eZSud6KYQAoKLsYB5QB-FI2zq5AZm6Apq1azKdupV3o,477
-project_llm_trainer-0.9.1.data/scripts/ds_train,sha256=41q4rOxwbvZDUY0FDdAIpG13PEaUWBpthhvFvww8uOc,388
-project_llm_trainer-0.9.1.data/scripts/plot_loss,sha256=O9ooioAJ-79-X06LosgqF8XOqQe-beRxYm3LsLunmoU,908
-project_llm_trainer-0.9.1.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.9.1.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.9.1.data/scripts/smart_train,sha256=N8dp2n7k6bghGczedBVwOdtf1O66oM_cNPh9QmZt0bM,914
-project_llm_trainer-0.9.1.dist-info/METADATA,sha256=1Rtg7wi3BLy-9cFnNouWF3bVkexiMHeFSyaSmvrOFpA,195
-project_llm_trainer-0.9.1.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.9.1.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.9.1.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,,