# Owner(s): ["oncall: distributed"] import copy import sys from collections import OrderedDict from typing import Dict, List, Optional, Tuple import torch from torch import distributed as dist from torch.distributed._tensor import ( DeviceMesh, distribute_module, DTensor, init_device_mesh, Replicate, Shard, ) from torch.distributed.fsdp.fully_sharded_data_parallel import ( CPUOffload, FullyShardedDataParallel as FSDP, ShardingStrategy, ) from torch.distributed.tensor.debug import CommDebugMode from torch.distributed.tensor.parallel import ( ColwiseParallel, parallelize_module, RowwiseParallel, ) from torch.testing._internal.common_distributed import skip_if_lt_x_gpu from torch.testing._internal.common_fsdp import FSDPTest from torch.testing._internal.common_utils import ( instantiate_parametrized_tests, run_tests, TEST_WITH_DEV_DBG_ASAN, ) from torch.testing._internal.distributed._tensor.common_dtensor import ( MLPModule, RMSNormPython, ) if not dist.is_available(): print("Distributed not available, skipping tests", file=sys.stderr) sys.exit(0) if TEST_WITH_DEV_DBG_ASAN: print( "Skip dev-asan as torch + multiprocessing spawn have known issues", file=sys.stderr, ) sys.exit(0) class SimpleModel(torch.nn.Module): def __init__(self) -> None: super().__init__() self.net1 = torch.nn.Linear(5, 8) self.relu = torch.nn.ReLU() self.net2 = torch.nn.Linear(8, 4) self.net3 = torch.nn.Linear(4, 12) def forward(self, x): return self.net3(self.net2(self.relu(self.net1(x)))) @staticmethod def get_sharded_param_names() -> List[str]: return ["net1.weight", "net1.bias", "net2.weight"] @staticmethod def get_non_sharded_param_names() -> List[str]: return ["net3.weight", "net3.bias"] def distribute_rmsnorm(module, device_mesh): def prepare_input_fn(mod, inputs, device_mesh): shard_tensor = DTensor.from_local(inputs[0], device_mesh, [Shard(0)]) return shard_tensor def prepare_output_fn(mod, outputs, device_mesh): return outputs.to_local() return distribute_module( module, device_mesh, input_fn=prepare_input_fn, output_fn=prepare_output_fn ) class TestTPFSDPIntegration(FSDPTest): def _get_params_and_sharding_info( self, model: SimpleModel, sharded_param_names: List[str], tensor_parallel_size: int, ) -> Tuple[Dict[str, int], Dict[str, Tuple[torch.Size, int]]]: """ """ assert ( type(model) is SimpleModel ), "Expects a `SimpleModel` since the sharding cases on the model definition" param_name_to_numel = OrderedDict() param_name_to_sharding_info = OrderedDict() for param_name, param in model.named_parameters(): if param_name not in sharded_param_names: param_name_to_numel[param_name] = param.numel() else: param_name_to_numel[param_name] = param.numel() // tensor_parallel_size param_name_to_sharding_info[param_name] = ( param.size(), 0 if "net1" in param_name else 1, ) return param_name_to_numel, param_name_to_sharding_info def _get_sub_pgs(self, tensor_parallel_size: int): """ Generates TP and FSDP subprocess groups. ``tensor_parallel_size`` gives the TP process group size. For example, if the global world size is 8 and the tensor parallel size is 2, then this creates: - 4 TP subprocess groups: [0, 1], [2, 3], [4, 5], [6, 7] - 2 FSDP subprocess groups: [0, 2, 4, 6], [1, 3, 5, 7] """ # 2-D mesh is [dp, tp] twod_mesh = DeviceMesh( device_type="cuda", mesh=torch.arange(0, self.world_size).view(-1, tensor_parallel_size), ) fsdp_pg = twod_mesh.get_group(mesh_dim=0) tp_pg = twod_mesh.get_group(mesh_dim=1) return twod_mesh, fsdp_pg, tp_pg def _sync_tp_grads( self, tp_fsdp_model: FSDP, tp_pg: dist.ProcessGroup, param_name_to_numel: Dict[str, int], non_sharded_param_names: List[str], ) -> None: """ Syncs the tensor parallel parameters' gradients following the data parallel paradigm where gradients are averaged over ranks (in this case, the ones in the tensor parallel process group). """ tp_world_size = tp_pg.size() fsdp_world_size = self.world_size // tp_world_size assert ( type(tp_fsdp_model) is FSDP and len([m for m in tp_fsdp_model.modules() if type(m) is FSDP]) == 1 ), ( "The following logic assumes a single top-level-only FSDP wrapping " "the model with TP already applied" ) for flat_param in tp_fsdp_model.params: splits = tuple(param_name_to_numel.values()) # Create a mask over the gradient elements to manually reduce unsharded_size = torch.Size([flat_param.numel() * fsdp_world_size]) unsharded_zeros = torch.zeros(unsharded_size, device=flat_param.device) per_param_masks = unsharded_zeros.split(splits) for param_idx, param_name in enumerate( param_name_to_numel.keys() ): # assumes fixed order if param_name not in non_sharded_param_names: per_param_masks[param_idx][:] = 1 unsharded_mask = ( torch.cat(per_param_masks).contiguous().type(torch.BoolTensor) ) sharded_mask = unsharded_mask.chunk(fsdp_world_size)[ self.rank // tp_world_size ] grad_device = flat_param.grad.device grad = flat_param.grad.detach().clone().cuda(self.rank) dist.all_reduce(grad, op=dist.ReduceOp.SUM, group=tp_pg) grad = grad.to(grad_device) flat_param.grad[~sharded_mask] = grad[~sharded_mask] # Average *all* gradient elements to match the FSDP only semantics flat_param.grad /= tp_world_size def _get_grads_as_flattened( self, model: FSDP, uses_tp: bool, param_name_to_numel: Dict[str, int], param_name_to_sharding_info: Dict[str, Tuple[torch.Size, int]], tp_pg: Optional[dist.ProcessGroup], fsdp_pg: Optional[dist.ProcessGroup], sharded_param_names: Optional[List[str]], ) -> torch.Tensor: """ Returns all unsharded gradients as a single flattened tensor. This returns the same value on all ranks. """ local_grads_as_flattened = ( torch.cat( [ torch.flatten(param.grad) if param.grad is not None else torch.zeros_like(torch.flatten(param)) for param in model.parameters() ] ) .contiguous() .cuda(self.rank) ) all_grads_as_flattened = torch.cat( [torch.empty_like(local_grads_as_flattened) for _ in range(fsdp_pg.size())] ).contiguous() dist.all_gather_into_tensor( all_grads_as_flattened, local_grads_as_flattened, group=fsdp_pg ) if not uses_tp: return all_grads_as_flattened splits = tuple(param_name_to_numel.values()) all_grads_per_param = list(all_grads_as_flattened.split(splits)) for param_idx, param_name in enumerate( param_name_to_numel.keys() ): # assumes fixed order if param_name in sharded_param_names: local_tensor_size = list(param_name_to_sharding_info[param_name][0]) sharding_dim = param_name_to_sharding_info[param_name][1] local_tensor_size[sharding_dim] //= tp_pg.size() local_tensor = all_grads_per_param[param_idx].view(*local_tensor_size) local_tensors = [ torch.empty_like(local_tensor) for _ in range(tp_pg.size()) ] dist.all_gather(local_tensors, local_tensor, group=tp_pg) all_grads_per_param[param_idx] = torch.cat( local_tensors, dim=sharding_dim ).reshape(-1) return torch.cat(all_grads_per_param).contiguous() @skip_if_lt_x_gpu(4) def test_fsdp_tp_integration(self): self.run_subtests( { "cpu_offload": [ CPUOffload(offload_params=False), CPUOffload(offload_params=True), ], "sharding_strategy": [None, ShardingStrategy.SHARD_GRAD_OP], "use_orig_params": [False, True], }, self._test_fsdp_tp_integration, ) def _test_fsdp_tp_integration( self, cpu_offload, sharding_strategy, use_orig_params ): """ Tests training for TP + FSDP integration by comparing an FSDP-only model with a TP + FSDP model. """ tensor_parallel_size = 2 LR = 3e-5 torch.manual_seed(0) model = SimpleModel().cuda(self.rank) tp_fsdp_model = copy.deepcopy(model) sharded_param_names = SimpleModel.get_sharded_param_names() non_sharded_param_names = SimpleModel.get_non_sharded_param_names() ( param_name_to_numel, param_name_to_sharding_info, ) = self._get_params_and_sharding_info( model, sharded_param_names, tensor_parallel_size, ) input_seed = self.rank torch.manual_seed(input_seed + 1) inp_size = [2, 3, 5] inp = torch.rand(*inp_size).cuda(self.rank) self.assertEqual(model(inp), tp_fsdp_model(inp)) # sanity check mesh_1d = init_device_mesh("cuda", (self.world_size,)) fsdp_model = FSDP( model, cpu_offload=cpu_offload, device_mesh=mesh_1d, sharding_strategy=sharding_strategy, use_orig_params=use_orig_params, ) mesh_2d = init_device_mesh( "cuda", (self.world_size // tensor_parallel_size, tensor_parallel_size), mesh_dim_names=["dp", "tp"], ) # Shard with TP and then wrap with FSDP sequence_parallelize_plan = { "net1": ColwiseParallel(input_layouts=Shard(0)), "net2": RowwiseParallel(output_layouts=Shard(0)), } tp_fsdp_model = parallelize_module( tp_fsdp_model, mesh_2d["tp"], sequence_parallelize_plan, ) tp_pg = mesh_2d["tp"].get_group(mesh_dim=0) assert isinstance(tp_fsdp_model.net1.weight, DTensor) assert isinstance(tp_fsdp_model.net2.weight, DTensor) tp_fsdp_model = FSDP( tp_fsdp_model, cpu_offload=cpu_offload, device_mesh=mesh_2d["dp"], sharding_strategy=sharding_strategy, use_orig_params=use_orig_params, ) fsdp_pg = mesh_2d["dp"].get_group(mesh_dim=0) # Check the forward by checking output equality fsdp_out = fsdp_model(inp) tp_fsdp_out = tp_fsdp_model(inp) self.assertEqual(fsdp_out, tp_fsdp_out) # Check the backward by checking gradient equality fsdp_out.sum().backward() tp_fsdp_out.sum().backward() self._sync_tp_grads( tp_fsdp_model, tp_pg, param_name_to_numel, non_sharded_param_names, ) model_grads = self._get_grads_as_flattened( fsdp_model, False, param_name_to_numel, param_name_to_sharding_info, None, self.process_group, None, ) model_tp_grads = self._get_grads_as_flattened( tp_fsdp_model, True, param_name_to_numel, param_name_to_sharding_info, tp_pg, fsdp_pg, sharded_param_names, ) self.assertEqual(model_grads, model_tp_grads) # Check the optimizer step by performing a second forward pass fsdp_optim = torch.optim.SGD(fsdp_model.parameters(), lr=LR) tp_fsdp_optim = torch.optim.SGD(tp_fsdp_model.parameters(), lr=LR) fsdp_optim.step() tp_fsdp_optim.step() torch.manual_seed(input_seed + 16) inp = torch.rand(*inp_size).cuda(self.rank) fsdp_out = fsdp_model(inp) tp_fsdp_out = tp_fsdp_model(inp) self.assertEqual(fsdp_out, tp_fsdp_out) @skip_if_lt_x_gpu(4) def test_fsdp_tp_extension_grad(self): """ Tests TP + FSDP extension with correct gradient (i.e. no ACT) """ mesh_2d = init_device_mesh( "cuda", (self.world_size // 2, 2), mesh_dim_names=["dp", "tp"] ) class TestModel(torch.nn.Module): def __init__(self) -> None: super().__init__() self.mlp = MLPModule("cuda") self.mlp_norm = RMSNormPython(10) def forward(self, x): return self.mlp(self.mlp_norm(x)) model = TestModel().cuda(self.rank) # Shard with TP and test gradient tp_mesh = mesh_2d["tp"] tp_model = parallelize_module( model, tp_mesh, { "mlp.net1": ColwiseParallel(input_layouts=Shard(0)), "mlp.net2": RowwiseParallel(output_layouts=Shard(0)), }, ) distribute_rmsnorm(tp_model.mlp_norm, tp_mesh) fsdp_2d_model = FSDP(tp_model, device_mesh=mesh_2d["dp"]) comm_mode = CommDebugMode() with comm_mode: fsdp_2d_model(torch.rand(2, 10).cuda(self.rank)).sum().backward() funcol = torch.ops.c10d_functional c10d_ops = torch.ops.c10d comm_counts = comm_mode.get_comm_counts() self.assertEqual(comm_mode.get_total_counts(), 7) # TP comms self.assertEqual(comm_counts[funcol.reduce_scatter_tensor], 2) self.assertEqual(comm_counts[funcol.all_gather_into_tensor], 2) self.assertEqual(comm_counts[funcol.all_reduce], 1) # FSDP comms self.assertEqual(comm_counts[c10d_ops._allgather_base_], 1) self.assertEqual(comm_counts[c10d_ops._reduce_scatter_base_], 1) grads = [p.grad for p in fsdp_2d_model.parameters() if p.grad is not None] for grad in grads: self.assertFalse(grad.isnan().any().item()) @skip_if_lt_x_gpu(4) def test_fsdp_tp_sync_module_state(self): mesh_2d = init_device_mesh( "cuda", (self.world_size // 2, 2), mesh_dim_names=["dp", "tp"] ) tp_mesh = mesh_2d["tp"] dp_mesh = mesh_2d["dp"] # set random seed for each rank torch.manual_seed(mesh_2d.get_rank()) class TestModel(torch.nn.Module): def __init__(self) -> None: super().__init__() replicated_dt = DTensor.from_local( torch.randn(8, 8), tp_mesh, [Replicate()], run_check=False ) replicated_buffer_dt = DTensor.from_local( torch.randn(8, 8), tp_mesh, [Replicate()], run_check=False ) self.param = torch.nn.Parameter(replicated_dt) self.buf = torch.nn.Buffer(replicated_buffer_dt) def forward(self, x): return self.param + self.buffer + 1 model = TestModel() def assert_local_shard_across_ranks(local_tensor, group, check_equal=True): gathered_tensors = [ torch.empty_like(local_tensor) for _ in range(group.size()) ] dist.all_gather(gathered_tensors, local_tensor, group=group) # on dp mesh dim local tensor does not equal tensor_to_compare = gathered_tensors[0] for tensor in gathered_tensors[1:]: if check_equal: self.assertTrue(torch.equal(tensor, tensor_to_compare)) else: self.assertFalse(torch.equal(tensor, tensor_to_compare)) dp_group = dp_mesh.get_group() # check on dp mesh dim param local tensor does not equal local_param = model.param.to_local() assert_local_shard_across_ranks(local_param, dp_group, check_equal=False) # check on dp mesh dim buffer local tensor does not equal local_buf = model.buf.to_local() assert_local_shard_across_ranks(local_buf, dp_group, check_equal=False) # wrap with fsdp sync param should sync dp mesh dim fsdp_mod = FSDP(model, device_mesh=dp_mesh, sync_module_states=True) with fsdp_mod.summon_full_params(fsdp_mod): # on dp mesh dim local param does equal after sync_module_states local_param = fsdp_mod.param.to_local() assert_local_shard_across_ranks(local_param, dp_group, check_equal=True) # on dp mesh dim local buf does equal after sync_module_states local_buf = fsdp_mod.buf.to_local() assert_local_shard_across_ranks(local_buf, dp_group, check_equal=True) instantiate_parametrized_tests(TestTPFSDPIntegration) if __name__ == "__main__": run_tests()