Fix ddp multi-output bug

[wyg1997]

修复 ddp 训练时，module 有多个 output 时出错的问题：Oneflow-Inc/models#187 (comment)

最小复现代码

import oneflow as flow
from oneflow.nn.parallel import DistributedDataParallel as ddp

train_x = [
    flow.tensor([[1, 2], [2, 3]], dtype=flow.float32),
    flow.tensor([[4, 6], [3, 1]], dtype=flow.float32),
]
train_y = [
    flow.tensor([[8], [13]], dtype=flow.float32),
    flow.tensor([[26], [9]], dtype=flow.float32),
]


class Model(flow.nn.Module):
    def __init__(self):
        super().__init__()
        self.lr = 0.01
        self.iter_count = 500
        self.w1 = flow.nn.Parameter(flow.tensor([[0], [0]], dtype=flow.float32))
        self.w2 = flow.nn.Parameter(flow.tensor([[0], [0]], dtype=flow.float32))

    def forward(self, x):
        x1 = flow.matmul(x, self.w1)
        x2 = flow.matmul(x, self.w2)
        return x1, x2


m = Model().to("cuda")
m = ddp(m)
loss = flow.nn.MSELoss(reduction="sum")
optimizer = flow.optim.SGD(m.parameters(), m.lr)

for i in range(0, m.iter_count):
    rank = flow.env.get_rank()
    x = train_x[rank].to("cuda")
    y = train_y[rank].to("cuda")

    y_pred, y2 = m(x)
    l = loss(y_pred, y)
    print(l)

    l.backward()
    optimizer.step()
    optimizer.zero_grad()

[CPFLAME]

bert下该问题已被解决。

这个分支跑起来很丝滑，
module 有多个 output 时出错的问题已被解决

[github-actions]

CI failed, removing label automerge

[github-actions]

Speed stats:

GPU Name: GeForce GTX 1080 

OneFlow resnet50 time: 128.1ms (= 6403.3ms / 50, input_shape=[16, 3, 224, 224])
PyTorch resnet50 time: 139.7ms (= 6984.2ms / 50, input_shape=[16, 3, 224, 224])
✔️ Relative speed: 1.09 (= 139.7ms / 128.1ms)

OneFlow resnet50 time: 74.6ms (= 3728.9ms / 50, input_shape=[8, 3, 224, 224])
PyTorch resnet50 time: 83.2ms (= 4159.9ms / 50, input_shape=[8, 3, 224, 224])
✔️ Relative speed: 1.12 (= 83.2ms / 74.6ms)

OneFlow resnet50 time: 51.3ms (= 2566.1ms / 50, input_shape=[4, 3, 224, 224])
PyTorch resnet50 time: 61.6ms (= 3078.1ms / 50, input_shape=[4, 3, 224, 224])
✔️ Relative speed: 1.20 (= 61.6ms / 51.3ms)

OneFlow resnet50 time: 59.1ms (= 2952.8ms / 50, input_shape=[2, 3, 224, 224])
PyTorch resnet50 time: 47.3ms (= 2364.4ms / 50, input_shape=[2, 3, 224, 224])
❌ Relative speed: 0.80 (= 47.3ms / 59.1ms)

OneFlow resnet50 time: 51.5ms (= 2575.6ms / 50, input_shape=[1, 3, 224, 224])
PyTorch resnet50 time: 47.4ms (= 2368.4ms / 50, input_shape=[1, 3, 224, 224])
✔️ Relative speed: 0.92 (= 47.4ms / 51.5ms)

OneFlow resnet50 time: 153.4ms (= 7671.7ms / 50, input_shape=[16, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 159.6ms (= 7982.4ms / 50, input_shape=[16, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.04 (= 159.6ms / 153.4ms)

OneFlow resnet50 time: 101.9ms (= 5096.2ms / 50, input_shape=[8, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 109.0ms (= 5450.9ms / 50, input_shape=[8, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.07 (= 109.0ms / 101.9ms)

OneFlow resnet50 time: 85.5ms (= 4273.1ms / 50, input_shape=[4, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 82.1ms (= 4106.6ms / 50, input_shape=[4, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 0.96 (= 82.1ms / 85.5ms)

OneFlow resnet50 time: 78.4ms (= 3919.2ms / 50, input_shape=[2, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 71.4ms (= 3567.7ms / 50, input_shape=[2, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 0.91 (= 71.4ms / 78.4ms)

OneFlow resnet50 time: 76.5ms (= 3823.6ms / 50, input_shape=[1, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 69.7ms (= 3485.1ms / 50, input_shape=[1, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 0.91 (= 69.7ms / 76.5ms)

[github-actions]

CI failed, removing label automerge

[github-actions]

Speed stats:

GPU Name: GeForce GTX 1080 

OneFlow resnet50 time: 127.0ms (= 6351.3ms / 50, input_shape=[16, 3, 224, 224])
PyTorch resnet50 time: 139.0ms (= 6952.3ms / 50, input_shape=[16, 3, 224, 224])
✔️ Relative speed: 1.09 (= 139.0ms / 127.0ms)

OneFlow resnet50 time: 73.9ms (= 3697.0ms / 50, input_shape=[8, 3, 224, 224])
PyTorch resnet50 time: 84.7ms (= 4236.5ms / 50, input_shape=[8, 3, 224, 224])
✔️ Relative speed: 1.15 (= 84.7ms / 73.9ms)

OneFlow resnet50 time: 49.2ms (= 2458.0ms / 50, input_shape=[4, 3, 224, 224])
PyTorch resnet50 time: 56.1ms (= 2802.8ms / 50, input_shape=[4, 3, 224, 224])
✔️ Relative speed: 1.14 (= 56.1ms / 49.2ms)

OneFlow resnet50 time: 38.8ms (= 1940.0ms / 50, input_shape=[2, 3, 224, 224])
PyTorch resnet50 time: 48.4ms (= 2420.4ms / 50, input_shape=[2, 3, 224, 224])
✔️ Relative speed: 1.25 (= 48.4ms / 38.8ms)

OneFlow resnet50 time: 35.1ms (= 1755.8ms / 50, input_shape=[1, 3, 224, 224])
PyTorch resnet50 time: 37.5ms (= 1875.5ms / 50, input_shape=[1, 3, 224, 224])
✔️ Relative speed: 1.07 (= 37.5ms / 35.1ms)

OneFlow resnet50 time: 147.1ms (= 7354.4ms / 50, input_shape=[16, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 160.0ms (= 7997.5ms / 50, input_shape=[16, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.09 (= 160.0ms / 147.1ms)

OneFlow resnet50 time: 95.8ms (= 4791.6ms / 50, input_shape=[8, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 103.5ms (= 5175.1ms / 50, input_shape=[8, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.08 (= 103.5ms / 95.8ms)

OneFlow resnet50 time: 75.6ms (= 3778.2ms / 50, input_shape=[4, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 77.3ms (= 3865.2ms / 50, input_shape=[4, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.02 (= 77.3ms / 75.6ms)

OneFlow resnet50 time: 68.5ms (= 3422.9ms / 50, input_shape=[2, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 71.7ms (= 3583.6ms / 50, input_shape=[2, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.05 (= 71.7ms / 68.5ms)

OneFlow resnet50 time: 67.0ms (= 3350.7ms / 50, input_shape=[1, 3, 224, 224], ddp, world size=2)
PyTorch resnet50 time: 67.0ms (= 3351.6ms / 50, input_shape=[1, 3, 224, 224], ddp, world size=2)
✔️ Relative speed: 1.00 (= 67.0ms / 67.0ms)