test_pipe.py 32 KB
Newer Older
Tom Birch's avatar
Tom Birch committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
# Copyright (c) Facebook, Inc. and its affiliates. All rights reserved.
#
# This source code is licensed under the BSD license found in the
# LICENSE file in the root directory of this source tree.

# Copyright 2019 Kakao Brain
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from collections import OrderedDict
from copy import deepcopy
import os
import time
24
from typing import Tuple
Tom Birch's avatar
Tom Birch committed
25
26
27
28
29
30

from packaging import version
import pytest
import torch
from torch import nn

31
32
33
34
35
36
37
from fairscale.nn.model_parallel.initialize import (
    destroy_model_parallel,
    get_pipeline_parallel_group,
    initialize_model_parallel,
)
from fairscale.nn.pipe import LazyModule, Pipe
from tests.nn.model_parallel.commons import get_worker_map, set_random_seed, torch_spawn
Tom Birch's avatar
Tom Birch committed
38
39
40


@torch_spawn([2])
41
42
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def parameters(pipeline_style):
Tom Birch's avatar
Tom Birch committed
43
    model = nn.Sequential(nn.Linear(1, 1))
44
    pipe = Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=1)
Tom Birch's avatar
Tom Birch committed
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
    if torch.distributed.get_rank() == 0:
        assert list(pipe.parameters()) != []
    else:
        assert list(pipe.parameters()) == []


@torch_spawn([2])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
def infiniband():
    if torch.distributed.get_rank() == 0:
        t = torch.Tensor(range(100)).cuda()
        torch.distributed.broadcast(t, 0)
    else:
        t = torch.empty(100).cuda()
        torch.distributed.broadcast(t, 0)

    assert torch.equal(t, torch.Tensor(range(100)).cuda())
    print(f"t on {torch.distributed.get_rank()} is {t}")


@torch_spawn([2])
@pytest.mark.skipif("OMPI_COMM_WORLD_RANK" not in os.environ, reason="mpi required")
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
def infiniband2():
    if torch.distributed.get_rank() == 0:
        t = torch.Tensor(range(100)).cuda()
71
        torch.distributed.send(t, 1, group=get_pipeline_parallel_group())
Tom Birch's avatar
Tom Birch committed
72
73
    else:
        t = torch.empty(100).cuda()
74
        torch.distributed.recv(t, 0, group=get_pipeline_parallel_group())
Tom Birch's avatar
Tom Birch committed
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111

    assert torch.equal(t, torch.Tensor(range(100)).cuda())
    print(f"t on {torch.distributed.get_rank()} is {t}")


@torch_spawn([2])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
def infiniband3():
    t = torch.Tensor(range(100)).cuda()
    torch.distributed.all_reduce(t, op=torch.distributed.ReduceOp.SUM)
    assert torch.equal(t, torch.Tensor(range(0, 200, 2)).cuda())


@torch_spawn([2])
@pytest.mark.skipif("OMPI_COMM_WORLD_RANK" not in os.environ, reason="mpi required")
def mpi():
    seed = 1234
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)

    torch.distributed.barrier()
    tensor_size = (1024, 1024, 10)
    torch.cuda.set_device(torch.distributed.get_rank())  # need to pin device or ucx gets unhappy

    if torch.distributed.get_rank() == 0:
        # t = torch.Tensor(range(10)).cuda(0)
        t = torch.rand(*tensor_size).cuda(0)
        torch.distributed.send(t, 1, tag=1234)
    else:
        t = torch.empty(*tensor_size).cuda(1)
        torch.distributed.recv(t, 0, tag=1234)
        t2 = torch.rand(*tensor_size).cuda(1)

        assert torch.equal(t, t2)


@torch_spawn([1])
112
113
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def public_attrs(pipeline_style):
Tom Birch's avatar
Tom Birch committed
114
115
116
117
118
119
120
121
122
123
124
125
    class MyString:
        def __init__(self, value):
            self.value = value

        def __str__(self):
            return self.value

    model = nn.Sequential(nn.Linear(1, 1))

    pipe = Pipe(
        model,
        balance=(1,),
126
        style=pipeline_style,
Tom Birch's avatar
Tom Birch committed
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
        worker_map=get_worker_map(),
        chunks=42.000,
        checkpoint=MyString("always"),
    )

    print(f"balance = {pipe.devices}")
    assert pipe.balance == [1]
    assert pipe.devices is None
    assert pipe.chunks == 42
    assert isinstance(pipe.chunks, int)
    assert pipe.checkpoint == "always"
    assert isinstance(pipe.checkpoint, str)


