test_oss.py 33.3 KB
Newer Older
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
1
2
3
4
5
# 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.

6
7
8
9
# pylint: disable=missing-module-docstring
# pylint: disable=missing-class-docstring
# pylint: disable=missing-function-docstring

10

11
12
import copy
from math import inf
13
import tempfile
14
from typing import Any, Dict, Type, cast
15
import unittest
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
16

17
import numpy as np
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
18
19
20
21
import pytest
import torch
import torch.distributed as dist
import torch.multiprocessing as mp
22
from torch.nn.parallel import DistributedDataParallel as DDP
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
23
24

import fairscale.optim as optim
25
from fairscale.utils.testing import check_same_model_params, skip_if_no_cuda, skip_if_py39_no_cuda, skip_if_single_gpu
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
26

27
28
BACKEND = dist.Backend.NCCL if torch.cuda.is_available() else dist.Backend.GLOO  # type: ignore
DEVICE = "cuda" if torch.cuda.is_available() else torch.device("cpu")
29
RECIPIENT_RANK = 1
30

31
32
33
34
35
36
37
38
39
try:
    from torch.distributed import broadcast_object_list  # noqa

    _torch_broadcast_object = True
except ImportError:
    from fairscale.optim.utils import broadcast_object  # noqa

    _torch_broadcast_object = False

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
40

41
42
43
def dist_init(rank, world_size, tempfile_name, backend=BACKEND):
    url = "file://" + tempfile_name
    dist.init_process_group(init_method=url, backend=backend, rank=rank, world_size=world_size)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
44
45


46
47
48
49
50
51
52
53
54
55
56
57
58
def sync_object_ranks(something_to_sync: Any, reference_rank: int, device: torch.device) -> Any:
    if _torch_broadcast_object:
        package = [something_to_sync]
        dist.broadcast_object_list(package, src=reference_rank, group=dist.group.WORLD)
        package_sync = package[0]
    else:
        package_sync = optim.utils.broadcast_object(
            something_to_sync, src_rank=reference_rank, group=dist.group.WORLD, dist_device=device
        )

    return package_sync


59
60
61
62
class TestSingleRank(unittest.TestCase):
    """
    All the following tests do not check for inter-process communication
    """
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
63

64
65
    def setUp(self):
        dist_init(0, 1, tempfile.mkstemp()[1])
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
66

67
68
    def tearDown(self):
        torch.distributed.destroy_process_group()
69

70
71
72
    def test_create(self):
        params = [torch.rand(1)]
        o = optim.OSS(params, lr=0.01)
73

74
75
76
    def test_state_dict(self):
        x = torch.tensor([1.0], device=DEVICE, requires_grad=True)
        o = optim.OSS([x], lr=0.1, momentum=0.9)
77
        x.backward()
78
79
80
        o.step()
        assert x == torch.tensor([0.9], device=DEVICE)
        assert o.optim.state[x]["momentum_buffer"] == torch.tensor([1.0], device=DEVICE)
81
        o.zero_grad()
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
        o.consolidate_state_dict()  # Sync state dict in between replicas - even if there are none
        state_dict = o.state_dict()

        # Check that the state dict is pytorch-compliant key wise
        assert "param_groups" in state_dict.keys()
        assert "state" in state_dict.keys()

        # Check that the pulled state is what we expect, and that we have all the expected keys
        assert state_dict["param_groups"][0]["lr"] == 0.1
        assert state_dict["param_groups"][0]["momentum"] == 0.9
        assert not state_dict["param_groups"][0]["nesterov"]
        assert state_dict["param_groups"][0]["weight_decay"] == 0.0
        assert state_dict["param_groups"][0]["dampening"] == 0.0

        # Check that the pulled state and the .param_groups attribute are in sync
        for k in state_dict["param_groups"][0].keys():
            if k != "params":
                assert state_dict["param_groups"][0][k] == o.param_groups[0][k]

        # Check that it's correctly loaded
        o = optim.OSS([x], lr=0.01)
        o.load_state_dict(state_dict)
        # Check that state is correct and on proper device
        assert o.optim.state[x]["momentum_buffer"] == torch.tensor([1.0], device=DEVICE)

        # We should now be using a lr of 0.1, both within the optimizer
        # and as exposed by the .param_groups attribute
        assert o.param_groups[0]["lr"] == 0.1
        x.backward()
