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

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(
@@ -203,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/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
             )
 

llm_trainer/parallel.py

@@ -140,7 +140,7 @@ class Parallel(ABC):
         return 1
 
     def wait(self, msg=None):
-        msg = f' for {msg}' if msg else 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
 

{project_llm_trainer-0.5.16.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.16
+Version: 0.5.17
 Summary: LLM and VLM trainer
 Author: qibin
 Author-email: qibin0506@gmail.com

{project_llm_trainer-0.5.16.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=gz31pZbbQvRTYrBhxV-MFaBAIFeqpe7rM6nFsjwT9lY,4328
 llm_trainer/dataset.py,sha256=4QlOo0SFB5816BUYegQjgobUqTUMQvdmZMM_OEAMSjE,4347
-llm_trainer/dpo_trainer.py,sha256=--ItH-rkkq24Da3M_Kf0VxpQ3t-k0fpZrzFGqkYsjks,12304
+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=g_ivzQop2SkvhlKAEWb0zUnIvNuHTfsOoIG6y29oTCw,16106
+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=j1L4n-JmDkDZblURrNKpEAWEqqGIAXAN9PT_fSS_OnE,4492
+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=YqWhD9jXbrUdm3KEjEHLyg_qHiXCy5R7PK-arCXxJ6M,26399
+llm_trainer/trainer.py,sha256=Q821nlLDKRZVpaRoiZ7DiJplpAJRRLtvR_33FbClGA0,26729
 llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
-project_llm_trainer-0.5.16.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.5.16.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
-project_llm_trainer-0.5.16.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
-project_llm_trainer-0.5.16.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.5.16.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.5.16.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.5.16.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
-project_llm_trainer-0.5.16.dist-info/METADATA,sha256=h0TMNrZMUU875tVasbuqt69EuPPMbo_nv6tHQLKeNbQ,196
-project_llm_trainer-0.5.16.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.5.16.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.5.16.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,,