[Bug][MT5] Throughput is unexpected #406
Comments
strint
changed the title
t5 单机4卡测试
t5 2机4卡测试
|
|
这里比较明显的问题是，我们 4 卡 2-D 并行是超过 Megatron 的，但是 两机8卡的吞吐比 单机四卡的还慢。而 Megatron 是一个线性的加速比。 |
|
这里有点问题，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的区别总结：
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 |
|
@xiezipeng-ML 这里说的hugging face版本的T5指的是 transformers 库的吗？如果是的话，直接支持transformers里面T5的oneflow后端之后，你觉得可以直接跑分布式训练吗？我上周移植了transformers的CLIP的infer，不知道训练会多多少东西。transformers的CLIP和t5应该会共用一些基础的模块吧。 |
是的 transformers仓库，我slack请教你 |
|
@xyn1201 这个的 nsys 结果是不是还没有 |
然后列一下dp4_mp2_pp1这组配置的对比结果，libai的是今天新跑的，megatron用的前面comment里的数据，两个模型的参数对齐了，但是数据集用的不一样，这个麻烦 @xiezipeng-ML 给说明一下 projects/T5 单机4卡测试
projects/T5 2机4卡测试
|
|
缺少了 Megatron 1n4d 2n4d 的 nsys，oneflow 1n4d nsys |
昨晚在libai的main分支里把IDEA的dataset换成了megatron的dataset测了下，两个datasets吞吐是一样的 |
|
单卡 mb4_gb32 |
初步分析结论之前两天的测试和本地测试受到 T5 （Megatron）和 MT5 （huggingface） 的区别，以及本地历史 LiBai 版本的影响，拖延了问题分析的进度。 目前的初步结论是 **2-D SBP 下， OneFlow T5 的 sbp infer 结果是不高效的，比 Megatron 多了几倍的通信开销，导致整体的吞吐慢了三倍。这个现象是随着 batch size 的增大而变得更差 ** 两机分析Megatron 2机 nsys 结果：主要看两个指标， 单个 iter 的前后向总耗时，以及 kernel 占比：
LiBai MT5 2机 nsys 结果：
单机4卡分析单机四卡的性能结果， oneflow 比 Megatron 快一些 ： oneflow 518ms vs Megatron 716ms 同时 Megatron 的 iter 间调度的间隔很大。 还有不少的优化空间。 oneflow 的调度是比较完美的。 Megatron 1n4d
OneFlow 1n4dOneFlow 的时间虽然比 Megatron 快，但是并不是最优的，仍有至少 12% 的冗余 send recv 通信，主要是在前向 fw encoder 部分包含大量的 send recv 通信。
单机单卡比较OneFlow 比 Megatron 优势非常明显： OneFlow 88ms vs Megatron 127ms
结论
|
SBP_INFER_RULE_TAG=2 和 自动并行 测试吞吐
|
|
Oneflow-Inc/oneflow#9288 另外它的op很多，不算variable快5000个了，所以初始化cost的时候很慢，需要20分钟，我优化了一下，估计能压缩一半。 在测试自动并行的同时，建议先看下半自动推导下mt5的boxing里面，哪里的sbp不符合预期，然后加上一些to_global来控制一下。 |
|
带 nccl logical op 和 sbp 的 op graph log：https://oneflow-test.oss-cn-beijing.aliyuncs.com/mt5_test/2n4g_log/output.log 搜索下 |
自动并行 2n4g 测试
|
|
看了一下 job/plan 发现了3个问题： 非预期的 SBP 变化会导致后续一系列 SBP 都乱掉，从而导致有多余的 nccl logical boxing (是不是还有其他影响有多余的 nccl logical boxing 还要测试) 原因是 但在2n4d情况下，该 reshape 前面被插入1个 1n4d job 片段2n4d job 片段1n4d 和 2n4d 他们两个的区别就是 batch size 变化了 (global)，猜测原因和 https://github.com/Oneflow-Inc/OneTeam/issues/1721 里面类似。但为什么 SBP_INFER_RULE_TAG=2 设置了后仍然不能阻止该非期望的 sbp 转换？需要再调试一下。 低效的 amp 转换代码这里的 broadcast_add 的左边 plan 片段低效冗余的 cast上述的 要进行这些计算，系统选择先把 bool cast to int64，该 cast 在每一层 transformer layer 都重复进行。 plan 片段同时 |
reshape只有一个输入的话，哪个sbp规则下都是match的，不可能发生改变吧？ |
|
|
|
emmm，等下我具体看看为什么推导出了不同的sbp |
|
不是 GreedilyFindMinCopyCostNdSbp 这个函数的问题，而是 GetValidNdSbpSignatureList 的问题。 |
哎，你也在看，这个今天修好了，在让 @xyn1201 测 |
|
我调试看起来不像是这个原因，而是 reshape 的 GetSbp 函数本身有问题。我再 debug 看看具体是什么。 |
哎，刚刚测试出修复失败了，我也再debug康康 |
|
reshape 的 sbp siganture list 在 1n4d 下正常，而 2n4d 下不正常的原因找到了： debug log 片段代码在: https://github.com/Oneflow-Inc/oneflow/blob/22eabed6a2432085cd4aa7bf7bf98464d30e9cba/oneflow/user/ops/reshape_user_op_util.cpp#L131-L132) 处判断当前 dimension 是否可以被 split 的时候是用
所以在 2n4d 下我们根据调试的信息可以看到 reshape 的 sbp signature list 里面没有 S(2) -> S(2) 这一项，但其实是可以 split 的，因为 4dp + 2mp，S(2) 要不切4份，要不切2份（取决于 S(2) 是 nd_sbp 的第1维还是第2维）， 所以出现了 #406 (comment) 里面所说的情况。这里的正确做法，应该根据 device mesh 的某一个维来判断是否能被 split，而不能只看 parallel_num。 但目前有一些困难，因为在 GetSbp 的时候，并不知晓推导的 sbp signature 将会被应用于 device mesh 的哪一维。这里只能添加上全部的 split(num_axes)，然后再到后面的 |
是的，原因就跟文晓讲的差不多。通过打印log可以看出来 备选策略里面没有S(2)。原因就是因为12不被8整除。reshape的sbp这部分是我之前重构的。为什么用的是parallel num，是因为reshape的get sbp函数只推导1d的sbp。而且一般这个1d的sbp推导是不涉及shape的，比如矩阵乘或者是加减这些op。在后面还有一个Filter，这个Filter做的才是根据shape筛选sbp。但是reshape本身跟shape又紧密关联，所以这里才必须要有这个filter。 2d sbp是根据1d sbp的直积得出，1d sbp把S(2) filter掉了，后面自然选不到 (S0, S2)。 昨天我做了一个修复尝试 Oneflow-Inc/oneflow@caf344f 只是测试结果并没有如愿修复bug。 点进吞吐可以看到log，S2还是没有出现。原因未知，不过今天稍微修复一下应该就行了。 |
debug_reshape_sbp_signature分支
refactor-GetSbpSignature分支
2个分支吞吐都有接近1倍的提升，但还低于megatron |
|
哎，refactor-GetSbpSignature 也测一下 2n4g mb16_gb512 康康是否会有内存暴涨的问题，然后输出一下boxing的log @xyn1201 |
|
refactor-GetSbpSignature分支 把S2加回来了，但是吞吐只有70%，还是需要定位一下其他的问题。 |
|
操作失误，上面测试的megatron数据是关掉zero的， 开zero测试
|
|
嗯嗯，这样容易接受多了 |
不启用activation checkpointing
|
|
如果把 parameter 的 zero 通信就地去做（在消费前去做）有1个问题，就是各张卡的计算节奏不是完全一致的（特别是不同的机器），那么就是频繁的出现你等我，我等你的情况，现象就是 send-recv 的 timeline 长度超出预期。send-recv 是需要同步的，他与计算交替执行，就会出现互等的情况。 而 megatron 里面把所有 parameter 的 zero 通信集中起来去做，这样互等的时间就会减少，注意不同卡上的 parameter 要按同样的顺序去进行通信交换。 这是我猜测的1个 send-recv 特别长的原因，就算换成 allgather, 这个问题应该也是存在的。 比如下图中特别长的 send-recv 在另外一台机器上就特别短：
|
|
我建议后面做这么几个测试：
|
关zero 关checkpointing测试 |
|
oneflow 补测一轮： 这个分支的结果吧
|
关zero 关checkpointing |
|
用这个分支再测试一次。关掉 ZeRO ，关掉 Checkpointing。 |
encoder layer profile对比 libai 和 megatron 的 encoder layer 的性能差异，但因为 libai t5 实际是 mt5，所以算子层面有一些差别，这个对比不是一个benchmark，只是为未来的优化做一个参考。 libai t5 encoder layer nsys snapshot
megatron t5 encoder layer nsys snapshot
我将上面 nsys 中所有的 kernel 都 dump 成表格 encoder_layer_pofile.xlsx，其中 sheet1 是 megatron 的，sheet2 是 libai 的。
下面会举一些例子如何通过上述表格观察到的差别： LayerNorm 差别libai t5 用的是 RMSLayerNorm，目前是用碎 op 组合的，耗时 69.32 μs。里面还有 cast_f2h/cast_h2f 等低效转换（因为有 two stage reduce）。 libai t5 layer norm stat snapshot
megatron 用的是 LayerNorm，耗时 36.83 μs。 megatron t5 layer norm stat snapshot
self_attention query_key_valuelibai t5 self_attention query_key_value 耗时 59.2 μs。 libai t5 self_attention query_key_value stat snapshot
megatron t5 self_attention query_key_value 耗时 59.59 μs。 megatron t5 self_attention query_key_value stat snapshot
matmul 是一样的耗时，比较符合预期。libai t5 里面没有 bias add，megatron 里面用的是 cutlass kernel。 softmaxlibai t5 self_attention softmax stat snapshot
megatron t5 self_attention softmax stat snapshot
两边的 mask 算法不一致。 dropoutlibai t5 self_attention softmax stat snapshot
megatron t5 self_attention softmax stat snapshot
view 操作libai t5 view stat snapshot
megatron t5 view stat snapshot
由于 oneflow 还不支持 lazy view，所以 tranpose, slice 都会带来 copy。而 megatron 里面只有唯一一次 contiguous 带来的 copy。 gelulibai t5 gelu_tanh stat snapshot
megatron t5 gelu stat snapshot
libai 里面的 gelu_tanh 是组合的，所以比较碎有较多开销。megatron 里面 gelu 是自定义的。 libai self_attention 还存在未知 SendRecvlibai t5 send recv stat snapshot
不知道是做什么的，需要进一步看 job/plan AllReduce对比不具有参考意义，因为单个 rank 上的 allreduce 需要等其他 rank 上一起执行，所以 timeline 上长也许只是等待时间长（因为各rank执行节奏不一致）。 |
debug_reshape_sbp_signature分支 关zero 关checkpointing 测试结果汇总
|
|
上文中说的 SelfAttention 中的未知 SendRecv 是必要的，它在代码这里，megatron 这里没有的原因是算法不一样，megatron 里面的 t5 没有这个 position_bias。 position_bias 这里 position_bias (S(0), B) 要与 attention_scores (S(0), S(1)) 做计算，需要做一个 (S(0), B) -> (S(0), S(1))，目前 2d SBP 里面是用 SendRecv 实现的，但可以用 SameDim0AllScatter 来实现（没有通信开销）。 但上述 (S(0), B) -> (S(0), S(1)) 的转换不用每一层 layer 都做，因为 position_bias 是在 layer 0 通过 compute_bias 计算出来的，后面的所有 layer 使用的都是 layer 0 的 position_bias，所以该转换只需要做一次。而 position_bias 在与 attention_scores 相加之前，需要先与 attention_mask (S(0), B) 相加（见这里），加完之后 position_bias sbp 也变为了 (S(0), B)。 我们只需要将 |
广义基础传输也没有通信开销吧，它是直接从本地拷数据 |
像这种场景，自动并行就能完美地处理。因为它知道后方有什么信息，知道在前面做通信会有多倍的通信代价。 所以关掉zero了以后，阔以直接打开自动并行，甚至原代码都不用动。 |
|
自动并行可以干预掉用户写的 to_global 吗？ |
这个忘记了，自动并行会移除 parallel cast 吗？ 我印象中某一档会移除 @Yipeng1994 @wyg1997 |
第二档会，可以选。 |
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
The text was updated successfully, but these errors were encountered: