Use fuse multi head att

[xiezipeng-ML]

No description provided.

[xiezipeng-ML]

batch size = 4, acc step = 8, amp, open Checkpointing

1n1g	use_fuse_multi_head_att = False	use_fuse_multi_head_att = True
Throughput	total_throughput: 151.70 samples/s	total_throughput: 155.41 samples/s
GPU Memory	3147MiB	3129MiB

在encoder和decoder中的self_att和cross_att中都使用了fuse_multihead_att.
在28号上简单测了一下，带来的提升有限，应该是transpose的使用次数太多，我下个commit准备把if，else直接取消，默认使用fuse_multihead_att来测一下.

@chengtbf @strint @ouyangyu @CPFLAME

[CPFLAME]

这里在进行cross_attention的时候, 每次都需要transpose一下

有没有办法只在外面transpose一次?

[chengtbf]

原本的 fuse multihead attention 就是在整个网络一开始做一次 transpose，使得整个 transformer layer 里的数据维度都是转置后的维度（到 loss 的部分可以考虑再转置回来，但是 Megatron 即使在 loss 部分貌似也是转置后的），这样可以省去每个 layer 内部的 transpose 操作。这样才能达到提升速度的目的。

[chengtbf]

参考代码 ：

• https://github.com/NVIDIA/Megatron-LM/blob/0bb597b42c53355a567aba2a1357cc34b9d99ddd/megatron/model/transformer.py#L395

默认进到 attention 里面的就是 [sq, b, h] ， batch size 在 dim 1

[chengtbf]

可以参考之前 程鹏 和 星宇 的代码：

• 16ca685#diff-0579a91f68617ee49042d24fa0fade3f3e1171363d3a14ae611309fa54601020R385

• libai/libai/models/gpt_model.py

  Line 311 in b1c7d32

  hidden_states = hidden_states.transpose(0, 1) # [seq, bs, dim]

[xiezipeng-ML]

• 这里魔改了一下，self_att和cross_att都使用了fuse_muti_head_att，attention层默认为fuse_multi_head_att，一共只多出3个必须的transpose：encode_embedding的输出进行一次transpose，decoder_embedding的输出进行一次transpose，loss接收的logits进行一次transpose
• 如果数据处理的时候直接处理成[seq_len, batch_size]的shape的话上述3个transpose可以取消
• 用这个pr下面的单测测过了修改后的模型和huggingface对齐：tests/model_utils/test_mt5_loader_2.py

@chengtbf @CPFLAME @strint @ouyangyu