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

llm_trainer/dpo_trainer.py

@@ -1,4 +1,3 @@
-import time
 from typing import Tuple, List, Optional
 import torch
 from torch.utils.data import Dataset
@@ -11,11 +10,11 @@ from .train_configs import TrainConfig
 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 .utils import (
+    autocastcontext,
+    get_dpo_collate_fn
 )
+from .partition_utils import sync_model_params
 
 from .checkpoint import (
     save_checkpoint,
@@ -37,7 +36,7 @@ class DPOTrainer(Trainer):
             eval_prompts=eval_prompts,
             eval_image_tags=eval_image_tags
         )
-
+        self.packed_sequences = False
         self.ref_model = self._init_ref_model()
 
     def _init_ref_model(self):
@@ -204,7 +203,7 @@ class DPOTrainer(Trainer):
                         if TrainerTools().parallel.parallel_train:
                             self.train_model.require_backward_grad_sync = need_update_grad
 
-                        with self.ctx:
+                        with autocastcontext(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)
                             aux_loss = policy_outputs.get('aux_loss')
@@ -217,7 +216,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/generate_utils.py

@@ -3,7 +3,10 @@ from contextlib import nullcontext
 import torch
 from llm_model import VlmModel, KVCache
 from .tools import TrainerTools
-from .utils import batch_repeat_image_tok
+from .utils import (
+    autocastcontext,
+    batch_repeat_image_tok
+)
 
 
 def _suppress_warper(logits: torch.Tensor, suppress_tokens: List[int]) -> torch.Tensor:
