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

llm_trainer/checkpoint.py

@@ -38,29 +38,32 @@ def save_best_checkpoint(
 ) -> bool:
     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_dir = os.environ.get('DIST_CHECKPOINT_DIR', 'checkpoint')
+        try:
+            if isinstance(TrainerTools().parallel, DsParallel):
+                checkpoint_dir = os.environ.get('DIST_CHECKPOINT_DIR', 'checkpoint')
 
-            if checkpoint_dir.endswith('/'):
-                best_checkpoint_dir = f'{checkpoint_dir[:-1]}_best'
-            else:
-                best_checkpoint_dir = f'{checkpoint_dir}_best'
+                if checkpoint_dir.endswith('/'):
+                    best_checkpoint_dir = f'{checkpoint_dir[:-1]}_best'
+                else:
+                    best_checkpoint_dir = f'{checkpoint_dir}_best'
 
-            if not os.path.exists(best_checkpoint_dir):
-                os.makedirs(best_checkpoint_dir)
+                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)
-            best_checkpoint_name = f'{checkpoint_name}_best'
+                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)
+                best_checkpoint_name = f'{checkpoint_name}_best'
 
-            if os.path.exists(checkpoint_name):
-                if os.path.exists(best_checkpoint_name):
-                    os.remove(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)
+                    shutil.copy2(checkpoint_name, best_checkpoint_name)
+        except:
+            pass
 
     TrainerTools().parallel.wait()
     return need_replace

llm_trainer/dpo_trainer.py

@@ -137,9 +137,10 @@ 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: Optional[float] = None
 
@@ -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/grpo_trainer.py

@@ -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,9 +465,10 @@ 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: Optional[float] = None
 
@@ -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.4.dist-info → project_llm_trainer-0.5.6.dist-info}/METADATA

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

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

@@ -1,11 +1,11 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=UVjOaDsiSIzRJ5VJZib6iXrdKv2A7K_gtJw3a9wNyoM,4293
+llm_trainer/checkpoint.py,sha256=GHRnPpvG0lz8mg2qv0itHr1rXLlj-itOqZWCtOV1IRU,4411
 llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
-llm_trainer/dpo_trainer.py,sha256=3hCjX06W6nt-8lio0YyGzXdI_saa5QlypXehQTtacO4,11871
+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=VltU9AngecJ5PEtpn_ToXfiLIc6VwwgOVwlLh-X2Je8,16001
+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=VHnuL8rZgxj2ewBVxmbN0jwVEPRVhbK2V2DYCtm_FxI,25886
+llm_trainer/trainer.py,sha256=YW59dJWTyQy77cLDGzBHhfinGyfkvmWCkl1SR9hM6a8,26071
 llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
-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,,
+project_llm_trainer-0.5.6.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.5.6.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.5.6.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.5.6.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.5.6.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.5.6.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.5.6.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.5.6.dist-info/METADATA,sha256=5JaiIS6GqDsm9o_6Zz-vkec_0rmENISg32Qld1Gn-u8,195
+project_llm_trainer-0.5.6.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.5.6.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.5.6.dist-info/RECORD,,