[Bug][MT5] Throughput is unexpected

[strint]

32m[10/19 15:36:23 lb.utils.events] [0 m eta: 7986 days, 7:30:37 iteration: 99/621340880 consumed_samples: 200 total_loss: 9.545 time: 1.1185 s/iter data_time: 0.0151 s/iter total_throughput: 1.79 samples/s lr: 1.02e-08

[xyn1201]

t5 单机4卡测试

• 机器：oneflow-25 单机4卡

• oneflow master Oneflow-Inc/oneflow@93d19f3

• libai main e9ca408

• 用例：t5_nl12_nah12_hs768_FP16_actrue_mp2_pp1_mb32_gb512_1n4g zero_stage=2

  • libai：4089 MiB /85.83 samples/s
    日志：oss://oneflow-test/libai_vs_megatron/1020_t5/25/libai_zero/
  • Megatron-deepspeed：4725 MiB /82.7 samples/s

t5 2机4卡测试

• 机器：oneflow-25 oneflow-28 2机一共8卡

• oneflow master Oneflow-Inc/oneflow@93d19f3

• libai main e9ca408

• 用例：t5_nl12_nah12_hs768_FP16_actrue_mp2_pp1_mb16_gb512_2n4g zero_stage=2

  • libai：4530MiB/50.92 samples/s
    日志：oss://oneflow-test/libai_vs_megatron/1019_t5/libai_zero/
  • Megatron-deepspeed：4783 MiB/164.3 samples/s

[chengtbf]

这里比较明显的问题是，我们 4 卡 2-D 并行是超过 Megatron 的，但是 两机8卡的吞吐比 单机四卡的还慢。而 Megatron 是一个线性的加速比。

[xiezipeng-ML]

这里有点问题，libai.models.T5Model是megatron的版本，IDEA需要的是huggingface版本的T5，也就是libai的projects下的T5（projects/T5是交付项目），这两个模型结构有区别，已经让yongning增加一份projects/T5的测试了，交付之前也是用projects/T5来和libai.model.T5Model来测的纯数据并行：here，两个模型不一样，感觉不能简单地去比较和megatron的性能，因为megatron实现的不是huggingface版本的T5

两个T5的区别总结：

• layernorm对应的算子不同（mt5用c++拼接算子:RMSLayernorm）

• decoder多一层embedding：

  libai/projects/T5/models/attention.py

  Lines 117 to 120 in b3c5ba2

  	if self.has_relative_attention_bias:
  	self.relative_attention_bias = Embedding(
  	self.relative_attention_num_buckets, self.num_heads, layer_idx=layer_idx
  	)

• dropout对应算子不同 (mt5使用的是：fix gelu nn.Module bug and support tanh mode. oneflow#8693)

• mt5（projects下的T5）的lm_head没有共享embedding的参数 （

  libai/projects/T5/models/t5_model.py

  Lines 129 to 134 in 9a4af26

  	if model_type == "mt5":
  	self.lm_head = Linear(
  	hidden_size, vocab_size, bias=False, layer_idx=2 * hidden_layers - 1
  	)
  	else:
  	self.lm_head = LMLogits(vocab_size, bias=False)

  ）

• mt5（projects下的T5）比t5（libai.models中的T5）少了position_embedding，但是mt5中的attention多出了position_bias的相关计算（

  libai/projects/T5/models/attention.py

  Line 272 in e9ca408

  def _relative_position_bucket(

  和

  libai/projects/T5/models/attention.py

  Line 320 in e9ca408

  def compute_bias(self, query_length, key_length, placement=None):

  ）

• mt5（projects下的T5）不包含任何bias. (Linear 和 LayerNorm)

• mt5（projects下的T5）因为要对齐huggingface的版本，没有用到t5（libai.models中的T5）当中的一些优化的地方，比如scale_mask_softmax_fusion，（mt5:

  libai/projects/T5/models/attention.py

  Lines 232 to 244 in 9a4af26

  	if attention_mask is not None:
  	attention_scores = flow.mul(attention_scores, attention_mask)
  	attention_scores = attention_scores - 10000.0 * (1 - attention_mask)
  	# TODO(xingyu.liao): graph will occur `where_scalar` errors
  	# when using `masked_fill`
  	# attention_scores = attention_scores.masked_fill(1 - attention_mask, -10000.0)
  	attention_weights = flow.softmax(attention_scores, dim=-1)
  	# [bsz, num_heads, tgt_len, src_len]
  	attention_weights = self.dropout(attention_weights)
  	else:
  	attention_weights = flow.softmax(attention_scores, dim=-1)
  	# [bsz, num_heads, tgt_len, src_len]
  	attention_weights = self.dropout(attention_weights)

  t5:

  libai/libai/layers/attention.py

  Lines 214 to 250 in 9a4af26

  	if attention_mask is not None:
  	if self.scale_mask_softmax_fusion:
  	if self.attn_mask_type == AttnMaskType.padding:
  	attention_mask = (
  	attention_mask.expand_as(attention_scores) if use_cache else attention_mask
  	)
  	attention_weights = flow._C.fused_scale_mask_softmax_dropout(
  	attention_scores,
  	attention_mask,
  	fill_value=-10000.0,
  	scale=self.coeff,
  	p=self.attention_dropout_prob,
  	)[0]
  	else:
  	if self.coeff is not None:
  	attention_scores *= self.coeff
  	attention_scores = flow.mul(attention_scores, attention_mask)
  	attention_scores = attention_scores - 10000.0 * (1 - attention_mask)
  	# TODO(xingyu.liao): graph will occur `where_scalar` errors
  	# when using `masked_fill`
  	# attention_scores = attention_scores.masked_fill(1 - attention_mask, -10000.0)
  	attention_weights = flow.softmax(attention_scores, dim=-1)
  	# [bsz, num_heads, tgt_len, src_len]
  	attention_weights = self.dropout(attention_weights)
  	else:
  	if self.scale_mask_softmax_fusion and self.attn_mask_type == AttnMaskType.causal:
  	attention_weights = flow._C.fused_scale_tril_softmax_mask_scale(
  	attention_scores,
  	p=self.attention_dropout_prob,
  	diagonal=0,
  	tril_scale_value=self.coeff,
  	tril_fill_value=-10000.0,
  	)[0]
  	else:
  	attention_weights = flow.softmax(attention_scores, dim=-1)
  	# [bsz, num_heads, tgt_len, src_len]
  	attention_weights = self.dropout(attention_weights)

  ）

• mt5（projects下的T5）的MLP层比t5是多出一层Linear的（https://github.com/Oneflow-Inc/libai/blob/main/projects/T5/models/mlp.py ）

• [form chengcheng] Attention 里的 FuseMultiHeadAttention 这个优化， megatron 将原本正常语义下的 batch size 转置到了 第一维 ， 这个从语义上是难以理解的，但是从性能上，可以只在 Transformer layer 之前做一次 transpose，内部的 matmul 可以使用 batch gemm 执行。 如果不这么做的话，需要在 每个 layer 内部都做 transpose ，单单这个优化就有 10% 的性能差距。 https://github.com/NVIDIA/Megatron-LM/blob/0bb597b42c53355a567aba2a1357cc34b9d99ddd/megatron/model/transformer.py#L395

mt5里没用到scale_mask_softmax_fusion，所以是走了t5的att中的else分支

def att_mt5(attention_scores, attention_mask):
    dropout = nn.Dropout(0)
    attention_scores = flow.mul(attention_scores, attention_mask)
    attention_scores = attention_scores - 10000.0 * (1 - attention_mask)
    attention_weights = flow.softmax(attention_scores, dim=-1)
    attention_weights = dropout(attention_weights)
    return attention_weights
def att_t5(attention_scores, attention_mask, scale_mask_softmax_fusion=True, coeff=None, attention_dropout_prob=0, use_cache=False):
    dropout = nn.Dropout(0)
    if scale_mask_softmax_fusion:
        if attn_mask_type == AttnMaskType.padding:
            attention_mask = (
                attention_mask.expand_as(attention_scores) if use_cache else attention_mask
            )
            attention_weights = flow._C.fused_scale_mask_softmax_dropout(
                attention_scores,
                attention_mask,
                fill_value=-10000.0,
                scale=coeff,
                p=attention_dropout_prob,
            )[0]
    else:
        if coeff is not None:
            attention_scores *= coeff
        attention_scores = flow.mul(attention_scores, attention_mask)
        attention_scores = attention_scores - 10000.0 * (1 - attention_mask)
        attention_weights = flow.softmax(attention_scores, dim=-1)
        attention_weights = dropout(attention_weights)
    return attention_weights

@chengtbf @strint @xyn1201

[jackalcooper]

@xiezipeng-ML 这里说的hugging face版本的T5指的是 transformers 库的吗？如果是的话，直接支持transformers里面T5的oneflow后端之后，你觉得可以直接跑分布式训练吗？我上周移植了transformers的CLIP的infer，不知道训练会多多少东西。transformers的CLIP和t5应该会共用一些基础的模块吧。

[xiezipeng-ML]

@xiezipeng-ML 这里说的hugging face版本的T5指的是 transformers 库的吗？如果是的话，直接支持transformers里面T5的oneflow后端之后，你觉得可以直接跑分布式训练吗？我上周移植了transformers的CLIP的infer，不知道训练会多多少东西。transformers的CLIP和t5应该会共用一些基础的模块吧。

是的 transformers仓库，我slack请教你

[chengtbf]

@xyn1201 这个的 nsys 结果是不是还没有

[xyn1201]

• 刚刚分别跑了dp4_mp2_pp1和dp2_mp4_pp1的2机4卡测试
  • dp4_mp2_pp1：吞吐是比较正常的
  • dp2_mp4_pp1：这个是IDEA给的配置，跑的很慢，15分钟第一个iter都没有跑完，后面就没再等了。

然后列一下dp4_mp2_pp1这组配置的对比结果，libai的是今天新跑的，megatron用的前面comment里的数据，两个模型的参数对齐了，但是数据集用的不一样，这个麻烦 @xiezipeng-ML 给说明一下

projects/T5 单机4卡测试

• 机器：oneflow-28 单机4卡

• oneflow master Oneflow-Inc/oneflow@f97f09f

• libai test_t5_time 0002b66

• 对比：

  • libai：mt5_pretrain.py mb16_gb256 dp2_mp2_pp1 zero_stage=2 如果跑mb32_gb512会OOM

    8913 MiB /68.00 samples/s nsys
    日志：oss://oneflow-test/mt5_test/1021/1n4g_28/libai/

  • Megatron-deepspeed：mb32_gb512 dp2_mp2_pp1 zero_stage=2

    4725 MiB /82.7 samples/s nsys

projects/T5 2机4卡测试

• 机器：oneflow-25 oneflow-28 2机一共8卡
• oneflow master Oneflow-Inc/oneflow@f97f09f
• libai test_t5_time 0002b66
• 用例：
  • libai：mt5_pretrain.py mb16_gb512 dp4_mp2_pp1 zero_stage=2
    8613 MiB/63.75 samples/s 25_nsys 28_nsys
    日志：oss://oneflow-test/mt5_test/1021/2n4g/libai/
  • Megatron-deepspeed：mb16_gb512 dp4_mp2_pp1 zero_stage=2
    4783 MiB/164.3 samples/s 25_nsys 28_nsys

[chengtbf]

缺少了 Megatron 1n4d 2n4d 的 nsys，oneflow 1n4d nsys

[xiezipeng-ML]

然后列一下dp4_mp2_pp1这组配置的对比结果，libai的是今天新跑的，megatron用的前面comment里的数据，两个模型的参数对齐了，但是数据集用的不一样，这个麻烦 @xiezipeng-ML 给说明一下

昨晚在libai的main分支里把IDEA的dataset换成了megatron的dataset测了下，两个datasets吞吐是一样的

[xyn1201]

单卡 mb4_gb32
libai_nsys
megatron_nsys
@chengtbf

[chengtbf]

初步分析结论

之前两天的测试和本地测试受到 T5 （Megatron）和 MT5 （huggingface） 的区别，以及本地历史 LiBai 版本的影响，拖延了问题分析的进度。

目前的初步结论是 **2-D SBP 下， OneFlow T5 的 sbp infer 结果是不高效的，比 Megatron 多了几倍的通信开销，导致整体的吞吐慢了三倍。这个现象是随着 batch size 的增大而变得更差 **

两机分析

Megatron 2机 nsys 结果：

主要看两个指标， 单个 iter 的前后向总耗时，以及 kernel 占比：

1. 单个 iter 的耗时是 357ms， 分为 fw encoder （78ms）+ fw decoder （24ms） + loss （9ms） + bw decoder （88ms） + bw encoder （150ms）
2. 其中，nccl 通信基本上都是 allreduce 通信，占总的计算时长比为： 66% （49.8% + 12.8% + 3.4%）

LiBai MT5 2机 nsys 结果：

1. 总耗时 1000ms，是 Megatron 的3倍，其中： fw encoder （360ms）+ fw decoder （94ms） + bw decoder （175ms）+ bw encoder （339ms）
2. nccl 占比： 占大头的不是 allreduce，而是 send recv（应该是 sbp 推导到了 bad case，send recv 很不高效） send recv 46.4% （应该全部都是多于的） + allreduce 17.9 % + allgather 12.9 %

单机4卡分析

单机四卡的性能结果， oneflow 比 Megatron 快一些 ： oneflow 518ms vs Megatron 716ms

同时 Megatron 的 iter 间调度的间隔很大。 还有不少的优化空间。 oneflow 的调度是比较完美的。

Megatron 1n4d

OneFlow 1n4d

OneFlow 的时间虽然比 Megatron 快，但是并不是最优的，仍有至少 12% 的冗余 send recv 通信，主要是在前向 fw encoder 部分包含大量的 send recv 通信。

单机单卡比较

OneFlow 比 Megatron 优势非常明显： OneFlow 88ms vs Megatron 127ms

结论

1. OneFlow 单卡速度领先 Megatron
2. 在 4 卡的 shape 下，2d sbp 的推导结果不是那么差，速度领先 Megatron
3. 在 8 卡的 shape 下，2d sbp 的推导结果非常差，多了几倍的通信开销，速度比 Megatron 慢 三倍

[xyn1201]

SBP_INFER_RULE_TAG=2 和 自动并行 测试吞吐

• 机器：oneflow-25 oneflow-28 2机一共8卡
• oneflow master Oneflow-Inc/oneflow@f97f09f
• libai main e9ca408
• libai吞吐数据 mt5_pretrain.py mb16_gb512 dp4_mp2_pp1 zero_stage=2
  • export SBP_INFER_RULE_TAG=2：11663 MiB/61.44 samples/s 吞吐和之前的数据持平
  • 自动并行：在跑，还没有输出，一会儿和一鹏看一下这个问题

[Yipeng1994]

Oneflow-Inc/oneflow#9288
在第一档允许了自动并行与ZeRO共存，但是实际效果没有测试过。我在16上跑，OOM了，毕竟自动并行还没有考虑内存，有些慌。
不过我看了一下，大的weight的sbp基本都是 (S0, S1)，并没有给出（B, B），所以功能上是符合预期了，就是有ZeRO，然后也有AutoParallel。

另外它的op很多，不算variable快5000个了，所以初始化cost的时候很慢，需要20分钟，我优化了一下，估计能压缩一半。
至于搜索算法就很快，半分钟就能出结果。

在测试自动并行的同时，建议先看下半自动推导下mt5的boxing里面，哪里的sbp不符合预期，然后加上一些to_global来控制一下。