# Owner(s): ["oncall: distributed"] import sys import unittest import torch from torch import distributed as dist from torch.distributed.fsdp.flat_param import ( FlatParamShardMetadata, HandleConfig, HandleShardingStrategy, ) from torch.distributed.fsdp.flatten_params_wrapper import FlattenParamsWrapper from torch.testing._internal.common_utils import TestCase, run_tests if not dist.is_available(): print("Distributed not available, skipping tests", file=sys.stderr) sys.exit(0) class TestFlattenParams(TestCase): """Base test class and used for CPU case.""" def _get_default_config(self): return HandleConfig(HandleShardingStrategy.FULL_SHARD, False, None, None) def _get_empty_module(self, seed=0): torch.manual_seed(seed) # keep everything deterministic class Test(torch.nn.Module): def forward(self, x): return x + 1 module = Test() def get_input(device, dtype): torch.manual_seed(1) # keep everything deterministic return torch.rand(1).to(device=device, dtype=dtype) module.get_input = get_input return module def _get_transformer(self, seed=0): torch.manual_seed(seed) # keep everything deterministic module = torch.nn.Transformer( d_model=32, num_encoder_layers=2, num_decoder_layers=2, dim_feedforward=128, dropout=0.1, ) module.register_buffer("dummy_buffer", torch.tensor(1.0)) def get_input(device, dtype): torch.manual_seed(1) # keep everything deterministic src = torch.rand(20, 8, 32).to(device=device, dtype=dtype) # T x B x C tgt = torch.rand(10, 8, 32).to(device=device, dtype=dtype) # T x B x C return (src, tgt) module.get_input = get_input return module def _get_shared_params_transformer(self, seed=0): module = self._get_transformer(seed=seed) # share the FFNs for enc_layer, dec_layer in zip(module.encoder.layers, module.decoder.layers): dec_layer.linear1.weight = enc_layer.linear1.weight dec_layer.linear2.weight = enc_layer.linear2.weight return module def _get_output(self, module): device = next(module.parameters()).device dtype = next(module.parameters()).dtype input = module.get_input(device, dtype) return module(*input) def _get_pnorm_after_step(self, module): optim = torch.optim.SGD(module.parameters(), lr=0.01) loss = self._get_output(module).sum() loss.backward() optim.step() return torch.norm(torch.stack([p.detach().norm() for p in module.parameters()])) def _test_num_params(self, module): ref_num_params = sum(p.numel() for p in module.parameters()) params_to_flatten = list(module.parameters()) flat_module = FlattenParamsWrapper( module, params_to_flatten, torch.device("cuda"), self._get_default_config(), ) flat_num_params = sum(p.numel() for p in flat_module.parameters()) self.assertEqual(ref_num_params, flat_num_params) self.assertEqual(flat_num_params, flat_module.flat_param.numel()) def _test_output(self, module): ref_output = self._get_output(module) params_to_flatten = list(module.parameters()) flat_module = FlattenParamsWrapper( module, params_to_flatten, torch.device("cuda"), self._get_default_config(), ) flat_output = self._get_output(flat_module) self.assertEqual(ref_output, flat_output) def test_partial_flattening(self): module = self._get_transformer() num_params = sum(p.numel() for p in module.parameters()) params_to_flatten = list(module.encoder.layers[1].parameters()) + list( module.decoder.layers[0].parameters() ) num_params_to_flatten = sum(p.numel() for p in params_to_flatten) module = FlattenParamsWrapper( module, params_to_flatten, torch.device("cuda"), self._get_default_config(), ) self.assertEqual(module.flat_param.numel(), num_params_to_flatten) self.assertEqual(sum(p.numel() for p in module.parameters()), num_params) # flattened parameters are removed self.assertEqual(len(list(module.encoder.layers[1].parameters())), 0) self.assertEqual(len(list(module.decoder.layers[0].parameters())), 0) # non-flattened parameters remain self.assertGreater(len(list(module.encoder.layers[0].parameters())), 0) self.assertGreater(len(list(module.decoder.layers[1].parameters())), 0) # test that changing the module dtype works properly orig_dtype = params_to_flatten[0].dtype new_dtype = torch.float32 if orig_dtype == torch.float16 else torch.float16 self.assertEqual(module.flat_param.dtype, orig_dtype) self.assertTrue( all(p.dtype == orig_dtype for p in module.encoder.layers[0].parameters()) ) module = module.to(dtype=new_dtype) self.assertEqual(module.flat_param.dtype, new_dtype) self.assertTrue( all(p.dtype == new_dtype for p in module.encoder.layers[0].parameters()) ) def test_flatten_nothing(self): module = self._get_transformer() module = FlattenParamsWrapper( module, [], torch.device("cuda"), self._get_default_config(), ) self.assertIsNone(module.flat_param) def test_empty_module(self): module = self._get_empty_module() in_data = torch.rand(1) ref_out = module(in_data) module = FlattenParamsWrapper( module, [], torch.device("cuda"), self._get_default_config(), ) self.assertEqual(len(list(module.parameters())), 0) self.assertIsNone(module.flat_param) fpw_out = module(in_data) self.assertEqual(ref_out, fpw_out) def test_num_params(self): module = self._get_transformer() self._test_num_params(module) def test_shared_params_num_params(self): module = self._get_shared_params_transformer() self._test_num_params(module) def test_output(self): module = self._get_transformer() self._test_output(module) def test_shared_params_output(self): module = self._get_shared_params_transformer() self._test_output(module) def test_shared_params_pnorm_after_step(self): # incorrect parameter sharing is likely to cause problems after an # optimization step module = self._get_shared_params_transformer() ref_pnorm_after_step = self._get_pnorm_after_step(module) module = self._get_shared_params_transformer() # recreate params_to_flatten = list(module.parameters()) flat_module = FlattenParamsWrapper( module, params_to_flatten, torch.device("cuda"), self._get_default_config(), ) flat_pnorm_after_step = self._get_pnorm_after_step(flat_module) self.assertEqual(ref_pnorm_after_step, flat_pnorm_after_step) def test_sharded_flat_param(self): module = torch.nn.Sequential( torch.nn.Linear(10, 10, bias=False), torch.nn.ReLU(), torch.nn.Linear(10, 10, bias=False), torch.nn.ReLU(), torch.nn.Linear(10, 10, bias=False), torch.nn.ReLU(), ) params_to_flatten = list(module.parameters()) flat_module = FlattenParamsWrapper( module, params_to_flatten, torch.device("cuda"), self._get_default_config(), ) flat_param_handle = flat_module.handle def _test(kwargs, expected): """ Tests the subroutine ``_get_shard_metadata()`` that computes shard metadata based on start and end indices in the unsharded flattened parameter. We manually set the relevant attributes on the flattened parameter to be able to check the effect of ``_get_shard_metadata()`` via ``shard_metadata()`` since normally the attributes are set in ``init_shard_info()`` with the start and end indices fixed based on rank and world size. """ flat_param = flat_module.flat_param flat_param._shard_param_offsets, flat_param._shard_indices = \ flat_param_handle._get_shard_metadata(kwargs["start"], kwargs["end"]) self.assertEqual( flat_param_handle.shard_metadata(), expected, msg=f"{flat_param_handle.shard_metadata()}, {expected}", ) _test( kwargs={"start": 0, "end": 0}, expected=FlatParamShardMetadata( param_names=["0.weight"], param_shapes=[(10, 10)], param_numels=[100], param_offsets=[(0, 0)], ), ) _test( kwargs={"start": 0, "end": 50}, expected=FlatParamShardMetadata( param_names=["0.weight"], param_shapes=[(10, 10)], param_numels=[100], param_offsets=[(0, 50)], ), ) _test( kwargs={"start": 0, "end": 99}, expected=FlatParamShardMetadata( param_names=["0.weight"], param_shapes=[(10, 10)], param_numels=[100], param_offsets=[(0, 99)], ), ) _test( kwargs={"start": 50, "end": 149}, expected=FlatParamShardMetadata( param_names=["0.weight", "2.weight"], param_shapes=[(10, 10), (10, 10)], param_numels=[100, 100], param_offsets=[(50, 99), (0, 49)], ), ) _test( kwargs={"start": 50, "end": 199}, expected=FlatParamShardMetadata( param_names=["0.weight", "2.weight"], param_shapes=[(10, 10), (10, 10)], param_numels=[100, 100], param_offsets=[(50, 99), (0, 99)], ), ) _test( kwargs={"start": 99, "end": 199}, expected=FlatParamShardMetadata( param_names=["0.weight", "2.weight"], param_shapes=[(10, 10), (10, 10)], param_numels=[100, 100], param_offsets=[(99, 99), (0, 99)], ), ) _test( kwargs={"start": 100, "end": 199}, expected=FlatParamShardMetadata( param_names=["2.weight"], param_shapes=[(10, 10)], param_numels=[100], param_offsets=[(0, 99)], ), ) _test( kwargs={"start": 100, "end": 299}, expected=FlatParamShardMetadata( param_names=["2.weight", "4.weight"], param_shapes=[(10, 10), (10, 10)], param_numels=[100, 100], param_offsets=[(0, 99), (0, 99)], ), ) _test( kwargs={"start": 100, "end": 1000}, expected=FlatParamShardMetadata( param_names=["2.weight", "4.weight"], param_shapes=[(10, 10), (10, 10)], param_numels=[100, 100], param_offsets=[(0, 99), (0, 99)], ), ) _test( kwargs={"start": 299, "end": 299}, expected=FlatParamShardMetadata( param_names=["4.weight"], param_shapes=[(10, 10)], param_numels=[100], param_offsets=[(99, 99)], ), ) @unittest.skipIf(not torch.cuda.is_available(), "test requires a GPU") class TestFlattenParamsCUDA(TestFlattenParams): def _get_transformer(self, seed=0): module = super()._get_transformer(seed=seed) return module.cuda() @unittest.skipIf(not torch.cuda.is_available(), "test requires a GPU") class TestFlattenParamsCUDAHalf(TestFlattenParams): def _get_transformer(self, seed=0): module = super()._get_transformer(seed=seed) return module.cuda().half() if __name__ == "__main__": run_tests()