pytorch
1105 строк · 40.9 Кб
1# Owner(s): ["oncall: distributed"]
2
3import functools
4import os
5import sys
6import warnings
7from collections import namedtuple
8from contextlib import nullcontext
9from copy import deepcopy
10from itertools import chain
11from typing import Any, Tuple
12
13import torch
14import torch.distributed as dist
15import torch.distributed.fsdp._traversal_utils as traversal_utils
16import torch.nn as nn
17from torch.distributed.fsdp import (
18CPUOffload,
19FlatParameter,
20FullyShardedDataParallel as FSDP,
21ShardingStrategy,
22)
23from torch.distributed.fsdp._flat_param import _FSDP_USE_UNSAFE_SETATTR
24from torch.distributed.fsdp._runtime_utils import HOMOGENEOUS_ATTR_NAMES
25from torch.distributed.fsdp.wrap import (
26always_wrap_policy,
27ModuleWrapPolicy,
28transformer_auto_wrap_policy,
29)
30from torch.distributed.optim import _apply_optimizer_in_backward
31from torch.nn import TransformerDecoderLayer, TransformerEncoderLayer
32from torch.nn.parallel import DistributedDataParallel as DDP
33from torch.testing._internal.common_distributed import skip_if_lt_x_gpu
34from torch.testing._internal.common_fsdp import (
35_assert_module_states,
36CUDAInitMode,
37FSDPInitMode,
38FSDPTest,
39FSDPTestMultiThread,
40NestedWrappedModule,
41TransformerWithSharedParams,
42)
43from torch.testing._internal.common_utils import (
44instantiate_parametrized_tests,
45parametrize,
46run_tests,
47TEST_WITH_DEV_DBG_ASAN,
48)
49
50if not dist.is_available():
51print("Distributed not available, skipping tests", file=sys.stderr)
52sys.exit(0)
53
54if TEST_WITH_DEV_DBG_ASAN:
55print(
56"Skip dev-asan as torch + multiprocessing spawn have known issues",
57file=sys.stderr,
58)
59sys.exit(0)
60
61
62class MyModel(nn.Module):
63def __init__(self):
64super().__init__()
65self.a = nn.Linear(2, 2)
66self.b = nn.Linear(2, 2)
67
68def forward(self, x, y):
69return self.b(self.a(x + y))
70
71
72class TestFSDPMiscMultiProcess(FSDPTest):
73@property
74def world_size(self):
75return 2
76
77@property
78def process_group(self):
79return dist.distributed_c10d._get_default_group()
80
81@skip_if_lt_x_gpu(2)
82@parametrize("use_index", [True, False])
83def test_fsdp_device_id(self, use_index):
84"""
85Tests the FSDP ``device_id`` argument:
86- Wrapping a CPU module should move the module to the GPU matching
87``device_id``
88- Wrapping a GPU module already on the GPU matching ``device_id``
89should not raise an error
90- Wrapping a GPU module already on GPU and passing a GPU device
91without specifying a device ID (i.e. ``torch.device("cuda")``) warns
92"""
93dev_id = (
94torch.cuda.current_device()
95if use_index
96else torch.device("cuda", torch.cuda.current_device())
97)
98
99def _check_device_matches(module, device_id):
100"""Checks that the ``FlatParameter``s in ``module`` have device
101matching ``device_id``."""
102devices = {
103p.device for p in module.parameters() if isinstance(p, FlatParameter)
104}
105assert len(devices) > 0
106self.assertEqual(1, len(devices))
107found_device = devices.pop()
108if use_index and not isinstance(device_id, torch.device):
109device = torch.device("cuda", device_id)
110else:
111device = device_id
112self.assertEqual(found_device, device)
113
114# Check that FSDP parameters are moved to `device_id` for a CPU module
115nested_wrapped_module = NestedWrappedModule.init(
116self.process_group,
117FSDPInitMode.RECURSIVE,
118CUDAInitMode.CUDA_NEVER,
119fsdp_kwargs={"device_id": dev_id},
120)
121_check_device_matches(nested_wrapped_module, dev_id)
122# Check that specifying `device_id` for a GPU module already on that
123# device does not raise an error
124nested_wrapped_module = NestedWrappedModule.init(
125self.process_group,
126FSDPInitMode.RECURSIVE,
127CUDAInitMode.CUDA_BEFORE,
128fsdp_kwargs={"device_id": dev_id},
129)
130_check_device_matches(nested_wrapped_module, dev_id)
131# Check that passing in `torch.device("cuda")` for a GPU module warns
132regex = "does not have an explicit index"
133context = self.assertWarnsRegex(
134expected_warning=UserWarning, expected_regex=regex
135)
136with context:
137nested_wrapped_module = NestedWrappedModule.init(
138self.process_group,
139FSDPInitMode.RECURSIVE,
140CUDAInitMode.CUDA_BEFORE,
141fsdp_kwargs={"device_id": torch.device("cuda")},
142)
143_check_device_matches(
144nested_wrapped_module, torch.device("cuda", torch.cuda.current_device())
145)
146
147@skip_if_lt_x_gpu(2)
148def test_fsdp_zero2_eval_with_prefetch(self):
149# Test FSDP validation with SHARD_GRAD_OP and forward_prefetch
150
151class Mnist(nn.Module):
152def __init__(self):
153super().__init__()
154self.conv1 = nn.Conv2d(1, 32, 3, 1)
155self.conv2 = nn.Conv2d(32, 64, 3, 1)
156self.dropout1 = nn.Dropout(0.25)
157self.dropout2 = nn.Dropout(0.5)
158self.fc1 = nn.Linear(9216, 128)
159self.fc2 = nn.Linear(128, 10)
160self.ln = nn.LayerNorm(9216)
161
162def forward(self, x, y):
163x = self.conv1(x)
164x = torch.nn.functional.relu(x)
165x = self.conv2(x)
166x = torch.nn.functional.relu(x)
167x = torch.nn.functional.max_pool2d(x, 2)
168x = self.dropout1(x)
169x = torch.flatten(x, 1)
170x = self.ln(x)
171x = self.fc1(x)
172x = torch.nn.functional.relu(x)
173x = self.dropout2(x)
174x = self.fc2(x)
175output = torch.nn.functional.log_softmax(x, dim=1)
176loss = torch.nn.functional.cross_entropy(output, y)
177return loss
178
179model = Mnist().cuda()
180model1 = Mnist().cuda()
181model1.load_state_dict(model.state_dict())
182fsdp_model = FSDP(
183model,
184sharding_strategy=ShardingStrategy.SHARD_GRAD_OP,
185forward_prefetch=True,
186use_orig_params=True,
187auto_wrap_policy=ModuleWrapPolicy([nn.Linear, nn.Conv2d]),
188)
189ddp_model = torch.nn.parallel.DistributedDataParallel(
190model1,
191)
192
193fsdp_opt = torch.optim.SGD(fsdp_model.parameters(), lr=1e-4)
194ddp_opt = torch.optim.SGD(ddp_model.parameters(), lr=1e-4)
195
196seed = self.rank + 20231010
197torch.manual_seed(seed)
198torch.cuda.manual_seed(seed)
199
200losses = []
201grads = []
202for i in range(5):
203x = torch.randn(8, 1, 28, 28, device="cuda").requires_grad_()
204y = torch.randint(low=0, high=9, size=(8,), device="cuda")
205for model, opt in ((fsdp_model, fsdp_opt), (ddp_model, ddp_opt)):
206seed = self.rank + i
207torch.manual_seed(seed)
208torch.cuda.manual_seed(seed)
209loss = model(x, y).sum()
210losses.append(loss)
211loss.backward()
212opt.step()
213grads.append(x.grad)
214opt.zero_grad()
215assert torch.allclose(losses[0], losses[1])
216assert torch.allclose(grads[0], grads[1])
217losses.clear()
218grads.clear()
219
220with torch.no_grad():
221fsdp_model.eval()
222ddp_model.eval()
223for _ in range(5):
224x = torch.randn(8, 1, 28, 28, device="cuda").requires_grad_()
225y = torch.randint(low=0, high=9, size=(8,), device="cuda")
226fsdp_loss = fsdp_model(x, y)
227ddp_loss = ddp_model(x, y)
228assert torch.allclose(fsdp_loss, ddp_loss)
229
230fsdp_model.train()
231ddp_model.train()
232for i in range(5):
233x = torch.randn(8, 1, 28, 28, device="cuda").requires_grad_()
234y = torch.randint(low=0, high=9, size=(8,), device="cuda")
235for model, opt in ((fsdp_model, fsdp_opt), (ddp_model, ddp_opt)):
236seed = self.rank + i
237torch.manual_seed(seed)
238torch.cuda.manual_seed(seed)
239loss = model(x, y).sum()
240losses.append(loss)
241loss.backward()
242opt.step()
243grads.append(x.grad)
244opt.zero_grad()
245assert torch.allclose(losses[0], losses[1])
246assert torch.allclose(grads[0], grads[1])
247losses.clear()
248grads.clear()
249
250@skip_if_lt_x_gpu(2)
251@parametrize("use_second_layer", [True, False])
252@parametrize("sharding_strategy", [ShardingStrategy.NO_SHARD, None])
253def test_fsdp_module_no_compute_grad(self, use_second_layer, sharding_strategy):
254# When use_second_layer=True, b is involved in forward computation but does
255# not receive grad in backward. Otherwise, b is not involved in forward
256# computation.
257
258class MyModel(nn.Module):
259def __init__(self):
260super().__init__()
261self.a = nn.Linear(10, 10)
262self.b = nn.Linear(10, 10)
263
264def forward(self, x, y):
265out1 = self.a(x)
266if use_second_layer:
267out2 = self.b(y)
268return out1, out2
269else:
270return out1
271
272fsdp = FSDP(
273MyModel().cuda(),
274sharding_strategy=sharding_strategy,
275auto_wrap_policy=always_wrap_policy,
276)
277x = torch.randn(10, 10, device="cuda")
278y = torch.randn(10, 10, device="cuda")
279for i in range(4):
280if use_second_layer:
281a, b = fsdp(x, y)
282else:
283a = fsdp(x, y)
284loss = a.sum()
285loss.backward()
286
287# self.a receives grad, self.b does not
288a_grad = fsdp.module.a._handle.flat_param.grad
289b_grad = fsdp.module.b._handle.flat_param.grad
290self.assertIsNotNone(a_grad)
291self.assertIsNone(b_grad)
292
293@skip_if_lt_x_gpu(2)
294def test_fsdp_not_all_outputs_used_in_loss(self):
295self.run_subtests(
296{
297"sharding_strategy": [
298ShardingStrategy.FULL_SHARD,
299ShardingStrategy.SHARD_GRAD_OP,
300ShardingStrategy.NO_SHARD,
301]
302},
303self._test_fsdp_not_all_outputs_used_in_loss,
304)
305
306def _test_fsdp_not_all_outputs_used_in_loss(
307self, sharding_strategy: ShardingStrategy
308):
309class MyModule(nn.Module):
310def __init__(self):
311super().__init__()
312self.lin1 = nn.Linear(4, 4)
313self.lin2 = nn.Linear(4, 4)
314
315def forward(self, x):
316a = self.lin1(x)
317b = self.lin2(x)
318return (a, b)
319
320def _check_resharded(fsdp_module):
321handle = fsdp_module._handle
322if not handle:
323return
324param = handle.flat_param
325if handle.uses_sharded_strategy:
326full_param = param._full_param_padded
327self.assertEqual(full_param.storage().size(), 0)
328
329self.assertEqual(param.data_ptr(), param._local_shard.data_ptr())
330
331def _check_equal(local, fsdp):
332with FSDP.summon_full_params(fsdp):
333for p1, p2 in zip(fsdp.parameters(), local.parameters()):
334torch.testing.assert_close(p1, p2)
335
336fsdp_ctor = functools.partial(FSDP, sharding_strategy=sharding_strategy)
337m = MyModule().cuda()
338m_local = deepcopy(m)
339local_m = m_local
340prev_params = [p.clone() for p in m_local.parameters()]
341
342m.lin1 = fsdp_ctor(m.lin1)
343m = fsdp_ctor(m)
344_check_equal(m_local, m)
345
346opt = torch.optim.SGD(m.parameters(), lr=1e-3)
347opt_local = torch.optim.SGD(local_m.parameters(), lr=1e-3)
348
349for i in range(6):
350t = torch.ones(4, device="cuda")
351a, b = m(t)
352local_a, local_b = local_m(t)
353if i < 2:
354# use both params in loss computation. Later,
355# b will go unused and we check grads are the
356# same as local training.
357loss = (a @ b).sum()
358loss_local = (local_a @ local_b).sum()
359else:
360loss = a.sum()
361loss_local = local_a.sum()
362
363loss.backward()
364loss_local.backward()
365_check_resharded(m)
366opt.step()
367opt_local.step()
368_check_equal(m_local, m)
369# Ensure at least some change from previous params, otherwise
370# above check would be vacuously true.
371self.assertTrue(
372any(
373not torch.equal(p1, p2)
374for p1, p2 in zip(prev_params, m_local.parameters())
375)
376)
377prev_params = [p.clone() for p in local_m.parameters()]
378opt.zero_grad()
379opt_local.zero_grad()
380
381dist.barrier()
382
383@skip_if_lt_x_gpu(2)
384def test_fsdp_optim_overlap_no_use_orig_params_error(self):
385fsdp_overlap = FSDP(
386MyModel().cuda(),
387auto_wrap_policy=always_wrap_policy,
388use_orig_params=False,
389)
390optim_cls = torch.optim.SGD
391optim_kwargs = {"lr": 0.03}
392_apply_optimizer_in_backward(
393optimizer_class=optim_cls,
394params=fsdp_overlap.parameters(),
395optimizer_kwargs=optim_kwargs,
396register_hook=False,
397)
398
399inp = torch.randn(10, 10, device="cuda")
400with self.assertRaisesRegex(
401RuntimeError, "only supported with use_orig_params=True"
402):
403fsdp_overlap(inp, inp)
404
405@skip_if_lt_x_gpu(2)
406def test_fsdp_optimizer_overlap(self):
407torch.manual_seed(0)
408for cpu_offload in [True, False]:
409offload = CPUOffload(offload_params=cpu_offload)
410model = MyModel().cuda()
411model_overlap = deepcopy(model)
412fsdp = FSDP(
413model.cuda(),
414auto_wrap_policy=always_wrap_policy,
415use_orig_params=True,
416cpu_offload=offload,
417)
418fsdp_overlap = FSDP(
419model_overlap.cuda(),
420auto_wrap_policy=always_wrap_policy,
421use_orig_params=True,
422cpu_offload=offload,
423)
424optim_cls = torch.optim.SGD
425optim_kwargs = {"lr": 0.03}
426_apply_optimizer_in_backward(
427optimizer_class=optim_cls,
428params=fsdp_overlap.parameters(),
429optimizer_kwargs=optim_kwargs,
430register_hook=False,
431)
432for p in fsdp_overlap.parameters():
433assert hasattr(p, "_in_backward_optimizers")
434optim = optim_cls(fsdp.parameters(), **optim_kwargs)
435
436# Verify params initially equal
437for p1, p2 in zip(fsdp.parameters(), fsdp_overlap.parameters()):
438self.assertEqual(p1, p2)
439
440with FSDP.summon_full_params(fsdp_overlap):
441fsdp_overlap_prev_params = [
442(n, p.clone()) for n, p in fsdp_overlap.named_parameters()
443]
444
445for i in range(6):
446inp = torch.randn(2, 2, device="cuda")
447with torch.no_grad():
448inp_clone = inp.clone()
449fsdp(inp, inp).sum().backward()
450fsdp_overlap(inp_clone, inp_clone).sum().backward()
451
452optim.step()
453optim.zero_grad()
454
455# Overlapped optimizer FSDP module should have sharded_grad as None.
456for fsdp_unit in FSDP.fsdp_modules(fsdp_overlap):
457handle = fsdp_unit._handle
458if handle:
459handle_grad = handle.sharded_grad
460self.assertEqual(
461None,
462handle_grad,
463"Overlapped FSDP sharded_grad is not None!",
464)
465
466# Note: FSDP without optimizer overlap won't set sharded_grad to None until the next
467# pre-forward since it needs to run FSDP specific logic that picks up that set_to_none=True
468# has been called (or that the gradients have been otherwise set to None)
469
470# Verify parameters are different than prev iteration
471with FSDP.summon_full_params(fsdp_overlap, with_grads=True):
472for (n, p), (n_prev, p_prev) in zip(
473fsdp_overlap.named_parameters(), fsdp_overlap_prev_params
474):
475self.assertNotEqual(
476p,
477p_prev,
478f"{n_prev} Params at iter {i} same as previous iter!",
479)
480
481# Verify overlap and non overlapped are the same
482with FSDP.summon_full_params(fsdp_overlap):
483with FSDP.summon_full_params(fsdp):
484for (n_overlap, p_overlap), (n, p) in zip(
485fsdp_overlap.named_parameters(), fsdp.named_parameters()
486):
487self.assertEqual(n_overlap, n)
488self.assertEqual(
489p,
490p_overlap,
491f"Rank {self.rank}: Params not equal at iteration {i}: {n_overlap} - {p} vs {p_overlap}",
492)
493self.assertEqual(
494None, p.grad, f"Expected param {n} grad to be None"
495)
496self.assertEqual(
497None,
498p_overlap.grad,
499f"Expected param {n_overlap} grad to be None",
500)
501
502fsdp_overlap_prev_params = [
503(n, p.clone()) for n, p in fsdp_overlap.named_parameters()
504]
505
506@skip_if_lt_x_gpu(2)
507def test_fsdp_cpu_training(self):
508"""Tests FSDP training on CPU."""
509gloo_pg = dist.new_group(backend="gloo")
510for ss in [
511ShardingStrategy.NO_SHARD,
512ShardingStrategy.FULL_SHARD,
513ShardingStrategy.SHARD_GRAD_OP,
514ShardingStrategy.HYBRID_SHARD,
515ShardingStrategy._HYBRID_SHARD_ZERO2,
516]:
517torch.manual_seed(42)
518model = MyModel()
519ref_model = DDP(deepcopy(model), process_group=gloo_pg)
520model = FSDP(
521model,
522auto_wrap_policy=always_wrap_policy,
523process_group=gloo_pg,
524device_id=torch.device("cpu"),
525)
526ref_optim = torch.optim.Adam(ref_model.parameters(), lr=1e-2)
527optim = torch.optim.Adam(model.parameters(), lr=1e-2)
528torch.manual_seed(42 + self.rank)
529inp = torch.randn(2, 2)
530for _ in range(10):
531losses = []
532for _model, _optim in ((ref_model, ref_optim), (model, optim)):
533loss = _model(inp, inp).sum()
534losses.append(loss)
535loss.backward()
536_optim.step()
537_optim.zero_grad()
538self.assertEqual(losses[0], losses[1])
539
540@skip_if_lt_x_gpu(2)
541def test_fsdp_cpu_init_stays_on_cpu(self):
542# Move me to MT test once warning logging and backward collective issue
543# is resolved.
544"""Tests that passing a CPU module to FSDP preserves that the wrapped
545module is on CPU after FSDP initialization, albeit after logging a
546warning, and that FSDP moves CPU input to GPU before the forward."""
547torch.cuda.set_device(self.rank)
548regex = "passed-in `module` is on CPU"
549context = self.assertWarnsRegex(
550expected_warning=UserWarning, expected_regex=regex
551)
552with context:
553nested_wrapped_module = NestedWrappedModule.init(
554self.process_group,
555FSDPInitMode.RECURSIVE,
556CUDAInitMode.CUDA_NEVER,
557)
558fsdp_model = FSDP(nested_wrapped_module, self.process_group)
559devices = {p.device for p in fsdp_model.parameters()}
560self.assertEqual(1, len(devices))
561self.assertEqual(torch.device("cpu"), devices.pop())
562fsdp_model = fsdp_model.cuda()
563# Ensure fwd + backward can be performed after moving to CUDA.
564# CPU input also tests that input is correctly moved to appropriate
565# CUDA device.
566inp = fsdp_model.module.get_input(device=torch.device("cpu"))
567fsdp_model(*inp).sum().backward()
568
569@skip_if_lt_x_gpu(2)
570def test_cpu_init_with_sync_module_states(self):
571"""
572Tests that passing ``sync_module_states=True`` raises an error for
573a CPU module since the synchronization requires GPU communication,
574while additionally passing ``device_id`` does not raise an error, even
575when the model has CPU buffers.
576"""
577
578def init_nested_wrapped_module():
579return NestedWrappedModule.init(
580self.process_group,
581FSDPInitMode.NO_FSDP,
582CUDAInitMode.CUDA_NEVER,
583)
584
585with self.assertRaisesRegex(
586ValueError,
587"The module has CPU parameters or buffers when `sync_module_states=True`",
588):
589FSDP(
590init_nested_wrapped_module(),
591self.process_group,
592sync_module_states=True,
593)
594
595# Check that `device_id` with `sync_module_states=True` works
596nested_wrapped_module = init_nested_wrapped_module()
597nested_wrapped_module.register_buffer(
598"buf", torch.ones((2, 2), device="cpu") * self.rank
599)
600nested_wrapped_module.module[0].register_buffer(
601"buf", torch.ones((3, 2), device="cpu") * self.rank
602)
603nested_wrapped_module = FSDP(
604nested_wrapped_module,
605self.process_group,
606auto_wrap_policy=ModuleWrapPolicy({nn.Linear}),
607device_id=torch.cuda.current_device(),
608sync_module_states=True,
609)
610# Each rank's buffers should be 0s since rank 0 is the source, and they
611# should be on GPU since we specified `device_id`
612self.assertEqual(
613nested_wrapped_module.buf.device,
614torch.device("cuda", torch.cuda.current_device()),
615)
616self.assertEqual(nested_wrapped_module.buf, torch.zeros((2, 2)))
617self.assertEqual(
618nested_wrapped_module.module.module[0].buf.device,
619torch.device("cuda", torch.cuda.current_device()),
620)
621self.assertEqual(
622nested_wrapped_module.module.module[0].buf, torch.zeros((3, 2))
623)
624
625
626class TestFSDPMiscMultiThread(FSDPTestMultiThread):
627@property
628def world_size(self):
629return 2
630
631@property
632def process_group(self):
633return dist.distributed_c10d._get_default_group()
634
635@skip_if_lt_x_gpu(2)
636def test_fsdp_namedtuple(self):
637class MyModule(nn.Module):
638def __init__(self):
639super().__init__()
640self.lin = nn.Linear(100, 100)
641
642def forward(self, x):
643return x
644
645m = MyModule().cuda()
646m = FSDP(m)
647t = torch.ones(1, device="cuda", requires_grad=True)
648
649MyOutputType = namedtuple(
650"MyOutputType", ["a", "b", "c", "d"], defaults=(t, t, t, t)
651)
652
653inp = MyOutputType()
654out = m(inp)
655# Ensure hooks are registered
656for x in out:
657self.assertNotEqual([], list(x._backward_hooks.values()))
658
659# TODO: we should check backward() and param is resharded
660# as well, but this is blocked by
661# https://github.com/pytorch/pytorch/issues/83107 and
662# https://github.com/pytorch/pytorch/issues/83129
663
664@skip_if_lt_x_gpu(2)
665def test_device_id_auto_wrap(self):
666"""Tests that ``auto_wrap_policy`` propagates ``device_id`` to all
667nested FSDP instances."""
668self.run_subtests(
669{"use_callable": [False, True]},
670self._test_device_id_auto_wrap,
671)
672
673def _test_device_id_auto_wrap(self, use_callable: bool):
674module_classes = {TransformerEncoderLayer, TransformerDecoderLayer}
675if use_callable:
676auto_wrap_policy = functools.partial(
677transformer_auto_wrap_policy,
678transformer_layer_cls=module_classes,
679)
680else:
681auto_wrap_policy = ModuleWrapPolicy(module_classes)
682fsdp_kwargs = {
683"auto_wrap_policy": auto_wrap_policy,
684"device_id": torch.cuda.current_device(),
685}
686fsdp_model = TransformerWithSharedParams.init(
687self.process_group,
688FSDPInitMode.RECURSIVE,
689CUDAInitMode.CUDA_BEFORE,
690fsdp_kwargs,
691)
692for fsdp_module in FSDP.fsdp_modules(fsdp_model):
693self.assertEqual(
694fsdp_module.compute_device,
695torch.device("cuda", torch.cuda.current_device()),
696)
697
698@skip_if_lt_x_gpu(2)
699def test_fsdp_device_id_cpu_offload(self):
700"""
701Tests FSDP when specifying both ``device_id`` and parameter CPU
702offloading.
703"""
704self.run_subtests(
705{"use_orig_params": [False, True]},
706self._test_fsdp_device_id_cpu_offload,
707)
708
709def _test_fsdp_device_id_cpu_offload(self, use_orig_params: bool):
710class MyModel(nn.Module):
711def __init__(self):
712super().__init__()
713self.seq = nn.Sequential(
714nn.Linear(10, 10),
715nn.Linear(10, 10),
716)
717self.lin = nn.Linear(10, 10)
718
719def forward(self, x):
720return self.lin(self.seq(x))
721
722model = MyModel()
723# Choose a wrapping policy such that there are (1) nested FSDP
724# instances and (2) the parent FSDP instance has managed parameters
725auto_wrap_policy = ModuleWrapPolicy({nn.Sequential})
726fsdp_model = FSDP(
727model,
728auto_wrap_policy=auto_wrap_policy,
729cpu_offload=CPUOffload(offload_params=True),
730device_id=torch.cuda.current_device(),
731use_orig_params=use_orig_params,
732)
733cpu_device = torch.device("cpu")
734for handle in traversal_utils._get_fsdp_handles(fsdp_model):
735self.assertEqual(handle.flat_param.device, cpu_device)
736
737@skip_if_lt_x_gpu(2)
738def test_module_device_mismatches_device_id(self):
739"""Tests that specifying a ``device_id`` argument to FSDP for a GPU
740module that does not match the GPU device ID raises an error."""
741# TODO: override FSDP MT Thread _run to set this instead of here for
742# every test.
743torch.cuda.set_device(self.rank)
744context = (
745self.assertRaisesRegex(ValueError, f"cuda:{self.rank} vs cuda:0")
746if self.rank != 0
747else nullcontext()
748)
749with context:
750NestedWrappedModule.init(
751self.process_group,
752FSDPInitMode.RECURSIVE,
753# Move wrapped modules to CUDA before wrapping with FSDP
754cuda_init_mode=CUDAInitMode.CUDA_BEFORE,
755# Should raise error since rank 1 is given `device_id=0` when
756# the model is on cuda:1
757fsdp_kwargs={"device_id": 0},
758)
759
760@skip_if_lt_x_gpu(2)
761def test_cpu_gpu_module(self):
762"""Tests a CPU + GPU module supported if device_id is passed
763in, errors if device_id is not.
764"""
765torch.cuda.set_device(self.rank)
766
767class CPUGPUModule(nn.Module):
768def __init__(self):
769super().__init__()
770self.a = nn.Linear(1, 1).cuda()
771self.b = nn.Linear(1, 1)
772
773cpu_gpu = CPUGPUModule()
774fsdp = FSDP(cpu_gpu, device_id=torch.cuda.current_device())
775for param in fsdp.parameters():
776self.assertEqual(param.device, torch.device(torch.cuda.current_device()))
777
778# without device_id, we hit an error
779with self.assertRaisesRegex(RuntimeError, "please pass in device_id"):
780FSDP(CPUGPUModule())
781
782@skip_if_lt_x_gpu(2)
783def test_fsdp_ignored_module_meta(self):
784torch.cuda.set_device(self.rank)
785
786class CPUGPUModule(nn.Module):
787def __init__(self):
788super().__init__()
789self.a = nn.Linear(1, 1)
790self.b = nn.Linear(1, 1)
791
792with torch.device("meta"):
793m = CPUGPUModule()
794m = FSDP(m, device_id=self.rank, ignored_modules=[m.a], use_orig_params=True)
795meta_device = torch.device("meta")
796self.assertEqual(meta_device, next(m.a.parameters()).device)
797
798# Test with param_init_fn
799with torch.device("meta"):
800m = CPUGPUModule()
801m = FSDP(
802m,
803device_id=torch.cuda.current_device(),
804ignored_modules=[m.a],
805use_orig_params=True,
806param_init_fn=lambda m: m.to_empty(
807device=torch.cuda.current_device(), recurse=False
808),
809)
810self.assertEqual(meta_device, next(m.a.parameters()).device)
811
812@skip_if_lt_x_gpu(2)
813def test_fsdp_device_id_no_move_ignored_params_and_bufs(self):
814class CPUGPUModule(nn.Module):
815def __init__(self):
816super().__init__()
817self.a = nn.Linear(1, 1)
818self.b = nn.Linear(1, 1)
819self.a.register_buffer("buf", torch.ones(1))
820
821m = CPUGPUModule()
822m = FSDP(m, device_id=self.rank, ignored_modules=[m.a], use_orig_params=True)
823ignored_params = m.a.parameters()
824ignored_bufs = m.a.buffers()
825for t in chain(ignored_params, ignored_bufs):
826self.assertEqual(torch.device("cpu"), t.device)
827
828@skip_if_lt_x_gpu(2)
829def test_multigpu_module(self):
830"""
831Module on multiple GPUs wrapped in FSDP should raise an error.
832"""
833
834class MultiGPUModule(nn.Module):
835def __init__(self, rank):
836super().__init__()
837self.rank = rank
838self.a = nn.Linear(1, 1).cuda(self.rank)
839self.b = nn.Linear(1, 1).cuda((self.rank + 1) % dist.get_world_size())
840
841with self.assertRaisesRegex(
842RuntimeError, "FSDP only supports single device modules"
843):
844FSDP(MultiGPUModule(self.rank))
845
846@skip_if_lt_x_gpu(2)
847def test_no_params(self):
848"""
849Test that device_id and cpu init work if module has no params
850(they are effective noops, but ensure FSDP does not assume module
851has parameters during init)
852"""
853# TODO: override FSDP MT Thread _run to set this instead of here for
854# every test.
855torch.cuda.set_device(self.rank)
856# Test CPU
857no_params = nn.ReLU()
858module = FSDP(no_params)
859# Test CUDA
860no_params = nn.ReLU().cuda()
861module = FSDP(no_params)
862# Test CPU + device_id
863no_params = nn.ReLU()
864module = FSDP(no_params, device_id=torch.cuda.current_device())
865# For modules with no params, wrong device_id will raise error about
866# inconsistency between compute_device and device_id, since compute_device
867# is computed as torch.cuda.current_device when there are no params.
868no_params = nn.ReLU().cuda()
869context = (
870(
871self.assertRaisesRegex(
872ValueError, f"Inconsistent.*cuda:{self.rank} vs cuda:0"
873)
874)
875if self.rank != 0
876else nullcontext()
877)
878with context:
879FSDP(no_params, device_id=0)
880
881@skip_if_lt_x_gpu(2)
882def test_fsdp_same_model_across_ranks(self):
883"""
884FSDP broadcasts model from rank 0 to ensure it starts off with the same
885values.
886"""
887
888class MyModel(nn.Module):
889def __init__(self, rank):
890super().__init__()
891# Seed via rank to make model different across ranks
892torch.manual_seed(rank)
893torch.cuda.manual_seed(rank)
894self.lin = nn.Linear(10, 10, bias=False)
895self.register_buffer("buffer", torch.ones(1) * rank)
896
897m = MyModel(self.rank).cuda()
898_assert_module_states(
899m, process_group=self.process_group, assert_fn=self.assertNotEqual
900)
901# Passing sync_module_states into FSDP makes model the same during init.
902fsdp = FSDP(m, sync_module_states=True)
903with fsdp.summon_full_params(fsdp):
904_assert_module_states(
905fsdp, process_group=self.process_group, assert_fn=self.assertEqual
906)
907
908# sync_module_states also works with CPU module with device_id passed in
909m = MyModel(self.rank)
910_assert_module_states(
911m, process_group=self.process_group, assert_fn=self.assertNotEqual
912)
913# Passing sync_module_states into FSDP makes model the same during init.
914fsdp = FSDP(m, device_id=torch.cuda.current_device(), sync_module_states=True)
915with fsdp.summon_full_params(fsdp):
916_assert_module_states(
917fsdp, process_group=self.process_group, assert_fn=self.assertEqual
918)
919
920@skip_if_lt_x_gpu(2)
921def test_homogeneous_attributes(self):
922"""
923Tests that passing heterogeneous values for attributes designated as
924homogeneous raises an error.
925"""
926# Manually construct this list but verify against the global list of
927# homogeneous attribute names
928all_attr_name_and_values = [
929("_use_orig_params", False, True),
930("limit_all_gathers", False, True),
931("_use_full_prec_in_eval", False, True),
932]
933self.assertEqual(
934[
935attr_name_and_values[0]
936for attr_name_and_values in all_attr_name_and_values
937],
938HOMOGENEOUS_ATTR_NAMES,
939)
940
941self.run_subtests(
942{"attr_name_and_values": all_attr_name_and_values},
943self._test_homogeneous_attributes,
944)
945
946def _test_homogeneous_attributes(self, attr_name_and_values: Tuple[str, Any, Any]):
947model = NestedWrappedModule.init(
948self.process_group,
949FSDPInitMode.NO_FSDP,
950CUDAInitMode.CUDA_BEFORE,
951{},
952)
953attr_name = attr_name_and_values[0]
954
955if "_use_full_prec_in_eval" == attr_name:
956model.module[1] = FSDP(model.module[1])
957os.environ["FSDP_USE_FULL_PREC_IN_EVAL"] = "1"
958fsdp_model = FSDP(model)
959else:
960fsdp_kwargs_inner = {attr_name.lstrip("_"): attr_name_and_values[1]}
961fsdp_kwargs_outer = {attr_name.lstrip("_"): attr_name_and_values[2]}
962model.module[1] = FSDP(model.module[1], **fsdp_kwargs_inner)
963fsdp_model = FSDP(model, **fsdp_kwargs_outer)
964
965# Run a forward to trigger lazy initialization and the error
966with self.assertRaisesRegex(
967ValueError, f"Expects one homogeneous value for {attr_name}"
968):
969inp = fsdp_model.module.get_input(torch.device("cuda"))
970fsdp_model(*inp)
971
972
973class TestFSDPMiscWorldSize1(FSDPTestMultiThread):
974@property
975def world_size(self) -> int:
976return 1
977
978@skip_if_lt_x_gpu(1)
979def test_world_size_1_sharding_strategy_warning(self):
980"""
981Tests that FSDP issues a warning when it switches to using ``NO_SHARD``
982when the world size is 1.
983"""
984warning_prefix = "FSDP is switching to use `NO_SHARD` instead of"
985# If the user already passes `NO_SHARD`, then there should not be a
986# warning
987with warnings.catch_warnings(record=True) as w:
988warnings.simplefilter("always") # trigger all warnings
989FSDP(nn.Linear(3, 3).cuda(), sharding_strategy=ShardingStrategy.NO_SHARD)
990for warning in w:
991self.assertTrue(
992warning.category != UserWarning
993or not str(warning.message).startswith(warning_prefix)
994)
995
996# Check that a warning is issued
997warning_suffix = " since the world size is 1."
998# - Pass `FULL_SHARD` or `None`
999expected_regex_full_shard = (
1000warning_prefix + " " + str(ShardingStrategy.FULL_SHARD) + warning_suffix
1001)
1002with self.assertWarnsRegex(UserWarning, expected_regex_full_shard):
1003FSDP(nn.Linear(3, 3).cuda(), sharding_strategy=ShardingStrategy.FULL_SHARD)
1004with self.assertWarnsRegex(UserWarning, expected_regex_full_shard):
1005FSDP(nn.Linear(3, 3).cuda())
1006# - Pass `SHARD_GRAD_OP`
1007expected_regex_shard_grad_op = (
1008warning_prefix + " " + str(ShardingStrategy.SHARD_GRAD_OP) + warning_suffix
1009)
1010with self.assertWarnsRegex(UserWarning, expected_regex_shard_grad_op):
1011FSDP(
1012nn.Linear(3, 3).cuda(), sharding_strategy=ShardingStrategy.SHARD_GRAD_OP
1013)
1014
1015@skip_if_lt_x_gpu(1)
1016def test_training_device_mismatch_errors(self):
1017"""
1018Tests that, when training starts, if FSDP parameters are not on the
1019expected device, then an informative error is raised. This applies for
1020both no parameter CPU offloading and parameter CPU offloading.
1021"""
1022# Incorrectly not moving from CPU -> GPU
1023model = torch.nn.Linear(10, 10)
1024fsdp_model = FSDP(model)
1025inp = torch.randn((2, 10))
1026with self.assertRaisesRegex(
1027RuntimeError,
1028"An FSDP-managed module unexpectedly has parameters on cpu. Make "
1029"sure to move the module to cuda:0 before training.",
1030):
1031fsdp_model(inp)
1032
1033# Incorrectly moving from CPU -> GPU
1034model = torch.nn.Linear(10, 10)
1035fsdp_model = FSDP(model, cpu_offload=CPUOffload(offload_params=True))
1036fsdp_model.to(torch.device("cuda"))
1037inp = torch.randn((2, 10))
1038with self.assertRaisesRegex(
1039RuntimeError,
1040"An FSDP-managed module with parameter CPU offloading enabled has "
1041"parameters on cuda:0. Make sure to not move the module from CPU "
1042"when offloading parameters.",
1043):
1044fsdp_model(inp)
1045
1046@skip_if_lt_x_gpu(2)
1047def test_unsafe_setattr(self):
1048"""
1049Tests that the environment variable for using unsafe setattr gates as
1050expected.
1051"""
1052self.run_subtests(
1053{"use_orig_params": [False, True]},
1054self._test_unsafe_setattr,
1055)
1056
1057def _test_unsafe_setattr(self, use_orig_params: bool):
1058called_setattr_override = False
1059
1060class SetattrLinear(nn.Module):
1061def __init__(self, in_dim: int, out_dim: int, device: torch.device) -> None:
1062super().__init__()
1063self.weight = nn.Parameter(
1064torch.randn((in_dim, out_dim), device=device)
1065)
1066
1067def forward(self, x: torch.Tensor) -> torch.Tensor:
1068return x @ self.weight
1069
1070def __setattr__(self, name: str, value: Any) -> None:
1071nonlocal called_setattr_override
1072called_setattr_override = True
1073return super().__setattr__(name, value)
1074
1075# Construct FSDP module without changing any environment variables and
1076# run forward, which triggers both unsharded and sharded view setting
1077module = SetattrLinear(5, 5, torch.device("cuda"))
1078fsdp_module = FSDP(module, use_orig_params=use_orig_params)
1079inp = torch.randn((8, 5), device=torch.device("cuda"))
1080called_setattr_override = False
1081fsdp_module(inp)
1082self.assertTrue(called_setattr_override)
1083
1084# Repeat with unsafe setattr explicitly enabled
1085os.environ[_FSDP_USE_UNSAFE_SETATTR] = "1"
1086module = SetattrLinear(5, 5, torch.device("cuda"))
1087fsdp_module = FSDP(module, use_orig_params=use_orig_params)
1088called_setattr_override = False
1089fsdp_module(inp)
1090self.assertFalse(called_setattr_override)
1091
1092# Repeat with unsafe setattr explicitly disabled
1093os.environ[_FSDP_USE_UNSAFE_SETATTR] = "0"
1094module = SetattrLinear(5, 5, torch.device("cuda"))
1095fsdp_module = FSDP(module, use_orig_params=use_orig_params)
1096called_setattr_override = False
1097fsdp_module(inp)
1098self.assertTrue(called_setattr_override)
1099
1100
1101instantiate_parametrized_tests(TestFSDPMiscMultiThread)
1102instantiate_parametrized_tests(TestFSDPMiscMultiProcess)
1103
1104if __name__ == "__main__":
1105run_tests()
1106