alibaba/Megatron-LLaMA

llama中decoder layer层里面的MLP问题

yuanzhoulvpi2017 opened this issue · 4 comments

yuanzhoulvpi2017 commented

开发者您好，对代码有一个疑问，想要请教一下。

情况描述

part 1

  1. 在debug代码的时候，发现你们的llama模型里面的decoder layer层的mlp部分，使用的是Megatron-lm的ParallelMLP。
  2. 但是这个算子，只有两个线性层。他们分别为
    1. self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear 、
    1. self.dense_4h_to_h = tensor_parallel.RowParallelLinear。
  1. 这个ParallelMLP 层的链接为

    Megatron-LLaMA/megatron/model/transformer.py

    Line 80 in 4da868f

    class ParallelMLP(MegatronModule):

part2

但是在hugginfgace的transformers代码中的llama的MLP部分、在llama官方的GitHub仓库提供的代码，llama的MLP部分（也叫feedward层），使用的是结构是有三个线性层。他们分别为:

    1. self.gate_proj = tensor_parallel.ColumnParallelLinear
    1. self.up_proj = tensor_parallel.ColumnParallelLinear
    1. self.down_proj = tensor_parallel.RowParallelLinear

hf的transformers的实现llama的MLP参考链接：https://github.com/huggingface/transformers/blob/4557a0dede92ce985576fac478b754d76bba3c18/src/transformers/models/llama/modeling_llama.py#L229

part3

  1. 上面也只是简单列举了贵团队和标准的llama的MLP部分的差异，计算方式也有比较大的差异，这里就不列举了。

我的问题：

想请教一下，你们的MLP计算方式和官方的计算方式差异是数学上等价的？还是bug？

感谢贵团队提供了这么优秀的开源作品，希望有同学可以解释一下这个问题, 谢谢，辛苦～

yuanzhoulvpi2017 commented

早上，脑子清晰，才搞懂这个问题：

问题：

  1. 为什么megaton训练llama的模型，在mlp部分只有dense_h_to_4h;dense_4h_to_h.
  2. 而hf版的llama的mlp部分是有gate_proj, up_proj, down_proj,三个线性层。

原因：

  1. 因为在llama里面，gate_proj, up_proj的线性层的input_size, output_size都是一样的
  2. 所以就把这两个矩阵叠在一起用，变成dense_h_to_4h.
  3. 然后把dense_h_to_4h输出的结果再chunk一下，就等价于原来的结果。

注意这里的dense_h_to_4h. 里面的4h只是名字上的替代，不是一定为h的4倍

花了一天时间，才搞懂这个东西[傻笑]，你们太6了

aiyinyuedejustin commented

6，这解释绝了

YLSnowy commented

早上，脑子清晰，才搞懂这个问题：

问题：

  1. 为什么megaton训练llama的模型，在mlp部分只有dense_h_to_4h;dense_4h_to_h.
  2. 而hf版的llama的mlp部分是有gate_proj, up_proj, down_proj,三个线性层。

原因：

  1. 因为在llama里面，gate_proj, up_proj的线性层的input_size, output_size都是一样的
  2. 所以就把这两个矩阵叠在一起用，变成dense_h_to_4h.
  3. 然后把dense_h_to_4h输出的结果再chunk一下，就等价于原来的结果。

注意这里的dense_h_to_4h. 里面的4h只是名字上的替代，不是一定为h的4倍

花了一天时间，才搞懂这个东西[傻笑]，你们太6了

MLP will take the input with h hidden state, project it to 4*h
hidden dimension, perform nonlinear transformation, and project the state back into h hidden dimension.
这个是megatron中的dense_h_to_4h的注释，所以，你是觉得这里的注释的4倍也是一个名字上的替代么

Hoonly commented

早上，脑子清晰，才搞懂这个问题：
问题：

  1. 为什么megaton训练llama的模型，在mlp部分只有dense_h_to_4h;dense_4h_to_h.
  2. 而hf版的llama的mlp部分是有gate_proj, up_proj, down_proj,三个线性层。

原因：

  1. 因为在llama里面，gate_proj, up_proj的线性层的input_size, output_size都是一样的
  2. 所以就把这两个矩阵叠在一起用，变成dense_h_to_4h.
  3. 然后把dense_h_to_4h输出的结果再chunk一下，就等价于原来的结果。

注意这里的dense_h_to_4h. 里面的4h只是名字上的替代，不是一定为h的4倍
花了一天时间，才搞懂这个东西[傻笑]，你们太6了

MLP will take the input with h hidden state, project it to 4*h hidden dimension, perform nonlinear transformation, and project the state back into h hidden dimension. 这个是megatron中的dense_h_to_4h的注释，所以，你是觉得这里的注释的4倍也是一个名字上的替代么

正好最近也在寻求这个问题的答案，4h确实只是指代，原始FFN是4h，但后来的演变主要看ffn_hidden_size具体是多少了