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

llm_trainer/dpo_trainer.py

@@ -12,10 +12,7 @@ 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,
-    unwrap_model_for_generation
-)
+from .partition_utils import sync_model_params
 
 from .checkpoint import (
     save_checkpoint,
@@ -217,7 +214,7 @@ class DPOTrainer(Trainer):
                             loss = self.criterion(policy_probs, ref_probs)
 
                             if aux_loss_coef and aux_loss:
-                                loss += aux_loss_coef *aux_loss
+                                loss += aux_loss_coef * aux_loss
 
                         if gradient_accumulation_steps > 1:
                             loss = loss / gradient_accumulation_steps

llm_trainer/grpo_trainer.py

@@ -66,8 +66,12 @@ class GRPOTrainer(Trainer):
 
     def _init_loss(self):
         criterion = GRPOLoss(
-            clip_eps=self.train_config.grpo_config.clip_eps,
-            kl_weight=self.train_config.grpo_config.kl_weight
+            beta=self.train_config.grpo_config.loss_beta,
+            clip_eps=self.train_config.grpo_config.loss_clip_eps,
+            delta=self.train_config.grpo_config.loss_delta,
+            importance_sampling_level=self.train_config.grpo_config.loss_importance_sampling_level,
+            loss_type=self.train_config.grpo_config.loss_type,
+            gen_max_new_tokens=self.train_config.grpo_config.gen_max_new_tokens
         )
 
         return criterion, None

llm_trainer/loss.py

@@ -2,6 +2,7 @@ from typing import List, Optional
 import torch
 from torch import nn
 import torch.nn.functional as F
+from .tools import TrainerTools
 
 
 class LMLoss(nn.Module):
@@ -115,6 +116,7 @@ class DPOLoss(nn.Module):
             )
 
         loss = losses.mean()
+
         # chosen_rewards = self.beta * (policy_chosen_probs - ref_chosen_probs).detach()
         # rejected_rewards = self.beta * (policy_reject_probs - ref_reject_probs).detach()
 
@@ -124,12 +126,21 @@ class DPOLoss(nn.Module):
 class GRPOLoss(nn.Module):
     def __init__(
             self,
+            beta: float,
             clip_eps: float,
-            kl_weight: float
+            delta: Optional[float] = None,
+            importance_sampling_level: str = 'token',
+            loss_type: str = 'grpo',
+            gen_max_new_tokens: Optional[float] = None
     ):
         super().__init__()
+
+        self.beta = beta
         self.clip_eps = clip_eps
-        self.kl_weight = kl_weight
+        self.delta = delta
+        self.importance_sampling_level = importance_sampling_level
+        self.loss_type = loss_type
+        self.gen_max_new_tokens = gen_max_new_tokens
 
     def forward(
             self,
@@ -139,33 +150,41 @@ class GRPOLoss(nn.Module):
             completion_mask: torch.Tensor,
             advantages: torch.Tensor
     ) -> torch.Tensor:
-        # Compute policy ratio
-        ratio = torch.exp(log_probs - old_log_probs)
 
-        # Compute surrogate loss with clipping
-        surrogate1 = ratio * advantages
-        surrogate2 = torch.clamp(ratio, 1 - self.clip_eps, 1 + self.clip_eps) * advantages
-        surrogate_loss = torch.min(surrogate1, surrogate2)
+        if self.beta != 0.0:
+            per_token_kl = torch.exp(ref_log_probs - log_probs) - (ref_log_probs - log_probs) - 1
+        else:
+            per_token_kl = None
+
+        log_ratio = log_probs - old_log_probs
+        if self.importance_sampling_level == "seq":
+            # GSPO
+            log_importance_weights = (log_ratio * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)
+            log_importance_weights = log_importance_weights.unsqueeze(-1)
+        else:
+            # GRPO
+            log_importance_weights = log_ratio
 
-        # Compute KL divergence penalty
-        kl_div = torch.exp(ref_log_probs - log_probs) - (ref_log_probs - log_probs) - 1
+        coef_1 = torch.exp(log_importance_weights)
+        coef_2 = torch.clamp(coef_1, 1 - self.clip_eps, 1 + self.clip_eps)
 
-        # Combine losses
-        per_token_loss = surrogate_loss - self.kl_weight * kl_div
-        loss = -((per_token_loss * completion_mask).sum(dim=1) / completion_mask.sum(dim=1)).mean()
+        # Two-sided clipping
+        if self.delta is not None:
+            coef_1 = torch.clamp(coef_1, max=self.delta)
 
-        return loss
+        per_token_loss1 = coef_1 * advantages
+        per_token_loss2 = coef_2 * advantages
+        per_token_loss = -torch.min(per_token_loss1, per_token_loss2)
+
+        if self.beta != 0.0:
+            per_token_loss = per_token_loss + self.beta * per_token_kl
 
