From a067e2512c1178e6ccca67f3e0a9433c4e2fe18d Mon Sep 17 00:00:00 2001 From: uphomer <1822358119@qq.com> Date: Thu, 15 Oct 2026 22:26:21 +0800 Subject: [PATCH] [fix] SFTDataset: keep assistant span when long prompt is truncated MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit 渲染长度超过 max_seq_len 时，原实现从右侧硬截断（input_ids[:max_length]）： 若 "system+tools+user" 前缀超过 max_seq_len，首个 <|im_start|>assistant 标记会被切掉/切残，generate_labels 搜不到 assistant 段，该样本 labels 全 -100， 静默丢失全部监督信号（仍做完整前向/反向，不报错、不影响 loss 曲线）。 以 sft_t2t_mini.jsonl + max_seq_len=340（蒸馏默认配置）全量实测： 905,718 条中 17,821 条（1.97%）labels 全空，其中 17,073 条是 tool-use 样本 （占全部 tool 样本的 20.1%）。 修复：仅当首个 assistant 标记会被截掉/切残，或回复恰好从窗口边界开始 （bos_pos + len(bos_id) >= max_length）时，左移窗口优先保留回复段； 其余样本行为与修复前完全一致（逐位相同）。 验证（全量）： - sft_t2t_mini：全空 17,821 -> 0 / 905,718；监督 token 均值 248.2 -> 250.9 - lora_medical：修复前后均 0 / 25,276（340 与 768，均值逐项一致） - 1500 条同种子新旧对比：1473 条原有监督样本逐位一致，0 回归 - train_distillation 冒烟 239 步，loss 有限、无 NaN --- dataset/lm_dataset.py | 13 ++++++++++++- 1 file changed, 12 insertions(+), 1 deletion(-) diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py index 6aa220b..d62ae3f 100644 --- a/dataset/lm_dataset.py +++ b/dataset/lm_dataset.py @@ -111,7 +111,18 @@ class SFTDataset(Dataset): conversations = pre_processing_chat(sample['conversations']) prompt = self.create_chat_prompt(conversations) prompt = post_processing_chat(prompt) - input_ids = self.tokenizer(prompt).input_ids[:self.max_length] + input_ids = self.tokenizer(prompt).input_ids + if len(input_ids) > self.max_length: + # 若右侧截断会把首个 assistant 头切掉/切残（导致 labels 全 -100、该样本无监督信号）， + # 则改为左移窗口，优先保留回复段；注意 bos_pos+len(bos_id)==max_length 时回复恰好 + # 从窗口边界开始、labels 全空，因此边界也需触发左移，条件取 >= + bos_pos = next((i for i in range(len(input_ids) - len(self.bos_id) + 1) + if input_ids[i:i + len(self.bos_id)] == self.bos_id), None) + if bos_pos is not None and bos_pos + len(self.bos_id) >= self.max_length: + start = min(bos_pos, len(input_ids) - self.max_length) + input_ids = input_ids[start:start + self.max_length] + else: + input_ids = input_ids[:self.max_length] input_ids += [self.tokenizer.pad_token_id] * (self.max_length - len(input_ids)) labels = self.generate_labels(input_ids) # # === 调试打印 ===