project-llm-trainer 0.7.8__py3-none-any.whl → 0.7.9__py3-none-any.whl

llm_trainer/dpo_trainer.py

@@ -12,7 +12,8 @@ from .loss import DPOLoss
 from .tools import TrainerTools
 from .utils import (
     autocast,
-    get_dpo_collate_fn
+    get_dpo_collate_fn,
+    fill_loss_mask
 )
 from .partition_utils import sync_model_params
 
@@ -84,7 +85,6 @@ class DPOTrainer(Trainer):
     def _calc_loss(self, inputs, attention_mask, logits, labels): ...
 
     def _log_probs_from_logits(self, logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
-        # https://github.com/OpenRLHF/OpenRLHF/pull/718#issuecomment-2641081881
         if logits.dtype in [torch.float32, torch.float64]:
             logits_labels = torch.gather(logits, dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)
             logsumexp_values = torch.stack(
@@ -102,25 +102,26 @@ class DPOTrainer(Trainer):
         return log_probs_labels
 
 
-    def _logprobs(self, logits, labels, mask):
+    def _logprobs(self, logits, labels, attention_mask):
         """
         Calculate the average log probabilities for a batch of sequences.
 
         Args:
             logits (torch.Tensor): Logits from the model with shape (B, T, V)
             labels (torch.Tensor): Ground truth labels with shape (B, T).
-            mask (torch.Tensor): Mask tensor with shape (B, T) indicating
+            attention_mask (torch.Tensor): Mask tensor with shape (B, T) indicating
                 which tokens are not padding (1 for valid tokens, 0 for padding).
 
         Returns:
             torch.Tensor: Average log probabilities for each sequence in the batch.
                           Shape is (B,) representing the mean log probability for each sequence.
         """
-        labels = labels[:, 1:].clone()
-        logits = logits[:, :-1, :]
+        loss_masks = attention_mask.clone().bool()
+        loss_masks = fill_loss_mask(loss_masks, labels)
 
-        # # Shift mask right by one to align with labels
-        mask = mask[:, 1:].clone()
+        logits = logits[:, :-1, :]
+        labels = labels[:, 1:].clone()
+        loss_masks = loss_masks[:, 1:]
 
         # dummy token; we'll ignore the losses on these tokens later
         labels[labels == -100] = 0
@@ -129,11 +130,10 @@ class DPOTrainer(Trainer):
         per_token_logps = self._log_probs_from_logits(logits, labels)
 
         # Apply the mask to set log-probs of padding tokens to 0
-        logprobs_sums = (per_token_logps * mask).sum(-1)
-
-        # logprobs_means = (per_token_logps * mask).sum(-1) / mask.sum(-1)
+        logprobs_sums = (per_token_logps * loss_masks).sum(-1)
+        logprobs_means = (per_token_logps * loss_masks).sum(-1) / loss_masks.sum(-1)
 
-        return logprobs_sums #, -logprobs_means.mean()
+        return logprobs_sums, logprobs_means
 
     def train(self):
         # 梯度累积步数
@@ -147,6 +147,7 @@ class DPOTrainer(Trainer):
         last_best_checkpoint_loss: Optional[float] = None
 
         aux_loss_coef = self.train_config.loss_config.aux_loss_coef
+        nll_loss_coef = self.train_config.dpo_config.nll_loss_coef
 
         for epoch in range(self.train_config.n_epochs):
             self.train_model.train()
@@ -188,36 +189,53 @@ class DPOTrainer(Trainer):
                     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)
+
                         rejected_inputs: torch.Tensor = batch_data['rejected_inputs'].to(TrainerTools().parallel.device)
                         rejected_labels: torch.Tensor = batch_data['rejected_labels'].to(TrainerTools().parallel.device)
 
-                        chosen_attention_mask: torch.Tensor = chosen_inputs != TrainerTools().tokenizer.pad
-                        rejected_attention_mask: torch.Tensor = rejected_inputs != TrainerTools().tokenizer.pad
+                        chosen_attention_masks: torch.Tensor = chosen_inputs != TrainerTools().tokenizer.pad
+                        rejected_attention_masks: torch.Tensor = rejected_inputs != TrainerTools().tokenizer.pad
 
                         # 在batch维度concat
                         # [chosen, chosen, reject, reject]
                         concat_inputs = torch.concat([chosen_inputs, rejected_inputs], dim=0)
                         concat_labels = torch.concat([chosen_labels, rejected_labels], dim=0)
-                        concat_mask = torch.concat([chosen_attention_mask, rejected_attention_mask], dim=0)
+                        concat_attention_masks = torch.concat([chosen_attention_masks, rejected_attention_masks], dim=0)
 
                         if TrainerTools().parallel.parallel_train:
                             self.train_model.require_backward_grad_sync = need_update_grad
 
                         with autocast(TrainerTools().parallel.device_type):
-                            policy_outputs = self.train_model(concat_inputs, attention_mask=concat_mask)
-                            policy_probs = self._logprobs(policy_outputs['logits'], concat_labels, concat_mask)
+                            policy_outputs = self.train_model(concat_inputs, attention_mask=concat_attention_masks)
+                            policy_logprobs_sums, policy_logprobs_means = self._logprobs(policy_outputs['logits'], concat_labels, concat_attention_masks)
                             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)