111
        o.step()
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
        assert x == torch.tensor([0.71], device=DEVICE)
        assert o.optim.state[x]["momentum_buffer"] == torch.tensor([1.9], device=DEVICE)

        # Check that the exposed param_groups are on the proper device
        assert o.param_groups[0]["params"][0].device == x.device

    def test_lr_scheduler(self):
        x = torch.tensor([1.0], device=DEVICE, requires_grad=True)
        x2 = torch.tensor([1.0], device=DEVICE, requires_grad=True)
        o = optim.OSS([x], lr=0.01)
        o2 = torch.optim.SGD([x2], lr=0.01)
        s = torch.optim.lr_scheduler.StepLR(o, 1)
        s2 = torch.optim.lr_scheduler.StepLR(o2, 1)
        for _ in range(5):
            x.backward()
            o.zero_grad()
            o.step()
            s.step()
            x2.backward()
            o2.zero_grad()
            o2.step()
            s2.step()
            assert x == x2

    def test_step_with_kwargs(self):
        class SGDWithStepKWArg(torch.optim.SGD):
            def step(self, closure=None, kwarg=[]):
                super().step()
                kwarg.append(5)

        kwarg = []
        x = torch.tensor([1.0], device=DEVICE, requires_grad=True)
        o = optim.OSS([x], SGDWithStepKWArg, lr=0.1)
        x.backward()
        o.step(0, kwarg=kwarg)
        assert kwarg == [5]
        assert x == torch.tensor([0.9], device=DEVICE)

    def test_step_with_extra_inner_key(self):
        class SGDWithNewKey(torch.optim.SGD):
            # Dummy optimizer which adds a new key to the param groups
            def step(self, closure=None):
                super().step()
                self.param_groups[0]["new_key"] = 0.1

        x = torch.tensor([1.0], device=DEVICE, requires_grad=True)
        o = optim.OSS([x], SGDWithNewKey, lr=0.1)
        x.backward()
        o.step()
        assert o.param_groups[0]["new_key"] == 0.1
        assert x == torch.tensor([0.9], device=DEVICE)
163

164
165
166
167
    def test_step_without_closure(self):
        class SGDWithoutClosure(torch.optim.SGD):
            def step(self):
                return super().step()
168

169
170
171
172
173
174
175
176
177
        x = torch.tensor([1.0], device=DEVICE, requires_grad=True)
        o = optim.OSS([x], SGDWithoutClosure, lr=0.1)
        x.backward()
        o.step()
        assert x == torch.tensor([0.9], device=DEVICE)

    def test_implicit_local_state_dict(self):
        x = torch.tensor([1.0], device=DEVICE, requires_grad=True)
        o = optim.OSS([x], lr=0.1)
178
179
        with pytest.raises(RuntimeError):
            _ = o.state_dict()
180
181


182
183
def run_test_add_param_group(rank, world_size, tempfile_name):
    dist_init(rank, world_size, tempfile_name)
184
185
186
187

    # Test with all parameters trainable to begin with
    def all_trainable():
        params = []
188
189
190
        sizes = [9, 7, 5, 3]
        sizes_world = sizes * world_size
        for size in sizes_world[:-1]:
191
192
193
194
195
196
197
198
199
200
201
202
            params.append(torch.rand(size, 1))

        # Make sure that the params are trainable, enforces size-based partitioning
        for p in params:
            p.requires_grad = True

        o = optim.OSS(params, lr=0.1)

        assert len(o.param_groups) == 1
        o.add_param_group({"params": [torch.rand(3, 1)]})

        assert len(o.param_groups) == 2
203
204
205

        # Verify that added group is added to the correct partition making all have the same number of elements
        assert sum([x.numel() for g in o.optim.param_groups for x in g["params"]]) == sum(sizes)
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
        assert len(o.optim.param_groups) == 2

    # Test a pathological config with a first big non-trainable param
    def some_trainable():
        params = []
        for size in [100, 3, 5, 2, 6, 4]:
            params.append(torch.rand(size, 1))

        # Make sure that the params are trainable, enforces size-based partitioning
        for p in params[1:]:
            p.requires_grad = True

        o = optim.OSS(params, lr=0.1)

        assert len(o.param_groups) == 1
        o.add_param_group({"params": [torch.rand(3, 1)]})

        assert len(o.param_groups) == 2
        assert len(o.optim.param_groups) == 2

    all_trainable()
    some_trainable()
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
228

229
230
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
231
232

def test_add_param_group():
233
    world_size = 4
234
    if torch.cuda.is_available() and torch.cuda.device_count() < world_size:
235
236
        world_size = min(world_size, torch.cuda.device_count())

237
    mp.spawn(run_test_add_param_group, args=(world_size, tempfile.mkstemp()[1]), nprocs=world_size, join=True)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
238
239


240
241
def run_test_zero_grad(rank, world_size, tempfile_name):
    dist_init(rank, world_size, tempfile_name)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
