swin 3d并行8卡训练会崩（数据+张量+接力）

[Ldpe2G]

现象描述

实验发现，swin 在 8卡 3d 并行的配置下，训练会出现训到某个点，开始loss会一直上升，精度会一直下降

实验配置:

train.dist.data_parallel_size = 2
train.dist.tensor_parallel_size = 2
train.dist.pipeline_parallel_size = 2

目前做过的实验

出问题的case

• 数据+模型+朴素接力（2，2， 2）， eager global, clip grad on ，local batch size = 128， 到 28% 往下掉成 1% 训崩了

• 数据+模型+朴素接力（2，2， 2）， eager global, clip grad off ，local batch size = 128， 到 28% 往下掉成 1% 训崩了

• 数据+模型+朴素接力（2，2， 2）， eager global, clip grad on ，local batch size = 32， 到 18% 往下掉成 1% 训崩了

没问题的case

• 纯数据并行 8卡, graph + amp + zero stage 1, clip grad on，local batch size = 32，最后收敛精度 > 75%

• 纯数据并行 8卡, eager global, clip grad on，local batch size = 32，最后收敛精度 > 75%

• 数据+模型并行（4，2）8卡, graph fp32, clip grad on，local batch size = 64，最后收敛精度 > 75%

其他不完全的实验

• 数据+朴素流水并行（2， 4）8卡, graph fp32, clip grad on，local batch size = 128，最后精度 ~69%，收敛不到指定精度
• 数据+模型并行（4，2）, eager global, clip grad on，local batch size = 64，慢的无法接受， 速度比纯数据并行 eager global 慢 ~7倍左右
• 数据+模型+朴素接力（2，2， 2）, graph fp32, clip grad on，local batch size = 128，训到 > 52% 就停了，但是在同样的迭代数下，比 纯数据并行精度要低很多

[Ldpe2G]

下一步实验计划，固定初始化，dataloader 去掉shuffle 和 随机 augment 保证输入一致，对比单卡和两卡接力，的前向输出和反向梯度

将 dataloader 的随机读取关掉，data augmentation 改成固定的方法去掉随机性，swin 中的 droppath 关掉，加载同样的初始化模型，发现单卡和两卡接力的第一个iter 前向输出一致 和 反向得到的权值的梯度也是一致的。

然后直接训练 ~7万个 iter ，得到的 loss 曲线

因为是每隔20个iter才打印一次 loss, 所以图上的iter数只有 3500 左右。loss曲线是基本重合的，而且接力并行也没有出现训崩的情况了，所以初步怀疑是上述提到的关掉的随机性中某个部分导致训崩的问题。

关掉的随机性和改动

• 训练集aug换成 test_aug
• 初始化模型加载的同一份固定权值（单卡初始化生成的）
• swin drop_path_rate 设成 0
• dataloader CyclicSampler, shuffle 设成 False
• mixup 设置为 none
• clip grad 直接注释掉

[Ldpe2G]

基于上面的实验将去掉的随机变量逐个加回去，然后经过反复的控制变量实验。

最终发现只需要在 libai 主函数入口处 train_net.py 实例化 DefaultTrainer 之前，设置一下 Numpy 的随机种子 numpy.random.seed(0)，2卡流水并行收敛慢+训崩的现象就消失了。

一开始感觉非常的诡异，经过仔细分析终于定位了，为什么 nlp 类的任务比如 bert 3d 并行是正常的，而图像分类的任务比如 swin 跑流水并行就会训崩。

原因分析：

简单来说就是，流水并行下，网络的输入和label没对上，这个问题是怎么产生的呢？

首先来看2卡流水并行配置下，rank 0 和 rank 1 各自的 dataloader 都是读所有的数据，且由于在 shuffle 训练集索引的时候，都是设置了同样的种子：samplers.py#L76

if self.shuffle:
    generator = flow.Generator()
    generator.manual_seed(self.seed + epoch)
    random_idx = flow.randperm(self.data_size_per_epoch, generator=generator).tolist()
    indices = [start_idx + x for x in random_idx[bucket_offset:]]

所以保证每个 rank 每个iter所读取到的 图像和标签都是一致的。

然后在 cifar100 数据集中会对 图像和标签预设placement_idx，如下所示：

def __getitem__(self, index: int):
    img, target = super().__getitem__(index)
    data_sample = Instance(
        images=DistTensorData(img, placement_idx=0),
        labels=DistTensorData(flow.tensor(target, dtype=flow.long), placement_idx=-1),
    )
    return data_sample

