Open xiezipeng-ML opened 2 years ago
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
可以取消tests/model_utils/test_mt5_loader_2.py
@chengtbf @CPFLAME @strint @ouyangyu
batch size = 4, acc step = 8, amp, open Checkpointing
在
encoder
和decoder
中的self_att
和cross_att
中都使用了fuse_multihead_att
. 在28号上简单测了一下,带来的提升有限,应该是transpose的使用次数太多,我下个commit准备把if,else直接取消,默认使用fuse_multihead_att
来测一下.@chengtbf @strint @ouyangyu @CPFLAME