@@ -124,13 +127,7 @@ def _generate(
     如果temperature很大但内容单一，需要增大k、p
     """
     use_kv_cache = True
-
-    ctx = torch.autocast(
-        device_type=device,
-        dtype=TrainerTools().dtype,
-        enabled=True,
-        cache_enabled=None
-    ) if TrainerTools().use_amp else nullcontext()
+    ctx = autocastcontext(device)
 
     if isinstance(model, VlmModel):
         tokens = batch_repeat_image_tok(tokens, tokens_per_image)
@@ -330,13 +327,7 @@ def batch_generate(
         device: Union[str, torch.device, int]
 ):
     use_kv_cache = True
-
-    ctx = torch.autocast(
-        device_type=device,
-        dtype=TrainerTools().dtype,
-        enabled=True,
-        cache_enabled=None
-    ) if TrainerTools().use_amp else nullcontext()
+    ctx = autocastcontext(device)
 
     if isinstance(model, VlmModel):
         tokens = batch_repeat_image_tok(tokens, tokens_per_image)

llm_trainer/grpo_trainer.py

@@ -1,4 +1,3 @@
-import time
 from typing import Tuple, List, Union, Callable, Optional
 import torch
 from torch.utils.data import Dataset
@@ -14,6 +13,7 @@ from .loss import GRPOLoss
 from .tools import TrainerTools
 from .generate_utils import batch_generate
 from .log import log
+from .utils import autocastcontext
 
 from .partition_utils import (
     sync_model_params,
@@ -41,6 +41,7 @@ class GRPOTrainer(Trainer):
             eval_image_tags=eval_image_tags
         )
 
+        self.packed_sequences = False
         self.reward_func = reward_func
         self.ref_model = self._init_ref_model()
 
@@ -66,8 +67,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
@@ -337,7 +342,7 @@ class GRPOTrainer(Trainer):
                             log(f'start train for batch {batch}/{batch_count_per_file}')
 
                         for grpo_step in range(self.train_config.grpo_config.grpo_steps):
-                            with self.ctx:
+                            with autocastcontext(TrainerTools().parallel.device_type):
                                 loss, aux_loss = self._maximize_grpo_objective(rollout_data)
                                 if aux_loss_coef and aux_loss:
                                     loss += aux_loss_coef * aux_loss

llm_trainer/log.py

@@ -7,6 +7,7 @@ def get_log_dir() -> str:
 
     return f'{log_dir}/' if not log_dir.endswith('/') else log_dir
 
+
 def log(msg: str, log_file=None):
     cur_time = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime())
     if not log_file:

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/parallel.py

@@ -140,6 +140,9 @@ class Parallel(ABC):
         return 1
 
     def wait(self, msg=None):
+        if self.world_size == 1:
+            return
+
         msg = f' for {msg}' if msg else ''
         log(f'wait at {self.device}{msg}')
         dist.barrier()

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/sft_trainer.py

@@ -21,6 +21,7 @@ class SFTTrainer(Trainer):
             eval_prompts=eval_prompts,
             eval_image_tags=eval_image_tags
         )
+        self.packed_sequences = False
 
     def _convert_train_args(self) -> Tuple[dict, dict, dict, bool]:
         sft_collate_fn = get_sft_collate_fn(self.train_config.mask_prompt)

llm_trainer/tools.py

@@ -31,15 +31,7 @@ class TrainerTools:
             self.tokenizer = Tokenizer(os.environ.get('TOKENIZERS_TYPE', 'zh_llama'))
             self.use_amp = 'cuda' in self.parallel.device and not isinstance(self.parallel, DsParallel)
 
-            dtype = os.environ.get('DTYPE', None)
-            self.dtype = dtypes[dtype] if dtype in dtypes else None
-
-            if not self.dtype:
-                self.dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
-
-            log(f'word_size={self.parallel.world_size},'
-                f' use_amp={self.use_amp},'
-                f' dtype={self.dtype}')
+            log(f'word_size={self.parallel.world_size}, use_amp={self.use_amp}')
 
     def _new_parallel(self):
         parallel_type = os.environ.get('PARALLEL_TYPE', 'none')

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

llm_trainer/trainer.py

@@ -1,4 +1,3 @@
-from contextlib import nullcontext
 from typing import Optional, Tuple, List, Dict, Any
 import copy
 
@@ -37,6 +36,9 @@ from .checkpoint import (
 
 from .utils import (
     set_seed,
+    autocastcontext,
+    create_doc_boundary_mask,
+    generate_position_ids,
     pretrain_collate_fn,
 )
 
@@ -55,6 +57,17 @@ class Trainer:
     ):
         set_seed()
 
+        # 是否打包序列，仅pretrain阶段需要打包序列，
+        # [[1, 1, eos, 2, 2, eos]]
+        #   doc_boundary_mask=[[[[0., 0., 0., 0., 0., 0.],
+        #           [0., 0., 0., 0., 0., 0.],
+        #           [0., 0., 0., 0., 0., 0.],
+        #           [-inf, -inf, -inf, 0., 0., 0.],
+        #           [-inf, -inf, -inf, 0., 0., 0.],
+        #           [-inf, -inf, -inf, 0., 0., 0.]]]]
+        #   position_ids=[[0, 1, 2, 0, 1, 2]]
+        self.packed_sequences = True
+
         self.train_config: TrainConfig = train_config
         self.eval_prompts = eval_prompts
         self.eval_image_tags = eval_image_tags
@@ -81,13 +94,6 @@ class Trainer:
 
         self.criterion, self.kd_loss = self._init_loss()
 
-        self.ctx = torch.autocast(
-            device_type=TrainerTools().parallel.device_type,
-            dtype=TrainerTools().dtype,
-            enabled=True,
-            cache_enabled=None
-        ) if TrainerTools().use_amp else nullcontext()
-
         load_checkpoint(
             self.train_model,
             optimizer=self.optimizer,
@@ -433,6 +439,14 @@ class Trainer:
 
         raise e
 
+    def _get_model_dtype(self):
+        if isinstance(TrainerTools().parallel, DsParallel):
+            import deepspeed
+            assert isinstance(self.train_model, deepspeed.DeepSpeedEngine)
+            return self.train_model.get_data_types()[0]
+        else:
+            return torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+
     def _eval(self, tag: str):
         with unwrap_model_for_generation(self.train_model) as generate_model:
             if TrainerTools().parallel.is_main_process:
@@ -526,8 +540,12 @@ class Trainer:
                         inputs, labels = inputs.to(TrainerTools().parallel.device), labels.to(TrainerTools().parallel.device)
                         attention_mask = inputs != TrainerTools().tokenizer.pad
 
-                        if TrainerTools().parallel.parallel_train:
-                            self.train_model.require_backward_grad_sync = need_update_grad
+                        if self.packed_sequences:
+                            doc_boundary_mask = create_doc_boundary_mask(inputs, self._get_model_dtype())
+                            position_ids = generate_position_ids(inputs)
+                        else:
+                            doc_boundary_mask = None
+                            position_ids = None
 
                         if self.pixel_values_provider and 'image_tags' in batch_data:
                             image_tags = batch_data['image_tags']
@@ -535,10 +553,15 @@ class Trainer:
                         else:
                             pixel_values = None
 
-                        with self.ctx:
+                        if TrainerTools().parallel.parallel_train:
+                            self.train_model.require_backward_grad_sync = need_update_grad
+
+                        with autocastcontext(TrainerTools().parallel.device_type):
                             result = self.train_model(
                                 inputs,
                                 attention_mask=attention_mask,
+                                doc_boundary_mask=doc_boundary_mask,
+                                position_ids=position_ids,
                                 pixel_values=pixel_values
                             )
 

llm_trainer/utils.py

@@ -1,4 +1,5 @@
 import random
+from contextlib import nullcontext
 import torch
 from torch.nn.utils.rnn import pad_sequence
 import torch.nn.functional as F
@@ -14,6 +15,115 @@ def set_seed(seed=42):
     torch.cuda.manual_seed_all(seed)
 
 
+def autocastcontext(device_type):
+    if TrainerTools().use_amp:
+        dtype = torch.bfloat16 if torch.cuda.is_available() and torch.cuda.is_bf16_supported() else torch.float16
+        return torch.autocast(
+            device_type=device_type,
+            dtype=dtype,
+            enabled=True,
+            cache_enabled=None
+        )
+    else:
+        return nullcontext()
+
+
+def create_doc_boundary_mask(
+        input_ids: torch.Tensor,
+        dtype: torch.dtype
+) -> torch.Tensor:
+    """
+    根据文档结束符 (eot) 的位置，创建一个 attention mask 来阻止跨文档的注意力。
+
+    这个函数生成的 mask 会阻止一个 token 关注 (attend to) 属于前面文档的 tokens。
+    例如，对于输入 `[[1, 2, eot, 3, 4, eot]]`，
+    tokens `3` 和 `4` 将无法关注 `1`, `2`, 和第一个 `eot`。
+
+    Args:
+        input_ids (torch.Tensor): 输入的 token ID 张量，形状为 (bsz, seq_len)。
+        dtype (torch.dtype): 数据类型。
+
+    Returns:
+        torch.Tensor: 符合 attention 机制要求的 mask 张量，
+                      形状为 (bsz, 1, seq_len, seq_len)。
+                      值为 -inf 的位置表示被屏蔽，值为 0 的位置表示允许注意力。
+    """
+    # 获取 batch size 和 sequence length
+    bsz, seq_len = input_ids.shape
+
+    # 1. 确定每个 eot_token 的位置
+    # is_eot 是一个布尔张量，形状为 (bsz, seq_len)
+    is_eot = (input_ids == TrainerTools().tokenizer.end)
+
+    # 2. 为每个 token 分配一个文档 ID
+    # 我们使用 cumsum (累加和) 来创建递增的文档 ID。一个 token 所属的文档 ID，
+    # 取决于它前面有多少个 eot。
+    # 示例:
+    # input_ids:        [[1, 2, 3, eot, 4, 5, eot]]
+    # is_eot:           [F, F, F, T, F, F, T] -> [0, 0, 0, 1, 0, 0, 1]
+    # doc_ids_ending:   [0, 0, 0, 1, 1, 1, 2] (cumsum 的结果)
+    # doc_ids:          [0, 0, 0, 0, 1, 1, 1] (向右移位后的结果)
+    # 这个结果正确地将文档 0 分配给了前四个 token，将文档 1 分配给了后三个 token。
+    doc_ids_ending = torch.cumsum(is_eot, dim=-1)
+    doc_ids = F.pad(doc_ids_ending[:, :-1], (1, 0), value=0)
+
+    # 3. 通过比较 query 和 key 的文档 ID 来创建 mask
+    # 我们的目标是：当 query token 所在的文档 ID 大于 key token 所在的文档 ID 时，进行屏蔽。
+    # query_doc_ids 形状: (bsz, seq_len, 1)
+    # key_doc_ids 形状:   (bsz, 1, seq_len)
+    query_doc_ids = doc_ids.unsqueeze(2)
+    key_doc_ids = doc_ids.unsqueeze(1)
+
+    # 利用 PyTorch 的广播机制，`query_doc_ids > key_doc_ids` 会创建一个
+    # 形状为 (bsz, seq_len, seq_len) 的布尔张量。
+    # 当 query 的文档 ID 大于 key 的文档 ID 时，值为 True，这正是我们需要屏蔽的位置。
+    boundary_mask = query_doc_ids > key_doc_ids
+
+    # 4. 将布尔 mask 转换为 attention 机制所需的浮点数 mask (-inf 和 0)
+    final_mask = torch.zeros(
+        (bsz, seq_len, seq_len), device=input_ids.device, dtype=dtype
+    )
+    final_mask.masked_fill_(boundary_mask, torch.finfo(dtype).min)
+
+    # 5. 增加一个维度以匹配 attention head 的输入要求 (bsz, num_heads, seq_len, seq_len)
+    #    这里我们只生成一个 mask，它可以被广播到所有的 head。
+    return final_mask.unsqueeze(1)
+
+
+def generate_position_ids(input_ids: torch.Tensor):
+    """
+    为打包序列生成 position_ids 张量。
+
+    参数:
+      input_ids (torch.Tensor): 输入的 token ID 张量 (batch_size, sequence_length)。
+      end_of_text_id (int): 代表文本结束的特殊 token ID。
+
+    返回:
+      torch.Tensor: 生成的 position_ids 张量。
+    """
+    # 获取输入张量的形状
+    batch_size, seq_length = input_ids.shape
+
+    # 创建一个与输入形状相同，全为0的张量来存储position_ids
+    # 第一个token的位置永远是0，所以这个初始化是正确的
+    position_ids = torch.zeros_like(input_ids, dtype=torch.long)
+
+    # 从第二个时间步 (t=1) 开始遍历整个序列
+    for t in range(1, seq_length):
+        # 检查前一个时间步 (t-1) 的token是否为 EOT token
+        # 这会为批次中的每个序列生成一个布尔值
+        is_reset_token = (input_ids[:, t - 1] == TrainerTools().tokenizer.end)
+
+        # 获取前一个时间步的位置ID
+        prev_position_ids = position_ids[:, t - 1]
+
+        # 如果前一个token是EOT，当前位置重置为0；否则，在前一个位置上加1
+        # torch.where 会根据 is_reset_token 的布尔值进行选择
+        position_ids[:, t] = torch.where(is_reset_token, 0, prev_position_ids + 1)
+
+    return position_ids
+
+
 def repeat_image_tok(
         tokens: torch.Tensor,
         tokens_per_image: int
@@ -43,43 +153,6 @@ def batch_repeat_image_tok(
     return torch.stack(new_tokens, dim=0)
 
 
-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)
-
-    return inputs, labels
-
-
-def _mask_prompt(labels):
-    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:
-                labels[batch, start_index:index + 1] = -100
-                start_index = -1
-
-    return labels
-
-
-def _zero_pad_sequences(
-    sequences: list[torch.Tensor], side: str = "left"
-) -> torch.Tensor:
-    assert side in ("left", "right")
-    max_len = max(seq.size(0) for seq in sequences)
-    padded_sequences = []
-    for seq in sequences:
-        pad_len = max_len - seq.size(0)
-        padding = (pad_len, 0) if side == "left" else (0, pad_len)
-        padded_sequences.append(F.pad(seq, padding))
-    return torch.stack(padded_sequences, dim=0)
-
-
 def pretrain_collate_fn(batch_data):
     inputs, labels = _pad_sequence(batch_data)
 
@@ -219,4 +292,41 @@ def join_batch(batch_data: list[dict]) -> dict:
             data = None
         result[key] = data
 
-    return result
+    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)
+
+    return inputs, labels
+
+
+def _mask_prompt(labels):
+    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:
+                labels[batch, start_index:index + 1] = -100
+                start_index = -1
+
+    return labels
+
+
+def _zero_pad_sequences(
+    sequences: list[torch.Tensor], side: str = "left"
+) -> torch.Tensor:
+    assert side in ("left", "right")
+    max_len = max(seq.size(0) for seq in sequences)
+    padded_sequences = []
+    for seq in sequences:
+        pad_len = max_len - seq.size(0)
+        padding = (pad_len, 0) if side == "left" else (0, pad_len)
+        padded_sequences.append(F.pad(seq, padding))
+    return torch.stack(padded_sequences, dim=0)

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

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

project_llm_trainer-0.7.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=_8ZwOKQH69c6Fa5Cey5hNep7XUoI4jPIXQaQcV3soGw,12367
+llm_trainer/ds_checkpoint.py,sha256=Wzy7PvVVWR794-BW4uragWFTAkkgDvjvkF-qMdyB4fc,2141
+llm_trainer/eval.py,sha256=ZyUfSo2Q8P-lrCdPEnGkoo5pGubd0AabREK5eMISRII,1109
+llm_trainer/generate_utils.py,sha256=zX5218RX4ltahCQCZVVCWQghCWhKslPk2NUnl_CakIE,15050
+llm_trainer/grpo_trainer.py,sha256=0iWvpuMI5CDNIjH08Dd1ihZFqDYenVnHACiMY2GLJtg,16449
+llm_trainer/log.py,sha256=XwychwKF6gvFPhthCIZCAEUZ0G3DY3fiQrOHqPWsxz0,463
+llm_trainer/loss.py,sha256=eYvOlCoguKnLvdGuqvQpGUoLVSADQ5coaU3DWYbJEdM,6811
+llm_trainer/parallel.py,sha256=yjStV21DJ26yM8-0O6GTMxdFAcyShY5GsQWSZmbI7HU,4543
+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=LudTRIaqLQYy6ym6jjMX7v9xtFBJelrR3nnPCwb48nM,1821
+llm_trainer/tokenizer.py,sha256=SSpgXtb0e1NtQqRW0gCq09TTZi47umggy-Fh5EMHKJg,6708
+llm_trainer/tools.py,sha256=5op5qrjjkK-Lr9oes5VxIVnOVYOYGoAdlIJq9mPUf64,2637
+llm_trainer/train_configs.py,sha256=U4hwXWKI6svDqiDOu6RPTitCzpxEYyjZUN6gwh_co8c,7510
+llm_trainer/trainer.py,sha256=2TC2GJeoGd0fDE6CFodk1chsSkk0v0yO0wrFYim5t4g,27938
+llm_trainer/utils.py,sha256=ox2fWtSOS7F2Nh7_FoHxuQgaps1jGW3q59VXz04wRuA,11491
+project_llm_trainer-0.7.0.data/scripts/calc_intermediate_size,sha256=AggpgNHokJiJMbEtVdOnolqr_4bH3i1UYuZNEAzC2Gc,460
+project_llm_trainer-0.7.0.data/scripts/ddp_train,sha256=x81AasaN2-9TwARFFF1l7iV1LmfMQ0bLw0i_CGbOwSw,299
+project_llm_trainer-0.7.0.data/scripts/ds_train,sha256=qL3qc3TcedBCw98UZUjW07ONcErRawLE1HymW2AmscA,265
+project_llm_trainer-0.7.0.data/scripts/plot_loss,sha256=MzFcdJESlVr1srj4Td6-AxPGUKkfB_QEcJwm0Bd-5fU,910
+project_llm_trainer-0.7.0.data/scripts/plot_lr,sha256=w_7XR_x3KYYyboeOVAeu_I4fveLFI-C0wBmRrNlmWUI,894
+project_llm_trainer-0.7.0.data/scripts/py_train,sha256=tOp9TquORQeU8XN5H7OVIk5O0Ypwi34p_GENxTwgwdk,265
+project_llm_trainer-0.7.0.data/scripts/smart_train,sha256=Pmt4Q0to4Hoz82iB9uFPZuz7uahNUbfE7FR1940EBy8,716
+project_llm_trainer-0.7.0.dist-info/METADATA,sha256=Q_UU9xBZIIBFOmfQJg1708lFfYn4bu5FA0fuxJCCcxQ,195
+project_llm_trainer-0.7.0.dist-info/WHEEL,sha256=Nw36Djuh_5VDukK0H78QzOX-_FQEo6V37m3nkm96gtU,91
+project_llm_trainer-0.7.0.dist-info/top_level.txt,sha256=LtRFg28i0QIG7iBCD2t095oSco99LCtkijibS9cMGik,12
+project_llm_trainer-0.7.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,,