比如在两卡流水并行配置下，placement_idx 就是相当于，图像 to_global 的时候 placement 参数为 ("cuda", ranks=[0])，标签 to_global 的时候 placement 参数为 ("cuda", ranks=[1])。

然后问题就在这里了，图像分类任务相比于 nlp 类的任务，在数据 to_global 前还多了个 data augmentation 的操作比如 mixup，里面包含了很多 numpy 的随机操作，如果我们不给每个 rank 设置一样的 numpy 随机种子，那么很可能虽然两个 rank 每个 iter 读取到的数据是一样的，但是在经过 mixup 之后生成的变换后的图片和softtarget 都是对不上的。

然后最关键的地方在于， 在做完 augmentation 之后，将 Local 的 tensor to global 的时候，由于是直接从 local to global，上面提到， 图像的 placement 参数是 ("cuda", ranks=[0])，所以直接保留 rank 0 的变换后的图像，rank 1丢弃。标签的 placement 参数是 ("cuda", ranks=[1])，则是直接保留 rank 1 的变换后的 softtarget，rank 0 的丢弃，这就是导致2卡流水并行收敛慢且训崩的原因，因为图像和标签没对应上。本来应该是 rank 0 的变换图像对应 rank 0 变换后的softtarget，现在确实 rank 0 图像对 rank 1 的 软标签

而为什么 nlp 没问题应该是 ，应该是因为用的 onehot label 所以没有 Mixup 因随机变换所导致的不同rank生成不一致label的问题。

然后在简单修复 DefaultTrainer get_batch 函数中对 图像类数据 local to global 的逻辑之后。简单来说就是对于 label 不直接 to global 到 rank 1 丢掉 rank 0的，而是先 to global 到 rank 0，就是保留 rank 0 的标签，然后再将 rank 0 的标签 to global 到 rank 1。

训练就正常了，下面是修复后的 Loss 曲线，~2万iter，每隔 20个iter打印一次 Loss：

绿色曲线是修复前两卡流水的Loss曲线，黄色曲线是修复后，可见和蓝色单卡曲线是接近的。

接下来的实验就是8卡3d并行上验证是否正常。

踩坑总结

感觉对于大部分从torch迁移过来的用户，在尝试 oneflow 的 global 来做一些图像类任务的混合并行实验很可能都会踩到的坑。

这个坑本质就是如果用户在使用 global 配置各种并行的时候，数据读取还是用的 torch 那套 dataloader 且 数据增强是在 local 上做的。然后在将数据 to global 的时候，就得要清楚各种配置下 local to global 会得到什么结果，且确认数据在 to global 之后是否符合预期。

这个bug我们内部的用户在写代码的时候都没有意识到这个问题，而且还得经过反复实验踩发现。

[yuanms2]

看看libai 之外有没有机制保证不出现这类问题

[Ldpe2G]

看看libai 之外有没有机制保证不出现这类问题

其实感觉这个问题跟任务有关，而且是要看用户意图，比如 nlp 类的任务做2卡流水，标签就是 Onehot 那既然每张卡都读了同样的数据，那 local to global 的时候，肯定是直接保留 rank 1 的标签就行了而不需要把 rank 0 的标签传输到 rank 1，还多了通信。

只是图像这边会对 标签 做变换，所以在不统一 随机种子的情况下，只能是把 rank 0 的标签传输到 rank1 来用。

感觉还得是要用户在写代码的时候要很清楚在各种并行配置下 local to global 会产生什么结果，然后是否符合预期。

[L1aoXingyu]

配置随机种子是不是在大规模下面效率更好一些

[Ldpe2G]

配置随机种子是不是在大规模下面效率更好一些

我感觉是的，目前的修复方式是，流水并行每个iter都需要额外传输标签，所以

• 要不就是加 随机种子
• 要不就是目前的改法，把 placement_idx 0 上的标签传到
  placment_idx -1, Dev fix swin pipline parallel train crash #255

但是有个地方是，目前我们是每个 rank 都会去读数据，所以加种子的方案是ok的，但是如果考虑上未来在数据读取上的改进，比如就简单的8卡接力并行，目前是每个卡都会去读数据，其实是可以改进成只需要 rank 0 读数据就行了， 如果这样改的话，就是应该要用目前的方案。

[Ldpe2G]

经过上诉几组实验，证明 pr #255 的修复是有效的，目前 数据+张量+朴素流水并行（2，2， 2）graph fp32 的收敛精度正常了。

但是实验发现貌似 eager global 的 2d sbp 还存在问题，一开始收敛正常，后期会 Loss 会上升。

下图均是 8卡 下各种并行配置的组合，可以看到 eager global 1d sbp 还有 2d sbp graph 的收敛都是正常的，最终收敛精度都能 > 76%