@torch_spawn([2])
@pytest.mark.parametrize("balance", [[2], [1, 1]])
143
144
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def sequential_like(balance, pipeline_style):
Tom Birch's avatar
Tom Birch committed
145
146
147
148
    a = nn.Linear(1, 1)
    b = nn.Linear(1, 1)

    model = nn.Sequential(a, b)
149
    model = Pipe(model, balance, style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181

    if balance == [2]:
        if torch.distributed.get_rank() == 0:
            assert len(model) == 2
            assert list(model) == [a, b]

            assert model[0] is a
            assert model[1] is b
            with pytest.raises(IndexError):
                _ = model[2]

            assert model[-1] is b
            assert model[-2] is a
        else:
            assert len(model) == 0
            assert list(model) == []
    else:
        assert len(model) == 1
        if torch.distributed.get_rank() == 0:
            assert list(model) == [a]
            assert model[0] is a
            assert model[-1] is a
        else:
            assert list(model) == [b]
            assert model[0] is b
            assert model[-1] is b

        with pytest.raises(IndexError):
            _ = model[1]


@torch_spawn([1])
182
183
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def balance_wrong_length(pipeline_style):
Tom Birch's avatar
Tom Birch committed
184
185
186
187
188
189
    a = nn.Linear(1, 1)
    b = nn.Linear(1, 1)

    model = nn.Sequential(a, b)

    with pytest.raises(ValueError):
190
        Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
191
192

    with pytest.raises(ValueError):
193
        Pipe(model, balance=[3], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
194
195
196


@torch_spawn([2])
197
198
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def balance_less_than_1(pipeline_style):
Tom Birch's avatar
Tom Birch committed
199
200
201
202
203
204
    a = nn.Linear(1, 1)
    b = nn.Linear(1, 1)

    model = nn.Sequential(a, b)

    with pytest.raises(ValueError):
205
        Pipe(model, balance=[0, 2], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
206
207

    with pytest.raises(ValueError):
208
        Pipe(model, balance=[-1, 3], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
209
210
211


@torch_spawn([1])
212
213
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def chunks_less_than_1(pipeline_style):
Tom Birch's avatar
Tom Birch committed
214
215
216
    model = nn.Sequential(nn.Linear(1, 1))

    with pytest.raises(ValueError):
217
        Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=0)
Tom Birch's avatar
Tom Birch committed
218
219

    with pytest.raises(ValueError):
220
        Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=-1)
Tom Birch's avatar
Tom Birch committed
221
222
223


@torch_spawn([1])
224
225
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def too_few_devices(pipeline_style):
Tom Birch's avatar
Tom Birch committed
226
227
228
229
    model = nn.Sequential(nn.Linear(1, 1), nn.Linear(1, 1), nn.Linear(1, 1), nn.Linear(1, 1))

    with pytest.raises(IndexError):
        # len(balance) > len(group.size())
230
        model = Pipe(model, balance=[1, 1, 1, 1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
231
232
233


@torch_spawn([1])
234
235
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def batch_size_indivisible(pipeline_style):
Tom Birch's avatar
Tom Birch committed
236
    model = nn.Sequential(nn.Linear(1, 1))
237
    model = Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=4)
Tom Birch's avatar
Tom Birch committed
238
239
240
241
242
243
244
245
246

    with pytest.warns(None) as record:
        model(torch.rand(7, 1))

    # Indivisible batch size is legal.
    assert not record


@torch_spawn([1])
247
248
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def batch_size_small(pipeline_style):
Tom Birch's avatar
Tom Birch committed
249
    model = nn.Sequential(nn.Linear(1, 1))
250
    model = Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=4)
Tom Birch's avatar
Tom Birch committed
251
252
253
254
255
256
257
258
259

    with pytest.warns(None) as record:
        model(torch.rand(2, 1))

    # Batch size smaller than chunks is legal.
    assert not record


@torch_spawn([1])
260
261
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def checkpoint_mode(pipeline_style):
Tom Birch's avatar
Tom Birch committed
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
    def count_grad_fn(grad_fn, name, visited=set()):
        if grad_fn in visited:
            return 0
        visited.add(grad_fn)

        if grad_fn is None:
            return 0
        if grad_fn.__class__.__name__ == name:
            return 1

        counter = 0
        for next_grad_fn, _ in grad_fn.next_functions:
            counter += count_grad_fn(next_grad_fn, name, visited=visited)
        return counter

    model = nn.Sequential(nn.Linear(1, 1))
    input = torch.rand(2, 1)

    always = Pipe(
        model,
        balance=[1],
283
        style=pipeline_style,
Tom Birch's avatar
Tom Birch committed
284
285
286
287
288
289
290
291
        worker_map=get_worker_map(),
        chunks=2,
        checkpoint="always",
        pipelined_backward=False,
    )
    except_last = Pipe(
        model,
        balance=[1],
292
        style=pipeline_style,
Tom Birch's avatar
Tom Birch committed
293
294
295
296
297
298
299
300
        worker_map=get_worker_map(),
        chunks=2,
        checkpoint="except_last",
        pipelined_backward=False,
    )
    never = Pipe(
        model,
        balance=[1],
301
        style=pipeline_style,
Tom Birch's avatar
Tom Birch committed
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
        worker_map=get_worker_map(),
        chunks=2,
        checkpoint="never",
        pipelined_backward=False,
    )

    always_output = always(input)
    except_last_output = except_last(input)
    never_output = never(input)

    assert count_grad_fn(always_output.grad_fn, "CheckpointBackward") == 2
    assert count_grad_fn(except_last_output.grad_fn, "CheckpointBackward") == 1
    assert count_grad_fn(never_output.grad_fn, "CheckpointBackward") == 0


@torch_spawn([1])
318
319
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def checkpoint_mode_invalid(pipeline_style):
Tom Birch's avatar
Tom Birch committed
320
321
322
323
324
325
    model = nn.Sequential(nn.Linear(1, 1))

    with pytest.raises(ValueError, match="checkpoint is not one of 'always', 'except_last', or 'never'"):
        Pipe(
            model,
            balance=[1],
326
            style=pipeline_style,
Tom Birch's avatar
Tom Birch committed
327
328
329
330
331
332
333
            worker_map=get_worker_map(),
            chunks=2,
            checkpoint="INVALID_CHECKPOINT",
        )


@torch_spawn([1])
334
335
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def checkpoint_mode_when_chunks_1(pipeline_style):
Tom Birch's avatar
Tom Birch committed
336
337
338
339
    model = nn.Sequential(nn.Linear(1, 1))

    # All checkpoint modes are fine.
    Pipe(
340
        model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=1, checkpoint="except_last",
Tom Birch's avatar
Tom Birch committed
341
    )
342
343
    Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=1, checkpoint="always")
    Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=1, checkpoint="never")
