test_oss.py 34.6 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
26
from fairscale.utils.testing import (
    check_same_model_params,
27
    check_same_models_across_ranks,
28
29
30
31
32
    skip_if_no_cuda,
    skip_if_py39_no_cuda,
    skip_if_single_gpu,
    torch_version,
)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
33

34
35
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")
36
RECIPIENT_RANK = 1
37

38
39
40
41
42
43
44
45
46
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
47

48
49
50
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
51
52


53
54
55
56
57
58
59
60
61
62
63
64
65
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


66
67
68
69
class TestSingleRank(unittest.TestCase):
    """
    All the following tests do not check for inter-process communication
    """
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
70

71
72
    def setUp(self):
        dist_init(0, 1, tempfile.mkstemp()[1])
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
73

74
75
    def tearDown(self):
        torch.distributed.destroy_process_group()
76

77
78
79
    def test_create(self):
        params = [torch.rand(1)]
        o = optim.OSS(params, lr=0.01)
80

81
82
83
    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)
84
        x.backward()
85
86
87
        o.step()
        assert x == torch.tensor([0.9], device=DEVICE)
        assert o.optim.state[x]["momentum_buffer"] == torch.tensor([1.0], device=DEVICE)
88
        o.zero_grad()
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
        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()
118
        o.step()
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
        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)

157
158
159
160
161
162
163
164
165
166
167
168
    @skip_if_no_cuda
    def test_device_change(self):
        x = torch.nn.Linear(1, 1).to("cpu")
        o = optim.OSS(x.parameters(), torch.optim.SGD, lr=0.1)

        # Move the model to device after OSS was constructed
        x.to(DEVICE)
        x(torch.zeros((1), device=DEVICE)).backward()

        # Check that OSS detects that the device changed
        o.step()

169
170
171
172
173
174
175
176
177
178
179
180
181
    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)
182

183
184
185
186
    def test_step_without_closure(self):
        class SGDWithoutClosure(torch.optim.SGD):
            def step(self):
                return super().step()
187

188
189
190
191
192
193
194
195
196
        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)
197
198
        with pytest.raises(RuntimeError):
            _ = o.state_dict()
199
200


201
202
def run_test_add_param_group(rank, world_size, tempfile_name):
    dist_init(rank, world_size, tempfile_name)
203
204
205
206

    # Test with all parameters trainable to begin with
    def all_trainable():
        params = []
207
208
209
        sizes = [9, 7, 5, 3]
        sizes_world = sizes * world_size
        for size in sizes_world[:-1]:
210
211
212
213
214
215
216
217
218
219
220
221
            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
222
223
224

        # 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)
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
        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
247

248
249
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
250
251

def test_add_param_group():
252
    world_size = 4
253
    if torch.cuda.is_available() and torch.cuda.device_count() < world_size:
254
255
        world_size = min(world_size, torch.cuda.device_count())

256
    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
257
258


259
260
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
261
262
263
264
265
266
267
268
269
270
271
    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

272
273
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
274
275
276

def test_zero_grad():
    world_size = 2
277
278
279
    if torch.cuda.is_available() and torch.cuda.device_count() < world_size:
        world_size = min(world_size, torch.cuda.device_count())

280
281
    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
282
283


284
285
def run_test_empty_shard(rank, world_size, tempfile_name, backend):
    dist_init(rank, world_size, tempfile_name, backend=backend)
286
    m = torch.nn.Linear(1, 1)
287
288
289
290
291
292
293
294
295
296
    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()
297
298
299
300

    dist.destroy_process_group()


301
302
@pytest.mark.parametrize("backend", ["gloo", "nccl"])
def test_empty_shard(backend):
303
    world_size = 4
304
305
306
307
308
    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)
309
310


311
312
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
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
    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)

328
329
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
330

Benjamin Lefaudeux's avatar
Benjamin Lefaudeux committed
331
@skip_if_single_gpu
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
332
def test_step():
Benjamin Lefaudeux's avatar
Benjamin Lefaudeux committed
333
    world_size = 2
