已开启
[RFC]多维并行的随机状态管理 #153
changzherui创建于  5月18日
changzherui
changzherui成员
5月18日 创建

当前发现的问题

当用户没有手动设置随机状态时(manual_seed()),DTensor的op_random操作会在op_dispatch里创建默认的OffsetBasedRNGTracker(显式执行),此时会执行platform.broadcast(rng_state, 0)操作,多个stage场景下会有问题。

业界情况调研

pytorch也有类似问题:在pp场景进行 lazy_init RNGTracker时,会出现卡死问题
具体说明:https://github.com/pytorch/pytorch/pull/141223
解决办法:手动调用manual_seed方法,不同stage设置不同的seed
image.png

例如torchtitan做法:初始化的时候会将不同的stage设置不同的seed
image.png

扩展

pytorch针对这个问题的讨论:https://github.com/pytorch/pytorch/issues/140301,
问题如下:

  1. 当前没有考虑PP场景
  2. 实现过于复杂
  3. 用户设置的seed有可能被覆盖(根据不同场景(TP\DP等)自动执行操作)

提出的解决办法:

  1. 扩展接口 torch.distributed.manual_seed(seed, sharded_groups=[pp, tp, fsdp]) 感知到DeviceMesh及并行分布
  2. 全局只创建一个RNG tracker对象,统一管理

结果:没有被采纳,设计初衷和DTensor违背

hyper解决办法

补齐Dtensor下的manual_seed方法,在pp场景,当前需要用户显式调用manual_seed方法,不同的stage设置不同的seed。
后续:统一在集群拉起阶段(init_dist)由框架自动拉起,类似torchtitan。

likedislike
changzheruichangzherui成员
5月18日 修改了issue 的描述
changzheruichangzherui成员
5月18日 修改了issue 的描述
changzheruichangzherui成员
5月18日 修改标题为 “[RFC]多维并行的随机状态管理”,原标题为“随机状态管理,在PP场景报错”
changzheruichangzherui成员
5月18日 issue类型由 Bug 改变为 RFC
changzheruichangzherui成员
5月18日 修改了issue 的描述
changzheruichangzherui成员
5月18日 修改了issue 的描述
changzheruichangzherui成员
5月18日 修改了issue 的描述
changzheruichangzherui成员
5月18日 修改了issue 的描述
changzheruichangzherui成员
5月18日 修改了issue 的描述
changzheruichangzherui成员
5月18日 修改了issue 的描述
changzheruichangzherui成员
5月19日 修改了issue 的描述
changzheruichangzherui成员
5月19日 将 changzherui1 设为负责人
Jjinxiaoxian
5月20日 关联了pull request:add ut for tensor parallel and collectives
Jjinxiaoxian
5月21日 关联了pull request:add ut for tensor parallel and collectives