project-llm-trainer 0.5.15__py3-none-any.whl → 0.5.17__py3-none-any.whl

llm_trainer/checkpoint.py

@@ -65,7 +65,7 @@ def save_best_checkpoint(
         except:
             pass
 
-    TrainerTools().parallel.wait()
+    TrainerTools().parallel.wait('save best checkpoint')
     return need_replace
 
 

llm_trainer/dpo_trainer.py

@@ -12,7 +12,10 @@ from .dataset import DPODataset
 from .loss import DPOLoss
 from .tools import TrainerTools
 from .utils import get_dpo_collate_fn
-from .partition_utils import sync_model_params
+from .partition_utils import (
+    sync_model_params,
+    unwrap_model_for_generation
+)
 
 from .checkpoint import (
     save_checkpoint,
@@ -35,28 +38,28 @@ class DPOTrainer(Trainer):
             eval_image_tags=eval_image_tags
         )
 
-        self.reference_model = self._init_reference_model()
+        self.ref_model = self._init_ref_model()
 
-    def _init_reference_model(self):
-        reference_model = self._new_model(self.train_config)
+    def _init_ref_model(self):
+        ref_model = self._new_model(self.train_config)
 
-        reference_model, _ = TrainerTools().parallel.process(
-            model=reference_model,
+        ref_model, _ = TrainerTools().parallel.process(
+            model=ref_model,
             optimizer=None,
-            kwargs=self._init_reference_args(),
+            kwargs=self._init_ref_model_args(),
             save_instance=False
         )
 
-        reference_model.eval()
-        for param in reference_model.parameters():
+        ref_model.eval()
+        for param in ref_model.parameters():
             param.requires_grad = False
 
         sync_model_params(
             _from=self.train_model,
-            _to=reference_model
+            _to=ref_model
         )
 
-        return reference_model
+        return ref_model
 
     def _init_loss(self):
         criterion = DPOLoss(
@@ -170,14 +173,19 @@ class DPOTrainer(Trainer):
                         skipping_train = True
                         continue
 
-                    skipping_train = False
-
                     # 是否需要更新梯度
-                    if gradient_accumulation_steps > 1:
+                    if skipping_train:
+                        need_update_grad = False
+                    elif gradient_accumulation_steps > 1:
                         need_update_grad = (batch + 1) % gradient_accumulation_steps == 0 or batch == batch_count_per_file - 1
                     else:
                         need_update_grad = True
 
+                    # 要放在need_update_grad赋值下面，解决在继续训练时未知原因的卡死现象
+                    if skipping_train:
+                        TrainerTools().parallel.wait('skip train')
+                        skipping_train = False
+
                     try:
                         chosen_inputs: torch.Tensor = batch_data['chosen_inputs'].to(TrainerTools().parallel.device)
                         chosen_labels: torch.Tensor = batch_data['chosen_labels'].to(TrainerTools().parallel.device)
@@ -198,17 +206,18 @@ class DPOTrainer(Trainer):
 
                         with self.ctx:
                             policy_outputs = self.train_model(concat_inputs, attention_mask=concat_mask)
-                            with torch.inference_mode():
-                                ref_outputs = self.reference_model(concat_inputs, attention_mask=concat_mask)
-
                             policy_probs = self._logprobs(policy_outputs['logits'], concat_labels, concat_mask)
-                            ref_probs = self._logprobs(ref_outputs['logits'], concat_labels, concat_mask)
+                            aux_loss = policy_outputs.get('aux_loss')
+
+                            with torch.no_grad():
+                                ref_outputs = self.ref_model(concat_inputs, attention_mask=concat_mask)
+                                ref_probs = self._logprobs(ref_outputs['logits'], concat_labels, concat_mask)
 
                             # calc loss
                             loss = self.criterion(policy_probs, ref_probs)
 
-                            if aux_loss_coef and policy_outputs['aux_loss']:
-                                loss += aux_loss_coef * policy_outputs['aux_loss']
+                            if aux_loss_coef and aux_loss:
+                                loss += aux_loss_coef *aux_loss
 
                         if gradient_accumulation_steps > 1:
                             loss = loss / gradient_accumulation_steps

llm_trainer/ds_checkpoint.py

@@ -37,7 +37,7 @@ def save_ds_checkpoint(model: nn.Module):
                 shutil.rmtree(oldest_ckpt)
             except: ...
 
-    TrainerTools().parallel.wait()
+    TrainerTools().parallel.wait('remove old ds checkpoint')
 
 
 def load_ds_checkpoint(

llm_trainer/grpo_trainer.py

@@ -42,27 +42,27 @@ class GRPOTrainer(Trainer):
         )
 
         self.reward_func = reward_func
-        self.reference_model = self._init_reference_model()
+        self.ref_model = self._init_ref_model()
 
         # 默认使用torch提供的pad_sequence
         # 如果pad_sequence不支持padding_side参数，则将改参数置为False，使用反转的方式
         self._use_origin_pad_sequence = True
 
-    def _init_reference_model(self):
-        reference_model = self._new_model(self.train_config)
+    def _init_ref_model(self):
+        ref_model = self._new_model(self.train_config)
 
-        reference_model, _ = TrainerTools().parallel.process(
-            model=reference_model,
+        ref_model, _ = TrainerTools().parallel.process(
+            model=ref_model,
             optimizer=None,
-            kwargs=self._init_reference_args(),
+            kwargs=self._init_ref_model_args(),
             save_instance=False
         )
 
-        reference_model.eval()
-        for param in reference_model.parameters():
+        ref_model.eval()
+        for param in ref_model.parameters():
             param.requires_grad = False
 
-        return reference_model
+        return ref_model
 
     def _init_loss(self):
         criterion = GRPOLoss(
@@ -225,7 +225,7 @@ class GRPOTrainer(Trainer):
             old_log_probs, _ = self._compute_log_probabilities(generate_model, input_ids, attention_mask, logits_to_keep)
 
             # Compute ref_log_probs from the reference model, which remains static.
-            ref_log_probs, _ = self._compute_log_probabilities(self.reference_model, input_ids, attention_mask, logits_to_keep)
+            ref_log_probs, _ = self._compute_log_probabilities(self.ref_model, input_ids, attention_mask, logits_to_keep)
 
         repeated_prompts = [p for p in prompts for _ in range(group_size)]
         repeated_answers = [a for a in answers for _ in range(group_size)]
@@ -290,7 +290,7 @@ class GRPOTrainer(Trainer):
         for epoch in range(self.train_config.n_epochs):
             sync_model_params(
                 _from=self.train_model,
-                _to=self.reference_model,
+                _to=self.ref_model,
                 mixup_alpha=self.train_config.grpo_config.mixup_alpha
             )
 
@@ -317,7 +317,9 @@ class GRPOTrainer(Trainer):
                         skipping_train = True
                         continue
 
-                    skipping_train = False
+                    if skipping_train:
+                        TrainerTools().parallel.wait('skip train')
+                        skipping_train = False
 
                     # start generate
                     if TrainerTools().parallel.is_main_process:

llm_trainer/parallel.py

@@ -139,9 +139,8 @@ class Parallel(ABC):
             return dist.get_world_size()
         return 1
 
-    def wait(self):
-        try:
-            log(f'wait at {self.device}')
-            dist.barrier()
-        except: ...
-        log(f'continue at {self.device}')
+    def wait(self, msg=None):
+        msg = f' for {msg}' if msg else ''
+        log(f'wait at {self.device}{msg}')
+        dist.barrier()
+        log(f'continue at {self.device}{msg}')

llm_trainer/trainer.py

@@ -1,6 +1,6 @@
-import time
 from contextlib import nullcontext
 from typing import Optional, Tuple, List, Dict, Any
+import copy
 
 import torch
 import torch.distributed as dist
@@ -65,6 +65,7 @@ class Trainer:
             assert len(self.eval_prompts) == len(self.eval_image_tags)
 
         parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim = self._convert_train_args()
+        self.parallel_kwargs = parallel_kwargs
         self.data_loader_kwargs: dict[str, Any] = data_loader_kwargs
         self.sampler_kwargs: dict[str, Any] = sampler_kwargs
 
@@ -323,8 +324,8 @@ class Trainer:
 
         return parallel_kwargs, data_loader_kwargs, sampler_kwargs, use_ds_optim
 
-    def _init_reference_args(self) -> dict:
-        parallel_kwargs, _, _, _ = self._convert_train_args()
+    def _init_ref_model_args(self) -> dict:
+        parallel_kwargs = copy.deepcopy(self.parallel_kwargs)
 
         if parallel_kwargs and isinstance(TrainerTools().parallel, DsParallel):
             # reference to https://github.com/huggingface/trl/blob/main/trl/models/utils.py:prepare_deepspeed
@@ -346,8 +347,13 @@ class Trainer:
             #             }
             #         )
 
-            if parallel_kwargs['zero_optimization']['stage'] != 3:
-                parallel_kwargs['zero_optimization']['stage'] = 0
+            parallel_kwargs.pop('activation_checkpointing', None)
+            parallel_kwargs.pop('gradient_clipping', None)
+
+            # ref_model暂时先使用stage 0, 解决训练卡住问题
+            parallel_kwargs["zero_optimization"] = {"stage": 0}
+            # if parallel_kwargs.get("zero_optimization", {}).get("stage", 0) != 3:
+            #     parallel_kwargs["zero_optimization"] = {"stage": 0}
 
         return parallel_kwargs
 
@@ -449,7 +455,7 @@ class Trainer:
                 )
                 generate_model.train()
 
-        TrainerTools().parallel.wait()
+        TrainerTools().parallel.wait('eval')
 
     def _on_batch_end(self, tag: str):
         self._eval(f'sign:batch/{tag}')
@@ -500,14 +506,19 @@ class Trainer:
                         skipping_train = True
                         continue
 
-                    skipping_train = False
-
                     # 是否需要更新梯度
-                    if gradient_accumulation_steps > 1:
+                    if skipping_train:
+                        need_update_grad = False
+                    elif gradient_accumulation_steps > 1:
                         need_update_grad = (batch + 1) % gradient_accumulation_steps == 0 or batch == batch_count_per_file - 1
                     else:
                         need_update_grad = True
 
+                    # 要放在need_update_grad赋值下面，解决在继续训练时未知原因的卡死现象
+                    if skipping_train:
+                        TrainerTools().parallel.wait('skip train')
+                        skipping_train = False
+
                     inputs = batch_data['inputs']
                     labels = batch_data['labels']
 

{project_llm_trainer-0.5.15.dist-info → project_llm_trainer-0.5.17.dist-info}/METADATA

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

{project_llm_trainer-0.5.15.dist-info → project_llm_trainer-0.5.17.dist-info}/RECORD

@@ -1,14 +1,14 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
-llm_trainer/checkpoint.py,sha256=Wh5CwceIajTgJ9i_mH3I1R9N2nOLFqVFmlEMkTiGcD4,4306
+llm_trainer/checkpoint.py,sha256=gz31pZbbQvRTYrBhxV-MFaBAIFeqpe7rM6nFsjwT9lY,4328
 llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
-llm_trainer/dpo_trainer.py,sha256=8LYRxviJKcB-rN_XprVsWr5YshU8KolggMm7irjbXvI,11990
-llm_trainer/ds_checkpoint.py,sha256=kM7--wZyo4WIg4C2xk3bwad-m3V8ICfNLF3aFKtvzSA,2115
+llm_trainer/dpo_trainer.py,sha256=pNJaXvk-g0lGkZoRhbODNH34hTiz8EdP4Z12ws4W0t8,12309
+llm_trainer/ds_checkpoint.py,sha256=Wzy7PvVVWR794-BW4uragWFTAkkgDvjvkF-qMdyB4fc,2141
 llm_trainer/eval.py,sha256=ZyUfSo2Q8P-lrCdPEnGkoo5pGubd0AabREK5eMISRII,1109
 llm_trainer/generate_utils.py,sha256=wrZoG2g7CsOyG4sb3px9vURHQFV6_9j5kQmpFc5A8yg,15335
-llm_trainer/grpo_trainer.py,sha256=PVTlKOEJpI0AMlh7Siw_MHpLm9CAZepCAMjjSZF6eRU,15996
+llm_trainer/grpo_trainer.py,sha256=tuzcSi1uBzUPVKojEheJ3-Tx8-g99mf6LYYxC5nsNiw,16040
 llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
 llm_trainer/loss.py,sha256=NZCQeUXnLSj__mmDflE8g89KgE0emAJXIab0IERCLno,6023
-llm_trainer/parallel.py,sha256=DQu8GqEFxD99HQ6hKuIxxyKi-05dMO33eMhImYlPuOI,4468
+llm_trainer/parallel.py,sha256=G9X0FddIJwd9j-5XOknB4AlBe4G2W6fUCaQH6ycC2Fo,4490
 llm_trainer/parallel_ddp.py,sha256=Pob9vUlBZnkL4oP1Re11kFob7nufMSE96pn7m7fuOEM,1345
 llm_trainer/parallel_ds.py,sha256=oy8RRxHud3rACWubFlJqqd0pjPEQhKeAPGPQUSdJX2c,1145
 llm_trainer/parallel_none.py,sha256=TG6Pm829Dg-yQu-97O-EHV3FCARBlNcP47KkGFAs16E,676
@@ -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=992wy0YhBG2WvxwdLEPL4_-JUl4NkwMPT-jj_BIHo6A,7347
-llm_trainer/trainer.py,sha256=FF75J-BRUp34No2TvQIgomvNozWYzVhDeOfaBgQLV9g,26079
+llm_trainer/trainer.py,sha256=Q821nlLDKRZVpaRoiZ7DiJplpAJRRLtvR_33FbClGA0,26729
 llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
-project_llm_trainer-0.5.15.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.5.15.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.5.15.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.5.15.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.5.15.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.5.15.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.5.15.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.5.15.dist-info/METADATA,sha256=7ObRAx3PO5Dn55rgnJRS-bXGp-NU-SHgoPKVdTUTGCc,196
-project_llm_trainer-0.5.15.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.5.15.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.5.15.dist-info/RECORD,,
+project_llm_trainer-0.5.17.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.5.17.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.5.17.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.5.17.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.5.17.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.5.17.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.5.17.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.5.17.dist-info/METADATA,sha256=BVzwe45PQXSE-f5-BCZulqWCK3PIpKzxv9z__moTEJY,196
+project_llm_trainer-0.5.17.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.5.17.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.5.17.dist-info/RECORD,,