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

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

39
40
41
42
43
try:
    from torch.distributed import broadcast_object_list  # noqa

    _torch_broadcast_object = True
except ImportError:
44
    from fairscale.utils.params import broadcast_object  # noqa
45
46
47

    _torch_broadcast_object = False

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
48

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


54
55
56
57
58
59
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:
60
        package_sync = utils.params.broadcast_object(
61
62
63
64
65
66
            something_to_sync, src_rank=reference_rank, group=dist.group.WORLD, dist_device=device
        )

    return package_sync


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

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

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

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

82
83
84
    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)
85
        x.backward()
86
87
88
        o.step()
        assert x == torch.tensor([0.9], device=DEVICE)
        assert o.optim.state[x]["momentum_buffer"] == torch.tensor([1.0], device=DEVICE)
89
        o.zero_grad()
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
118
        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()
119
        o.step()
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
        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)

158
159
160
161
162
163
164
165
166
167
168
169
    @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()

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

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

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


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

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

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

249
250
    dist.destroy_process_group()

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

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

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


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

273
274
    dist.destroy_process_group()

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

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

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


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

    dist.destroy_process_group()


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


312
313
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
314
315
316
317
318
319
320
321
322
323
324
325
    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()
326
327
    assert m.weight == torch.tensor([[0.75]], device=rank), f"{rank}: {m.weight.item()}, 0.75 expected"
    assert m.bias == torch.tensor([1.85], device=rank), f"{rank}: {m.bias.item()}, 1.85 expected"
Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
328

329
330
    dist.destroy_process_group()

Mandeep Singh Baines's avatar
Mandeep Singh Baines committed
331

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


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

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

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

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

377
378
    dist.destroy_process_group()

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

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


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

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

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

404
405
    dist.destroy_process_group()

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

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

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


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

    # Run a dummy step so that the optimizer state dict exists
422
    batch, input_width, hidden, target_width = 3, 3, 3, 5
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
451
    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()
452
        assert len(optimizer_state_dict["state"]) == len(list(model.parameters()))
453
454
455
    else:
        optimizer_state_dict = {}

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

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

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

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

474
475
476
477
478
    # Check that if the model is moved to cpu, the optimizer consolidation still works
    model.cpu()
    optimizer = optim.OSS(model.parameters(), lr=0.1, momentum=0.99)
    optimizer.consolidate_state_dict(recipient_rank=reference_rank)

479
    dist.destroy_process_group()
480
481


482
@skip_if_single_gpu
483
def test_collect_shards():
484
    world_size = 2
485
    temp_file_name = tempfile.mkstemp()[1]
486
487
488
    reference_rank = 0

    mp.spawn(
489
        run_test_collect_shards, args=(world_size, reference_rank, temp_file_name), nprocs=world_size, join=True,
490
    )
491
492


493
def run_test_reproducibility(rank, world_size, tempfile_name, broadcast_fp16):
494
495
    dist_init(rank, world_size, tempfile_name)
    device = torch.device(rank) if torch.cuda.device_count() > 1 else DEVICE
496
    torch.cuda.set_device(rank)
497
498
499
500
501
502
503
504

    # 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)
505
    model = DDP(model, device_ids=[device])
506
507
508
509

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

510
    optimizer = optim.OSS(model.parameters(), optim=torch.optim.RMSprop, lr=0.1, broadcast_fp16=broadcast_fp16)
511
512
513
514
515
516
517
518
519
520

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

    _ = optimizer.step(closure=closure)

521
522
523
    # 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())
524
525
526
527
528
529
530

    # 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)
531
    model.load_state_dict(model_state_dict)
532
533
534
535
536

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

537
    assert torch.allclose(reference_loss, test_loss), f"{reference_loss} vs {test_loss}. Reproducibility is broken"
538
539
540
541

    dist.destroy_process_group()