242
243
244
245
246
247
248
249
250
251
252
    x = torch.rand(1)
    m = torch.nn.Linear(1, 1)
    o = optim.OSS(m.parameters(), lr=0.1)
    y = m(x)
    y.backward(x)
    assert m.weight.grad
    assert m.bias.grad
    o.zero_grad()
    assert not m.weight.grad
    assert not m.bias.grad

253
254
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
255
256
257

def test_zero_grad():
    world_size = 2
258
259
260
    if torch.cuda.is_available() and torch.cuda.device_count() < world_size:
        world_size = min(world_size, torch.cuda.device_count())

261
262
    temp_file_name = tempfile.mkstemp()[1]
    mp.spawn(run_test_zero_grad, args=(world_size, temp_file_name), nprocs=world_size, join=True)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
263
264


265
266
def run_test_empty_shard(rank, world_size, tempfile_name, backend):
    dist_init(rank, world_size, tempfile_name, backend=backend)
267
    m = torch.nn.Linear(1, 1)
268
269
270
271
272
273
274
275
276
277
    x = torch.rand(20, 1)

    if torch.cuda.is_available():
        m = m.to(rank)
        x = x.to(rank)

    o = optim.OSS(m.parameters(), lr=0.1)
    y = m(x).sum()
    y.backward()
    o.step()
278
279
280
281

    dist.destroy_process_group()


282
283
@pytest.mark.parametrize("backend", ["gloo", "nccl"])
def test_empty_shard(backend):
284
    world_size = 4
285
286
287
288
289
    if torch.cuda.is_available() and torch.cuda.device_count() < world_size:
        world_size = min(world_size, torch.cuda.device_count())
    if world_size == 1 or (backend == "nccl" and not torch.cuda.is_available()):
        pytest.skip("Not enough GPUs to test with NCCL, or CUDA not present")
    mp.spawn(run_test_empty_shard, args=(world_size, tempfile.mkstemp()[1], backend), nprocs=world_size, join=True)
290
291


292
293
def run_test_step(rank, world_size, tempfile_name):
    dist_init(rank, world_size, tempfile_name, backend="gloo")
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
    x = torch.tensor([float(rank + 1)], device=rank)
    m = torch.nn.Linear(1, 1)
    m.weight.data = torch.tensor([[1.0]])
    m.bias.data = torch.tensor([2.0])
    m.to(rank)
    o = optim.OSS(m.parameters(), lr=0.1)
    y = m(x)
    y.backward(x)
    for p in m.parameters():
        dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
        p.grad.data /= world_size
    o.step()
    assert m.weight == torch.tensor([[0.75]], device=rank)
    assert m.bias == torch.tensor([1.85], device=rank)

309
310
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
311

Benjamin Lefaudeux's avatar
Benjamin Lefaudeux committed
312
@skip_if_single_gpu
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
313
def test_step():
Benjamin Lefaudeux's avatar
Benjamin Lefaudeux committed
314
    world_size = 2
315
316
317
    temp_file_name = tempfile.mkstemp()[1]

    mp.spawn(run_test_step, args=(world_size, temp_file_name), nprocs=world_size, join=True)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
318
319


320
321
def run_test_step_with_closure(rank, world_size, tempfile_name, optimizer=None):
    dist_init(rank, world_size, tempfile_name)
322

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
323
324
325
326
327
328
329
330
331
332
333
334
    x_val = rank + 1
    weight = 1.0
    bias = 2.0
    error = 1.0
    target = torch.tensor([x_val * weight + bias + error], device=rank)
    loss_fn = torch.nn.L1Loss()

    x = torch.tensor([float(x_val)], device=rank)
    m = torch.nn.Linear(1, 1)
    m.weight.data = torch.tensor([[weight]])
    m.bias.data = torch.tensor([bias])
    m.to(rank)
335

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
336
    o = optim.OSS(m.parameters(), lr=0.1)
337

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
    y = m(x)
    y.backward(x)
    for p in m.parameters():
        dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
        p.grad.data /= world_size

    def closure():
        o.zero_grad()
        output = m(x)
        loss = loss_fn(output, target)
        loss.backward()
        return loss

    loss = o.step(closure=closure)

    assert loss == torch.tensor(error, device=rank)
    assert m.weight == torch.tensor([[1.1]], device=rank)
    assert m.bias == torch.tensor([2.1], device=rank)

357
358
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
359
360
361

@skip_if_no_cuda
def test_step_with_closure():
362
    world_size = min(2, torch.cuda.device_count())
363
364
365
    temp_file_name = tempfile.mkstemp()[1]

    mp.spawn(run_test_step_with_closure, args=(world_size, temp_file_name), nprocs=world_size, join=True)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