-        # kl = self._approx_kl_divergence(
-        #     log_probs=log_probs,
-        #     ref_log_probs=ref_log_probs,
-        #     mask=mask,
-        # )
-        #
-        # ratio = (log_probs - old_log_probs).exp()
-        # surr1 = ratio * advantages
-        # surr2 = ratio.clamp(1 - self.clip_eps, 1 + self.clip_eps) * advantages
-        # loss = -torch.min(surr1, surr2) + self.kl_weight * kl
-        #
-        # loss = self._masked_mean(loss, mask, dim=-1).mean()
-        # return loss, kl.mean()
+        if self.loss_type == "bnpo":
+            loss = (per_token_loss * completion_mask).sum() / completion_mask.sum().clamp(min=1.0)
+        elif self.loss_type == "dr_grpo":
+            assert self.gen_max_new_tokens is not None
+            loss = (per_token_loss * completion_mask).sum() / (per_token_loss.size(0) * self.gen_max_new_tokens)
+        else:
+            loss = ((per_token_loss * completion_mask).sum(-1) / completion_mask.sum(-1).clamp(min=1.0)).mean()
+
+        return loss

llm_trainer/partition_utils.py

@@ -4,6 +4,7 @@ import itertools
 from packaging import version
 from torch import nn
 from torch.nn.parallel import DistributedDataParallel as DDP
+import torch.distributed as dist
 
 from .tools import TrainerTools
 from .parallel_ds import DsParallel
@@ -45,12 +46,40 @@ def unwrap_model_for_generation(model: nn.Module):
 
 
 def sync_model_params(_from: nn.Module, _to: Optional[nn.Module], mixup_alpha: float = 1.0):
+    """
+        必须在所有rank上调用，非rank0, _to 可以设置为None.
+        当前函数不适用于_to是一个zero3模型
+    """
     if isinstance(TrainerTools().parallel, DsParallel):
-        _sync_ds_model_params(_from, _to, mixup_alpha)
-    elif isinstance(TrainerTools().parallel, DdpParallel):
-        _sync_ddp_model_params(_from, _to, mixup_alpha)
+        state_dict = _get_ds_model_params(_from, only_rank0=_to is None)
+    elif isinstance(_from, DDP):
+        state_dict = _from.module.state_dict()
+    else:
+        state_dict = _from.state_dict()
+
+    if not _to or not state_dict:
+        return
+
+    unwrap_to_model = unwrap_model(_to)
+    if mixup_alpha == 1.0:
+        # 直接覆盖
+        unwrap_to_model.load_state_dict(state_dict)
     else:
-        _copy_params(_from, _to, mixup_alpha)
+        # 混合参数
+        for param_name, target_param in unwrap_to_model.named_parameters():
+            if param_name in state_dict:
+                from_param_tensor = state_dict[param_name]
+                target_param.data.mul_(1.0 - mixup_alpha).add_(
+                    from_param_tensor.data.to(target_param.device),
+                    alpha=mixup_alpha
+                )
+
+    # if isinstance(TrainerTools().parallel, DsParallel):
+    #     _sync_ds_model_params(_from, _to, mixup_alpha)
+    # elif isinstance(TrainerTools().parallel, DdpParallel):
+    #     _sync_ddp_model_params(_from, _to, mixup_alpha)
+    # else:
+    #     _copy_params(_from, _to, mixup_alpha)
 
 
 def unwrap_model(model) -> nn.Module:
@@ -66,6 +95,57 @@ def unwrap_model(model) -> nn.Module:
     return model
 
 
+def _get_ds_full_state_dict_on_rank0(model: nn.Module) -> Optional[dict]:
+    """
+        需要在所有rank上调用，然后只有rank0有值
+    """
+    import deepspeed
+    assert isinstance(model, deepspeed.DeepSpeedEngine)
+
+    if model.zero_optimization_stage() != 3:
+        if TrainerTools().parallel.is_main_process:
+            return {k: v.cpu().clone() for k, v in model.module.state_dict().items()}
+        return None
+
+    # --- ZeRO-3 ---
+    # 只调用一次 GatheredParameters，传入所有参数
+    with deepspeed.zero.GatheredParameters(model.parameters(), modifier_rank=0):
+        if TrainerTools().parallel.is_main_process:
+            # 在这个 'with' 代码块内，rank 0 上的 model.module 拥有完整的参数
+            # 所以我们可以像操作普通模型一样直接调用 state_dict()
+            full_state_dict = model.module.state_dict()
+
+            # 将其克隆到 CPU 并返回
+            return {k: v.cpu().clone() for k, v in full_state_dict.items()}
+
+    # 其他 rank 执行到这里时，上下文结束，直接返回 None
+    return None
+
+
+def _get_ds_model_params(model: nn.Module, only_rank0=False):
+    """
+        从一个正在运行的 DeepSpeedEngine 中高效地提取完整的 FP32 state_dict，
+        兼容 ZeRO Stages 0, 1, 2, 3。
+        包含了对 ZeRO-3 中分片参数的正确处理。
+    """
+
+    import deepspeed
+    assert isinstance(model, deepspeed.DeepSpeedEngine)
+    state_dict = _get_ds_full_state_dict_on_rank0(model)
+
+    # 现在，只有 rank 0 上的 state_dict 是一个有效的字典，其他 rank 上是 None。
+    # 我们需要将其广播给所有进程。
+    if not only_rank0 and TrainerTools().parallel.world_size > 1:
+        # 准备一个列表，rank 0 有数据，其他 rank 是占位符
+        object_list = [state_dict] if TrainerTools().parallel.is_main_process else [None]
+        # 执行广播，这个操作是阻塞的，会同步所有进程
+        dist.broadcast_object_list(object_list, src=0)
+        # 所有进程从列表中获取广播后的 state_dict 副本
+        state_dict = object_list[0]
+
+    return state_dict
+
+
 def _copy_params(model, target_model, mixup_alpha):
     for target_param, copy_param in zip(target_model.parameters(), model.parameters()):
         target_param.data.mul_(1.0 - mixup_alpha).add_(copy_param.data, alpha=mixup_alpha)