+                                ref_outputs = self.ref_model(concat_inputs, attention_mask=concat_attention_masks)
+                                ref_logprobs_sums, _ = self._logprobs(ref_outputs['logits'], concat_labels, concat_attention_masks)
+
+                            policy_chosen_logps = policy_logprobs_sums[:chosen_inputs.shape[0]]
+                            policy_rejected_logps = policy_logprobs_sums[chosen_inputs.shape[0]:]
+
+                            ref_chosen_logps = ref_logprobs_sums[:chosen_inputs.shape[0]]
+                            ref_rejected_logps = ref_logprobs_sums[chosen_inputs.shape[0]:]
+
+                            nll_loss = -policy_logprobs_means[:chosen_inputs.shape[0]].mean()
 
                             # calc loss
-                            loss = self.criterion(policy_probs, ref_probs)
+                            loss = self.criterion(
+                                policy_chosen_logps,
+                                policy_rejected_logps,
+                                ref_chosen_logps,
+                                ref_rejected_logps
+                            )
 
                             if aux_loss_coef and aux_loss:
                                 loss += aux_loss_coef * aux_loss
 
+                            if nll_loss_coef and nll_loss:
+                                loss += nll_loss_coef * nll_loss
+
                         if gradient_accumulation_steps > 1:
                             loss = loss / gradient_accumulation_steps
 

llm_trainer/loss.py

@@ -92,17 +92,13 @@ class DPOLoss(nn.Module):
 
     def forward(
             self,
-            policy_logps: torch.Tensor,
-            reference_logps: torch.Tensor,
+            policy_chosen_logps: torch.Tensor,
+            policy_reject_logps: torch.Tensor,
+            ref_chosen_logps: torch.Tensor,
+            ref_reject_logps: torch.Tensor
     ) -> torch.Tensor:
