PyTorch因果语言模型中对数概率求和不一致问题排查

我使用Hugging Face Transformers库与PyTorch开展项目，计算因果语言模型的对数概率时发现：同一句子单独计算的对数概率总和，与将其放入批量（和另一句长句子一起）计算时的对应结果不一致。

简化复现代码 #

import transformers
import torch
import random
import numpy as np

torch.backends.cudnn.deterministic = True

def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)

set_seed(42)
hf_name = "cerebras/Cerebras-GPT-111M"
tokenizer_arguments = {"truncation": True,
                       "max_length": 1096,
                       "padding_side": "left",
                       "add_special_tokens": False}
tokenizer = transformers.AutoTokenizer.from_pretrained(hf_name, **tokenizer_arguments)
tokenizer.pad_token = tokenizer.eos_token
model = transformers.AutoModelForCausalLM.from_pretrained(hf_name, torch_dtype="bfloat16")

one_sentence = "Hello this is John"
longer_sentence = "Hello guys, my name is John, nice to meet you"
no_batch = [one_sentence]
batch = [one_sentence, longer_sentence]

def tokenize(content):
    return tokenizer(content,
                     return_tensors="pt",
                     truncation=False,
                     padding="longest",
                     add_special_tokens=True)

input_no_batch = tokenize(no_batch)
input_batch = tokenize(batch)

with torch.no_grad():
    outputs_no_batch = model(**input_no_batch)
    outputs_batch = model(**input_batch)

def get_log_probs(logits, _masks):
    logits = logits * _masks.unsqueeze(-1)
    log_probs = torch.log_softmax(logits, dim=-1)
    return log_probs

att_masks = input_no_batch['attention_mask']
att_masks2 = input_batch['attention_mask']

no_batch_probs = get_log_probs(outputs_no_batch.logits, att_masks)
batch_probs = get_log_probs(outputs_batch.logits, att_masks2)

after_masks_no_batch = no_batch_probs * att_masks.unsqueeze(-1)
after_masks_batch = batch_probs * att_masks2.unsqueeze(-1)

print(torch.sum(after_masks_no_batch).item())  # Sum of log probabilities for single sentence
print(torch.sum(after_masks_batch[0, :, :]).item())  # Sum of log probabilities for the same sentence in batch

预期行为 #

torch.sum(after_masks_no_batch).item()与torch.sum(after_masks_batch[0, :, :]).item()应输出相等数值，因为二者针对的是同一句子的对数概率求和。

已尝试的排查步骤 #

• 确保模型处于评估模式
• 核对分词器参数一致性
• 打印中间形状与数值追踪差异

差异产生原因 #

1. bfloat16精度限制：模型使用torch_dtype="bfloat16"，该精度的有效位数远低于float32，批量与单样本计算时的数值舍入误差会被累积放大，导致最终结果偏差。
2. 左padding的计算路径差异：分词器设置了padding_side="left"，单样本输入无padding，而批量输入中短句子左侧会填充pad token。即使注意力掩码屏蔽了pad token，模型内部的矩阵运算（如层归一化、注意力得分计算）在两种场景下的数值计算路径不同，低精度进一步放大了这种差异。
3. 对数概率计算逻辑问题：当前get_log_probs先将logits与掩码相乘（把pad位置logits置0）再做log_softmax，这会改变softmax的计算分布，引入额外误差。
4. CuDNN确定性设置的局限性：虽然开启了torch.backends.cudnn.deterministic = True，但bfloat16的部分运算可能未完全支持确定性实现，不同批量大小的优化路径也可能引入非确定性误差。

解决思路 #

• 改用float32精度：将模型加载时的torch_dtype改为torch.float32，牺牲部分显存换取更高精度，可大幅减少计算误差。
• 统一输入padding：单样本输入时手动添加与批量场景一致的左侧padding，确保输入张量形状完全匹配，消除计算路径差异。例如：

  # 单样本tokenize后手动pad到批量最大长度
  max_len = input_batch['input_ids'].shape[1]
  input_no_batch = tokenizer.pad(input_no_batch, padding='max_length', max_length=max_len, return_tensors='pt')
  

• 修正对数概率计算逻辑：先计算log_softmax再用掩码过滤pad位置，避免干扰softmax的分布计算：

  def get_log_probs(logits, _masks):
      log_probs = torch.log_softmax(logits, dim=-1)
      # 仅保留有效token的对数概率，pad位置置0不影响求和
      return log_probs * _masks.unsqueeze(-1)
  

• 显式设置模型评估模式：推理前调用model.eval()，确保dropout等训练相关层被关闭（部分模型默认可能未开启，但显式设置更稳妥）。

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