Tom Birch's avatar
Tom Birch committed
344
345
346


@torch_spawn([1])
347
348
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def checkpoint_eval(pipeline_style):
Tom Birch's avatar
Tom Birch committed
349
350
    model = nn.Sequential(nn.Linear(1, 1))
    model = Pipe(
351
        model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=2, pipelined_backward=False,
Tom Birch's avatar
Tom Birch committed
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
    )
    input = torch.rand(2, 1)

    def find_grad_fn(grad_fn, name):
        if grad_fn is None:
            return False
        if grad_fn.__class__.__name__ == name:
            return True
        for next_grad_fn, _ in grad_fn.next_functions:
            if find_grad_fn(next_grad_fn, name):
                return True
        return False

    model.train()
    train_output = model(input)
    assert find_grad_fn(train_output.grad_fn, "CheckpointBackward")
    assert find_grad_fn(train_output.grad_fn, "RecomputeBackward")

    model.eval()
    eval_output = model(input)
    assert not find_grad_fn(eval_output.grad_fn, "CheckpointBackward")
    assert not find_grad_fn(eval_output.grad_fn, "RecomputeBackward")


376
377
378
379
380
381
def torch_version() -> Tuple[int, ...]:
    result = version.parse(torch.__version__).release
    assert result
    return result


Tom Birch's avatar
Tom Birch committed
382
@torch_spawn([2])
383
384
385
@pytest.mark.xfail(torch_version() < (1, 6, 0), reason="Doesn't work on torch < 1.6.0", strict=True)
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def checkpoint_non_float_input(pipeline_style):
Tom Birch's avatar
Tom Birch committed
386
387
388
389
390
391
392
393
394
395
396
397
    class ForkNonFloat(nn.Module):
        def forward(self, input):
            return (input * 2, torch.tensor([False]))

    class JoinNonFloat(nn.Module):
        def forward(self, input):
            return input[0] * 2

    model = nn.Sequential(ForkNonFloat(), JoinNonFloat())
    model = Pipe(
        model,
        balance=[1, 1],
398
        style=pipeline_style,
Tom Birch's avatar
Tom Birch committed
399
400
401
402
403
404
405
406
407
408
409
        worker_map=get_worker_map(),
        chunks=1,
        checkpoint="always",
        pipelined_backward=False,
    )

    input = torch.rand(1, requires_grad=True)
    output = model(input)
    if model.group.rank() == 1:
        # with torch.autograd.detect_anomaly():
        output.backward()
