Balance Module

🐞 Bug

Balance by time didn’t ignore the time for CUDA initialization. I think adding a warm-up run can eliminate the effect of that. Balance by size actually accumulate the size so that the latter layer always consumes more memory than the previous layers.

Code that reproduces

Change the test to the following can reproduce for balance by size.

def test_balance_by_size():
    class Expand(nn.Module):
        def __init__(self, times):
            super().__init__()
            self.times = times

        def forward(self, x):
            for i in range(self.times):
                x = x + torch.rand_like(x, requires_grad=True)
            return x

    model = nn.Sequential(*[Expand(i) for i in [6, 5, 4, 3, 2, 1]])
    sample = torch.rand(10, 100, 100)
    balance = balance_by_size(model, sample, partitions=2, device='cuda')
    assert balance == [2, 4]

[842974287]

That makes sense, thanks!

[sublee]

I fixed this issue with 3011ff0613fc91606c637d5df139e44b2d3bfaa5. If you still have the same issue, please reopen it.