project-llm-trainer 0.5.2__py3-none-any.whl → 0.5.4__py3-none-any.whl

llm_trainer/checkpoint.py

@@ -14,17 +14,14 @@ DEFAULT_CHECKPOINT_NAME = "checkpoint.pth"
 
 def save_checkpoint(
         model: nn.Module,
-        optimizer: Optional[Optimizer] = None,
-        suffix: Optional[str] = None
+        optimizer: Optional[Optimizer] = None
 ):
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import save_ds_checkpoint
-        save_ds_checkpoint(model, suffix)
+        save_ds_checkpoint(model)
     else:
         if TrainerTools().parallel.is_main_process:
             checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
-            if suffix:
-                checkpoint_name = f"{checkpoint_name}_{suffix}"
 
             raw_model = model if not isinstance(model, DDP) else model.module
             ckpt = {'model_state_dict': raw_model.state_dict()}
@@ -37,29 +34,27 @@ def save_checkpoint(
 
 def save_best_checkpoint(
         current_loss: float,
-        last_best_checkpoint_loss: float,
-        suffix: Optional[str] = None
+        last_best_checkpoint_loss: Optional[float] = None
 ) -> bool:
-    need_replace = current_loss <= last_best_checkpoint_loss
+    need_replace = not last_best_checkpoint_loss or current_loss <= last_best_checkpoint_loss
     if need_replace and TrainerTools().parallel.is_main_process:
         if isinstance(TrainerTools().parallel, DsParallel):
-            checkpoint_name = os.environ.get('DIST_CHECKPOINT_DIR', 'checkpoint')
-            if suffix:
-                checkpoint_name = f"{checkpoint_name}_{suffix}"
+            checkpoint_dir = os.environ.get('DIST_CHECKPOINT_DIR', 'checkpoint')
 
-            best_checkpoint_name = f'best_{checkpoint_name}'
-            if not os.path.exists(best_checkpoint_name):
-                os.makedirs(best_checkpoint_name)
+            if checkpoint_dir.endswith('/'):
+                best_checkpoint_dir = f'{checkpoint_dir[:-1]}_best'
+            else:
+                best_checkpoint_dir = f'{checkpoint_dir}_best'
 
-            if os.path.exists(checkpoint_name):
-                shutil.rmtree(best_checkpoint_name)
-                shutil.copytree(checkpoint_name, best_checkpoint_name)
+            if not os.path.exists(best_checkpoint_dir):
+                os.makedirs(best_checkpoint_dir)
+
+            if os.path.exists(checkpoint_dir):
+                shutil.rmtree(best_checkpoint_dir)
+                shutil.copytree(checkpoint_dir, best_checkpoint_dir)
         else:
             checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
-            if suffix:
-                checkpoint_name = f"{checkpoint_name}_{suffix}"
-
-            best_checkpoint_name = f'best_{checkpoint_name}'
+            best_checkpoint_name = f'{checkpoint_name}_best'
 
             if os.path.exists(checkpoint_name):
                 if os.path.exists(best_checkpoint_name):
@@ -75,16 +70,13 @@ def load_checkpoint(
         model: nn.Module,
         optimizer: Optional[Optimizer] = None,
         device: Optional[Union[torch.device, str]] = None,
-        load_module_only: bool = False,
-        suffix: Optional[str] = None
+        load_module_only: bool = False
 ):
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import load_ds_checkpoint
-        load_ds_checkpoint(model, load_module_only=load_module_only, suffix=suffix)
+        load_ds_checkpoint(model, load_module_only=load_module_only)
     else:
         checkpoint_name = os.environ.get('CHECKPOINT_NAME', DEFAULT_CHECKPOINT_NAME)
-        if suffix:
-            checkpoint_name = f"{checkpoint_name}_{suffix}"
 
         state_dict = torch.load(checkpoint_name, weights_only=True, map_location=device)
         raw_model = model.module if isinstance(model, DDP) else model
@@ -96,14 +88,13 @@ def load_checkpoint(
 
 def load_checkpoint_for_eval(
         model: nn.Module,
-        device: Optional[Union[torch.device, str]] = None,
-        suffix: Optional[str] = None
+        device: Optional[Union[torch.device, str]] = None
 ):
     if isinstance(TrainerTools().parallel, DsParallel):
         from .ds_checkpoint import load_ds_checkpoint_for_eval
         load_ds_checkpoint_for_eval(model)
     else:
-        load_checkpoint(model, None, device, suffix=suffix)
+        load_checkpoint(model, None, device)
 
 
 def save_steps(global_steps: int, lr_scheduler: Optional[LRScheduler] = None):

llm_trainer/dpo_trainer.py

@@ -141,7 +141,7 @@ class DPOTrainer(Trainer):
         skipping_train = False
 
         current_loss: float = 0.0
-        last_best_checkpoint_loss: float = 0.0
+        last_best_checkpoint_loss: Optional[float] = None
 
         aux_loss_coef = self.train_config.loss_config.aux_loss_coef
 

llm_trainer/ds_checkpoint.py

@@ -1,5 +1,4 @@
 import os
-from typing import Optional
 from glob import glob
 import shutil
 from torch import nn
@@ -17,14 +16,9 @@ load_state_dict_from_zero_checkpoint	从 ZeRO 检查点加载模型和优化器
 convert_zero_checkpoint_to_fp32_state_dict	将 ZeRO 检查点转换为独立的 FP32 状态字典文件	否	是	创建可移植的 FP32 权重文件，用于部署、分享等
 """
 
-def save_ds_checkpoint(
-        model: nn.Module,
-        suffix: Optional[str] = None
-):
+def save_ds_checkpoint(model: nn.Module):
     assert isinstance(model, DeepSpeedEngine)
     ckpt_dir = os.environ.get('DIST_CHECKPOINT_DIR', 'checkpoint')
-    if suffix:
-        ckpt_dir = f"{ckpt_dir}_{suffix}"
 
     try:
         # 包括model、optimizer等状态
@@ -44,13 +38,10 @@ def save_ds_checkpoint(
 
 def load_ds_checkpoint(
         model: nn.Module,
-        load_module_only: bool = False,
-        suffix: Optional[str] = None
+        load_module_only: bool = False
 ):
     assert isinstance(model, DeepSpeedEngine)
     ckpt_dir = os.environ.get('DIST_CHECKPOINT_DIR', 'checkpoint')
-    if suffix:
-        ckpt_dir = f"{ckpt_dir}_{suffix}"
 
     # 包括model、optimizer等状态
     if os.path.exists(ckpt_dir):

llm_trainer/grpo_trainer.py

@@ -283,7 +283,7 @@ class GRPOTrainer(Trainer):
         skipping_train = False
 
         current_loss: float = 0.0
-        last_best_checkpoint_loss: float = 0.0
+        last_best_checkpoint_loss: Optional[float] = None
 
         aux_loss_coef = self.train_config.loss_config.aux_loss_coef
 

llm_trainer/trainer.py

@@ -469,7 +469,7 @@ class Trainer:
         skipping_train = False
 
         current_loss: float = 0.0
-        last_best_checkpoint_loss: float = 0.0
+        last_best_checkpoint_loss: Optional[float] = None
 
         for epoch in range(self.train_config.n_epochs):
             self.train_model.train()

{project_llm_trainer-0.5.2.dist-info → project_llm_trainer-0.5.4.dist-info}/METADATA

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

{project_llm_trainer-0.5.2.dist-info → project_llm_trainer-0.5.4.dist-info}/RECORD

@@ -1,11 +1,11 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=wC4GdIY2HAnxGHzUND5Yq-J_ynhDPT7A_2sXlWRElBc,4647
+llm_trainer/checkpoint.py,sha256=UVjOaDsiSIzRJ5VJZib6iXrdKv2A7K_gtJw3a9wNyoM,4293
 llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
-llm_trainer/dpo_trainer.py,sha256=xfYXlLA5TbqPKCUbk5_V79TreEh-dnLMaN72a3-Tdzg,11860
-llm_trainer/ds_checkpoint.py,sha256=wz48HoLBBt8QGO1tXfvJwrXoiGtPG_gjwHfEqARllso,2175
+llm_trainer/dpo_trainer.py,sha256=3hCjX06W6nt-8lio0YyGzXdI_saa5QlypXehQTtacO4,11871
+llm_trainer/ds_checkpoint.py,sha256=D092fkS1Up4QmpV9YCpqbSzfX_caCAeX-UiOrhOE1I8,1947
 llm_trainer/eval.py,sha256=fjASCILU3fSPJxo9cP3rIXEEnkc5ZlUyHqXlZtUiHrw,888
 llm_trainer/generate_utils.py,sha256=CbJ3mfAD6DkQ0GUHcJQ1AK02m-ocwmd-BPXEpiwvNNQ,14933
-llm_trainer/grpo_trainer.py,sha256=vTNi3n6R4NbwFh_s8LYN1TWEJm8AW2F5NVJlT5MHxKk,15990
+llm_trainer/grpo_trainer.py,sha256=VltU9AngecJ5PEtpn_ToXfiLIc6VwwgOVwlLh-X2Je8,16001
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
 llm_trainer/loss.py,sha256=NZCQeUXnLSj__mmDflE8g89KgE0emAJXIab0IERCLno,6023
 llm_trainer/parallel.py,sha256=DQu8GqEFxD99HQ6hKuIxxyKi-05dMO33eMhImYlPuOI,4468
@@ -18,16 +18,16 @@ llm_trainer/sft_trainer.py,sha256=gxQA7T1o1QGUsHp2CX1Qb_fO5LppBJuNbc0H4ixCYUA,17
 llm_trainer/tokenizer.py,sha256=SSpgXtb0e1NtQqRW0gCq09TTZi47umggy-Fh5EMHKJg,6708
 llm_trainer/tools.py,sha256=yF17lp6oOfLe2XJeKDQ1juZcbv-6vFamJSLwEeArduA,2975
 llm_trainer/train_configs.py,sha256=c6bgivkkWRYcPD3NzI5uRItAUhZiIBgKVMuMgVFRnFo,7336
-llm_trainer/trainer.py,sha256=g8YUP0FmBP3MGwewyoyOW35p9CY98rS62pzjnOMiWvE,25875
+llm_trainer/trainer.py,sha256=VHnuL8rZgxj2ewBVxmbN0jwVEPRVhbK2V2DYCtm_FxI,25886
 llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
-project_llm_trainer-0.5.2.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.5.2.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.5.2.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.5.2.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.5.2.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.5.2.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.5.2.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.5.2.dist-info/METADATA,sha256=AP3lS957E984t8klMW_Z6VmrIU7-sBtrzAszA6V-KcQ,195
-project_llm_trainer-0.5.2.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.5.2.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.5.2.dist-info/RECORD,,
+project_llm_trainer-0.5.4.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.5.4.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.5.4.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.5.4.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.5.4.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.5.4.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.5.4.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.5.4.dist-info/METADATA,sha256=lZWPvJQFiqZTl9b1FUSFl1Fl6berO7DKGZ-A_7ZkidE,195
+project_llm_trainer-0.5.4.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.5.4.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.5.4.dist-info/RECORD,,