410
    elif pipeline_style == Pipe.MultiProcess:
Tom Birch's avatar
Tom Birch committed
411
412
        model.back_helper(output)

413
414
    torch.distributed.barrier()

Tom Birch's avatar
Tom Birch committed
415
416

@torch_spawn([1])
417
418
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def no_grad(pipeline_style):
Tom Birch's avatar
Tom Birch committed
419
    model = nn.Sequential(nn.Linear(1, 1))
420
    model = Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=2)
Tom Birch's avatar
Tom Birch committed
421
422
423
424
425
426
427
428
429
430
431
    input = torch.rand(2, 1)

    latent = None

    def hook(module, input, output):
        _ = module
        _ = input

        nonlocal latent
        latent = output

432
433
    partition = model.mp_partitions[0]
    partition.module.register_forward_hook(hook)
Tom Birch's avatar
Tom Birch committed
434
435
436
437
438
439
440
441

    with torch.no_grad():
        model(input)

    assert latent.grad_fn is None


@torch_spawn([1])
442
443
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def exception(pipeline_style):
Tom Birch's avatar
Tom Birch committed
444
445
446
447
448
449
450
451
    class ExpectedException(Exception):
        pass

    class Raise(nn.Module):
        def forward(self, *_):
            raise ExpectedException()

    model = nn.Sequential(Raise())
452
    model = Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=1)
Tom Birch's avatar
Tom Birch committed
453
454
455
456
457
458
459
460

    with pytest.raises(ExpectedException):
        model(torch.rand(1))


# FIXME(tom) should probably signal to all hosts in group to stop
@torch_spawn([4])
@pytest.mark.xfail(strict=True)
461
462
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def exception_early_stop_asap(pipeline_style):
Tom Birch's avatar
Tom Birch committed
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
    """Even the first partitions have finished to process, the partition before
    the failed partition hould be killed as soon as possible.
    """

    class ExpectedExceptio(Exception):
        pass

    class Pass(nn.Module):
        def forward(self, x):
            return x

    counter = 0

    class Counter(nn.Module):
        def forward(self, x):
            time.sleep(0.1)

            nonlocal counter
            counter += 1

            return x

    class Raise(nn.Module):
        def forward(self, x):
            raise ExpectedException()

    model = nn.Sequential(Pass(), Pass(), Counter(), Raise())
490
    model = Pipe(model, [1, 1, 1, 1], style=pipeline_style, worker_map=get_worker_map(), chunks=3)
Tom Birch's avatar
Tom Birch committed
491
492
493
494
495
496
497
498
499

    with pytest.raises(ExpectedException):
        model(torch.rand(3))

    # If the early stop doesn't work, it would be 3 instead.
    assert counter == 2


@torch_spawn([1])
500
501
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def input_pair(pipeline_style):
Tom Birch's avatar
Tom Birch committed
502
503
504
505
506
507
508
509
510
511
512
513
    class Two(nn.Module):
        def __init__(self):
            super().__init__()
            self.fc_a = nn.Linear(1, 1)
            self.fc_b = nn.Linear(1, 1)

        def forward(self, a_and_b):
            a, b = a_and_b
            return (self.fc_a(a), self.fc_b(b))

    model = nn.Sequential(Two())
    model = Pipe(
514
        model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=2, pipelined_backward=False,
Tom Birch's avatar
Tom Birch committed
515
516
517
518
519
520
521
522
523
524
525
526
527
528
    )

    a = torch.rand(10, 1, requires_grad=True)
    b = torch.rand(10, 1, requires_grad=True)

    a_out, b_out = model((a, b))
    loss = (a_out + b_out).mean()
    loss.backward()

    assert a.grad is not None
    assert b.grad is not None


@torch_spawn([1])
529
530
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def input_singleton(pipeline_style):
Tom Birch's avatar
Tom Birch committed
531
532
533
534
535
536
537
538
539
540
541
    class One(nn.Module):
        def __init__(self):
            super().__init__()
            self.fc = nn.Linear(1, 1)

        def forward(self, only_a):
            (a,) = only_a
            return (self.fc(a),)

    model = nn.Sequential(One())
    model = Pipe(
542
        model, balance=[1], style=pipeline_style, worker_map=get_worker_map(), chunks=2, pipelined_backward=False,
Tom Birch's avatar
Tom Birch committed
543
544
545
546
547
548
549
550
551
552
553
554
555
    )

    a = torch.rand(10, 1, requires_grad=True)

    (a_out,) = model((a,))
    loss = a_out.mean()
    loss.backward()

    assert all(p.grad is not None for p in model.parameters())
    assert a.grad is not None