但是只要是 eager global 2d sbp 的组合，就会训到后期开始训崩（精度到50%左右就会往下掉），比如图中的绿色和紫色线。

每隔 20 个 iter 打印一次 loss，所以图中的数据实际是跑了 ~1.5 万个iter。

接下来的实验还是先保证 graph 在开启各种优化下的收敛精度，eager global 2d sbp 估计还有啥隐藏的bug，之前也没有足够的模型的测试。

接下来的实验

• 数据+张量+朴素流水并行（2，2， 2）graph amp 2d sbp, checkpointing on, clip grad on ，local batch size 128,
  • 收敛正常，最终精度 > 76%
• 数据+朴素流水并行（2，4）graph amp 2d sbp, zero stage 1, checkpointing on, clip grad on ，local batch size 128,
  • 收敛正常，最终精度 > 75.6%

[Ldpe2G]

经过实验定位，发现只需要在 droppath 模块中，调用 flow.rand 的时候，传入 generator 参数，generator 的 seed 所有rank一致 ，然后 swin 2d sbp eager global 训崩的问题就消失了。

相关pr: #268

修复后 loss 曲线

图中可见添加 generator 之后 3d 并行 2d sbp eager（蓝色曲线）的loss曲线和纯数据并行的曲线基本一致，而且还做了纯数据并行下，加不加generator的对比实验，实验表明 generator 对纯数据并行影响不大。

一些分析

简单做了些实验分析，在 3d 并行的配置下，在运行至 droppath 模块中的 flow.rand 的时候， sbp 可能会是 [S(0), S(0)] 和 [S(0), B]，下面看下加不加 generator 的结果：

不加 generator

import oneflow as flow
sbp = [flow.sbp.split(0), flow.sbp.broadcast]
# sbp = [flow.sbp.split(0), flow.sbp.split(0)]
tensor = flow.rand(4, generator=None, sbp=sbp, placement=flow.placement("cuda", ranks=[[0,1],[2,3]]))
print(flow.env.get_rank(), tensor.to_local().numpy())

# sbp = [S(0), S(0)]，4个rank上的分量均不一样。
1 [0.28674227]
0 [0.91543716]
2 [0.05125458]
3 [0.89985836]

# sbp = [S(0), B]，4个rank上的分量均不一样，但是按照sbp的设置，
# 0和1要一样，2和3要一样，然后 01组和23组之间要不一样。
0 [0.30961376 0.08173643]
1 [0.7047522 0.7899459]
2 [0.86135066 0.37061813]
3 [0.10754773 0.7707022 ]

添加 generator

所有rank种子一致

import oneflow as flow
generator = flow.Generator()
generator.manual_seed(0)
sbp = [flow.sbp.split(0), flow.sbp.broadcast]
# sbp = [flow.sbp.split(0), flow.sbp.split(0)]
tensor = flow.rand(4, generator=generator, sbp=sbp, placement=flow.placement("cuda", ranks=[[0,1],[2,3]]))
print(flow.env.get_rank(), tensor.to_local().numpy())

# sbp = [S(0), S(0)]，4个rank上的分量变成一样了。
0 [0.54892594]
1 [0.54892594]
2 [0.54892594]
3 [0.54892594]


# sbp = [S(0), B]，4个rank上的分量一样，按照sbp的设置，
# 虽然符合了0和1一样，2和3一样，但是 01组和23组之间是要不一样的。
0 [0.54892594 0.21360275]
1 [0.54892594 0.21360275]
2 [0.54892594 0.21360275]
3 [0.54892594 0.21360275]

手动配置各个rank的种子

# sbp = [S(0), S(0)]
# 每个rank的种子设置为 generator.manual_seed(flow.env.get_rank())
# sbp = [S(0), S(0)]，4个rank上的分量就不一样了。
0 [0.54892594]
1 [0.4787291]
2 [0.20269513]
3 [0.1734449]

# sbp = [S(0), B]
# 每个rank的种子设置为 generator.manual_seed(flow.env.get_rank() // 2)
# 这就符合sbp了
0 [0.54892594 0.21360275]
1 [0.54892594 0.21360275]
2 [0.4787291  0.02171235]
3 [0.4787291  0.02171235]

踩坑总结

目前还没想明白，为啥给 flow.rand 传个所有rank种子一致的 generator 就能解决 2d sbp eagar 训崩的问题，因为 droppath 就类似 dropout 一样的模块，按道理即使随机出来的值不符合 sbp 的设置应该也不会有这么大的影响，而且设置了统一种子之后，在sbp包含 s 的情况下其实也不符合 sbp，但是却不会有影响，不过炼丹就是这么玄学。