@@ -79,6 +159,7 @@ def _sync_ds_model_params(_from: nn.Module, _to: Optional[nn.Module], mixup_alph
 
     if _from.zero_optimization_stage() == 3:
         with deepspeed.zero.GatheredParameters(list(origin_from.parameters()) + list(_to.parameters()), modifier_rank=0):
+            # why only rank 0?
             if TrainerTools().parallel.is_main_process:
                 _copy_params(origin_from, _to, mixup_alpha)
     else:

llm_trainer/train_configs.py

@@ -136,10 +136,13 @@ class DPOConfig:
 @dataclass(kw_only=True)
 class GRPOConfig:
     grpo_steps: int = 1
-    clip_eps: float = 0.1
-    kl_weight: float = 0.04
     group_size: int = 12
     mixup_alpha: float = 1.0
+    loss_beta: float = 0.04
+    loss_clip_eps: float = 0.1
+    loss_delta: Optional[float] = None
+    loss_importance_sampling_level: str = 'token' # token or seq
+    loss_type: str = 'grpo' # grpo or bnpo or dr_grpo
     gen_max_new_tokens: Optional[int] = None
     gen_temperature: Optional[float] = None
     gen_k: Optional[int] = None

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

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

project_llm_trainer-0.6.0.dist-info/RECORD

@@ -0,0 +1,33 @@
+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=mETXpU1ZSasg1UM72wnh9NaoTuXBibuNuodfuW7u8Iw,12269
+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=-wbozslll_bcGUMqrbS0a73jhosyjc3oC3PHLSev6lw,16344
+llm_trainer/log.py,sha256=LxqTGRNZUGMTSQCePRpk-rYyxSnSIbT4kOdP8Fbzr0M,462
+llm_trainer/loss.py,sha256=eYvOlCoguKnLvdGuqvQpGUoLVSADQ5coaU3DWYbJEdM,6811
+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
+llm_trainer/partition_utils.py,sha256=eEYNhfEIF4hGzZ3OLa6sEBIECz261drptEz_n7fZYtk,8396
+llm_trainer/scheduler.py,sha256=LAI_0VxClsIQkix0bRoduRD4vPfVuIZDhZgTAT_KK8k,4901
+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=U4hwXWKI6svDqiDOu6RPTitCzpxEYyjZUN6gwh_co8c,7510
+llm_trainer/trainer.py,sha256=Q821nlLDKRZVpaRoiZ7DiJplpAJRRLtvR_33FbClGA0,26729
+llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
+project_llm_trainer-0.6.0.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.6.0.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.6.0.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.6.0.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.6.0.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.6.0.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.6.0.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.6.0.dist-info/METADATA,sha256=_F0QQHrdQNGXG8eDGRDsgEvdX6fYWXSDg5Ad089CXHk,195
+project_llm_trainer-0.6.0.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.6.0.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.6.0.dist-info/RECORD,,

project_llm_trainer-0.5.17.dist-info/RECORD

@@ -1,33 +0,0 @@
-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=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=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=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
-llm_trainer/partition_utils.py,sha256=xzv8kwlbKp3dai2pBwX89gN5ymeHk1bGbTkGru5H-UM,5167
-llm_trainer/scheduler.py,sha256=LAI_0VxClsIQkix0bRoduRD4vPfVuIZDhZgTAT_KK8k,4901
-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=992wy0YhBG2WvxwdLEPL4_-JUl4NkwMPT-jj_BIHo6A,7347
-llm_trainer/trainer.py,sha256=Q821nlLDKRZVpaRoiZ7DiJplpAJRRLtvR_33FbClGA0,26729
-llm_trainer/utils.py,sha256=LWNhyQ0NDEZ9mZtk2Ryvh6EulvHIaUGIflugSpqmeFI,6791
-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,,