@torch_spawn([1])
556
557
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def input_varargs(pipeline_style):
Tom Birch's avatar
Tom Birch committed
558
    model = nn.Sequential(nn.Linear(1, 1))
559
    model = Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
560
561
562
563
564
565
566
567
568
569

    a = torch.rand(1)
    b = torch.rand(1)

    # TypeError: forward() takes 2 positional arguments but 3 were given
    with pytest.raises(TypeError):
        model(a, b)


@torch_spawn([1])
570
571
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def non_tensor(pipeline_style):
Tom Birch's avatar
Tom Birch committed
572
573
574
575
576
    class NonTensor(nn.Module):
        def forward(self, _):
            return "hello"

    model = nn.Sequential(NonTensor())
577
    model = Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
578
579
580
581
582
583
584
585
586
587
588
589
    x = torch.rand(1)

    # TypeError: expected Tensor as element 0 in argument 0, but got str
    with pytest.raises(TypeError):
        model(x)

    # TypeError: expected Tensor to scatter, but got str
    with pytest.raises(TypeError):
        model("hello")


@torch_spawn([1])
590
591
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def non_tensor_tuple(pipeline_style):
Tom Birch's avatar
Tom Birch committed
592
593
594
595
596
    class NonTensorTuple(nn.Module):
        def forward(self, x):
            return (x, "hello")

    model = nn.Sequential(NonTensorTuple())