-        batch_size = reference_logps.shape[0]
-        ref_chosen_probs = reference_logps[:batch_size//2]
-        ref_reject_probs = reference_logps[batch_size//2:]
-        policy_chosen_probs = policy_logps[:batch_size//2]
-        policy_reject_probs = policy_logps[batch_size//2:]
-
-        pi_logratios = policy_chosen_probs - policy_reject_probs
-        ref_logratios = ref_chosen_probs - ref_reject_probs
+        pi_logratios = policy_chosen_logps - policy_reject_logps
+        ref_logratios = ref_chosen_logps - ref_reject_logps
         logits = pi_logratios - ref_logratios
 
         if self.ipo:

llm_trainer/tokenizer.py

@@ -3,7 +3,7 @@ import warnings
 from typing import List, Dict, Union
 from transformers import Qwen2TokenizerFast
 from transformers import AddedToken
-from transformers import LlamaTokenizer, LlamaTokenizerFast
+from transformers import LlamaTokenizerFast
 import torch
 
 TOKEN_TYPE_QWEN = 'qwen'
@@ -164,3 +164,18 @@ class Tokenizer:
 
         return chat_template
 
+    def get_special_tokens_dict(self):
+        return {
+            self.text_end: self.end,
+            self.text_pad: self.pad,
+            self.text_unk: self.unk,
+            self.text_user: self.user,
+            self.text_assistant: self.assistant,
+            self.text_think_start: self.think_start,
+            self.text_think_end: self.think_end,
+            self.text_answer_start: self.answer_start,
+            self.text_answer_end: self.answer_end,
+            self.text_system: self.system,
+            self.text_image: self.image,
+        }
+

llm_trainer/utils.py

@@ -154,16 +154,22 @@ def batch_repeat_image_tok(
 
 
 def pretrain_collate_fn(batch_data):
-    inputs, labels = _pad_sequence(batch_data)
+    # [[x,x,x], [y,y,y]]
+    inputs = pad_sequence(batch_data, batch_first=True, padding_value=TrainerTools().tokenizer.pad)
+    # crossEntropy默认的ignore_index是-100
+    labels = pad_sequence(batch_data, batch_first=True, padding_value=-100)
 
     # inputs, labels
-    return {'inputs': inputs, 'labels': labels}
+    return {
+        'inputs': inputs,
+        'labels': labels
+    }
 
 
 def get_sft_collate_fn(mask_prompt: bool):
     def sft_collate_fn(batch_data):
         """
-         如果是sft，则不计算prompt部分的loss, 例如：
+        如果是sft，则不计算prompt部分的loss, 例如：
         logits: [USER]你好[BOT]我好[SEP]
         labels: [USER]你好[BOT]我好[SEP]
 
@@ -184,11 +190,19 @@ def get_sft_collate_fn(mask_prompt: bool):
             batch_train_data.append(item['inputs'])
             image_tags.append(item['image_tag'])
 
-        inputs, labels = _pad_sequence(batch_train_data)
+        # [[x,x,x], [y,y,y]]
+        inputs = pad_sequence(batch_train_data, batch_first=True, padding_value=TrainerTools().tokenizer.pad)
+        # crossEntropy默认的ignore_index是-100
+        labels = pad_sequence(batch_train_data, batch_first=True, padding_value=-100)
+
         if mask_prompt:
             labels = _mask_prompt(labels)
 
-        return {'inputs': inputs, 'labels': labels, 'image_tags': image_tags}
+        return {
+            'inputs': inputs,
+            'labels': labels,
+            'image_tags': image_tags
+        }
 
     return sft_collate_fn
 
@@ -295,13 +309,24 @@ def join_batch(batch_data: list[dict]) -> dict:
     return result
 
 
-def _pad_sequence(batch_data):
-    # [[x,x,x], [y,y,y]]
-    inputs = pad_sequence(batch_data, batch_first=True, padding_value=TrainerTools().tokenizer.pad)
-    # crossEntropy默认的ignore_index是-100
-    labels = pad_sequence(batch_data, batch_first=True, padding_value=-100)
+def fill_loss_mask(loss_masks, labels):
+    """
+    将loss_mask中prompt部分强制设置为False
+    loss_masks: shape  (B, T)
+    labels: shape (B, T)
+    """
+    tokenizer = TrainerTools().tokenizer
+    # 支持多轮会话的mask
+    for batch, label in enumerate(labels):
+        start_index = -1
+        for index, token in enumerate(label):
+            if token == tokenizer.system or token == tokenizer.user:
+                start_index = index
+            elif token == tokenizer.end and start_index != -1:
+                loss_masks[batch, start_index:index + 1] = False
+                start_index = -1
 
-    return inputs, labels
+    return loss_masks
 
 
 def _mask_prompt(labels):

{project_llm_trainer-0.7.8.data → project_llm_trainer-0.7.9.data}/scripts/ds_train

@@ -10,14 +10,15 @@ if __name__ == '__main__':
     if len(arguments) > 1:
         # 0,1,2,3
         cuda_visible_devive = arguments[1]
-    else:
-        cuda_visible_devive = None
 
-    # cuda location
-    if len(arguments) > 2:
-        cuda_loc = arguments[2]
+        # cuda location
+        if len(arguments) > 2:
+            cuda_loc = arguments[2]
+        else:
+            cuda_loc = 'localhost'
     else:
-        cuda_loc = 'localhost'
+        cuda_visible_devive = None
+        cuda_loc = None
 
     os.environ['PARALLEL_TYPE'] = 'ds'
 

{project_llm_trainer-0.7.8.dist-info → project_llm_trainer-0.7.9.dist-info}/METADATA

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

{project_llm_trainer-0.7.8.dist-info → project_llm_trainer-0.7.9.dist-info}/RECORD

@@ -1,13 +1,13 @@
 llm_trainer/__init__.py,sha256=HWgtTEVeQSnZmEyYQm2K6eFEG4X2QAoigMlB5Z2tcXE,260
 llm_trainer/checkpoint.py,sha256=X5ZeUtJlxVz7pnWQLaS-y7UIZOaOAnZTt2L8rSAPzUs,4428
 llm_trainer/dataset.py,sha256=UL3fGeM4XSlyNQRZH-139u3LujqAQx3YyaxNRewk6LE,8935
-llm_trainer/dpo_trainer.py,sha256=RMfbTsl3eav4yTJ2PK59mi6a0ECVOg8WwYVsHvMbNUE,12353
+llm_trainer/dpo_trainer.py,sha256=Bgds18UWFhzf_UNCFN-iBCdhKf9pcXJBFPEc32oJeXA,13354
 llm_trainer/ds_checkpoint.py,sha256=X2IWgpgi0yOtogph7n6DEwvK_0Ceb7juu1WMutv3HSk,2270
 llm_trainer/eval.py,sha256=ZyUfSo2Q8P-lrCdPEnGkoo5pGubd0AabREK5eMISRII,1109
 llm_trainer/generate_utils.py,sha256=8K3YFbp7IF_lCkmkzjHhqTW26EBFb2AilQmarVcfMvs,15001
 llm_trainer/grpo_trainer.py,sha256=MXnP8Kc9CQJw0CB3uMbHxIYwvpuujai4hgbbpUut_K4,16808
 llm_trainer/log.py,sha256=XwychwKF6gvFPhthCIZCAEUZ0G3DY3fiQrOHqPWsxz0,463
-llm_trainer/loss.py,sha256=glf4IeDWHvA2cJo-QKLRL8P6OxK4QjRJGrYJWOZiTPQ,6929
+llm_trainer/loss.py,sha256=RhTxftLMj1Tqc5pkUvJiZumfbMEPWL8GBGxdTfQggmk,6744
 llm_trainer/parallel.py,sha256=yjStV21DJ26yM8-0O6GTMxdFAcyShY5GsQWSZmbI7HU,4543
 llm_trainer/parallel_ddp.py,sha256=Pob9vUlBZnkL4oP1Re11kFob7nufMSE96pn7m7fuOEM,1345
 llm_trainer/parallel_ds.py,sha256=oy8RRxHud3rACWubFlJqqd0pjPEQhKeAPGPQUSdJX2c,1145
@@ -15,19 +15,19 @@ llm_trainer/parallel_none.py,sha256=TG6Pm829Dg-yQu-97O-EHV3FCARBlNcP47KkGFAs16E,
 llm_trainer/partition_utils.py,sha256=eEYNhfEIF4hGzZ3OLa6sEBIECz261drptEz_n7fZYtk,8396
 llm_trainer/scheduler.py,sha256=LAI_0VxClsIQkix0bRoduRD4vPfVuIZDhZgTAT_KK8k,4901
 llm_trainer/sft_trainer.py,sha256=LudTRIaqLQYy6ym6jjMX7v9xtFBJelrR3nnPCwb48nM,1821
-llm_trainer/tokenizer.py,sha256=SSpgXtb0e1NtQqRW0gCq09TTZi47umggy-Fh5EMHKJg,6708
+llm_trainer/tokenizer.py,sha256=0-xQCMz1xiPTDAZiYsVsiECSoZ_1eIvW9XsZOoFfakQ,7250
 llm_trainer/tools.py,sha256=5op5qrjjkK-Lr9oes5VxIVnOVYOYGoAdlIJq9mPUf64,2637
 llm_trainer/train_configs.py,sha256=N3ykM1uaLHcSNRC8ErYIxp9VYhSP7voJyAP-2D4ZJe0,7574
 llm_trainer/trainer.py,sha256=jS31zEXIIj9BoPTPlmaGYq61x72HGCjKfS2u3_gOkDk,27924
-llm_trainer/utils.py,sha256=xcdzpvPvXRKqsOK2yB7PZ9GmOvZMDFcglDPUZY2hJTY,11484
-project_llm_trainer-0.7.8.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
-project_llm_trainer-0.7.8.data/scripts/ddp_train,sha256=Z-309mM56CN0m3bxoeC5us4LUuwuNnoiOm3-fDdLMjQ,566
-project_llm_trainer-0.7.8.data/scripts/ds_train,sha256=3nXNNKmYI7miqyBdf-Ijl_rW1cGIKrAMZ1CSswN_gGo,665
-project_llm_trainer-0.7.8.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
-project_llm_trainer-0.7.8.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
-project_llm_trainer-0.7.8.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
-project_llm_trainer-0.7.8.data/scripts/smart_train,sha256=3oLIDuuqb4U4TU1lXy9V8lw_0gIf7i8tGsxlQ_s6bro,1220
-project_llm_trainer-0.7.8.dist-info/METADATA,sha256=rSYUrEkdjPCyYUqT2SOw3-hzT40wU3AwEw-ouHh1rBY,195
-project_llm_trainer-0.7.8.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
-project_llm_trainer-0.7.8.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
-project_llm_trainer-0.7.8.dist-info/RECORD,,
+llm_trainer/utils.py,sha256=xC5plG-8-_Al5yIF5xIU5lroOcBBk98TEhtUJrazZPE,12305
+project_llm_trainer-0.7.9.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.7.9.data/scripts/ddp_train,sha256=Z-309mM56CN0m3bxoeC5us4LUuwuNnoiOm3-fDdLMjQ,566
+project_llm_trainer-0.7.9.data/scripts/ds_train,sha256=tME0xmMdX1D9XuVo07D9dilW5VIWavBS3UK9DoY67WI,709
+project_llm_trainer-0.7.9.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.7.9.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.7.9.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.7.9.data/scripts/smart_train,sha256=3oLIDuuqb4U4TU1lXy9V8lw_0gIf7i8tGsxlQ_s6bro,1220
+project_llm_trainer-0.7.9.dist-info/METADATA,sha256=mDGLc1BjmIlOPz85JYB5bFnlXJgJ5VaNesW4z0HDZCA,195
+project_llm_trainer-0.7.9.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.7.9.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.7.9.dist-info/RECORD,,