366
367


368
369
def run_test_sharding(rank, world_size, tempfile_name):
    dist_init(rank, world_size, tempfile_name)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
370
    params = []
371
372
373
374
    sizes = [9, 7, 5, 3]
    sizes_world = sizes * world_size

    for size in sizes_world:
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
375
        params.append(torch.rand(size, 1))
376
377
378
379
380

    # Make sure that the params are trainable, enforces size-based partitioning
    for p in params:
        p.requires_grad = True

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
381
    o = optim.OSS(params, lr=0.1)
382
    assert sum([x.numel() for x in o.optim.param_groups[0]["params"]]) == sum(sizes)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
383

384
385
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
386
387

def test_sharding():
388
389
390
    world_size = 4
    if torch.cuda.is_available():
        world_size = min(world_size, torch.cuda.device_count())
391

392
    _, temp_file_name = tempfile.mkstemp()
393
    mp.spawn(run_test_sharding, args=(world_size, temp_file_name), nprocs=world_size, join=True)
394
395


396
397
def run_test_collect_shards(rank, world_size, reference_rank, tempfile_name):
    dist_init(rank, world_size, tempfile_name)
398
399
400
    device = torch.device(rank) if torch.cuda.device_count() > 1 else DEVICE

    # Run a dummy step so that the optimizer state dict exists
401
    batch, input_width, hidden, target_width = 3, 3, 3, 5
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
    target = torch.rand((batch, target_width), device=device)
    inputs = torch.rand((batch, input_width), device=device)

    model = torch.nn.Sequential(torch.nn.Linear(input_width, hidden), torch.nn.Linear(hidden, target_width))
    model.to(device)

    loss_fn = torch.nn.L1Loss()
    loss_fn.to(device)

    # With SGD, Momentum is required to get a state to shard
    optimizer = optim.OSS(model.parameters(), lr=0.1, momentum=0.99)

    def closure():
        optimizer.zero_grad()
        output = model(inputs)
        loss = loss_fn(output, target)
        loss.backward()
        return loss

    _ = optimizer.step(closure=closure)

    # Update the optimizer state on the reference rank
    optimizer.consolidate_state_dict(recipient_rank=reference_rank)

    # Fetch the state on the reference rank
    # - check that it has the correct size
    # - load it again
    if rank == reference_rank:
        optimizer_state_dict = optimizer.state_dict()
431
        assert len(optimizer_state_dict["state"]) == len(list(model.parameters()))
432
433
434
    else:
        optimizer_state_dict = {}

435
436
    # distribute to the other ranks
    optimizer_state_dict = sync_object_ranks(optimizer_state_dict, reference_rank, device)
437
438

    # Load the optimizer state dict
439
    optimizer.load_state_dict(optimizer_state_dict)
440
    dist.destroy_process_group()
441
442
443
444


def test_collect_shards():
    world_size = 3
445
446
    temp_file_name = tempfile.mkstemp()[1]

447
448
    if torch.cuda.is_available():
        world_size = min(world_size, torch.cuda.device_count())
449
450
451
    reference_rank = 0

    mp.spawn(
452
        run_test_collect_shards, args=(world_size, reference_rank, temp_file_name), nprocs=world_size, join=True,
453
    )
454
455


456
457
458
459
460
461
462
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
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
def run_test_reproducibility(rank, world_size, reference_rank, tempfile_name):
    dist_init(rank, world_size, tempfile_name)
    device = torch.device(rank) if torch.cuda.device_count() > 1 else DEVICE

    # Run a dummy step so that the optimizer state dict exists
    batch, input_width, hidden, target_width = 3, 3, 3, 5
    target = torch.rand((batch, target_width), device=device)
    inputs = torch.rand((batch, input_width), device=device)

    model = torch.nn.Sequential(torch.nn.Linear(input_width, hidden), torch.nn.Linear(hidden, target_width))
    model.to(device)

    loss_fn = torch.nn.L1Loss()
    loss_fn.to(device)

    optimizer = optim.OSS(model.parameters(), optim=torch.optim.RMSprop, lr=0.1)

    def closure():
        optimizer.zero_grad()
        output = model(inputs)
        loss = loss_fn(output, target)
        loss.backward()
        return loss

    _ = optimizer.step(closure=closure)

    # Update the optimizer state on the reference rank
    optimizer.consolidate_state_dict(recipient_rank=reference_rank)

    # Fetch the state on the reference rank, broadcast to the other ones
    if rank == reference_rank:
        optimizer_state_dict = optimizer.state_dict()
    else:
        optimizer_state_dict = {}

    # Run two steps, log the loss
    _ = optimizer.step(closure=closure)
    reference_loss = optimizer.step(closure=closure)

    # Load the optimizer state dict, rewind the state two steps back
    optimizer.load_state_dict(optimizer_state_dict)

    # Run two new steps, log the loss again and check that we get the same
    _ = optimizer.step(closure=closure)
    test_loss = optimizer.step(closure=closure)

    assert torch.allclose(reference_loss, test_loss)

    dist.destroy_process_group()