542
@skip_if_single_gpu
543
544
@pytest.mark.parametrize("broadcast_fp16", [False, True])
def test_reproducibility(broadcast_fp16: bool):
545
546
547
548
    world_size = 2
    temp_file_name = tempfile.mkstemp()[1]

    mp.spawn(
549
        run_test_reproducibility, args=(world_size, temp_file_name, broadcast_fp16), nprocs=world_size, join=True,
550
551
552
    )


553
def run_test_multiple_groups(rank, world_size, tempfile_name):
554
    # Only work with the even ranks, to check that the global_rank indexing is properly used
555
    dist_init(rank=rank, world_size=world_size, tempfile_name=tempfile_name, backend="gloo")
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
587
588
589
590
591
592
593
    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:]:
594
595
596
597
598
                            assert torch.all(
                                torch.eq(receptacle[0], sync_p)
                            ), "Models differ in between ranks {} - {}".format(
                                torch.norm(receptacle[0]), torch.norm(sync_p)
                            )
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620

    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)

621
622
    dist.destroy_process_group(process_group)

623

624
@skip_if_py39_no_cuda
625
626
def test_multiple_groups():
    world_size = 6
627
    temp_file_name = tempfile.mkstemp()[1]
628
629

    mp.spawn(
630
        run_test_multiple_groups, args=(world_size, temp_file_name), nprocs=world_size, join=True,
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
670
671
672
673
674
675
676


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()
677
        torch.testing.assert_allclose(loss_oss, loss)
678
679
680
681
682
683
684

        # 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
685
        for params in sharded_optimizer._per_device_params.values():
686
687
688
689
690
691
692
            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
693
694
695
        # Check twice, catch an hypothetic iterator dumb mistake
        check(norm)

696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
    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,
    )
711
712
713
714


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

716
717
718
    device = torch.device(rank)
    torch.manual_seed(rank)  # make sure that the different rank get different data

719
720
    # 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
721
722
723
724
    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)

725
    model_oss1 = torch.nn.Sequential(torch.nn.Linear(input_width, hidden), torch.nn.Linear(hidden, hidden)).to(device)
726
727
    head_oss1 = torch.nn.Linear(hidden, target_width).to(device)

728
    model_oss2 = copy.deepcopy(model_oss1)
729
    head_oss2 = copy.deepcopy(head_oss1)
730
731
732
733
734
735
736

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

739
740
    model_oss2 = DDP(module=model_oss2, device_ids=[rank],)
    sharded_optimizer2 = optim.OSS(model_oss2.parameters(), lr=0.1, momentum=0.99)
741
    sharded_optimizer2.add_param_group({"params": head_oss2.parameters()})
742

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

745
    def run_grad_step(model, head, optimizer):
746
        model.zero_grad()
747
        outputs = head(model(inputs))
748

749
750
751
752
    # 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)
753

754
755
    # 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)
756
    sharded_optimizer2.add_param_group({"params": head_oss2.parameters()})
757
    sharded_optimizer2.load_state_dict(state_dict2)
758
759
760
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (before any steps)"
    )
761
762

    # now take a step and check that parameters are equal
763
764
    run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
    run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
765
766
767
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (after stepping)"
    )
768

769
770
771
772
    # 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)
773
774
775
776
777
778
779
780

    # 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
781
782
    run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
    run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
783
784
785
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (after consolidating)"
    )
786

787
788
789
790
    # 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)
791
792
793

    # reload the state_dict
    sharded_optimizer2 = optim.OSS(model_oss2.parameters(), lr=0.1, momentum=0.99)
794
    sharded_optimizer2.add_param_group({"params": head_oss2.parameters()})
795
796
797
    sharded_optimizer2.load_state_dict(state_dict2)

    # take a step
798
799
    run_grad_step(model_oss1, head_oss1, sharded_optimizer1)
    run_grad_step(model_oss2, head_oss2, sharded_optimizer2)
800
801
802
    check_same_model_params(
        model_oss1, model_oss2, "parameters of the two identical models have diverged (after reloading)"
    )
