Swin 数据并行从单进程到8进程，显存暴涨 about libai

[leaves-zwx]

ddp 正常而数据并行不正常，这个现象很像是有多余的 boxing，可以先看看 boxing log 里面有哪些 boxing。ddp 只有最后模型更新前有1个 allreduce 的 boxing，数据并行理论上也应该是这样。

[strint]

graph.debug(1) 会默认在0卡打印graph内的显存开销；

graph执行第一次后，print(graph)可以看到 graph 内创建的 op，及其输入、输出的 SBP；

[Ldpe2G]

ddp 正常而数据并行不正常，这个现象很像是有多余的 boxing，可以先看看 boxing log 里面有哪些 boxing。ddp 只有最后模型更新前有1个 allreduce 的 boxing，数据并行理论上也应该是这样。

发现不少 p->s oneflow.INFO.LOG

graph.debug(1) 会默认在0卡打印graph内的显存开销；

graph执行第一次后，print(graph)可以看到 graph 内创建的 op，及其输入、输出的 SBP；

设置了 graph.debug(1) 没看到有打印显存开销，print(graph) 看内部的op sbp都是broadcast，但是也没看到输入和输出的sbp

[Ldpe2G]

设置了 graph.debug(1) 没看到有打印显存开销，print(graph) 看内部的op sbp都是broadcast，但是也没看到输入和输出的sbp

要首次运行之后才有，之前是没运行graph就打印了，发现中间某些层的输入是 P，网络输入和输出都是 S(0)

[Ldpe2G]

Swin 数据并行下包含的除了 P2B 之外的 boxing

P2S

Swin 中包含一个 WindowAttention 模块，下面在该模块 forward 函数中加入注释，解释SBP的转换过程：

class WindowAttention(nn.Module):
	# ......
    def forward(self, x, mask):
        B_, N, C = x.shape
        # x.sbp = S(0)
        qkv = (
            self.qkv(x) # S(0) 
            .reshape(B_, N, 3, self.num_heads, C // self.num_heads) # S(0)
            .permute(2, 0, 3, 1, 4) # S(1)
        )
        # qkv.sbp = S(1)
        q, k, v = qkv[0], qkv[1], qkv[2]
        # q.sbp = P, k.sbp = P, v.sbp = P
	    # S(1) 的tensor经过slice操作之后就变成了 P
        q = q * self.scale
        # q.sbp = P
        attn = flow.matmul(q, k, transpose_b=True)
        # attn.sbp = P
        relative_position_bias = self.relative_position_bias_table[
            self.relative_position_index.view(-1)
        ].view(
            self.window_size[0] * self.window_size[1],
            self.window_size[0] * self.window_size[1],
            -1,
        )  
        # relative_position_bias.sbp = 
        relative_position_bias = relative_position_bias.permute(
            2, 0, 1
        ).contiguous()  # nH, Wh*Ww, Wh*Ww
        unsqueeze_relative_position_bias = relative_position_bias.unsqueeze(0)
        # unsqueeze_relative_position_bias.sbp = B
        
        # attn.sbp = P
        # unsqueeze_relative_position_bias.sbp = B
        # 这里加操作推导出 attn 要是 S(0)，所以插入了一个 P2S 的 boxing
        attn = attn + unsqueeze_relative_position_bias
	   # 结果 attn.sbp = S(0)
        
        attn = self.softmax(attn)
        # 结果 attn.sbp = S(0)

        attn = self.attn_drop(attn)
        # 结果 attn.sbp = S(0)

        # batch_matmul，推导出 attn 和 v 的 sbp 都要是 S(0)
        # 所以差一个 P2S 的 boxing，v要从 P 转 S(0)
        x = flow.matmul(attn, v).transpose(1, 2).reshape(B_, N, C)
       	# x.sbp = S(0)
        
        x = self.proj(x)
        x = self.proj_drop(x)
        # x.sbp = S(0)
        return x

S2S

S2B

(Send)2(Recv)

[chengtbf]

S(1) 的tensor经过slice操作之后就变成了 P

这个用的是哪个 oneflow commit？ 这个在：

• Oneflow-Inc/oneflow#8521

修复了。 昨天合并的

[Ldpe2G]

S(1) 的tensor经过slice操作之后就变成了 P

这个用的是哪个 oneflow commit？ 这个在：

• SliceOp support S->S, not support S->P(TODO) oneflow#8521

修复了。 昨天合并的

用的29号的版本，拿我更新一下oneflow再测一下

[Ldpe2G]

S(1) 的tensor经过slice操作之后就变成了 P

这个用的是哪个 oneflow commit？ 这个在：

• SliceOp support S->S, not support S->P(TODO) oneflow#8521

修复了。 昨天合并的

用的29号的版本，拿我更新一下oneflow再测一下

更新了oneflow之后重新测了下显存

libai eager global fp32 batch 32

1进程 0号卡显存： 5722 M
8进程 0号卡显存： 6474 M

libai graph fp32 batch 32

1进程 0号卡显存：5098 M
8进程 0号卡显存：5986 M

可以看到相对于之前，8卡下显存占用明显少了很多