def test_reproducibility():
    world_size = 2
    temp_file_name = tempfile.mkstemp()[1]

    if torch.cuda.is_available() and torch.cuda.device_count() < world_size:
        # Bail out if not enough devices
        return

    reference_rank = 0

    mp.spawn(
        run_test_collect_shards, args=(world_size, reference_rank, temp_file_name), nprocs=world_size, join=True,
    )


522
def run_test_multiple_groups(rank, world_size, tempfile_name):
523
    # Only work with the even ranks, to check that the global_rank indexing is properly used
524
    dist_init(rank=rank, world_size=world_size, tempfile_name=tempfile_name, backend="gloo")
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
    sub_group_ranks = [0, 2, 4]
    process_group = torch.distributed.new_group(ranks=sub_group_ranks, backend="gloo")

    # Make sure that all the ranks get different training data
    # So that the sync check in between their models is meaningful
    torch.manual_seed(rank)
    np.random.seed(rank)

    # Standard deep learning setup
    device = "cpu"
    epochs, batch, input_width, hidden, target_width = 5, 3, 20, 10, 5
    loss_fn = torch.nn.L1Loss().to(device)

    def check(optimizer):
        # Just run a couple of epochs, check that the model is properly updated
        for _ in range(epochs):
            target = torch.rand((batch, target_width), device=device)
            inputs = torch.rand((batch, input_width), device=device)

            def closure():
                optimizer.zero_grad()
                output = model(inputs)
                loss = loss_fn(output, target)
                loss /= world_size
                loss.backward()
                dist.all_reduce(loss, group=process_group)  # Not strictly needed for the test below

                return loss

            _ = optimizer.step(closure=closure)

            # Check that all the params are the same on all ranks
            for pg in optimizer.param_groups:
                for p in pg["params"]:
                    receptacle = [p.clone() for _ in sub_group_ranks] if rank == 0 else []
                    dist.gather(p, receptacle, dst=0, group=process_group)
                    if rank == 0:
                        for sync_p in receptacle[1:]:
563
564
565
566
567
                            assert torch.all(
                                torch.eq(receptacle[0], sync_p)
                            ), "Models differ in between ranks {} - {}".format(
                                torch.norm(receptacle[0]), torch.norm(sync_p)
                            )
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589

    if rank in sub_group_ranks:
        # Model fitting in the broadcast bucket
        model = torch.nn.Sequential(torch.nn.Linear(input_width, hidden), torch.nn.Linear(hidden, target_width)).to(
            device
        )

        # With SGD, Momentum is required to get a state to shard
        optimizer = optim.OSS(
            model.parameters(), lr=0.1, momentum=0.99, group=process_group, broadcast_buffer_size=2 ** 20
        )
        check(optimizer)

        # Model not-fitting in the broadcast bucket
        model = torch.nn.Sequential(torch.nn.Linear(input_width, hidden), torch.nn.Linear(hidden, target_width)).to(
            device
        )

        # With SGD, Momentum is required to get a state to shard
        optimizer = optim.OSS(model.parameters(), lr=0.1, momentum=0.99, group=process_group, broadcast_buffer_size=0)
        check(optimizer)

590
591
    dist.destroy_process_group(process_group)

592

593
@skip_if_py39_no_cuda
594
595
def test_multiple_groups():
    world_size = 6
596
    temp_file_name = tempfile.mkstemp()[1]
597
598

    mp.spawn(
599
        run_test_multiple_groups, args=(world_size, temp_file_name), nprocs=world_size, join=True,
600
    )
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645