597
    model = Pipe(model, balance=[1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
598
599
600
601
602
603
604
605
606
607
608
609
610
611
    x = torch.rand(1)

    # TypeError: CheckpointBackward.forward: expected Variable (got str) for return value 1
    with pytest.raises(TypeError):
        model(x)

    # TypeError: expected Tensor to scatter, but got str
    with pytest.raises(TypeError):
        model((x, "hello"))


@torch_spawn([1])
@pytest.mark.parametrize("checkpoint", ["never", "always", "except_last"])
@pytest.mark.parametrize("lazy", [True, False])
612
613
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def deferred_batch_norm(checkpoint, lazy, pipeline_style):
Tom Birch's avatar
Tom Birch committed
614
615
616
617
    bn = nn.BatchNorm2d(3)
    pipe_bn = deepcopy(bn)
    pipe_fn = lambda: pipe_bn  # noqa: E731
    if lazy:
618
        model = [LazyModule(pipe_fn)]
Tom Birch's avatar
Tom Birch committed
619
620
621
622
623
    else:
        model = nn.Sequential(pipe_bn)
    pipe = Pipe(
        model,
        balance=[1],
624
        style=pipeline_style,
Tom Birch's avatar
Tom Birch committed
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
        worker_map=get_worker_map(),
        chunks=2,
        checkpoint=checkpoint,
        deferred_batch_norm=True,
    )

    x = torch.rand(4, 3, 10, 10)
    pipe(x).mean().backward()
    bn(x).mean().backward()

    assert torch.allclose(pipe[0].running_mean, bn.running_mean, atol=1e-4)
    assert torch.allclose(pipe[0].running_var, bn.running_var, atol=1e-4)


@torch_spawn([1])
@pytest.mark.parametrize("checkpoint", ["never", "always"])
@pytest.mark.parametrize("lazy", [True, False])
642
643
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def deferred_batch_norm_params(checkpoint, lazy, pipeline_style):
Tom Birch's avatar
Tom Birch committed
644
645
646
647
    bn = nn.BatchNorm2d(3)
    pipe_bn = deepcopy(bn)
    pipe_fn = lambda: pipe_bn  # noqa: E731
    if lazy:
648
        model = [LazyModule(pipe_fn)]
Tom Birch's avatar
Tom Birch committed
649
650
651
652
653
    else:
        model = nn.Sequential(pipe_bn)
    pipe = Pipe(
        model,
        balance=[1],
654
        style=pipeline_style,
Tom Birch's avatar
Tom Birch committed
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
        worker_map=get_worker_map(),
        chunks=1,
        checkpoint=checkpoint,
        deferred_batch_norm=True,
    )

    x = torch.rand(4, 3, 10, 10)
    pipe(x).mean().backward()
    bn(x).mean().backward()

    assert pipe[0].weight.grad is not None
    assert pipe[0].bias.grad is not None

    assert torch.allclose(pipe[0].weight.grad, bn.weight.grad, atol=1e-4)
    assert torch.allclose(pipe[0].bias.grad, bn.bias.grad, atol=1e-4)


672
@torch_spawn([4])
673
674
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def devices(pipeline_style):
Tom Birch's avatar
Tom Birch committed
675
676
677
678
679
680
    a = nn.Linear(1, 1)
    b = nn.Linear(1, 1)
    c = nn.Linear(1, 1)

    # There are extra two ranks.
    model = nn.Sequential(a, b, c)
681
    model = Pipe(model, [1, 1, 1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
682
683

    # Extra devices must be discarded.
684
    if model.group.rank() == 3:
Tom Birch's avatar
Tom Birch committed
685
686
687
688
        assert model.pipeline is None


@torch_spawn([2])
689
690
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def partitions(pipeline_style):
Tom Birch's avatar
Tom Birch committed
691
692
693
694
    a = nn.Linear(1, 1)
    b = nn.Linear(1, 1)

    model = nn.Sequential(a, b)
695
    model = Pipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
696

697
    assert isinstance(model.mp_partitions, list)
Tom Birch's avatar
Tom Birch committed
698
    assert len(model) == 1
699
    assert isinstance(model.mp_partitions[0].module, nn.Sequential)
Tom Birch's avatar
Tom Birch committed
700

701
702
703
704
    if model.group.rank() == 0:
        assert "0.0.weight" in model.state_dict()
    else:
        assert "0.1.weight" in model.state_dict()
Tom Birch's avatar
Tom Birch committed
705
706
707
708


@torch_spawn([2])
@pytest.mark.skipif(not torch.cuda.is_available(), reason="cuda required")
709
710
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def deny_moving(pipeline_style):
Tom Birch's avatar
Tom Birch committed
711
712
713
714
    a = nn.Linear(1, 1)
    b = nn.Linear(1, 1)

    model = nn.Sequential(a, b)
715
    model = Pipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732

    model.cuda()
    model.cpu()
    model.to(torch.device("cuda"))
    model.to(0)
    model.to("cuda")
    model.to(device=0)
    model.to(torch.rand(1))
    model.to(tensor=torch.rand(1))

    # Casting is allowed.
    model.half()
    model.to(torch.double)
    model.to(dtype=torch.float)


@torch_spawn([1])
733
734
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def empty_module(pipeline_style):
Tom Birch's avatar
Tom Birch committed
735
736
    # Empty sequential module is not illegal.
    model = nn.Sequential()
737
    model = Pipe(model, [], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
738
739
740
741
742
743
744
745
746
747
748

    assert model(torch.tensor([42])) == torch.tensor([42])
    assert model((torch.tensor([42]),)) == (torch.tensor([42]),)

    # But only tensor or tensors is legal in Pipe.

    with pytest.raises(TypeError):
        model(42)


@torch_spawn([2])
749
750
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def named_children(pipeline_style):
Tom Birch's avatar
Tom Birch committed
751
752
753
754
    a = nn.Linear(1, 1)
    b = nn.Linear(1, 1)

    model = nn.Sequential(OrderedDict([("a", a), ("b", b)]))
755
    model = Pipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
756
757

    names = set(n for n, _ in model.named_modules())
758
759
760
761
    if model.group.rank() == 0:
        assert "0.a" in names
    else:
        assert "0.b" in names
Tom Birch's avatar
Tom Birch committed
762
763
764
765
766
767
768
769

    # Pipe doesn't support __getattr__. Unlike nn.Sequential, Pipe requires
    # several methods in its namespace.
    with pytest.raises(AttributeError):
        model.a


@torch_spawn([1])
770
771
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def recommend_auto_balance(pipeline_style):
Tom Birch's avatar
Tom Birch committed
772
773
774
775
776
777
778
779
780
781
782
783
784
785
    with pytest.raises(ValueError, match="fairscale.nn.pipe.balance"):
        # balance is required
        Pipe(nn.Sequential())

    with pytest.raises(ValueError, match="fairscale.nn.pipe.balance"):
        # module and sum of balance have differen length (module: 0, sum of balance: 1)
        Pipe(nn.Sequential(), [1])

    with pytest.raises(ValueError, match="fairscale.nn.pipe.balance"):
        # module and sum of balance have different length (module: 2, sum of balance: 1)
        Pipe(nn.Sequential(nn.Linear(1, 1), nn.Linear(1, 1)), [1])


@torch_spawn([2])
786
787
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def lazy_construction(pipeline_style):
Tom Birch's avatar
Tom Birch committed
788
789
790
791
792
793
794
795
796
797
798
799
    init_count = 0

    class Custom(nn.Module):
        def __init__(self):
            super(Custom, self).__init__()
            nonlocal init_count
            init_count += 1

        def forward(self, x):
            return x

    model = [
800
801
802
803
        LazyModule(lambda: Custom()),
        LazyModule(lambda: Custom()),
        LazyModule(lambda: Custom()),
        LazyModule(lambda: Custom()),
Tom Birch's avatar
Tom Birch committed
804
805
    ]

806
    pipe = Pipe(model, balance=[2, 2], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
807
808
809
810
811
812
813
814

    assert isinstance(pipe[0], Custom)
    assert isinstance(pipe[1], Custom)
    assert len(pipe) == 2
    assert init_count == 2


@torch_spawn([2])
815
816
817
@pytest.mark.skipif("OMPI_COMM_WORLD_RANK" in os.environ, reason="doesn't apply to mpi")
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def missing_worker_map(pipeline_style):
Tom Birch's avatar
Tom Birch committed
818
819
    model = nn.Sequential(nn.ReLU(), nn.ReLU())

820
821
    with pytest.raises(ValueError, match="'RpcTransport' requires 'worker_map' to be set"):
        Pipe(model, [1, 1], style=pipeline_style)
Tom Birch's avatar
Tom Birch committed
822
823
824
825


@torch_spawn([2])
@pytest.mark.skip(reason="currently broken")
826
827
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def verify_module_duplicate_parameters_on_distinct_partitions(pipeline_style):
Tom Birch's avatar
Tom Birch committed
828
829
830
831
832
833
834
835
836
837
    class Surrogate(nn.Module):
        def __init__(self, module):
            super().__init__()
            self.module = module

    conv = nn.Conv2d(3, 3, 1)
    model = nn.Sequential(Surrogate(conv), Surrogate(conv))

    # FIXME(tom) can't have duplicate params with separate processes
    with pytest.raises(ValueError, match="module with duplicate parameters on distinct devices is not supported"):
838
        Pipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
839
840
841


@torch_spawn([4])
842
843
@pytest.mark.parametrize("pipeline_style", [Pipe.MultiProcess, Pipe.AsyncSchedule])
def pipelined_backward(pipeline_style):
Tom Birch's avatar
Tom Birch committed
844
845
846
847
    model = nn.Sequential(nn.ReLU(), nn.ReLU())

    destroy_model_parallel()
    initialize_model_parallel(1, 4)
848
    pipe = Pipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
849
850
851
852
853

    assert pipe.pipelined_backward is False

    destroy_model_parallel()
    initialize_model_parallel(2, 2)
854
    pipe = Pipe(model, [1, 1], style=pipeline_style, worker_map=get_worker_map())
Tom Birch's avatar
Tom Birch committed
855
856

    assert pipe.pipelined_backward is True
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
1002
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022
1023
1024
1025
1026
1027
1028
1029
1030
1031
1032
1033
1034
1035
1036
1037
1038
1039


@torch_spawn([4])
def async_event_loop():

    model = nn.Sequential(nn.Linear(10, 10), nn.ReLU(), nn.Linear(10, 10), nn.ReLU())
    pipe = Pipe(model, [1, 1, 1, 1], style=Pipe.AsyncSchedule, worker_map=get_worker_map(), chunks=10)

    inputs = torch.rand(100, 10)

    output = pipe(inputs)
    if pipe.final_stage:
        loss = output.mean()
        loss.backward()


@torch_spawn([4])
def reuse_lazy():
    if False:  # speed
        reused = LazyModule(lambda: nn.Linear(10, 10))
        model = [reused, nn.Linear(10, 10), nn.ReLU(), reused, nn.ReLU(), reused, nn.ReLU()]
        # model = [reused, reused, nn.Linear(10, 10), nn.ReLU(), reused, reused, nn.ReLU(), reused, reused, nn.ReLU()]
        pipe = Pipe(model, [3, 1, 1], style=Pipe.AsyncSchedule, worker_map=get_worker_map())
        pipe.eval()
        output = pipe(torch.rand(10))

        print(f"output on {pipe.group.rank()}, {output}")
        torch.distributed.barrier()

    set_random_seed(1234)
    # test both foward
    reused = nn.Linear(10, 10)
    layers = [reused, nn.Linear(10, 10), nn.ReLU(), reused, nn.ReLU(), reused, nn.ReLU()]
    model = nn.Sequential(*layers)
    model.eval()

    set_random_seed(1234)
    # ensure identical weights but no sharing between model and pipe
    reused = nn.Linear(10, 10)
    layers = [reused, nn.Linear(10, 10), nn.ReLU(), reused, nn.ReLU(), reused, nn.ReLU()]
    pipe = Pipe(layers, [3, 1, 1], style=Pipe.AsyncSchedule, worker_map=get_worker_map())
    pipe.eval()
    model_optimizer = torch.optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
    pipe_optimizer = torch.optim.SGD(pipe.parameters(), lr=0.01, momentum=0.9) if len(list(pipe.parameters())) else None
    inputs = torch.rand(10)
    if False:  # speed
        model_out = model(inputs)
        pipe_out = pipe(inputs)

        torch.distributed.barrier()

        if pipe.final_stage:
            assert torch.equal(model_out, pipe_out)

    model.train()
    pipe.train()
    model_out = model(inputs)
    pipe_out = pipe(inputs)
    if pipe.final_stage:
        pipe_loss = pipe_out.mean()
        pipe_loss.backward()

    model_loss = model_out.mean()
    model_loss.backward()

    model_optimizer.step()
    if pipe_optimizer:
        pipe_optimizer.step()

    model.eval()
    pipe.eval()
    model_out = model(inputs)
    pipe_out = pipe(inputs)

    print(f"before barrier on {torch.distributed.get_rank()}")
    torch.distributed.barrier()
    print(f"after barrier on {torch.distributed.get_rank()}")

    if pipe.final_stage:
        assert torch.equal(model_out, pipe_out)


def test_instantiate_partition():
    from fairscale.nn.pipe.async_schedule import Location
    from fairscale.nn.pipe.pipe import instantiate_partition

    class FakeGroup:
        def __init__(self, rank, size):
            self._rank = rank
            self._size = size

        def rank(self):
            return self._rank

        def size(self):
            return self._size

    def check_partitions(model, balance, expected_order, expected_ranks):
        """Check the instantiated model matches expectation of order and rank

        model: a list of modules or an nn.Sequential
        balance: the balance argument to Pipe
        expected_order: the index of modules in `model` in the order they will
            be executed, grouped by nn.Sequential
        expected_rank: the rank that each module will be executed on
        """

        invocations = []
        invocation_wrapper = dict()

        # Collect `Invocation` and `Invocation` -> `ModuleWrapper` mapping from
        # instantiated model
        for rank in range(len(balance)):
            instantiated = instantiate_partition(model, balance, FakeGroup(rank, len(balance)), Pipe.AsyncSchedule)
            for part in instantiated:
                assert isinstance(part.module, nn.Sequential)
                for inv in part.invocations:
                    invocations.append(inv)
                    invocation_wrapper[inv] = part

        modules = []
        prev = None
        current = Location(0, 0)
        ranks = []

        for order, inv in enumerate(sorted(invocations, key=lambda x: x.order)):
            # Check integrity of Location chain
            assert inv.order == order
            assert inv.source == prev
            assert inv.this == current
            prev = inv.this
            current = inv.dest
            modules.append(list(invocation_wrapper[inv].module.children()))
            ranks.append(inv.this.stage)

        # assert len(modules) == len(expected_order)
        for left, right in zip(modules, expected_order):
            assert len(left) == len(right), f"{right}"
            assert list(map(id, left)) == list(map(id, (model[e] for e in right))), f"{right}"

        assert ranks == expected_ranks

    reused = nn.Linear(20, 20)
    model = [reused, nn.Linear(10, 10), nn.ReLU(), reused, nn.ReLU(), reused, nn.ReLU()]
    balance = [3, 1, 1]

    check_partitions(
        model, balance, expected_order=[[0], [1, 2], [0], [4], [0], [6]], expected_ranks=[0, 0, 0, 1, 0, 2]
    )

    reused2 = nn.Linear(5, 5)
    model = [reused, reused2, nn.Linear(10, 10), nn.ReLU(), reused, reused2, nn.ReLU(), reused, reused2, nn.ReLU()]
    balance = [4, 1, 1]

    check_partitions(
        model,
        balance,
        expected_order=[[0], [1], [2, 3], [0], [1], [6], [0], [1], [9]],
        expected_ranks=[0, 0, 0, 0, 0, 1, 0, 0, 2],
    )

    reused2 = nn.Linear(5, 5)
    model = [
        nn.Linear(10, 10),
        reused,
        nn.Linear(10, 10),
        nn.ReLU(),
        reused,
        reused2,
        nn.ReLU(),
        reused,
        reused2,
        nn.ReLU(),
    ]
    # 0 1 2 3 1 5 6 1 5 9
    balance = [4, 2, 1]

    check_partitions(
        model,
        balance,
        expected_order=[[0], [1], [2, 3], [1], [5], [6], [1], [5], [9]],
        expected_ranks=[0, 0, 0, 0, 1, 1, 0, 1, 2],
    )