BioGPT因果语言模型调用遇Attention mask尺寸不匹配错误求助

问题 #

尝试使用BioGPT的因果语言模型时出现异常，第一次迭代正常，第二次调用模型触发以下错误：

ValueError: Attention mask should be of size (1, 1, 0, 12), but is torch.Size([1, 1, 1, 1])

操作步骤：

1. 安装依赖：

!pip install transformers sacremoses -q

2. 执行生成代码：

input_sequence = "Hello, I'm a language model,"

inputs = torch.as_tensor(tokenizer.encode(input_sequence)).unsqueeze(0).to(device)
past_key_values = None

count = 0
complete_token = []
with torch.no_grad():
    while count<10:
        count += 1
        print("Iteration no.: " + str(count))
        if count > 1:
            inputs = input_token

        model_out = model(input_ids=inputs.to(device), past_key_values=past_key_values)
        logits = model_out.logits[:, -1, :]
        past_key_values = model_out.past_key_values

        topk_values, topk_indices = torch.topk(logits, 5)

        log_probs = F.softmax(topk_values, dim=-1)
        inputs_in_topk = torch.multinomial(log_probs, num_samples=1, replacement=True)
        input_token = torch.gather(topk_indices, 1, inputs_in_topk)
        complete_token.append(input_token)

错误原因 #

第二次迭代时，仅传入单个新生成的token作为input_ids，但past_key_values中保存了第一次迭代时的完整上下文长度（原输入的12个token）。BioGPT结合past_key_values使用时，会自动计算需要的attention mask维度，此时传入的input_ids长度为1，导致模型预期的mask维度和实际传入的不匹配（实际mask只覆盖当前1个token，模型需要覆盖历史12个token+当前token的长度）。

解决方法 #

提供三种可行的修正方式：

方式一：放弃使用past_key_values，每次拼接生成的token到输入序列 #

这种方式更直观，无需处理复杂的attention mask：

input_sequence = "Hello, I'm a language model,"

# 初始输入编码
inputs = torch.as_tensor(tokenizer.encode(input_sequence)).unsqueeze(0).to(device)
complete_token = []

with torch.no_grad():
    for count in range(10):
        print("Iteration no.: " + str(count+1))
        model_out = model(input_ids=inputs.to(device))
        logits = model_out.logits[:, -1, :]
        
        topk_values, topk_indices = torch.topk(logits, 5)
        log_probs = F.softmax(topk_values, dim=-1)
        inputs_in_topk = torch.multinomial(log_probs, num_samples=1, replacement=True)
        input_token = torch.gather(topk_indices, 1, inputs_in_topk)
        
        # 将新生成的token拼接到输入序列末尾
        inputs = torch.cat([inputs, input_token], dim=-1)
        complete_token.append(input_token)

# 解码最终生成的文本
generated_text = tokenizer.decode(inputs[0], skip_special_tokens=True)
print(generated_text)

方式二：正确传递attention_mask（适合需要用past_key_values优化速度的场景） #

如果要保留past_key_values来加速生成，需要手动构建匹配历史长度的attention mask：

input_sequence = "Hello, I'm a language model,"

inputs = torch.as_tensor(tokenizer.encode(input_sequence)).unsqueeze(0).to(device)
past_key_values = None
# 初始attention mask：全1，覆盖所有初始输入token
attention_mask = torch.ones_like(inputs)

count = 0
complete_token = []
with torch.no_grad():
    while count<10:
        count += 1
        print("Iteration no.: " + str(count))
        if count > 1:
            inputs = input_token
            # 更新attention mask：拼接1，新增对当前token的掩码
            attention_mask = torch.cat([attention_mask, torch.ones_like(inputs)], dim=-1)

        model_out = model(
            input_ids=inputs.to(device), 
            past_key_values=past_key_values,
            attention_mask=attention_mask.to(device)
        )
        logits = model_out.logits[:, -1, :]
        past_key_values = model_out.past_key_values

        topk_values, topk_indices = torch.topk(logits, 5)
        log_probs = F.softmax(topk_values, dim=-1)
        inputs_in_topk = torch.multinomial(log_probs, num_samples=1, replacement=True)
        input_token = torch.gather(topk_indices, 1, inputs_in_topk)
        complete_token.append(input_token)

方式三：使用transformers内置的generate方法（推荐） #

transformers库提供了封装好的generate方法，无需手动实现循环，更简洁且不易出错：

input_sequence = "Hello, I'm a language model,"
inputs = tokenizer(input_sequence, return_tensors="pt").to(device)

# 生成10个新token
outputs = model.generate(
    **inputs,
    max_new_tokens=10,
    do_sample=True,
    top_k=5,
    temperature=1.0
)

generated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
print(generated_text)

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