|
| 1 | +# Copyright 2019 Kakao Brain |
| 2 | +# |
| 3 | +# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved. |
| 4 | +# |
| 5 | +# This source code is licensed under the BSD license found in the |
| 6 | +# LICENSE file in the root directory of this source tree. |
| 7 | +import pytest |
| 8 | +import torch |
| 9 | +from torch import nn |
| 10 | + |
| 11 | +from torch.distributed._pipeline.sync import Pipe |
| 12 | +from torch.distributed._pipeline.sync.skip import pop, skippable, stash |
| 13 | +from torch.distributed._pipeline.sync.skip.portal import PortalBlue, PortalCopy, PortalOrange |
| 14 | + |
| 15 | + |
| 16 | +@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required") |
| 17 | +@pytest.mark.parametrize("balance", [[3], [1, 2], [2, 1], [1, 1, 1]], ids=["3", "1:2", "2:1", "1:1:1"]) |
| 18 | +@pytest.mark.parametrize("checkpoint", ["never", "always", "except_last"]) |
| 19 | +def test_1to3(balance, checkpoint): |
| 20 | + if torch.cuda.device_count() < len(balance): |
| 21 | + pytest.skip("at least %d cuda devices required" % len(balance)) |
| 22 | + |
| 23 | + @skippable(stash=["1to3"]) |
| 24 | + class Layer1(nn.Module): |
| 25 | + def __init__(self): |
| 26 | + super().__init__() |
| 27 | + self.conv = nn.Conv2d(3, 3, 1) |
| 28 | + |
| 29 | + def forward(self, input): |
| 30 | + yield stash("1to3", input) |
| 31 | + output = self.conv(input) |
| 32 | + return output # noqa |
| 33 | + |
| 34 | + class Layer2(nn.Module): |
| 35 | + def __init__(self): |
| 36 | + super().__init__() |
| 37 | + self.conv = nn.Conv2d(3, 3, 1) |
| 38 | + |
| 39 | + def forward(self, input): |
| 40 | + output = self.conv(input) |
| 41 | + return output |
| 42 | + |
| 43 | + @skippable(pop=["1to3"]) |
| 44 | + class Layer3(nn.Module): |
| 45 | + def __init__(self): |
| 46 | + super().__init__() |
| 47 | + self.conv = nn.Conv2d(3, 3, 1) |
| 48 | + |
| 49 | + def forward(self, input): |
| 50 | + skip_1to3 = yield pop("1to3") |
| 51 | + output = self.conv(input) + skip_1to3 |
| 52 | + return output |
| 53 | + |
| 54 | + model = nn.Sequential(Layer1(), Layer2(), Layer3()) |
| 55 | + model = Pipe(model, balance, chunks=3, checkpoint=checkpoint) |
| 56 | + |
| 57 | + in_device = model.devices[0] |
| 58 | + out_device = model.devices[-1] |
| 59 | + |
| 60 | + input = torch.rand(30, 3, 224, 224, device=in_device, requires_grad=True) |
| 61 | + output = model(input) |
| 62 | + loss = output.mean() |
| 63 | + loss.backward() |
| 64 | + |
| 65 | + assert torch.allclose(output.norm(), torch.tensor(1039.0, device=out_device), atol=6e-1) |
| 66 | + assert torch.allclose(input.grad.norm(), torch.tensor(0.0004533053, device=in_device)) |
| 67 | + |
| 68 | + |
| 69 | +def test_none_skip(): |
| 70 | + @skippable(stash=["none"]) |
| 71 | + class Stash(nn.Module): |
| 72 | + def forward(self, input): |
| 73 | + yield stash("none", None) |
| 74 | + return input # noqa |
| 75 | + |
| 76 | + @skippable(pop=["none"]) |
| 77 | + class Pop(nn.Module): |
| 78 | + def forward(self, input): |
| 79 | + none = yield pop("none") |
| 80 | + assert none is None |
| 81 | + return input |
| 82 | + |
| 83 | + model = nn.Sequential(Stash(), Pop()) |
| 84 | + model = Pipe(model, [1, 1], devices=["cpu", "cpu"], chunks=5) |
| 85 | + |
| 86 | + input = torch.rand(10, requires_grad=True) |
| 87 | + output = model(input) |
| 88 | + |
| 89 | + def assert_grad_fn_is_not_portal(grad_fn, visited=None): |
| 90 | + if visited is None: |
| 91 | + visited = set() |
| 92 | + if grad_fn in visited or grad_fn is None: |
| 93 | + return |
| 94 | + |
| 95 | + assert not isinstance(grad_fn, PortalBlue._backward_cls) |
| 96 | + assert not isinstance(grad_fn, PortalCopy._backward_cls) |
| 97 | + assert not isinstance(grad_fn, PortalOrange._backward_cls) |
| 98 | + |
| 99 | + visited.add(grad_fn) |
| 100 | + for next_grad_fn, _ in grad_fn.next_functions: |
| 101 | + assert_grad_fn_is_not_portal(next_grad_fn, visited) |
| 102 | + |
| 103 | + assert_grad_fn_is_not_portal(output.grad_fn) |
| 104 | + |
| 105 | + output.sum().backward() |
| 106 | + assert input.grad.mean().item() == 1 |
0 commit comments