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

llm_trainer/checkpoint.py

@@ -1,5 +1,6 @@
 import os
 from typing import Optional, Union, Tuple
+import shutil
 import torch
 from torch import nn
 from torch.optim import Optimizer
@@ -34,6 +35,42 @@ def save_checkpoint(
             torch.save(ckpt, checkpoint_name)
 
 
+def save_best_checkpoint(
+        current_loss: float,
+        last_best_checkpoint_loss: float,
+        suffix: Optional[str] = None
+) -> bool:
+    need_replace = 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}"
+
+            best_checkpoint_name = f'best_{checkpoint_name}'
+            if not os.path.exists(best_checkpoint_name):
+                os.makedirs(best_checkpoint_name)
+
+            if os.path.exists(checkpoint_name):
+                shutil.rmtree(best_checkpoint_name)
+                shutil.copytree(checkpoint_name, best_checkpoint_name)
+        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}'
+
+            if os.path.exists(checkpoint_name):
+                if os.path.exists(best_checkpoint_name):
+                    os.remove(best_checkpoint_name)
+
+                shutil.copy2(checkpoint_name, best_checkpoint_name)
+
+    TrainerTools().parallel.wait()
+    return need_replace
+
+
 def load_checkpoint(
         model: nn.Module,
         optimizer: Optional[Optimizer] = None,

llm_trainer/dpo_trainer.py

@@ -16,6 +16,7 @@ from .partition_utils import sync_model_params
 
 from .checkpoint import (
     save_checkpoint,
+    save_best_checkpoint,
     save_steps,
 )
 
@@ -139,6 +140,9 @@ class DPOTrainer(Trainer):
         loss_accumulation = 0.0
         skipping_train = False
 
+        current_loss: float = 0.0
+        last_best_checkpoint_loss: float = 0.0
+
         aux_loss_coef = self.train_config.loss_config.aux_loss_coef
 
         for epoch in range(self.train_config.n_epochs):
@@ -243,6 +247,9 @@ class DPOTrainer(Trainer):
 
                             if (batch - last_ckpt_batch) >= self.train_config.eval_batch_interval:
                                 save_checkpoint(model=self.train_model, optimizer=self.optimizer)
+                                if save_best_checkpoint(current_loss, last_best_checkpoint_loss):
+                                    last_best_checkpoint_loss = current_loss
+
                                 last_ckpt_batch = batch
                                 self._on_batch_end(tag=f'epoch:{epoch}/batch:{batch}')
 
@@ -252,8 +259,12 @@ class DPOTrainer(Trainer):
 
             # end epoch
             if not skipping_train:
-                save_checkpoint(model=self.train_model, optimizer=self.optimizer)
                 save_steps(global_steps=global_steps, lr_scheduler=self.lr_scheduler)
+
+                save_checkpoint(model=self.train_model, optimizer=self.optimizer)
+                if save_best_checkpoint(current_loss, last_best_checkpoint_loss):
+                    last_best_checkpoint_loss = current_loss
+
                 TrainerTools().parallel.on_epoch_end(epoch)
                 self._on_epoch_end(tag=f'epoch:{epoch}')
 

llm_trainer/grpo_trainer.py

@@ -22,6 +22,7 @@ from .partition_utils import (
 
 from .checkpoint import (
     save_checkpoint,
+    save_best_checkpoint,
     save_steps,
 )
 
@@ -280,6 +281,10 @@ class GRPOTrainer(Trainer):
     def train(self):
         global_steps = 0
         skipping_train = False
+
+        current_loss: float = 0.0
+        last_best_checkpoint_loss: float = 0.0
+
         aux_loss_coef = self.train_config.loss_config.aux_loss_coef
 
         for epoch in range(self.train_config.n_epochs):
@@ -361,6 +366,9 @@ class GRPOTrainer(Trainer):
 
                         if (batch - last_ckpt_batch) >= self.train_config.eval_batch_interval:
                             save_checkpoint(model=self.train_model, optimizer=self.optimizer)
+                            if save_best_checkpoint(current_loss, last_best_checkpoint_loss):
+                                last_best_checkpoint_loss = current_loss
+
                             last_ckpt_batch = batch
                             self._on_batch_end(tag=f'epoch:{epoch}/batch:{batch}')
 
@@ -370,8 +378,12 @@ class GRPOTrainer(Trainer):
 
             # end epoch
             if not skipping_train:
-                save_checkpoint(model=self.train_model, optimizer=self.optimizer)
                 save_steps(global_steps=global_steps, lr_scheduler=self.lr_scheduler)
+
+                save_checkpoint(model=self.train_model, optimizer=self.optimizer)
+                if save_best_checkpoint(current_loss, last_best_checkpoint_loss):
+                    last_best_checkpoint_loss = current_loss
+
                 TrainerTools().parallel.on_epoch_end(epoch)
                 self._on_epoch_end(tag=f'epoch:{epoch}')
 

llm_trainer/scheduler.py

@@ -30,12 +30,12 @@ class WarmupCosineAnnealingLRScheduler(LRScheduler):
             self,
             *,
             optimizer: torch.optim.Optimizer,
+            warmup_iters: int,
             initial_lr: float,
             min_lr: float,
             max_lr: float,
-            warmup_iters: int,
-            period: int, # 每个周期的步数
-            period_mul: int = 1, # 周期长度的倍数
+            cosine_annealing_period: int, # 每个周期的步数
+            cosine_annealing_period_mul: int = 0, # 周期长度的倍数
             need_log: bool = False
     ):
         super().__init__()
@@ -46,8 +46,8 @@ class WarmupCosineAnnealingLRScheduler(LRScheduler):
         self._max_lr = max_lr
         self._warmup_iters = warmup_iters
 
-        self._period = period
-        self._period_mul = period_mul
+        self._cosine_annealing_period = cosine_annealing_period
+        self._cosine_annealing_period_mul = cosine_annealing_period_mul
 
         self.T_cur = 0  # 当前周期内已走过的步数
         self.cycle = 0  # 当前周期编号
@@ -85,7 +85,12 @@ class WarmupCosineAnnealingLRScheduler(LRScheduler):
         return self._steps > self._warmup_iters
 
     def _update_lr(self):
-        if self._steps <= self._warmup_iters:
+        # 如果period_mul是0，则认为没有周期，超过余弦退火总步数，则一直保持最小lr
+        if self._cosine_annealing_period_mul == 0 and self._steps >= self._cosine_annealing_period + self._warmup_iters:
+            lr = self._min_lr
+            for param_group in self._optimizer.param_groups:
+                param_group['lr'] = lr
+        elif self._steps <= self._warmup_iters:
             # Warmup: adjust learning rate linearly
             # (max_lr - initial_lr) / warmup_iters
             lr = self._initial_lr + self._steps * self._lr_increment
@@ -97,7 +102,7 @@ class WarmupCosineAnnealingLRScheduler(LRScheduler):
 
             """每步更新学习率"""
             # 计算当前周期的最大步数
-            T_max = self._period * (self._period_mul ** self.cycle)
+            T_max = self._cosine_annealing_period * (max(self._cosine_annealing_period_mul, 1) ** self.cycle)
 
             # 更新周期状态
             self.T_cur += 1

llm_trainer/train_configs.py

@@ -109,13 +109,13 @@ class DataLoaderConfig:
 @dataclass(kw_only=True)
 class LrConfig:
     enable_lr_scheduler: bool = False
-    initial_lr: Optional[float] = None
+    initial_lr: float
     weight_decay: float = 0.1
+    warmup_iters: Optional[int] = None
     max_lr: Optional[float] = None
     min_lr: Optional[float] = None
-    period: Optional[int] = None
-    period_mul: Optional[int] = None
-    warmup_iters: Optional[int] = None
+    cosine_annealing_period: Optional[int] = None
+    cosine_annealing_period_mul: int = 0
 
 
 @dataclass(kw_only=True)

llm_trainer/trainer.py

@@ -30,6 +30,7 @@ from .scheduler import (
 from .checkpoint import (
     load_checkpoint,
     save_checkpoint,
+    save_best_checkpoint,
     load_steps,
     save_steps,
 )
@@ -172,20 +173,20 @@ class Trainer:
 
     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
-            warmup_iters = self.train_config.lr_config.warmup_iters
-            period = self.train_config.lr_config.period
-            period_mul = self.train_config.lr_config.period_mul
+            cosine_annealing_period = self.train_config.lr_config.cosine_annealing_period
+            cosine_annealing_period_mul = self.train_config.lr_config.cosine_annealing_period_mul
 
             return WarmupCosineAnnealingLRScheduler(
                 optimizer=self.optimizer,
+                warmup_iters=warmup_iters,
                 initial_lr=initial_lr,
                 min_lr=min_lr,
                 max_lr=max_lr,
-                warmup_iters=warmup_iters,
-                period=period,
-                period_mul=period_mul,
+                cosine_annealing_period=cosine_annealing_period,
+                cosine_annealing_period_mul=cosine_annealing_period_mul,
                 need_log=TrainerTools().parallel.is_main_process
             )
 
@@ -467,6 +468,9 @@ class Trainer:
         loss_accumulation = 0.0
         skipping_train = False
 
+        current_loss: float = 0.0
+        last_best_checkpoint_loss: float = 0.0
+
         for epoch in range(self.train_config.n_epochs):
             self.train_model.train()
             file_count = len(self.train_config.file_dataset)
@@ -539,7 +543,7 @@ class Trainer:
                             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():
@@ -553,7 +557,7 @@ class 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
@@ -565,6 +569,9 @@ class Trainer:
 
                             if (batch - last_ckpt_batch) >= self.train_config.eval_batch_interval:
                                 save_checkpoint(model=self.train_model, optimizer=self.optimizer)
+                                if save_best_checkpoint(current_loss, last_best_checkpoint_loss):
+                                    last_best_checkpoint_loss = current_loss
+
                                 last_ckpt_batch = batch
                                 self._on_batch_end(tag=f'epoch:{epoch}/batch:{batch}')
 
@@ -574,8 +581,12 @@ class Trainer:
 
             # end epoch
             if not skipping_train:
-                save_checkpoint(model=self.train_model, optimizer=self.optimizer)
                 save_steps(global_steps=global_steps, lr_scheduler=self.lr_scheduler)
+
+                save_checkpoint(model=self.train_model, optimizer=self.optimizer)
+                if save_best_checkpoint(current_loss, last_best_checkpoint_loss):
+                    last_best_checkpoint_loss = current_loss
+
                 TrainerTools().parallel.on_epoch_end(epoch)
                 self._on_epoch_end(tag=f'epoch:{epoch}')
 

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

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

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

@@ -1,11 +1,11 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=xTmmQSJ_jQDVSTT3km1p_8eRrc7yE_dEsi92z9OX5ec,3251
+llm_trainer/checkpoint.py,sha256=wC4GdIY2HAnxGHzUND5Yq-J_ynhDPT7A_2sXlWRElBc,4647
 llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
-llm_trainer/dpo_trainer.py,sha256=wMREatLt0I8Ajdm_sI2U8Zj-IN1L6txP9s_tH1oI3-s,11431
+llm_trainer/dpo_trainer.py,sha256=xfYXlLA5TbqPKCUbk5_V79TreEh-dnLMaN72a3-Tdzg,11860
 llm_trainer/ds_checkpoint.py,sha256=wz48HoLBBt8QGO1tXfvJwrXoiGtPG_gjwHfEqARllso,2175
 llm_trainer/eval.py,sha256=fjASCILU3fSPJxo9cP3rIXEEnkc5ZlUyHqXlZtUiHrw,888
 llm_trainer/generate_utils.py,sha256=CbJ3mfAD6DkQ0GUHcJQ1AK02m-ocwmd-BPXEpiwvNNQ,14933
-llm_trainer/grpo_trainer.py,sha256=qiC3KwxYPSB9UKqyk4eSRvORP3b6GM-2ozqI8u3QvI0,15568
+llm_trainer/grpo_trainer.py,sha256=vTNi3n6R4NbwFh_s8LYN1TWEJm8AW2F5NVJlT5MHxKk,15990
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
 llm_trainer/loss.py,sha256=NZCQeUXnLSj__mmDflE8g89KgE0emAJXIab0IERCLno,6023
 llm_trainer/parallel.py,sha256=DQu8GqEFxD99HQ6hKuIxxyKi-05dMO33eMhImYlPuOI,4468
@@ -13,21 +13,21 @@ llm_trainer/parallel_ddp.py,sha256=Pob9vUlBZnkL4oP1Re11kFob7nufMSE96pn7m7fuOEM,1
 llm_trainer/parallel_ds.py,sha256=oy8RRxHud3rACWubFlJqqd0pjPEQhKeAPGPQUSdJX2c,1145
 llm_trainer/parallel_none.py,sha256=TG6Pm829Dg-yQu-97O-EHV3FCARBlNcP47KkGFAs16E,676
 llm_trainer/partition_utils.py,sha256=xzv8kwlbKp3dai2pBwX89gN5ymeHk1bGbTkGru5H-UM,5167
-llm_trainer/scheduler.py,sha256=Xz8HhwoRMjRe41sf_NHhpZfkTlEs0I2MYusvMY6hCVw,3531
+llm_trainer/scheduler.py,sha256=lyC9TFuF_y8EXYq9d-WAqN4CSaq_w9kSKeh_BOo3EpI,4039
 llm_trainer/sft_trainer.py,sha256=gxQA7T1o1QGUsHp2CX1Qb_fO5LppBJuNbc0H4ixCYUA,1783
 llm_trainer/tokenizer.py,sha256=SSpgXtb0e1NtQqRW0gCq09TTZi47umggy-Fh5EMHKJg,6708
 llm_trainer/tools.py,sha256=yF17lp6oOfLe2XJeKDQ1juZcbv-6vFamJSLwEeArduA,2975
-llm_trainer/train_configs.py,sha256=m57W71SI5VCCU9aJ_nJkB-3AJrSGiNXmV28rdpuYmLg,7332
-llm_trainer/trainer.py,sha256=zTJVyY1cAjJdTkyXCOy2ZPVP18SOMLdWhD54Mz2JRe4,25314
+llm_trainer/train_configs.py,sha256=c6bgivkkWRYcPD3NzI5uRItAUhZiIBgKVMuMgVFRnFo,7336
+llm_trainer/trainer.py,sha256=g8YUP0FmBP3MGwewyoyOW35p9CY98rS62pzjnOMiWvE,25875
 llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
-project_llm_trainer-0.5.1.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.5.1.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.5.1.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.5.1.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.5.1.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.5.1.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.5.1.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.5.1.dist-info/METADATA,sha256=x-Bobn0EH7wyKznJydUeVLK9sdIrkBmDYDbEpyG4pKc,195
-project_llm_trainer-0.5.1.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.5.1.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.5.1.dist-info/RECORD,,
+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,,