From d6401cbb92c773584df7182252efc4c439f79a2e Mon Sep 17 00:00:00 2001 From: Linxiushen <1787979356@qq.com> Date: Sun, 20 Sep 2026 10:36:32 +0800 Subject: [PATCH] =?UTF-8?q?[fix]=20PretrainDataset=20=E6=8C=89=20token=20i?= =?UTF-8?q?d=20=E5=B1=8F=E8=94=BD=20padding=EF=BC=8C=E4=BC=9A=E8=AF=AF?= =?UTF-8?q?=E4=BC=A4=E6=AD=A3=E6=96=87=E9=87=8C=E5=AD=97=E9=9D=A2=E7=9A=84?= =?UTF-8?q?=20<|endoftext|>?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit pad_token 是 <|endoftext|>(id 0)。正文里字面出现这个串时会被编成同一个 id， 而 labels[input_ids == pad_token_id] = -100 是按 token id 屏蔽的， 于是这个真实的正文 token 被当成 padding 从 loss 里丢掉。 改为按位置屏蔽 labels[n_real:] = -100，n_real = len(bos + 正文 + eos)。 padding 只可能出现在序列尾部，所以按位置屏蔽在语义上更准确，也更便宜。 SFTDataset 不受影响：它的 generate_labels 从全 -100 出发、扫描 assistant 区间来填，不依赖 pad id。 --- dataset/lm_dataset.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/dataset/lm_dataset.py b/dataset/lm_dataset.py index bbc4197..d02ba4e 100644 --- a/dataset/lm_dataset.py +++ b/dataset/lm_dataset.py @@ -48,10 +48,15 @@ class PretrainDataset(Dataset): sample = self.samples[index] tokens = self.tokenizer(str(sample['text']), add_special_tokens=False, max_length=self.max_length - 2, truncation=True).input_ids tokens = [self.tokenizer.bos_token_id] + tokens + [self.tokenizer.eos_token_id] - input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - len(tokens)) + n_real = len(tokens) # bos + 正文 + eos，后面才是 padding + input_ids = tokens + [self.tokenizer.pad_token_id] * (self.max_length - n_real) input_ids = torch.tensor(input_ids, dtype=torch.long) labels = input_ids.clone() - labels[input_ids == self.tokenizer.pad_token_id] = -100 + # 按位置屏蔽，而不是按 token id。pad_token 是 <|endoftext|>(id 0)， + # 正文里字面出现这个串时会被编成同一个 id + # （tokenizer('前面<|endoftext|>后面') -> [938, 605, 0, 701, 605]）， + # 用 `input_ids == pad_token_id` 会把这个真实的正文 token 一起屏蔽掉。 + labels[n_real:] = -100 return input_ids, labels