PR Category

Performance Optimization

PR Types

Performance

Description

在动半 sharding optmize 时，反向计算的 grad 是 partial 状态的，在做 grad clip 会分别有：

1. grad(partial) -> allreduce(replicate) -> l2_norm
2. grad(partial) -> allreduce(replicate) -> multiply

在 grad clip 开始时需要融合冗余的 allreduce 通信：

1. 方法1 fuse: 修改 grad 是 replicate 的，先进行 allreduce 后再进入 grad clip，通过开关 FLAGS_fuse_allreduce_in_opt 控制
2. 方法2 fuse: 修改 grad 是 shard 的，通过 reduce_scatter 尽可能在 sharding mesh纬度切分，一般是 shard(0)，其他 mesh 纬度尽可能对 tensor 进行 shard，不能切分的保持replicate，通过开关 FLAGS_fuse_reducescatter_in_opt 控制

在 llama7B sd8 测试吞吐(tokens/s/card)情况：
base：6153.6
开启FLAGS_fuse_allreduce_in_opt：6909.8(+12.2%)
开启FLAGS_fuse_reducescatter_in_opt：7372.6(+19.7%)

PCard-90131