大模型解码过程中计算顺序对结果的影响

这个问题是在研究样本打包的过程中引起的，但是其影响并不限于样本打包，一切影响计算顺序的操作都可能影响计算结果 。

一段脚本

model = AutoModelForCausalLM.from_pretrained(model_path, device_map=device, torch_dtype=torch.float32)
model = model.eval()

sample1 = [0,1,2,3,4]
sample2 = [0,1,2]
input_ids = [
    sample1 + sample2,
    sample1 + sample2,
    sample2 + sample1,
    sample2 + sample1
]
attention_mask = [
    [1,1,1,1,1,0,0,0],  # no/fake packing
    [1,1,1,1,1,1,1,1],  # real packing
    [0,0,0,1,1,1,1,1],
    [0,0,0,1,1,1,1,1]
]
position_ids = [
    [0,1,2,3,4,5,6,7], 
    [0,1,2,3,4,5,6,7], 
    [0,1,2,3,4,5,6,7],  # without reset position id
    [0,1,2,0,1,2,3,4]   # reset position id
]
logits = model(torch.tensor(input_ids), attention_mask=torch.tensor(attention_mask), position_ids=torch.tensor(position_ids)).logits

# expect to be the same, but not
print((logits[0][0:5] == logits[1][0:5]).all())
print((logits[0][0:5] == logits[3][3:8]).all())
print((logits[0][0:5] == logits[2][3:8]).all())
print((logits[3][3:8] == logits[2][3:8]).all())

几个实验现象：

  • 由于ROPE具有相对位置编码的特性，重置position id其实是不必要的，但是无论是否重置和怎么mask，计算结果都无法完全一致，除了logits[0]和logits[1]。
  • 不管是四条数据一个batch推理或者是分开推理，上面的误差都存在
  • fp32下，logits的最大差异在$10^{-5}$量级，bf16下，能达到$10^{-1}$量级

影响范围

  • 样本打包和分块下三角的attention mask
    • 关于打包的概念，看这个图 https://huggingface.co/blog/sirluk/llm-sequence-packing 数据打包是需要mask掉不相关的跨文档影响的，这是普遍的操作
      • 那么根据上面的脚本，数据放在前面和后面，其实推理结果有差异，那训练过程也有差异，社区没有看到关于这个问题的讨论，很奇怪（？？？）
  • 题外话：社区对于打包研究是很多的
    • 也有不普遍的操作，就是meta的 In-context Pretraining: Language Modeling Beyond Document Boundaries，不mask跨文档影响，只是打包时候要用相关的数据
    • 智谱也考虑了打包时候对长度不同的样本的loss设计 LongAlign: A Recipe for Long Context Alignment of Large Language Models
    • 智谱甚至考虑了长度不均下的训练负载均衡 GLM Long: Scaling Pre-trained Model Contexts to Millions
    • remark: 长度不均下的pretrain仅仅是效率问题，但是SFT下，长度不均会影响模型的学习效果，这个问题目前还没有看到相关的研究，但是去年在 llama2 的技术报告刚发布的时候，meta 说他们的 sft 阶段和 pretrain 阶段一样把数据 concat 成 4K 长度做训练，且不做 attention_mask。我们觉着很合理，不同的数据 concat 在一起，可以培养模型学习 session 是否和当前 query 有关的能力，因此就去复现了，结果发现“测试数据中，分类任务的 ACC 有明显下降”。经过一通分析后，我们发现，新的训练方式改变了短 answer 数据的 loss 占比，毕竟模型在计算 loss 的时候，是先算一个句子内每个 token 的 平均 loss，再算一个 batch_size 内的平均 loss。分类任务的 answer 通常只有 1 个 token：不 concat 的时候，它的 loss 贡献就是 1 / batch_size；concat 的时候，它就需要先和别的 answer 的 token 算平均 loss，再贡献 1 / batch_size。这也就是说，采用 llama2 提到的 先 concat 语料再做 sft 训练，会对短 answer 数据很不公平，也就更容易造成短 answer 数据的欠拟合，pretrain 由于所有 token 都算 loss 则没有这个现象。最终，我们通过上采样短 answer 数据，成功的避免了分类任务的效果下滑。

  • 一个不太好的消息：llama factory不支持FA2下的packing分块注意力 代码片段1：https://github.com/hiyouga/LLaMA-Factory/blob/f54733460427cf2b126b8b8737fdd732c0e19d9c/src/llamafactory/data/collator.py#L41 代码片段2：https://github.com/hiyouga/LLaMA-Factory/blob/f54733460427cf2b126b8b8737fdd732c0e19d9c/src/llamafactory/data/collator.py#L224

    • 如果sft数据一条都很长，不涉及packing，这个对我们来说也无所谓了。以后要用到的人可能要留意这里
    • FA2+packing 业界也有了。。https://huggingface.co/blog/zh/packing-with-FA2 给大佬磕头
  • 推理也受影响
    • Continuous batching
    • Beam search Ref: https://huggingface.co/blog/poedator/4d-masks

原理猜想

需要更多实验进行验证
  • 计算顺序带来的累积误差
    • transformer的计算中，有很多累积操作，比如attention的加权求和，residual连接的累加
    • 这些累积操作中，计算顺序不同会带来不同的舍入误差，从而影响最终结果
  • 非线性操作的敏感性
    • transformer中有很多非线性操作，比如softmax，gelu等
    • 这些非线性操作对输入的微小变化可能会放大，从而导致输出的显著差异
  • 归一化层的影响
    • transformer中有归一化层，比如layer norm
    • 归一化层会根据输入的分布进行调整，计算顺序不同可能会影响输入的分布，从而影响归一化的结果

💬 发表新评论