803
804
805
806
807
808

    dist.destroy_process_group()


@skip_if_no_cuda
def test_state_dict_distributed():
809
    world_size = 2
810
811
812
    temp_file_name = tempfile.mkstemp()[1]

    if torch.cuda.is_available():
813
        world_size = max(world_size, torch.cuda.device_count())
814
815
816
817

    mp.spawn(
        run_state_dict_distributed, args=(world_size, temp_file_name), nprocs=world_size, join=True,
    )
818
819


820
def run_ddp_parity(rank, world_size, backend, temp_file_name, change_train_graph, broadcast_fp16):
821
822
823
824
825
826
827
    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)
828
829
830
831
    hidden = 5
    in_channels = 3
    out_channels = 3
    batch = 64
832

833
    def check_optimizer_equivalence(optimizer: Type[torch.optim.Optimizer], change_train_graph: bool = False):
834
        # Any model works. Add one different buffer per rank
835
836
837
        trunk = torch.nn.Sequential(
            torch.nn.Linear(in_channels, hidden), torch.nn.Linear(hidden, hidden), torch.nn.Linear(hidden, hidden)
        )
838
839
840
841
842
843
        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
844
        oss_module = torch.nn.Sequential(trunk, head)
845
846

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

852
853
        optimizer_settings: Dict[Any, Any] = {}
        if isinstance(optimizer, torch.optim.SGD):
854
855
856
857
858
859
860
861
862
863
            optimizer_settings["momentum"] = 0.9

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

864
        oss_ddp_model = DDP(module=oss_module, device_ids=[rank], broadcast_buffers=True, find_unused_parameters=True)
865

866
867
868
869
        # 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)
870

871
        ddp_trainable_params = [
872
873
            {"params": list(ddp_trunk.parameters())[:-1] + list(ddp_head.parameters()), "lr": 1e-5},
            {"params": list(ddp_trunk.parameters())[-1], "lr": 1e-4},
874
875
        ]
        ddp_optimizer = optimizer(ddp_trainable_params, **optimizer_settings)  # type: ignore
876
        ddp_model = DDP(module=ddp_module, device_ids=[rank], broadcast_buffers=True, find_unused_parameters=True)
877

878
879
        def check_step():
            input_tensor = torch.rand((batch, in_channels)).to(device)
880
881
882
883
884
885
886
887
888

            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()
889
                sharded_loss = oss_ddp_model(input_tensor).abs().sum()
890
891
892
893
894
895
896
                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(
897
898
                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}"
899

900
901
            check_same_model_params(oss_ddp_model, ddp_model)

902
        # The model should be synchronized in between the ranks at construction time, check that
903
        check_same_model_params(oss_ddp_model, ddp_model)
904
905
906
907

        # The models should stay the same in between ddp and sharded optimizer
        for i in range(5):
            check_step()
908
909
910
911
912
913
914

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

916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
        # 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()
936

937
    for opt in [torch.optim.Adam, torch.optim.SGD]:
938
        check_optimizer_equivalence(opt, change_train_graph=change_train_graph)
939
940
941
942
943
944

    dist.destroy_process_group()


@skip_if_no_cuda
@skip_if_single_gpu
945
946
@pytest.mark.parametrize("change_train_graph", [True, False])
@pytest.mark.parametrize("backend", [dist.Backend.NCCL, dist.Backend.GLOO])
947
948
@pytest.mark.parametrize("broadcast_fp16", [False, True])
def test_ddp_parity(change_train_graph: bool, backend: dist.Backend, broadcast_fp16: bool):
949
950
    temp_file_name = tempfile.mkstemp()[1]
    world_size = torch.cuda.device_count()
951
    mp.spawn(
952
953
954
955
        run_ddp_parity,
        args=(world_size, backend, temp_file_name, change_train_graph, broadcast_fp16),
        nprocs=world_size,
        join=True,
956
    )