def run_gradient_clipping(rank, world_size, tempfile_name):
    dist_init(rank, world_size, tempfile_name, backend="gloo")
    device = torch.device(rank)
    torch.manual_seed(rank)  # make sure that the different rank get different data

    # Run a dummy step so that the optimizer state dict exists
    batch, input_width, hidden, target_width = 3, 20, 10, 5
    target = torch.rand((batch, target_width), device=device)
    inputs = torch.rand((batch, input_width), device=device)
    NORMS = [1.0, 2.0, 1, 2, inf]
    CLIP_NORM = 0.3

    def check(norm):
        model_oss = torch.nn.Sequential(
            torch.nn.Linear(input_width, hidden),
            torch.nn.Linear(hidden, hidden),
            torch.nn.Linear(hidden, target_width),
        ).to(device)
        model = copy.deepcopy(model_oss)

        # For this test the gradients are (all) reduced in the same way in between the torch reference and fairscale.
        # Normally OSS would use ShardedDDP and only reduce to the proper rank, but this does not change the
        # gradient norm computation from OSS and adds a dependency.
        # to keep the comparison apples-to-apples DDP is used in both cases
        model_oss = DDP(module=model_oss, device_ids=[rank],)
        sharded_optimizer = optim.OSS(model_oss.parameters(), lr=0.1, momentum=0.99)

        model = DDP(model, device_ids=[rank],)

        loss_fn = torch.nn.L1Loss()
        loss_fn.to(device)

        model.zero_grad()
        model_oss.zero_grad()

        outputs = model(inputs)
        outputs_oss = model_oss(inputs)

        loss = loss_fn(outputs, target)
        loss.backward()

        loss_oss = loss_fn(outputs_oss, target)
        loss_oss.backward()
646
        torch.testing.assert_allclose(loss_oss, loss)
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661

        # Check the equivalence with the non-sharded optim
        oss_total_norm = sharded_optimizer.clip_grad_norm(CLIP_NORM, norm_type=norm)
        total_norm = torch.nn.utils.clip_grad_norm_(model.parameters(), CLIP_NORM, norm_type=norm)
        assert torch.allclose(oss_total_norm, total_norm), "torch and fairscale should return the same grad norm"

        # Check that the params have indeed been clipped
        for params in sharded_optimizer.per_device_params.values():
            for param in filter(lambda x: x.grad is not None, params[rank]):
                assert torch.norm(param.grad, p=norm) < CLIP_NORM, f"param grad norm above clip : {param.grad}"

    for norm in NORMS:
        print(f"Checking norm {norm}")
        check(norm)

Benjamin Lefaudeux's avatar
Benjamin Lefaudeux committed
662
663
664
        # Check twice, catch an hypothetic iterator dumb mistake
        check(norm)

665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
    dist.destroy_process_group()


@skip_if_no_cuda
def test_gradient_clipping():
    world_size = 3
    temp_file_name = tempfile.mkstemp()[1]

    if torch.cuda.is_available():
        world_size = min(world_size, torch.cuda.device_count())
    reference_rank = 0

    mp.spawn(
        run_gradient_clipping, args=(world_size, temp_file_name), nprocs=world_size, join=True,
    )
680
681
682
683


def run_state_dict_distributed(rank, world_size, tempfile_name):
    dist_init(rank, world_size, tempfile_name, backend="gloo")
684

685
686
687
    device = torch.device(rank)
    torch.manual_seed(rank)  # make sure that the different rank get different data

688
689
    # Setup two problems in parallel, we'll make sure that the second track (with save/load) follows the first one(untouched)
    # We split the model in two to test the multiple param groups support
690
691
692
693
    batch, input_width, hidden, target_width = 3, 20, 10, 5
    target = torch.rand((batch, target_width), device=device)
    inputs = torch.rand((batch, input_width), device=device)

694
    model_oss1 = torch.nn.Sequential(torch.nn.Linear(input_width, hidden), torch.nn.Linear(hidden, hidden)).to(device)
695
696
    head_oss1 = torch.nn.Linear(hidden, target_width).to(device)

697
    model_oss2 = copy.deepcopy(model_oss1)
698
    head_oss2 = copy.deepcopy(head_oss1)
699
700
701
702
703
704
705

    # For this test the gradients are (all) reduced in the same way in between the torch reference and fairscale.
    # Normally OSS would use ShardedDDP and only reduce to the proper rank, but this does not change the
    # gradient norm computation from OSS and adds a dependency.
    # to keep the comparison apples-to-apples DDP is used in both cases
    model_oss1 = DDP(module=model_oss1, device_ids=[rank],)
    sharded_optimizer1 = optim.OSS(model_oss1.parameters(), lr=0.1, momentum=0.99)
706
707
    sharded_optimizer1.add_param_group({"params": head_oss1.parameters()})

708
709
    model_oss2 = DDP(module=model_oss2, device_ids=[rank],)
    sharded_optimizer2 = optim.OSS(model_oss2.parameters(), lr=0.1, momentum=0.99)
710
    sharded_optimizer2.add_param_group({"params": head_oss2.parameters()})
711

712
    loss_fn = torch.nn.L1Loss().to(device)
713

714
    def run_grad_step(model, head, optimizer):
715
        model.zero_grad()
716
        outputs = head(model(inputs))
717

