project-llm-trainer 0.5.3__py3-none-any.whl → 0.5.5__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,28 +34,26 @@ 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'{checkpoint_name}_best'
-            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'{checkpoint_name}_best'
 
             if os.path.exists(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

@@ -137,11 +137,12 @@ class DPOTrainer(Trainer):
         # 梯度累积步数
         gradient_accumulation_steps = self.train_config.gradient_accumulation_steps
         global_steps = 0
-        loss_accumulation = 0.0
         skipping_train = False
 
+        loss_accumulation = 0.0
+        batches_accumulated = 0
         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
 
@@ -214,14 +215,15 @@ class DPOTrainer(Trainer):
 
                         loss_accumulation += loss.detach().item()
                         self._backward_loss(loss)
+                        batches_accumulated += 1
 
                         if need_update_grad:
-                            loss_tensor = torch.tensor(loss_accumulation, device=TrainerTools().parallel.device)
+                            loss_tensor = torch.tensor(loss_accumulation * gradient_accumulation_steps / batches_accumulated, device=TrainerTools().parallel.device)
 
                             if TrainerTools().parallel.parallel_train:
                                 dist.all_reduce(loss_tensor, dist.ReduceOp.AVG)
 
-                            final_log_loss = loss_tensor.item()
+                            current_loss = loss_tensor.item()
 
                             # ds模式已经集成gradient_clipping
                             if not isinstance(TrainerTools().parallel, DsParallel) and self.lr_scheduler.can_clip_grad():
@@ -235,10 +237,11 @@ class DPOTrainer(Trainer):
                                 epoch_tag=f'epoch: {epoch}',
                                 file_tag=f'file: {file_idx + 1}/{file_count}',
                                 batch_tag=f'batch: {batch}/{batch_count_per_file}',
-                                loss=final_log_loss
+                                loss=current_loss
                             )
                             # reset to default
                             loss_accumulation = 0.0
+                            batches_accumulated = 0
                     except Exception as e:
                         self._on_exception(e, epoch, batch)
                     finally:

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
 
@@ -345,6 +345,8 @@ class GRPOTrainer(Trainer):
                             if TrainerTools().parallel.parallel_train:
                                 dist.all_reduce(loss, dist.ReduceOp.AVG)
 
+                            current_loss = loss.detach().item()
+
                             # ds模式已经集成gradient_clipping
                             if not isinstance(TrainerTools().parallel, DsParallel) and self.lr_scheduler.can_clip_grad():
                                 # clip grad
@@ -357,7 +359,7 @@ class GRPOTrainer(Trainer):
                                 epoch_tag=f'epoch: {epoch}',
                                 file_tag=f'file: {file_idx + 1}/{file_count}',
                                 batch_tag=f'batch: {batch}/{batch_count_per_file}',
-                                loss=loss.detach().item()
+                                loss=current_loss
                             )
                     except Exception as e:
                         self._on_exception(e, epoch, batch)

llm_trainer/trainer.py

@@ -465,11 +465,12 @@ class Trainer:
         # 梯度累积步数
         gradient_accumulation_steps = self.train_config.gradient_accumulation_steps
         global_steps = 0
-        loss_accumulation = 0.0
         skipping_train = False
 
+        loss_accumulation = 0.0
+        batches_accumulated = 0
         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()
@@ -536,9 +537,10 @@ class Trainer:
 
                         loss_accumulation += loss.detach().item()
                         self._backward_loss(loss)
+                        batches_accumulated += 1
 
                         if need_update_grad:
-                            loss_tensor = torch.tensor(loss_accumulation, device=TrainerTools().parallel.device)
+                            loss_tensor = torch.tensor(loss_accumulation * gradient_accumulation_steps / batches_accumulated, device=TrainerTools().parallel.device)
 
                             if TrainerTools().parallel.parallel_train:
                                 dist.all_reduce(loss_tensor, dist.ReduceOp.AVG)
@@ -561,6 +563,7 @@ class Trainer:
                             )
                             # reset to default
                             loss_accumulation = 0.0
+                            batches_accumulated = 0
                     except Exception as e:
                         self._on_exception(e, epoch, batch)
                     finally:

{project_llm_trainer-0.5.3.dist-info → project_llm_trainer-0.5.5.dist-info}/METADATA

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

{project_llm_trainer-0.5.3.dist-info → project_llm_trainer-0.5.5.dist-info}/RECORD

@@ -1,11 +1,11 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=67q1zDYvcbS2zce1PVt3nmsPzqniHu0f2pI-cyyCkng,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=1A_4QP2_xqM_YeqdXy-0RaMvEL80gim-pgnPQyHww9U,12052
+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=sCYjvksdm9f7TpN23KXuCmua_8VFTZEfVEcflL89P_I,16058
 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=YW59dJWTyQy77cLDGzBHhfinGyfkvmWCkl1SR9hM6a8,26071
 llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
-project_llm_trainer-0.5.3.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.5.3.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.5.3.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.5.3.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.5.3.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.5.3.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.5.3.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.5.3.dist-info/METADATA,sha256=iDB3C1trVLQsnwsRxeFm7Oi2YpNevuX3XO2WZFlL7wg,195
-project_llm_trainer-0.5.3.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.5.3.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.5.3.dist-info/RECORD,,
+project_llm_trainer-0.5.5.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.5.5.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.5.5.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.5.5.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.5.5.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.5.5.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.5.5.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.5.5.dist-info/METADATA,sha256=ajxfapuo4Q2xfdJ3kjZoCzs7Q5ynGp6BssXRFOIbF7Y,195
+project_llm_trainer-0.5.5.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.5.5.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.5.5.dist-info/RECORD,,