334
335
336
    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
337
338


339
340
def run_test_step_with_closure(rank, world_size, tempfile_name, optimizer=None):
    dist_init(rank, world_size, tempfile_name)
341

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
342
343
344
345
346
347
348
349
350
351
352
353
    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)
354

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

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
    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)

376
377
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
378
379
380

@skip_if_no_cuda
def test_step_with_closure():
381
    world_size = min(2, torch.cuda.device_count())
382
383
384
    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
385
386


387
388
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
389
    params = []
390
391
392
393
    sizes = [9, 7, 5, 3]
    sizes_world = sizes * world_size

    for size in sizes_world:
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
394
        params.append(torch.rand(size, 1))
395
396
397
398
399

    # 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
400
    o = optim.OSS(params, lr=0.1)
401
    assert sum([x.numel() for x in o.optim.param_groups[0]["params"]]) == sum(sizes)
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
402

403
404
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
405
406

def test_sharding():
407
408
409
    world_size = 4
    if torch.cuda.is_available():
        world_size = min(world_size, torch.cuda.device_count())
410

411
    _, temp_file_name = tempfile.mkstemp()
412
    mp.spawn(run_test_sharding, args=(world_size, temp_file_name), nprocs=world_size, join=True)
413
414


415
416
def run_test_collect_shards(rank, world_size, reference_rank, tempfile_name):
    dist_init(rank, world_size, tempfile_name)
417
    device = torch.device(rank) if torch.cuda.device_count() > 1 else DEVICE
418
    torch.cuda.set_device(rank)
419
420

    # Run a dummy step so that the optimizer state dict exists
421
    batch, input_width, hidden, target_width = 3, 3, 3, 5
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
    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()
451
        assert len(optimizer_state_dict["state"]) == len(list(model.parameters()))
452
453
454
    else:
        optimizer_state_dict = {}

455
456
    # distribute to the other ranks
    optimizer_state_dict = sync_object_ranks(optimizer_state_dict, reference_rank, device)
457
458

    # Load the optimizer state dict
459
    optimizer.load_state_dict(optimizer_state_dict)
460
461
462
463
464

    # Check that the states are not None, but {}
    for state in optimizer.state.values():
        for _, _ in state.items():
            pass
465
466
467
468
469
470
471
472

    # Test the state dict materialization on all ranks
    _ = optimizer.step(closure=closure)
    optimizer_state_dict = optimizer.state_dict(all_ranks=True)  # one per rank
    optimizer.load_state_dict(optimizer_state_dict)
    _ = optimizer.step(closure=closure)
    check_same_models_across_ranks(model, dist.group.WORLD, params_should_be_equal=True, check_broadcast_buffers=False)

473
    dist.destroy_process_group()
474
475


476
@skip_if_single_gpu
477
def test_collect_shards():
478
    world_size = 2
479
    temp_file_name = tempfile.mkstemp()[1]
480
481
482
    reference_rank = 0

    mp.spawn(
483
        run_test_collect_shards, args=(world_size, reference_rank, temp_file_name), nprocs=world_size, join=True,
484
    )
485
486


487
def run_test_reproducibility(rank, world_size, tempfile_name):
488
489
    dist_init(rank, world_size, tempfile_name)
    device = torch.device(rank) if torch.cuda.device_count() > 1 else DEVICE
490
    torch.cuda.set_device(rank)
491
492
493
494
495
496
497
498

    # 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)
499
    model = DDP(model, device_ids=[device])
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514

    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)

515
516
517
    # Get a snapshot of the state at this point
    optimizer_state_dict = copy.deepcopy(optimizer.state_dict(all_ranks=True))
    model_state_dict = copy.deepcopy(model.state_dict())
518
519
520
521
522
523
524

    # 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)
525
    model.load_state_dict(model_state_dict)
526
527
528
529
530

    # 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)

531
    assert torch.allclose(reference_loss, test_loss), f"{reference_loss} vs {test_loss}. Reproducibility is broken"