718
719
720
721
    # pull the current state, broadcast it to all ranks
    sharded_optimizer2.consolidate_state_dict(recipient_rank=RECIPIENT_RANK)  # all ranks
    state_dict2 = sharded_optimizer2.state_dict() if rank == RECIPIENT_RANK else {}
    state_dict2 = sync_object_ranks(state_dict2, RECIPIENT_RANK, device)
722

723
724
    # re-create a new optimizer from scratch with absurd values, load the previous state
    sharded_optimizer2 = optim.OSS(model_oss2.parameters(), lr=1e6, momentum=0.0001)
725
    sharded_optimizer2.add_param_group({"params": head_oss2.parameters()})
726
    sharded_optimizer2.load_state_dict(state_dict2)
727
728
729
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (before any steps)"
    )
730
731

    # now take a step and check that parameters are equal
732
733
    run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
    run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
734
735
736
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (after stepping)"
    )
737

738
739
740
741
    # save the state dict for one model only, then distribute to the other ranks
    sharded_optimizer2.consolidate_state_dict(recipient_rank=RECIPIENT_RANK)  # all ranks
    state_dict2 = sharded_optimizer2.state_dict() if rank == RECIPIENT_RANK else {}
    state_dict2 = sync_object_ranks(state_dict2, RECIPIENT_RANK, device)
742
743
744
745
746
747
748
749

    # Check that the pulled state and the .param_groups attribute are in sync
    for replica in range(len(state_dict2["param_groups"])):
        for k in state_dict2["param_groups"][replica].keys():
            if k != "params":
                assert state_dict2["param_groups"][replica][k] == sharded_optimizer2.param_groups[0][k]

    # take a step
750
751
    run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
    run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
752
753
754
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (after consolidating)"
    )
755

756
757
758
759
    # save again for one rank, then distribute to the others
    sharded_optimizer2.consolidate_state_dict(recipient_rank=RECIPIENT_RANK)  # all ranks
    state_dict2 = sharded_optimizer2.state_dict() if rank == RECIPIENT_RANK else {}
    state_dict2 = sync_object_ranks(state_dict2, RECIPIENT_RANK, device)
760
761
762

    # reload the state_dict
    sharded_optimizer2 = optim.OSS(model_oss2.parameters(), lr=0.1, momentum=0.99)
763
    sharded_optimizer2.add_param_group({"params": head_oss2.parameters()})
764
765
766
    sharded_optimizer2.load_state_dict(state_dict2)

    # take a step
767
768
    run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
    run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
769
770
771
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (after reloading)"
    )
772
773
774
775
776
777

    dist.destroy_process_group()


@skip_if_no_cuda
def test_state_dict_distributed():
778
    world_size = 2
779
780
781
    temp_file_name = tempfile.mkstemp()[1]

    if torch.cuda.is_available():
782
        world_size = max(world_size, torch.cuda.device_count())
783
784
785
786

    mp.spawn(
        run_state_dict_distributed, args=(world_size, temp_file_name), nprocs=world_size, join=True,
    )
787
788
789
790
791
792
793
794
795
796


def run_ddp_parity(rank, world_size, backend, temp_file_name):
    url = "file://" + temp_file_name
    dist.init_process_group(init_method=url, backend=backend, rank=rank, world_size=world_size)

    device = torch.device("cuda")
    torch.cuda.set_device(rank)
    torch.manual_seed(rank)
    np.random.seed(rank)
797
798
799
800
    hidden = 5
    in_channels = 3
    out_channels = 3
    batch = 64
801

802
    def check_optimizer_equivalence(optimizer: Type[torch.optim.Optimizer], change_train_graph: bool = False):
803
        # Any model works. Add one different buffer per rank
804
805
806
        trunk = torch.nn.Sequential(
            torch.nn.Linear(in_channels, hidden), torch.nn.Linear(hidden, hidden), torch.nn.Linear(hidden, hidden)
        )
807
808
809
810
811
812
        trunk.register_buffer("test_buffer", torch.ones((1)) * rank)
        trunk.to(device)

        head = torch.nn.Linear(hidden, out_channels).to(device)

        # Define a model to be trained by OSS
813
        oss_module = torch.nn.Sequential(trunk, head)
814
815
816
817
818
        oss_trainable_params = [
            {"params": trunk.parameters(), "lr": 1e-5},
            {"params": head.parameters(), "lr": 1e-4},
        ]

819
820
        optimizer_settings: Dict[Any, Any] = {}
        if isinstance(optimizer, torch.optim.SGD):
821
822
823
824
825
826
827
828
829
830
            optimizer_settings["momentum"] = 0.9

        sharded_optimizer = optim.OSS(
            params=oss_trainable_params,
            optim=optimizer,
            group=None,
            broadcast_buffer_size=2 ** 10,
            **optimizer_settings,
        )

