[fix] SFTDataset: keep assistant span when long prompt is truncated

渲染长度超过 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
pull/888/head
uphomer 2026-10-15 22:26:21 +08:00
parent f659b55761
commit a067e2512c
1 changed files with 12 additions and 1 deletions

View File

@ -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)
# # === 调试打印 ===