532
533
534
535

    dist.destroy_process_group()


536
@skip_if_single_gpu
537
538
539
540
541
def test_reproducibility():
    world_size = 2
    temp_file_name = tempfile.mkstemp()[1]

    mp.spawn(
542
        run_test_reproducibility, args=(world_size, temp_file_name), nprocs=world_size, join=True,
543
544
545
    )


546
def run_test_multiple_groups(rank, world_size, tempfile_name):
547
    # Only work with the even ranks, to check that the global_rank indexing is properly used
548
    dist_init(rank=rank, world_size=world_size, tempfile_name=tempfile_name, backend="gloo")
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
    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:]:
587
588
589
590
591
                            assert torch.all(
                                torch.eq(receptacle[0], sync_p)
                            ), "Models differ in between ranks {} - {}".format(
                                torch.norm(receptacle[0]), torch.norm(sync_p)
                            )
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613

    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)

614
615
    dist.destroy_process_group(process_group)

616

617
@skip_if_py39_no_cuda
618
619
def test_multiple_groups():
    world_size = 6
620
    temp_file_name = tempfile.mkstemp()[1]
621
622

    mp.spawn(
623
        run_test_multiple_groups, args=(world_size, temp_file_name), nprocs=world_size, join=True,
624
    )
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669


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()
670
        torch.testing.assert_allclose(loss_oss, loss)
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685

        # 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
686
687
688
        # Check twice, catch an hypothetic iterator dumb mistake
        check(norm)

689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
    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,
    )
704
705
706
707


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

709
710
711
    device = torch.device(rank)
    torch.manual_seed(rank)  # make sure that the different rank get different data

712
713
    # 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
714
715
716
717
    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)

718
    model_oss1 = torch.nn.Sequential(torch.nn.Linear(input_width, hidden), torch.nn.Linear(hidden, hidden)).to(device)
719
720
    head_oss1 = torch.nn.Linear(hidden, target_width).to(device)

721
    model_oss2 = copy.deepcopy(model_oss1)
722
    head_oss2 = copy.deepcopy(head_oss1)
723
724
725
726
727
728
729

    # 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)
730
731
    sharded_optimizer1.add_param_group({"params": head_oss1.parameters()})

732
733
    model_oss2 = DDP(module=model_oss2, device_ids=[rank],)
    sharded_optimizer2 = optim.OSS(model_oss2.parameters(), lr=0.1, momentum=0.99)
734
    sharded_optimizer2.add_param_group({"params": head_oss2.parameters()})
735

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

738
    def run_grad_step(model, head, optimizer):
739
        model.zero_grad()
740
        outputs = head(model(inputs))
741

742
743
744
745
    # 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)
746

747
748
    # 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)
749
    sharded_optimizer2.add_param_group({"params": head_oss2.parameters()})
750
    sharded_optimizer2.load_state_dict(state_dict2)
751
752
753
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (before any steps)"
    )
754
755

    # now take a step and check that parameters are equal
756
757
    run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
    run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
758
759
760
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (after stepping)"
    )
761

762
763
764
765
    # 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)
766
767
768
769
770
771
772
773

    # 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
774
775
    run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
    run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
776
777
778
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (after consolidating)"
    )
779

780
781
782
783
    # 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)
784
785
786

    # reload the state_dict
    sharded_optimizer2 = optim.OSS(model_oss2.parameters(), lr=0.1, momentum=0.99)
787
    sharded_optimizer2.add_param_group({"params": head_oss2.parameters()})
788
789
790
    sharded_optimizer2.load_state_dict(state_dict2)

    # take a step
791
792
    run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
    run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
793
794
795
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (after reloading)"
    )
796
797
798
799
800
801

    dist.destroy_process_group()


@skip_if_no_cuda
def test_state_dict_distributed():
802
    world_size = 2
803
804
805
    temp_file_name = tempfile.mkstemp()[1]

    if torch.cuda.is_available():
806
        world_size = max(world_size, torch.cuda.device_count())
