使用DataCollatorForLanguageModeling时mlm=True且batch_size>1触发RuntimeError

问题原因 #

报错RuntimeError: output with shape [3, 1, 32] doesn't match the broadcast shape [3, 3, 32]的核心原因是自定义的SmilesDataset返回的样本张量多了一个额外维度。

当mlm=True时，DataCollatorForLanguageModeling会生成掩码并对输入进行掩码替换，这个过程需要对批量张量做广播操作。如果每个样本的input_ids是二维张量（比如[1, 32]），批量拼接后会变成[3, 1, 32]，而掩码操作需要的形状是[3, 32]（对应批量中每个样本的序列），广播时就会出现维度不匹配问题。

而mlm=False时，collator仅做简单拼接和padding，不会触发掩码广播逻辑；batch_size=1时，[1,1,32]的形状不会引发广播冲突，所以能正常运行。

解决方案 #

检查并修改SmilesDataset的__getitem__方法，确保返回的input_ids、attention_mask等张量是一维的，而非二维。

错误示例（导致额外维度） #

class SmilesDataset(Dataset):
    def __init__(self, data, tokenizer):
        self.data = data
        self.tokenizer = tokenizer
    
    def __getitem__(self, idx):
        text = self.data[idx]
        tokenized = self.tokenizer(text, padding='max_length', truncation=True, max_length=32)
        # 错误：用[]包裹生成了二维张量
        return {
            'input_ids': torch.tensor([tokenized['input_ids']]),
            'attention_mask': torch.tensor([tokenized['attention_mask']])
        }

修正后的代码 #

class SmilesDataset(Dataset):
    def __init__(self, data, tokenizer):
        self.data = data
        self.tokenizer = tokenizer
    
    def __getitem__(self, idx):
        text = self.data[idx]
        tokenized = self.tokenizer(text, padding='max_length', truncation=True, max_length=32)
        # 正确：直接转换为一维张量
        return {
            'input_ids': torch.tensor(tokenized['input_ids']),
            'attention_mask': torch.tensor(tokenized['attention_mask'])
        }

修改后，批量拼接后的input_ids形状会变成[3, 32]，DataCollatorForLanguageModeling的掩码广播逻辑就能正常工作，不会再触发维度不匹配的错误。

内容的提问来源于stack exchange，提问作者AnonHymn