然后就是对于类似 flow.rand 这类的source op，是否需要再想一下怎么对用户更加的友好，因为从上面的实验结果看，如果要生成符合sbp的随机值，就需要用户手动根据 sbp 给每个 rank 设置 generator 的种子，感觉心智负担有点大了。

目前实验结果表明是对于 swin 模型，所有 rank 全设置成一致是可以的，但是可能换个其他网络，用了这个 droppath 模块又会出其他问题，或者就需要每个rank严格按照 sbp 来设置种子。

而且为什么 graph 下不需要设置 generator 2d sbp 不会训崩，而 eager 下不设置 generator 2d sbp 就会训崩，两者能否统一一下？

[yuanms2]

不错，继续研究研究，把原因搞清楚

[yuanms2]

一个猜测，不同rank 上处理的数据不一样，如果随机数一样，就会使得“相同的随机数” 遇见了“不同的数据和反馈信号”，让系统陷入“自相矛盾”和“冲突”的境地，如果随机数不同，系统就不需要面对这种纠结了。

[lixinqi]

nn.Graph的随机数种子会固定下来的吧？ @strint @liujuncheng @hjchen2

[Ldpe2G]

这里更新一下 droppath

一个猜测，不同rank 上处理的数据不一样，如果随机数一样，就会使得“相同的随机数” 遇见了“不同的数据和反馈信号”，让系统陷入“自相矛盾”和“冲突”的境地，如果随机数不同，系统就不需要面对这种纠结了。

重新看了下 droppath 的实现，这个模块和 dropout 不一样的地方是，它是随机drop掉整个样本，其实现就是生成 batch_size 个随机数，加个阈值然后二值化，乘以原来的输入张量，来drop掉一些样本。

那么在 2d sbp 的设置下，当输入张量是 [S(0), B] 的时候，在 B 这个维度，所有 rank 拿到的样本是一样的，那是否是要求生成的随机数要一致，也就是drop掉的batch索引要一致？

然后对于纯数据并行没影响，因为本来每个rank拿到的样本就不同，那是否drop掉同样的batch索引应该没影响？。

[strint]

graph的种子会随着 job 的所有rank同步，从而所有rank一致。

eager的种子各个rank独立生成，所以默认是不同的。eager global 的种子是否默认提供同步

• 上周完善graph 支持时，顺便聊到eager了，但是还没达成一致
• houjiang建议使用 generator.manual_seed
• 看到这个隐晦的问题，看起来还是保证 global 语义，自动帮忙同步比较好

参考资料：

• rand source op global 支持的总issue，https://github.com/Oneflow-Inc/OneTeam/issues/1072
• eager global 种子的问题，source op support s and fixed generator bug oneflow#7571 (comment)

[Ldpe2G]

再思考了下，对于 [S(0), B] ，生成的随机数，应该是 S(0) 维度是否一样没所谓，但是B维度一定要一致？

[strint]

再思考了下，对于 [S(0), B] ，生成的随机数，应该是 S(0) 维度是否一样没所谓，但是B维度一定要一致？

是的。现在有个处理方法，可以保证S/P维度的种子不一样，B维度的一样。

[strint]

eager global 的种子是否默认提供同步

ND 种子的生成

方法1

rand seq: a b c d e f g

[B, S]
[[0, 1], [2, 3]]

rank 0: a b

rank 1: a c

rank 2: a b

rank 3: a c

方法2

• 每一层传递种子，层内用generator；
• 种子是否一样来表达B/S/P；
• 迭代过程

数值的偏移

rankpermute arrange
[0, n]

[strint]

方案1

generator = flow.Generator()
generator.manual_seed(0)
sbp = [flow.sbp.split(0), flow.sbp.broadcast]
# sbp = [flow.sbp.split(0), flow.sbp.split(0)]
tensor = flow.rand(4, generator=generator, sbp=sbp, placement=flow.placement("cuda", ranks=[[0,1],[2,3]]))

1、rank 0 generator seed to all rank in placement -> rank_0_seed
2、inner_generator(rank_0_seed)

[strint]

问题1、一个op内的seed序列生成方案

• 效率优先：生成不同的seed；
• 正确性优先：一个seed，slice出数据；

或者两种模式都支持，做成可配置的。

问题2、eager global 多op共享generator的问题

• op间共享generator
• op内再用建一个内部generator

TODO

结合具体case + 具体API，写一下暴露的问题、处理方案 。

• 支持save generator state，以支持训练的resume

[Ldpe2G]

这个issue主要的问题都解决了，generator相关的讨论可以移到 oneteam 。
@strint