807
808
809
810

    mp.spawn(
        run_state_dict_distributed, args=(world_size, temp_file_name), nprocs=world_size, join=True,
    )
811
812


813
def run_ddp_parity(rank, world_size, backend, temp_file_name, change_train_graph):
814
815
816
817
818
819
820
    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)
821
822
823
824
    hidden = 5
    in_channels = 3
    out_channels = 3
    batch = 64
825

826
    def check_optimizer_equivalence(optimizer: Type[torch.optim.Optimizer], change_train_graph: bool = False):
827
        # Any model works. Add one different buffer per rank
828
829
830
        trunk = torch.nn.Sequential(
            torch.nn.Linear(in_channels, hidden), torch.nn.Linear(hidden, hidden), torch.nn.Linear(hidden, hidden)
        )
831
832
833
834
835
836
        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
837
        oss_module = torch.nn.Sequential(trunk, head)
838
839

        # Make sure that the param groups are interleaved, to catch an ordering bug in the state dict
840
        oss_trainable_params = [
841
842
            {"params": list(trunk.parameters())[:-1] + list(head.parameters()), "lr": 1e-5},
            {"params": list(trunk.parameters())[-1], "lr": 1e-4},
843
844
        ]

845
846
        optimizer_settings: Dict[Any, Any] = {}
        if isinstance(optimizer, torch.optim.SGD):
847
848
849
850
851
852
853
854
855
856
            optimizer_settings["momentum"] = 0.9

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

857
        oss_ddp_model = DDP(module=oss_module, device_ids=[rank], broadcast_buffers=True, find_unused_parameters=True)
858

859
860
861
862
        # 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)
863

864
        ddp_trainable_params = [
865
866
            {"params": list(ddp_trunk.parameters())[:-1] + list(ddp_head.parameters()), "lr": 1e-5},
            {"params": list(ddp_trunk.parameters())[-1], "lr": 1e-4},
867
868
        ]
        ddp_optimizer = optimizer(ddp_trainable_params, **optimizer_settings)  # type: ignore
869
        ddp_model = DDP(module=ddp_module, device_ids=[rank], broadcast_buffers=True, find_unused_parameters=True)
870

871
872
        def check_step():
            input_tensor = torch.rand((batch, in_channels)).to(device)
873
874
875
876
877
878
879
880
881

            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()
882
                sharded_loss = oss_ddp_model(input_tensor).abs().sum()
883
884
885
886
887
888
889
                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(
890
891
                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}"
892

893
894
            check_same_model_params(oss_ddp_model, ddp_model)

895
        # The model should be synchronized in between the ranks at construction time, check that
896
        check_same_model_params(oss_ddp_model, ddp_model)
897
898
899
900

        # The models should stay the same in between ddp and sharded optimizer
        for i in range(5):
            check_step()
901
902
903
904
905
906
907

            # 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()
908

909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
        # Check that the checkpoints are compatible (post pytorch 1.5)
        if torch_version()[1] > 5:
            # - 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
            # 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
            ddp_optimizer.load_state_dict(sharded_optim_state_dict)  # mixup on purpose !
            sharded_optimizer.load_state_dict(ddp_state_dict)
            check_step()

            #  - 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)
            check_step()
929

930
    for opt in [torch.optim.Adam, torch.optim.SGD]:
931
        check_optimizer_equivalence(opt, change_train_graph=change_train_graph)
932
933
934
935
936
937

    dist.destroy_process_group()


@skip_if_no_cuda
@skip_if_single_gpu
938
939
940
@pytest.mark.parametrize("change_train_graph", [True, False])
@pytest.mark.parametrize("backend", [dist.Backend.NCCL, dist.Backend.GLOO])
def test_ddp_parity(change_train_graph: bool, backend: dist.Backend):
941
942
    temp_file_name = tempfile.mkstemp()[1]
    world_size = torch.cuda.device_count()
943
944
945
    mp.spawn(
        run_ddp_parity, args=(world_size, backend, temp_file_name, change_train_graph), nprocs=world_size, join=True
    )