831
        oss_ddp_model = DDP(module=oss_module, device_ids=[rank], broadcast_buffers=True, find_unused_parameters=True)
832

833
834
835
836
        # Define a model to be trained by normal pytorch + DDP
        ddp_trunk = copy.deepcopy(trunk)
        ddp_head = copy.deepcopy(head)
        ddp_module = torch.nn.Sequential(ddp_trunk, ddp_head)
837

838
839
840
841
842
        ddp_trainable_params = [
            {"params": ddp_trunk.parameters(), "lr": 1e-5},
            {"params": ddp_head.parameters(), "lr": 1e-4},
        ]
        ddp_optimizer = optimizer(ddp_trainable_params, **optimizer_settings)  # type: ignore
843
        ddp_model = DDP(module=ddp_module, device_ids=[rank], broadcast_buffers=True, find_unused_parameters=True)
844

845
846
        def check_step():
            input_tensor = torch.rand((batch, in_channels)).to(device)
847
848
849
850
851
852
853
854
855

            def closure_ddp(input_tensor=input_tensor):
                ddp_optimizer.zero_grad()
                ddp_loss = ddp_model(input_tensor).abs().sum()
                ddp_loss.backward()
                return ddp_loss

            def closure_sharded(input_tensor=input_tensor):
                sharded_optimizer.zero_grad()
856
                sharded_loss = oss_ddp_model(input_tensor).abs().sum()
857
858
859
860
861
862
863
                sharded_loss.backward()
                return sharded_loss

            loss_ddp = cast(torch.Tensor, ddp_optimizer.step(closure=closure_ddp))
            loss_sharded_optim = cast(torch.Tensor, sharded_optimizer.step(closure=closure_sharded))

            assert torch.allclose(
864
865
                loss_ddp, loss_sharded_optim, rtol=1e-3
            ), f"Losses differ in between Pytorch optim and OSS\n {loss_ddp.item()} - {loss_sharded_optim.item()} - world size {world_size}"
866

867
868
            check_same_model_params(oss_ddp_model, ddp_model)

869
        # The model should be synchronized in between the ranks at construction time, check that
870
        check_same_model_params(oss_ddp_model, ddp_model)
871
872
873
874

        # The models should stay the same in between ddp and sharded optimizer
        for i in range(5):
            check_step()
875
876
877
878
879
880
881

            # Check that altering the trainable parameters does not cause DDP and OSS to diverge
            if change_train_graph:
                # Flip the first parameter from trainable to non-trainable and vice-versa
                next(ddp_module.parameters()).requires_grad = not next(ddp_module.parameters()).requires_grad
                next(oss_module.parameters()).requires_grad = not next(oss_module.parameters()).requires_grad
                # sharded_optimizer.refresh_trainable()
882

883
884
885
886
887
888
889
890
        # Check that the checkpoints are compatible
        # - get states
        ddp_state_dict = ddp_optimizer.state_dict()
        sharded_optimizer.consolidate_state_dict(recipient_rank=RECIPIENT_RANK)
        sharded_optim_state_dict = sharded_optimizer.state_dict() if rank == RECIPIENT_RANK else {}
        sharded_optim_state_dict = sync_object_ranks(sharded_optim_state_dict, RECIPIENT_RANK, device)

        # - cross load the states
891
892
        # run one step and check that the models are still the same
        ddp_state_dict_ref = copy.deepcopy(ddp_state_dict)  # OSS will remove some states
893
894
        ddp_optimizer.load_state_dict(sharded_optim_state_dict)  # mixup on purpose !
        sharded_optimizer.load_state_dict(ddp_state_dict)
895
        check_step()
896

897
898
899
900
        #  - self load, rewind, check no problem
        # run one step and check that the models are still the same
        ddp_optimizer.load_state_dict(ddp_state_dict_ref)
        sharded_optimizer.load_state_dict(sharded_optim_state_dict)
901
902
        check_step()

903
904
905
    for opt in [torch.optim.Adam, torch.optim.SGD]:
        check_optimizer_equivalence(opt, change_train_graph=False)
        check_optimizer_equivalence(opt, change_train_graph=True)
906
907
908
909
910
911
912
913
914
915
916

    dist.destroy_process_group()


@skip_if_no_cuda
@skip_if_single_gpu
def test_ddp_parity():
    temp_file_name = tempfile.mkstemp()[1]
    world_size = torch.cuda.device_count()
    backend = dist.Backend.NCCL
    mp.spawn(run_ddp_parity, args=(world_size, backend, temp_file_name), nprocs=world_size, join=True)